[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

55–58

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

39

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

48

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

133

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

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

39

ConvertControlFlowToLLVM is already a Module pass.

[rriddle]

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

39

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

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

476

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

133

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

133

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

39

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

 
 #define GEN_PASS_DECL_CONVERTCONTROLFLOWTOLLVM
 #include "mlir/Conversion/Passes.h.inc"
 
 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);
+///
+/// If `abortOnFailedAssert` is false, messages of failed assertions are
+/// printed, but program execution continues. This is useful for testing
+/// asserts.
+void populateControlFlowToLLVMConversionPatterns(
+    LLVMTypeConverter &converter, RewritePatternSet &patterns,
+    bool abortOnFailedAssert = 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/Conversion/Passes.td

let description = [{
     already present in the IR will be kept as is.
   }];
   let constructor = "mlir::cf::createConvertControlFlowToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,
+    Option<"abortOnFailedAssert", "abort-on-failed-assert", "bool",
+           /*default=*/"true",
+           "Abort program when an assertion is failing">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // ControlFlowToSPIRV
 //===----------------------------------------------------------------------===//
 
 def ConvertControlFlowToSPIRV : Pass<"convert-cf-to-spirv"> {

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/include/mlir/ExecutionEngine/CRunnerUtils.h

 extern "C" MLIR_CRUNNERUTILS_EXPORT void printF32(float f);
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printF64(double d);
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printOpen();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printClose();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printComma();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printNewline();
 
 //===----------------------------------------------------------------------===//
+// Small runtime support library for debug assertions.
+//===----------------------------------------------------------------------===//
+extern "C" MLIR_CRUNNERUTILS_EXPORT void printStr(int8_t *str);
+
+//===----------------------------------------------------------------------===//
 // Small runtime support library for timing execution and printing GFLOPS
 //===----------------------------------------------------------------------===//
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printFlops(double flops);
 extern "C" MLIR_CRUNNERUTILS_EXPORT double rtclock();
 
 #endif // MLIR_EXECUTIONENGINE_CRUNNERUTILS_H

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 
 #define PASS_NAME "convert-cf-to-llvm"
 
 namespace {
+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.
+void printMsg(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('\n');
+  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);
+}
+
 /// 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());
+    printMsg(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, llvm::None);
       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,
rewriter.replaceOpWithNewOp<LLVM::SwitchOp>(
         op, adaptor.getOperands(), op->getSuccessors(), op->getAttrs());
     return success();
   }
 };
 
 } // namespace
 
 void mlir::cf::populateControlFlowToLLVMConversionPatterns(
-    LLVMTypeConverter &converter, RewritePatternSet &patterns) {
+    LLVMTypeConverter &converter, RewritePatternSet &patterns,
+    bool abortOnFailedAssert) {
   // clang-format off
   patterns.add<
-      AssertOpLowering,
       BranchOpLowering,
       CondBranchOpLowering,
       SwitchOpLowering>(converter);
   // clang-format on
+  patterns.add<AssertOpLowering>(converter, abortOnFailedAssert);
 }
 
 //===----------------------------------------------------------------------===//
 // Pass Definition
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A pass converting MLIR operations into the LLVM IR dialect.
 struct ConvertControlFlowToLLVM
     : public impl::ConvertControlFlowToLLVMBase<ConvertControlFlowToLLVM> {
   ConvertControlFlowToLLVM() = default;
 
   /// Run the dialect converter on the module.
   void runOnOperation() override {
     LLVMConversionTarget target(getContext());
     RewritePatternSet patterns(&getContext());
 
     LowerToLLVMOptions options(&getContext());
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
 
     LLVMTypeConverter converter(&getContext(), options);
-    mlir::cf::populateControlFlowToLLVMConversionPatterns(converter, patterns);
+    mlir::cf::populateControlFlowToLLVMConversionPatterns(converter, patterns,
+                                                          abortOnFailedAssert);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::cf::createConvertControlFlowToLLVMPass() {
   return std::make_unique<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 = "printStr";
 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/lib/ExecutionEngine/CRunnerUtils.cpp

 extern "C" void printU64(uint64_t u) { fprintf(stdout, "%" PRIu64, u); }
 extern "C" void printF32(float f) { fprintf(stdout, "%g", f); }
 extern "C" void printF64(double d) { fprintf(stdout, "%lg", d); }
 extern "C" void printOpen() { fputs("( ", stdout); }
 extern "C" void printClose() { fputs(" )", stdout); }
 extern "C" void printComma() { fputs(", ", stdout); }
 extern "C" void printNewline() { fputc('\n', stdout); }
 
+// For debug assertions only: Print the given string.
+extern "C" void printStr(int8_t *str) {
+  fputs(reinterpret_cast<char *>(str), stderr);
+}
+
 extern "C" void memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
                            UnrankedMemRefType<char> *dstArg) {
   DynamicMemRefType<char> src(*srcArg);
   DynamicMemRefType<char> dst(*dstArg);
 
   int64_t rank = src.rank;
   MLIR_MSAN_MEMORY_IS_INITIALIZED(src.sizes, rank * sizeof(int64_t));
 

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

+// RUN: mlir-opt %s -convert-cf-to-llvm="abort-on-failed-assert=0" \
+// RUN:     -convert-func-to-llvm | \
+// RUN: mlir-cpu-runner -e main -entry-point-result=void \
+// RUN:     -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext 2>&1 | \
+// 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
+}