[mlir][LLVM] Add runtime debug assertions for memref::CollapseShapeOp

[springerm]

Summary

This change adds two helper functions: runtimeAbort/runtimeAssert. These can be used to generate runtime assertions.

This change adds runtime asserts to check for invalid memref::CollapseShapeOps. In the static case, these are already detected during op verification. With this change, the program execution crashes explictly at runtime when the assertion is violated. For now, this assertion is activated only on an opt-in basis (to avoid performance regressions).

[qcolombet]

LGTM, couple of nits below.

Thanks for doing this Matthias!

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1552

Do you expect the following runtimeAssert to be replaced by "nothing"?
Put differently do we need to keep this cast?

1578

Should this be part of this patch?

[springerm]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1552

Just a leftover from debugging

1578

Yes, I just moved it a few lines further down, so that the insertion point is correct when this function is left.

[springerm]

address comments

[springerm]

clang-format

[springerm]

fix tests

[nicolasvasilache]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1636

Big -1 on removing insertion guard here: there exist code paths (i.e. if ref is empty) that change the insertion point and leak across function boundary. This is a huge footgun.

[nicolasvasilache]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1636

Note: fix is as simple as adding a guard just before or, better, at the top of the function to be globally safe.

[springerm]

address comments

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1636

I expect fillInStridesForCollapsedMemDescriptor to actually change the insertion point. This function is creating new ops and when leaving this function it should be set to right after these new ops. (So that runtimeVerifyCollapseShapeContiguity can read from the descriptor that we just filled.)

This is just like reifyResultShapes (and various other helpers that we have); these are also incrementing the insertion point. (It is a bit more complex here because we are splitting blocks, but that is an implementation detail of this function and the caller does not have to worry about that.)

That being said, the rewriter.setInsertionPoint(op); was unnecessarily strict, assuming that this function is only called from the CollapseShapeOp lowering pattern. I fixed that.

mlir/include/mlir/Conversion/LLVMCommon/LoweringOptions.h

 /// options between passes, patterns, and type converter.
 class LowerToLLVMOptions {
 public:
   explicit LowerToLLVMOptions(MLIRContext *ctx);
   LowerToLLVMOptions(MLIRContext *ctx, const DataLayout &dl);
 
   bool useBarePtrCallConv = false;
 
+  bool enableDebugAssertions = false;
+
   enum class AllocLowering {
     /// Use malloc for for heap allocations.
     Malloc,
 
     /// Use aligned_alloc for heap allocations.
     AlignedAlloc,
 
     /// Do not lower heap allocations. Users must provide their own patterns for

mlir/include/mlir/Conversion/Passes.td

Option<"useAlignedAlloc", "use-aligned-alloc", "bool", /*default=*/"false",
            "Use aligned_alloc in place of malloc for heap allocations">,
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,
     Option<"useGenericFunctions", "use-generic-functions",
            "bool",
            /*default=*/"false",
            "Use generic allocation and deallocation functions instead of the "
-           "classic 'malloc', 'aligned_alloc' and 'free' functions">
+           "classic 'malloc', 'aligned_alloc' and 'free' functions">,
+    Option<"enableDebugAssertions", "enable-debug-assertions", "bool",
+           /*default=*/"false",
+           "Emit assertions around certain ops for easier debugging">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefToSPIRV
 //===----------------------------------------------------------------------===//
 
 def MapMemRefStorageClass : Pass<"map-memref-spirv-storage-class"> {

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

 namespace detail {
 struct LLVMTypeStorage;
 struct LLVMDialectImpl;
 } // namespace detail
 } // namespace LLVM
 } // namespace mlir
 
 namespace mlir {
+class RewriterBase;
+
 namespace LLVM {
 template <typename Values>
 class GEPIndicesAdaptor;
 
 /// Bit-width of a 'GEPConstantIndex' within GEPArg.
 constexpr int kGEPConstantBitWidth = 29;
 /// Wrapper around a int32_t for use in a PointerUnion.
 using GEPConstantIndex =
public:
   /// Constructs a GEPArg with a constant index.
   /*implicit*/ GEPArg(int32_t integer) : BaseT(integer) {}
 
   /// Constructs a GEPArg with a dynamic index.
   /*implicit*/ GEPArg(Value value) : BaseT(value) {}
 
   using BaseT::operator=;
 };
+
+/// Emit IR that aborts the program at runtime.
+void runtimeAbort(OpBuilder &b, Location loc);
+
+/// C/C++ style runtime assert that aborts the program if the given condition
+/// is "false".
+void runtimeAssert(RewriterBase &rewriter, Location loc, Value condition);
 } // namespace LLVM
 } // namespace mlir
 
 ///// Ops /////
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMOps.h.inc"
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMIntrinsicOps.h.inc"

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

 struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> {
   using ConvertOpToLLVMPattern<cf::AssertOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(cf::AssertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
-    // 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);
+    LLVM::runtimeAbort(rewriter, loc);
 
     // Generate assertion test.
     rewriter.setInsertionPointToEnd(opBlock);
     rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
         op, adaptor.getArg(), continuationBlock, failureBlock);
 
     return success();
   }

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

for (auto dstIndex : llvm::reverse(en.value())) {
       dstDesc.setStride(b, loc, dstIndex, currentStrideToExpand);
       Value size = dstDesc.size(b, loc, dstIndex);
       currentStrideToExpand =
           b.create<LLVM::MulOp>(loc, size, currentStrideToExpand);
     }
   }
 }
 
+// Note: Check memref::CollapseShapeOp verifier for details. This function is
+// implementing the same contiguity check, but at runtime.
+static void runtimeVerifyCollapseShapeContiguity(
+    RewriterBase &rewriter, Location loc, TypeConverter *typeConverter,
+    MemRefType srcType, MemRefDescriptor &srcDesc, MemRefDescriptor &dstDesc,
+    ArrayRef<ReassociationIndices> reassociation) {
+  auto llvmIndexType = typeConverter->convertType(rewriter.getIndexType());
+  unsigned resultStrideIndex = reassociation.size() - 1;
+  for (const ReassociationIndices &reassoc : llvm::reverse(reassociation)) {
+    auto trailingReassocs = ArrayRef<int64_t>(reassoc).drop_front();
+    Value stride = dstDesc.stride(rewriter, loc, resultStrideIndex--);
+    for (int64_t idx : llvm::reverse(trailingReassocs)) {
+      Value srcSize;
+      if (srcType.isDynamicDim(idx)) {
+        srcSize = srcDesc.size(rewriter, loc, idx);
+      } else {
+        srcSize = rewriter.create<LLVM::ConstantOp>(
+            loc, llvmIndexType, rewriter.getIndexAttr(srcType.getDimSize(idx)));
+      }
+      stride = rewriter.create<LLVM::MulOp>(loc, stride, srcSize);
+
+      // Both source and computed stride must have the same value. In that
+      // case, we can be sure, that the dimensions are collapsible (because they
+      // are contiguous).
+      Value srcStride = srcDesc.stride(rewriter, loc, idx - 1);
+      Value sameStride = rewriter.create<LLVM::ICmpOp>(
+          loc, LLVM::ICmpPredicate::eq, stride, srcStride);
+      LLVM::runtimeAssert(rewriter, loc, sameStride);
+    }
+  }
+}
+
 static void fillInStridesForCollapsedMemDescriptor(
-    ConversionPatternRewriter &rewriter, Location loc, Operation *op,
+    ConversionPatternRewriter &rewriter, Location loc,
     TypeConverter *typeConverter, MemRefType srcType, MemRefDescriptor &srcDesc,
     MemRefDescriptor &dstDesc, ArrayRef<ReassociationIndices> reassociation) {
+  auto ip = rewriter.saveInsertionPoint();
   auto llvmIndexType = typeConverter->convertType(rewriter.getIndexType());
   // See comments for computeCollapsedLayoutMap for details on how the strides
   // are calculated.
   auto srcShape = srcType.getShape();
   for (auto &en : llvm::enumerate(reassociation)) {
-    rewriter.setInsertionPoint(op);
+    rewriter.restoreInsertionPoint(ip);
     auto dstIndex = en.index();
     ArrayRef<int64_t> ref = llvm::makeArrayRef(en.value());
     while (srcShape[ref.back()] == 1 && ref.size() > 1)
       ref = ref.drop_back();
     if (!ShapedType::isDynamic(srcShape[ref.back()]) || ref.size() == 1) {
       dstDesc.setStride(rewriter, loc, dstIndex,
                         srcDesc.stride(rewriter, loc, ref.back()));
     } else {
if (!ShapedType::isDynamic(srcShape[ref.back()]) || ref.size() == 1) {
       //            |   ...                       |   |            |
       //            +--------------+--------------+   |   +--------+
       //                           |                  |   |
       //                           v                  v   v
       //   +--------------------------------------------------+
       //   | continue(%newStride):                            |
       //   |   %newMemRefDes = setStride(%newStride,dstIndex) |
       //   +--------------------------------------------------+
-      OpBuilder::InsertionGuard guard(rewriter);
       Block *initBlock = rewriter.getInsertionBlock();
       Block *continueBlock =
           rewriter.splitBlock(initBlock, rewriter.getInsertionPoint());
       continueBlock->insertArgument(unsigned(0), srcDesc.getIndexType(), loc);
-      rewriter.setInsertionPointToStart(continueBlock);
-      dstDesc.setStride(rewriter, loc, dstIndex, continueBlock->getArgument(0));
 
       Block *curEntryBlock = initBlock;
       Block *nextEntryBlock;
       for (auto srcIndex : llvm::reverse(ref)) {
         if (srcShape[srcIndex] == 1 && srcIndex != ref.front())
           continue;
         rewriter.setInsertionPointToEnd(curEntryBlock);
         Value srcStride = srcDesc.stride(rewriter, loc, srcIndex);
if (!ShapedType::isDynamic(srcShape[ref.back()]) || ref.size() == 1) {
           OpBuilder::InsertionGuard guard(rewriter);
           nextEntryBlock = rewriter.createBlock(
               initBlock->getParent(), Region::iterator(continueBlock), {});
         }
         rewriter.create<LLVM::CondBrOp>(loc, predNeOne, continueBlock,
                                         srcStride, nextEntryBlock, llvm::None);
         curEntryBlock = nextEntryBlock;
       }
+
+      rewriter.setInsertionPointToStart(continueBlock);
+      dstDesc.setStride(rewriter, loc, dstIndex, continueBlock->getArgument(0));
     }
   }
 }
 
 static void fillInDynamicStridesForMemDescriptor(
     ConversionPatternRewriter &b, Location loc, Operation *op,
-    TypeConverter *typeConverter, MemRefType srcType, MemRefType dstType,
+    LLVMTypeConverter *typeConverter, MemRefType srcType, MemRefType dstType,
     MemRefDescriptor &srcDesc, MemRefDescriptor &dstDesc,
     ArrayRef<ReassociationIndices> reassociation) {
-  if (srcType.getRank() > dstType.getRank())
-    fillInStridesForCollapsedMemDescriptor(b, loc, op, typeConverter, srcType,
+  if (srcType.getRank() > dstType.getRank()) {
+    fillInStridesForCollapsedMemDescriptor(b, loc, typeConverter, srcType,
                                            srcDesc, dstDesc, reassociation);
-  else
+    if (typeConverter->getOptions().enableDebugAssertions)
+      runtimeVerifyCollapseShapeContiguity(b, loc, typeConverter, srcType,
+                                           srcDesc, dstDesc, reassociation);
+  } else {
     fillInStridesForExpandedMemDescriptor(b, loc, srcType, srcDesc, dstDesc,
                                           reassociation);
   }
+}
 
 // ReshapeOp creates a new view descriptor of the proper rank.
-// For now, the only conversion supported is for target MemRef with static sizes
-// and strides.
 template <typename ReshapeOp>
 class ReassociatingReshapeOpConversion
     : public ConvertOpToLLVMPattern<ReshapeOp> {
 public:
   using ConvertOpToLLVMPattern<ReshapeOp>::ConvertOpToLLVMPattern;
   using ReshapeOpAdaptor = typename ReshapeOp::Adaptor;
 
   LogicalResult
if (llvm::all_of(strides, isStaticStrideOrOffset)) {
            llvm::reverse(llvm::seq<int64_t>(0, dstShape.size()))) {
         dstDesc.setStride(rewriter, loc, dimIndex, stride);
         stride = rewriter.create<LLVM::MulOp>(loc, dstShape[dimIndex], stride);
       }
     } else {
       // There could be mixed static/dynamic strides. For simplicity, we
       // recompute all strides if there is at least one dynamic stride.
       fillInDynamicStridesForMemDescriptor(
-          rewriter, loc, reshapeOp, this->typeConverter, srcType, dstType,
+          rewriter, loc, reshapeOp, this->getTypeConverter(), srcType, dstType,
           srcDesc, dstDesc, reshapeOp.getReassociationIndices());
     }
     rewriter.replaceOp(reshapeOp, {dstDesc});
     return success();
   }
 };
 
 /// Conversion pattern that transforms a subview op into:
void runOnOperation() override {
     const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
     LowerToLLVMOptions options(&getContext(),
                                dataLayoutAnalysis.getAtOrAbove(op));
     options.allocLowering =
         (useAlignedAlloc ? LowerToLLVMOptions::AllocLowering::AlignedAlloc
                          : LowerToLLVMOptions::AllocLowering::Malloc);
 
     options.useGenericFunctions = useGenericFunctions;
+    options.enableDebugAssertions = enableDebugAssertions;
 
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
 
     LLVMTypeConverter typeConverter(&getContext(), options,
                                     &dataLayoutAnalysis);
     RewritePatternSet patterns(&getContext());
     populateMemRefToLLVMConversionPatterns(typeConverter, patterns);
     LLVMConversionTarget target(getContext());
     target.addLegalOp<func::FuncOp>();
     if (failed(applyPartialConversion(op, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace

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

 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
 
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Type.h"
return builder.create<LLVM::GEPOp>(
       loc, LLVM::LLVMPointerType::get(IntegerType::get(ctx, 8)), globalPtr,
       ArrayRef<GEPArg>{0, 0});
 }
 
 bool mlir::LLVM::satisfiesLLVMModule(Operation *op) {
   return op->hasTrait<OpTrait::SymbolTable>() &&
          op->hasTrait<OpTrait::IsIsolatedFromAbove>();
 }
+
+void mlir::LLVM::runtimeAbort(OpBuilder &b, Location loc) {
+  // Insert the `abort` declaration if necessary.
+  auto module =
+      b.getInsertionBlock()->getParentOp()->getParentOfType<ModuleOp>();
+  auto abortFunc = module.lookupSymbol<LLVM::LLVMFuncOp>("abort");
+  if (!abortFunc) {
+    OpBuilder::InsertionGuard guard(b);
+    b.setInsertionPointToStart(module.getBody());
+    auto abortFuncTy = LLVM::LLVMFunctionType::get(
+        LLVM::LLVMVoidType::get(b.getContext()), {});
+    abortFunc =
+        b.create<LLVM::LLVMFuncOp>(b.getUnknownLoc(), "abort", abortFuncTy);
+  }
+  b.create<LLVM::CallOp>(loc, abortFunc, llvm::None);
+  b.create<LLVM::UnreachableOp>(loc);
+}
+
+void mlir::LLVM::runtimeAssert(RewriterBase &rewriter, Location loc,
+                               Value condition) {
+  // Split block.
+  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());
+  runtimeAbort(rewriter, loc);
+
+  // Generate assertion branch.
+  rewriter.setInsertionPointToEnd(opBlock);
+  rewriter.create<LLVM::CondBrOp>(loc, condition, continuationBlock,
+                                  failureBlock);
+
+  // Reset insertion point.
+  rewriter.setInsertionPointToStart(continuationBlock);
+}

mlir/test/Conversion/MemRefToLLVM/debug-assertions.mlir

+// RUN: mlir-opt %s -convert-memref-to-llvm="enable-debug-assertions" -allow-unregistered-dialect | \
+// RUN:   FileCheck %s
+
+module {
+  llvm.func @test_collapse_shape() {
+    %m = "test.dummy"() : () -> (memref<?x?xf32, strided<[?, ?], offset: ?>>)
+    %0 = memref.collapse_shape %m [[0, 1]] : memref<?x?xf32, strided<[?, ?], offset: ?>> into memref<?xf32, strided<[?], offset: ?>>
+    llvm.return
+  }
+}
+
+// CHECK: llvm.func @abort()
+// CHECK: llvm.func @test_collapse_shape
+// CHECK:   "test.dummy"()
+// CHECK: ^bb1:
+// CHECK:   llvm.extractvalue
+// CHECK:   llvm.br ^bb2
+// CHECK: ^bb2(%[[stride:.*]]: i64):
+// CHECK:   %[[stride_inserted:.*]] = llvm.insertvalue %[[stride]]
+// CHECK:   %[[stride_extracted:.*]] = llvm.extractvalue
+// CHECK:   %[[src_dim_size:.*]] = llvm.extractvalue
+// CHECK:   %[[computed_stride:.*]] = llvm.mul %[[stride_extracted]], %[[src_dim_size]]
+// CHECK:   %[[src_stride:.*]] = llvm.extractvalue
+// CHECK:   %[[cmp:.*]] = llvm.icmp "eq" %[[computed_stride]], %[[src_stride]]
+// CHECK:   llvm.cond_br %[[cmp]], ^bb3, ^bb4
+// CHECK: ^bb3:
+// CHECK:   llvm.return
+// CHECK: ^bb4:
+// CHECK:   llvm.call @abort()