[mlir][cf] Print message in cf.assert to LLVM lowering

[springerm]

Summary

The assert message was previously ignored. This change adds a new helper function printStr to the runtime library and calls it in case of a failed assertion.

Depends On D138576

[springerm]

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

97–100

Note: This (copying StringRef into uint8_t vector) is copied from LLVM: llvm/lib/IR/Constants.cpp (ConstantDataArray::getString).

[tpopp]

I'm leaving a few high level comments first. I think this abortOnFailedAssert=false case should be handled by a separate test pass instead.

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

81

If we do this, we need to limit the pass to only running on Modules or risk race conditions.

90

I would rename this to something like createPrintStr. Otherwise, I thought this was immediately printing something rather than generating a print statement.

132

This is convenient, but I don't think it's a great idea to provide this flag that changes the meanings of assertions in the primary populateControlFlowToLLVM. I think there should be a separate test pass that creates prints instead if you want this.

mlir/lib/ExecutionEngine/CRunnerUtils.cpp

50 ↗

(On Diff #477718)

I'm just explicitly mentioning that I don't know the consequences of creating this functionality when reviewing the rest of this code.

[springerm]

address comments

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

81

ConvertControlFlowToLLVM is already a Module pass.

[rriddle]

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

81

Only classes should be in anonymous namespaces, pull these out and mark static.

[springerm]

fix cmake

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

476 ↗

(On Diff #477788)

Can't you just call libc in the lowering directly?

476 ↗

(On Diff #477788)

That is: I'm concerned about any general lowering to LLVM that would depend on the runner support library right now.

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

132

I agree with this, proliferations of "options" is something I'm concerned about in general and better be avoided when we can.

[springerm]

address comments

[springerm]

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

132

Kept the flag here but it is no longer exposed to populateControlFlowToLLVM. There's a special populate just for this pattern and a new test pass.

[springerm]

rebase

[springerm]

minor update

[springerm]

rebase

[mehdi_amini]

mlir/include/mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h

34

set to false

mlir/test/Integration/Dialect/ControlFlow/assert.mlir

5

Do we need any utils here?

[springerm]

address comments

[mehdi_amini]

The commit message probably needs an update.

mlir/include/mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h

 
 namespace cf {
 /// Collect the patterns to convert from the ControlFlow dialect to LLVM. The
 /// conversion patterns capture the LLVMTypeConverter by reference meaning the
 /// references have to remain alive during the entire pattern lifetime.
 void populateControlFlowToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                  RewritePatternSet &patterns);
 
+/// Populate the cf.assert to LLVM conversion pattern. If `abortOnFailure` is
+/// set to false, the program execution continues when a condition is
+/// unsatisfied.
+void populateAssertToLLVMConversionPattern(LLVMTypeConverter &converter,
+                                           RewritePatternSet &patterns,
+                                           bool abortOnFailure = true);
+
 /// Creates a pass to convert the ControlFlow dialect into the LLVMIR dialect.
 std::unique_ptr<Pass> createConvertControlFlowToLLVMPass();
 } // namespace cf
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_CONTROLFLOWTOLLVM_CONTROLFLOWTOLLVM_H

mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h

 /// Helper functions to lookup or create the declaration for commonly used
 /// external C function calls. The list of functions provided here must be
 /// implemented separately (e.g. as part of a support runtime library or as part
 /// of the libc).
 LLVM::LLVMFuncOp lookupOrCreatePrintI64Fn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintU64Fn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintF32Fn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintF64Fn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreatePrintStrFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintOpenFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintCloseFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintCommaFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintNewlineFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreateMallocFn(ModuleOp moduleOp, Type indexType);
 LLVM::LLVMFuncOp lookupOrCreateAlignedAllocFn(ModuleOp moduleOp,
                                               Type indexType);
 LLVM::LLVMFuncOp lookupOrCreateFreeFn(ModuleOp moduleOp);

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

 #define GEN_PASS_DEF_CONVERTCONTROLFLOWTOLLVM
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 
 #define PASS_NAME "convert-cf-to-llvm"
 
+static std::string generateGlobalMsgSymbolName(ModuleOp moduleOp) {
+  std::string prefix = "assert_msg_";
+  int counter = 0;
+  while (moduleOp.lookupSymbol(prefix + std::to_string(counter)))
+    ++counter;
+  return prefix + std::to_string(counter);
+}
+
+/// Generate IR that prints the given string to stderr.
+static void createPrintMsg(OpBuilder &builder, Location loc, ModuleOp moduleOp,
+                           StringRef msg) {
+  auto ip = builder.saveInsertionPoint();
+  builder.setInsertionPointToStart(moduleOp.getBody());
+  MLIRContext *ctx = builder.getContext();
+
+  // Create a zero-terminated byte representation and allocate global symbol.
+  SmallVector<uint8_t> elementVals;
+  elementVals.append(msg.begin(), msg.end());
+  elementVals.push_back(0);
+  auto dataAttrType = RankedTensorType::get(
+      {static_cast<int64_t>(elementVals.size())}, builder.getI8Type());
+  auto dataAttr =
+      DenseElementsAttr::get(dataAttrType, llvm::makeArrayRef(elementVals));
+  auto arrayTy =
+      LLVM::LLVMArrayType::get(IntegerType::get(ctx, 8), elementVals.size());
+  std::string symbolName = generateGlobalMsgSymbolName(moduleOp);
+  auto globalOp = builder.create<LLVM::GlobalOp>(
+      loc, arrayTy, /*constant=*/true, LLVM::Linkage::Private, symbolName,
+      dataAttr);
+
+  // Emit call to `printStr` in runtime library.
+  builder.restoreInsertionPoint(ip);
+  auto msgAddr = builder.create<LLVM::AddressOfOp>(
+      loc, LLVM::LLVMPointerType::get(arrayTy), globalOp.getName());
+  SmallVector<LLVM::GEPArg> indices(1, 0);
+  Value gep = builder.create<LLVM::GEPOp>(
+      loc, LLVM::LLVMPointerType::get(builder.getI8Type()), msgAddr, indices);
+  Operation *printer = LLVM::lookupOrCreatePrintStrFn(moduleOp);
+  builder.create<LLVM::CallOp>(loc, TypeRange(), SymbolRefAttr::get(printer),
+                               gep);
+}
+
 namespace {
 /// Lower `cf.assert`. The default lowering calls the `abort` function if the
 /// assertion is violated and has no effect otherwise. The failure message is
 /// ignored by the default lowering but should be propagated by any custom
 /// lowering.
 struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> {
-  using ConvertOpToLLVMPattern<cf::AssertOp>::ConvertOpToLLVMPattern;
+  explicit AssertOpLowering(LLVMTypeConverter &typeConverter,
+                            bool abortOnFailedAssert = true)
+      : ConvertOpToLLVMPattern<cf::AssertOp>(typeConverter, /*benefit=*/1),
+        abortOnFailedAssert(abortOnFailedAssert) {}
 
   LogicalResult
   matchAndRewrite(cf::AssertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
+    auto module = op->getParentOfType<ModuleOp>();
 
+    // Split block at `assert` operation.
+    Block *opBlock = rewriter.getInsertionBlock();
+    auto opPosition = rewriter.getInsertionPoint();
+    Block *continuationBlock = rewriter.splitBlock(opBlock, opPosition);
+
+    // Failed block: Generate IR to print the message and call `abort`.
+    Block *failureBlock = rewriter.createBlock(opBlock->getParent());
+    createPrintMsg(rewriter, loc, module, op.getMsg());
+    if (abortOnFailedAssert) {
       // Insert the `abort` declaration if necessary.
-    auto module = op->getParentOfType<ModuleOp>();
       auto abortFunc = module.lookupSymbol<LLVM::LLVMFuncOp>("abort");
       if (!abortFunc) {
         OpBuilder::InsertionGuard guard(rewriter);
         rewriter.setInsertionPointToStart(module.getBody());
         auto abortFuncTy = LLVM::LLVMFunctionType::get(getVoidType(), {});
         abortFunc = rewriter.create<LLVM::LLVMFuncOp>(rewriter.getUnknownLoc(),
                                                       "abort", abortFuncTy);
       }
-
-    // Split block at `assert` operation.
-    Block *opBlock = rewriter.getInsertionBlock();
-    auto opPosition = rewriter.getInsertionPoint();
-    Block *continuationBlock = rewriter.splitBlock(opBlock, opPosition);
-
-    // Generate IR to call `abort`.
-    Block *failureBlock = rewriter.createBlock(opBlock->getParent());
       rewriter.create<LLVM::CallOp>(loc, abortFunc, std::nullopt);
       rewriter.create<LLVM::UnreachableOp>(loc);
+    } else {
+      rewriter.create<LLVM::BrOp>(loc, ValueRange(), continuationBlock);
+    }
 
     // Generate assertion test.
     rewriter.setInsertionPointToEnd(opBlock);
     rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
         op, adaptor.getArg(), continuationBlock, failureBlock);
 
     return success();
   }
+
+private:
+  /// If set to `false`, messages are printed but program execution continues.
+  /// This is useful for testing asserts.
+  bool abortOnFailedAssert = true;
 };
 
 /// The cf->LLVM lowerings for branching ops require that the blocks they jump
 /// to first have updated types which should be handled by a pattern operating
 /// on the parent op.
 static LogicalResult verifyMatchingValues(ConversionPatternRewriter &rewriter,
                                           ValueRange operands,
                                           ValueRange blockArgs, Location loc,
void mlir::cf::populateControlFlowToLLVMConversionPatterns(
   patterns.add<
       AssertOpLowering,
       BranchOpLowering,
       CondBranchOpLowering,
       SwitchOpLowering>(converter);
   // clang-format on
 }
 
+void mlir::cf::populateAssertToLLVMConversionPattern(
+    LLVMTypeConverter &converter, RewritePatternSet &patterns,
+    bool abortOnFailure) {
+  patterns.add<AssertOpLowering>(converter, abortOnFailure);
+}
+
 //===----------------------------------------------------------------------===//
 // Pass Definition
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A pass converting MLIR operations into the LLVM IR dialect.
 struct ConvertControlFlowToLLVM
     : public impl::ConvertControlFlowToLLVMBase<ConvertControlFlowToLLVM> {

mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp

 /// Helper functions to lookup or create the declaration for commonly used
 /// external C function calls. The list of functions provided here must be
 /// implemented separately (e.g. as  part of a support runtime library or as
 /// part of the libc).
 static constexpr llvm::StringRef kPrintI64 = "printI64";
 static constexpr llvm::StringRef kPrintU64 = "printU64";
 static constexpr llvm::StringRef kPrintF32 = "printF32";
 static constexpr llvm::StringRef kPrintF64 = "printF64";
+static constexpr llvm::StringRef kPrintStr = "puts";
 static constexpr llvm::StringRef kPrintOpen = "printOpen";
 static constexpr llvm::StringRef kPrintClose = "printClose";
 static constexpr llvm::StringRef kPrintComma = "printComma";
 static constexpr llvm::StringRef kPrintNewline = "printNewline";
 static constexpr llvm::StringRef kMalloc = "malloc";
 static constexpr llvm::StringRef kAlignedAlloc = "aligned_alloc";
 static constexpr llvm::StringRef kFree = "free";
 static constexpr llvm::StringRef kGenericAlloc = "_mlir_memref_to_llvm_alloc";

 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintF64Fn(ModuleOp moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintF64,
                           Float64Type::get(moduleOp->getContext()),
                           LLVM::LLVMVoidType::get(moduleOp->getContext()));
 }
 
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintStrFn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(
+      moduleOp, kPrintStr,
+      LLVM::LLVMPointerType::get(IntegerType::get(moduleOp->getContext(), 8)),
+      LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintOpenFn(ModuleOp moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintOpen, {},
                           LLVM::LLVMVoidType::get(moduleOp->getContext()));
 }
 
 LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintCloseFn(ModuleOp moduleOp) {
   return lookupOrCreateFn(moduleOp, kPrintClose, {},
                           LLVM::LLVMVoidType::get(moduleOp->getContext()));

mlir/test/Integration/Dialect/ControlFlow/assert.mlir

+// RUN: mlir-opt %s -test-cf-assert \
+// RUN:     -convert-func-to-llvm | \
+// RUN: mlir-cpu-runner -e main -entry-point-result=void | \
+// RUN: FileCheck %s
+
+func.func @main() {
+  %a = arith.constant 0 : i1
+  %b = arith.constant 1 : i1
+  // CHECK: assertion foo
+  cf.assert %a, "assertion foo"
+  // CHECK-NOT: assertion bar
+  cf.assert %b, "assertion bar"
+  return
+}

mlir/test/lib/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
 add_subdirectory(Arith)
 add_subdirectory(Bufferization)
+add_subdirectory(ControlFlow)
 add_subdirectory(DLTI)
 add_subdirectory(Func)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVM)
 add_subdirectory(Math)
 add_subdirectory(MemRef)
 add_subdirectory(NVGPU)

mlir/test/lib/Dialect/ControlFlow/CMakeLists.txt

+# Exclude tests from libMLIR.so
+add_mlir_library(MLIRArithTestPasses
+  TestAssert.cpp
+
+  EXCLUDE_FROM_LIBMLIR
+
+  LINK_LIBS PUBLIC
+  MLIRControlFlowToLLVM
+  MLIRFuncDialect
+  MLIRLLVMCommonConversion
+  MLIRLLVMDialect
+  MLIRPass
+  MLIRTransforms
+)

mlir/test/lib/Dialect/ControlFlow/TestAssert.cpp

+//===- TestAssert.cpp - Test cf.assert Lowering  ----------------*- c++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass for integration testing of wide integer
+// emulation patterns. Applies conversion patterns only to functions whose
+// names start with a specified prefix.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
+#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+struct TestAssertPass
+    : public PassWrapper<TestAssertPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAssertPass)
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<cf::ControlFlowDialect, LLVM::LLVMDialect>();
+  }
+  StringRef getArgument() const final { return "test-cf-assert"; }
+  StringRef getDescription() const final {
+    return "Function pass to test cf.assert lowering to LLVM without abort";
+  }
+
+  void runOnOperation() override {
+    LLVMConversionTarget target(getContext());
+    RewritePatternSet patterns(&getContext());
+
+    LLVMTypeConverter converter(&getContext());
+    mlir::cf::populateAssertToLLVMConversionPattern(converter, patterns,
+                                                    /*abortOnFailure=*/false);
+
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+namespace mlir::test {
+void registerTestCfAssertPass() { PassRegistration<TestAssertPass>(); }
+} // namespace mlir::test

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerMemRefBoundCheck();
 void registerPatternsTestPass();
 void registerSimpleParametricTilingPass();
 void registerTestAffineLoopParametricTilingPass();
 void registerTestArithEmulateWideIntPass();
 void registerTestAliasAnalysisPass();
 void registerTestBuiltinAttributeInterfaces();
 void registerTestCallGraphPass();
+void registerTestCfAssertPass();
 void registerTestConstantFold();
 void registerTestControlFlowSink();
 void registerTestGpuSerializeToCubinPass();
 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataLayoutPropagation();
 void registerTestDataLayoutQuery();
 void registerTestDeadCodeAnalysisPass();
 void registerTestDecomposeCallGraphTypes();
void registerTestPasses() {
   mlir::test::registerMemRefBoundCheck();
   mlir::test::registerPatternsTestPass();
   mlir::test::registerSimpleParametricTilingPass();
   mlir::test::registerTestAffineLoopParametricTilingPass();
   mlir::test::registerTestAliasAnalysisPass();
   mlir::test::registerTestArithEmulateWideIntPass();
   mlir::test::registerTestBuiltinAttributeInterfaces();
   mlir::test::registerTestCallGraphPass();
+  mlir::test::registerTestCfAssertPass();
   mlir::test::registerTestConstantFold();
   mlir::test::registerTestControlFlowSink();
   mlir::test::registerTestDiagnosticsPass();
   mlir::test::registerTestDialectConversionPasses();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":Support",
         ":Transforms",
         "//llvm:AllTargetsCodeGens",
         "//llvm:Support",
         "//mlir/test:TestAffine",
         "//mlir/test:TestAnalysis",
         "//mlir/test:TestArith",
         "//mlir/test:TestBufferization",
+        "//mlir/test:TestControlFlow",
         "//mlir/test:TestDLTI",
         "//mlir/test:TestDialect",
         "//mlir/test:TestFunc",
         "//mlir/test:TestFuncToLLVM",
         "//mlir/test:TestGPU",
         "//mlir/test:TestIR",
         "//mlir/test:TestLLVM",
         "//mlir/test:TestLinalg",

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

deps = [
         "//mlir:BufferizationDialect",
         "//mlir:BufferizationTransforms",
         "//mlir:IR",
         "//mlir:Pass",
     ],
 )
 
 cc_library(
+    name = "TestControlFlow",
+    srcs = glob(["lib/Dialect/ControlFlow/*.cpp"]),
+    includes = ["lib/Dialect/Test"],
+    deps = [
+        "//mlir:ControlFlowDialect",
+        "//mlir:ControlFlowToLLVM",
+        "//mlir:FuncDialect",
+        "//mlir:LLVMCommonConversion",
+        "//mlir:LLVMDialect",
+        "//mlir:Pass",
+        "//mlir:Transforms",
+    ],
+)
+
+cc_library(
     name = "TestShapeDialect",
     srcs = [
         "lib/Dialect/Shape/TestShapeFunctions.cpp",
         "lib/Dialect/Shape/TestShapeMappingAnalysis.cpp",
     ],
     deps = [
         "//llvm:Support",
         "//mlir:FuncDialect",