[MLIR][Standard] Add default lowering for `assert`

[frgossen]

Summary

The default lowering of assert calls abort in case the assertion is
violated. The failure message is ignored but should be used by custom lowerings
that can assume more about their environment.

[rriddle]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1429

nit: remove the newline here.

1452

nit: Can you use Block instead of auto here and below?

1463

Why not use LLVM::CondBranchOp directly here?

1463

This seems incorrect. I would expect the continuation block to be destination if the condition is true. (The true block is first is cond_br)

[bondhugula]

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

86

The targets are switched.

[frgossen]

Address comments

[silvas]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1443

Mutating the module like this is not allowed. This conversion pattern could be applied inside a function pass which will be running in parallel, and this mutation of the module body will trigger a race condition very easily. It's likely that if you craft a test case that simply has two std.assert ops, each in a different function, you will cause a ThreadSanitizer failure.

What is currently being done for creating the malloc/free declarations?

[silvas]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1443

Or do we have some sort of guarantee that this is running single-threaded on the whole module?

[silvas]

Thanks for adding this!

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1443

I see that we are already doing this kind of "insert into the mdoule from deep inside a conversion pattern" in a number of places. Sigh. I guess we are assuming that these patterns are being applied serially for now :/

Sorry for the noise.

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1443

Conversion to LLVM is a module pass. It replaces standard functions with llvm functions, so it cannot be a function pass anyway. If it becomes a performance issue, we can potentially split it into two parts: rewrite functions serially, then rewrite function bodies in parallel. This will need some infrastructural support to work around type mismatches (many patterns rely on knowing the original type) and to avoid splitting this into two passes.

[frgossen]

@bondhugula, can I land this?

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 using UnsignedDivIOpLowering =
     VectorConvertToLLVMPattern<UnsignedDivIOp, LLVM::UDivOp>;
 using UnsignedRemIOpLowering =
     VectorConvertToLLVMPattern<UnsignedRemIOp, LLVM::URemOp>;
 using UnsignedShiftRightOpLowering =
     OneToOneConvertToLLVMPattern<UnsignedShiftRightOp, LLVM::LShrOp>;
 using XOrOpLowering = VectorConvertToLLVMPattern<XOrOp, LLVM::XOrOp>;
 
+/// Lower `std.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<AssertOp> {
+  using ConvertOpToLLVMPattern<AssertOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op->getLoc();
+    AssertOp::Adaptor transformed(operands);
+
+    // 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::LLVMType::getFunctionTy(getVoidType(), {}, /*isVarArg=*/false);
+      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);
+
+    // Generate assertion test.
+    rewriter.setInsertionPointToEnd(opBlock);
+    rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
+        op, transformed.arg(), continuationBlock, failureBlock);
+
+    return success();
+  }
+};
+
 // Lowerings for operations on complex numbers.
 
 struct CreateComplexOpLowering
     : public ConvertOpToLLVMPattern<CreateComplexOp> {
   using ConvertOpToLLVMPattern<CreateComplexOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
void mlir::populateStdToLLVMNonMemoryConversionPatterns(
   // clang-format off
   patterns.insert<
       AbsFOpLowering,
       AddCFOpLowering,
       AddFOpLowering,
       AddIOpLowering,
       AllocaOpLowering,
       AndOpLowering,
+      AssertOpLowering,
       AtomicRMWOpLowering,
       BranchOpLowering,
       CallIndirectOpLowering,
       CallOpLowering,
       CeilFOpLowering,
       CmpFOpLowering,
       CmpIOpLowering,
       CondBranchOpLowering,

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

func @rsqrt_multidim_vector(%arg0 : vector<4x3xf32>) {
   // CHECK: %[[EXTRACT:.*]] = llvm.extractvalue %arg0[0] : !llvm<"[4 x <3 x float>]">
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(dense<1.000000e+00> : vector<3xf32>) : !llvm<"<3 x float>">
   // CHECK: %[[SQRT:.*]] = "llvm.intr.sqrt"(%[[EXTRACT]]) : (!llvm<"<3 x float>">) -> !llvm<"<3 x float>">
   // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : !llvm<"<3 x float>">
   // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[DIV]], %0[0] : !llvm<"[4 x <3 x float>]">
   %0 = rsqrt %arg0 : vector<4x3xf32>
   std.return
 }
+
+// -----
+
+// Lowers `assert` to a function call to `abort` if the assertion is violated.
+// CHECK: llvm.func @abort()
+// CHECK-LABEL: @assert_test_function
+// CHECK-SAME:  (%[[ARG:.*]]: !llvm.i1)
+func @assert_test_function(%arg : i1) {
+  // CHECK: llvm.cond_br %[[ARG]], ^[[CONTINUATION_BLOCK:.*]], ^[[FAILURE_BLOCK:.*]]
+  // CHECK: ^[[CONTINUATION_BLOCK]]:
+  // CHECK: llvm.return
+  // CHECK: ^[[FAILURE_BLOCK]]:
+  // CHECK: llvm.call @abort() : () -> ()
+  // CHECK: llvm.unreachable
+  assert %arg, "Computer says no"
+  return
+}
+