[mlir][memref] Add runtime verification for memref::CastOp

[springerm]

Summary

Verify unranked -> ranked casts and casts of dynamic sizes/offset/strides to static ones.

Depends On D138647

[springerm]

update

[springerm]

update

[springerm]

rebase

[springerm]

minor update

[springerm]

rebase

[springerm]

minor update

[guraypp]

Looks good to me. But can we add more tests. For example, the case below, here c2 should be verified by your code but c3 should fail, right?

func.func @c1(%m: memref<?x?xf32>) -> memref<*xf32> {
  %0 = memref.cast %m : memref<?x?xf32> to memref<*xf32>
  return %0 : memref<*xf32>
}
func.func @c2(%m: memref<*xf32>) -> memref<?x?xf32> {
  %0 = memref.cast %m : memref<*xf32> to memref<?x?xf32>
  return %0 : memref<?x?xf32>
}
func.func @c3(%m: memref<*xf32>) -> memref<?xf32> {
  %0 = memref.cast %m : memref<*xf32> to memref<?xf32>
  return %0 : memref<?xf32>
}
func.func @main() {
  %a = memref.alloc() : memref<5x5xf32>
  %c = memref.cast %a : memref<5x5xf32> to memref<?x?xf32>
  %res = func.call @c1(%c) : (memref<?x?xf32>) -> memref<*xf32> 
  func.call @c2(%res) : (memref<*xf32>) -> (memref<?x?xf32>)
  func.call @c3(%res) : (memref<*xf32>) -> (memref<?xf32>)
  return
}

Also, runtime error text does not tell me which code/location fails. It prints the following text for example. I am not sure you can print the location correctly, but can we at least print the problematic IR?
memref::CastOp: rank mismatch

[springerm]

address comments

[guraypp]

Error messages are now very clear, thanks!

mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Interfaces/RuntimeVerifiableOpInterface.h"
 
 namespace mlir {
 namespace memref {
 namespace {
+struct CastOpInterface
+    : public RuntimeVerifiableOpInterface::ExternalModel<CastOpInterface,
+                                                         CastOp> {
+  void generateRuntimeVerification(Operation *op, OpBuilder &builder,
+                                   Location loc) const {
+    auto castOp = cast<CastOp>(op);
+    auto srcType = castOp.getSource().getType().cast<BaseMemRefType>();
+
+    // Nothing to check if the result is an unranked memref.
+    auto resultType = castOp.getType().dyn_cast<MemRefType>();
+    if (!resultType)
+      return;
+
+    if (srcType.isa<UnrankedMemRefType>()) {
+      // Check rank.
+      Value srcRank = builder.create<RankOp>(loc, castOp.getSource());
+      Value resultRank =
+          builder.create<arith::ConstantIndexOp>(loc, resultType.getRank());
+      Value isSameRank = builder.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::eq, srcRank, resultRank);
+      builder.create<cf::AssertOp>(loc, isSameRank,
+                                   "memref::CastOp: rank mismatch");
+    }
+
+    // Check dimension sizes.
+    for (const auto &it : llvm::enumerate(resultType.getShape())) {
+      // Static dim size -> static/dynamic dim size does not need verification.
+      if (auto rankedSrcType = srcType.dyn_cast<MemRefType>())
+        if (!rankedSrcType.isDynamicDim(it.index()))
+          continue;
+
+      // Static/dynamic dim size -> dynamic dim size does not need verification.
+      if (resultType.isDynamicDim(it.index()))
+        continue;
+
+      Value srcDimSz =
+          builder.create<DimOp>(loc, castOp.getSource(), it.index());
+      Value resultDimSz =
+          builder.create<arith::ConstantIndexOp>(loc, it.value());
+      Value isSameSz = builder.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::eq, srcDimSz, resultDimSz);
+      builder.create<cf::AssertOp>(loc, isSameSz,
+                                   "memref::CastOp: size mismatch of dim " +
+                                       std::to_string(it.index()));
+    }
+
+    // Get source offset and strides. We do not have an op to get extract
+    // offsets and strides from unranked memrefs, so cast the source to a type
+    // with fully dynamic layout, from which we can then extract the offset and
+    // strides. (Rank was already verified.)
+    int64_t dynamicOffset = ShapedType::kDynamic;
+    SmallVector<int64_t> dynamicShape(resultType.getRank(),
+                                      ShapedType::kDynamic);
+    auto stridedLayout = StridedLayoutAttr::get(builder.getContext(),
+                                                dynamicOffset, dynamicShape);
+    auto dynStridesType =
+        MemRefType::get(dynamicShape, resultType.getElementType(),
+                        stridedLayout, resultType.getMemorySpace());
+    Value helperCast =
+        builder.create<CastOp>(loc, dynStridesType, castOp.getSource());
+    auto metadataOp = builder.create<ExtractStridedMetadataOp>(loc, helperCast);
+
+    // Get result offset and strides.
+    int64_t resultOffset;
+    SmallVector<int64_t> resultStrides;
+    if (failed(getStridesAndOffset(resultType, resultStrides, resultOffset)))
+      return;
+
+    // Check offset.
+    if (resultOffset != ShapedType::kDynamic) {
+      // Static/dynamic offset -> dynamic offset does not need verification.
+      Value srcOffset = metadataOp.getResult(1);
+      Value resultOffsetVal =
+          builder.create<arith::ConstantIndexOp>(loc, resultOffset);
+      Value isSameOffset = builder.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::eq, srcOffset, resultOffsetVal);
+      builder.create<cf::AssertOp>(loc, isSameOffset,
+                                   "memref::CastOp: offset mismatch");
+    }
+
+    // Check strides.
+    for (const auto &it : llvm::enumerate(resultStrides)) {
+      // Static/dynamic stride -> dynamic stride does not need verification.
+      if (it.value() == ShapedType::kDynamic)
+        continue;
+
+      Value srcStride =
+          metadataOp.getResult(2 + resultType.getRank() + it.index());
+      Value resultStrideVal =
+          builder.create<arith::ConstantIndexOp>(loc, it.value());
+      Value isSameStride = builder.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::eq, srcStride, resultStrideVal);
+      builder.create<cf::AssertOp>(loc, isSameStride,
+                                   "memref::CastOp: stride mismatch of dim " +
+                                       std::to_string(it.index()));
+    }
+  }
+};
+
 struct ExpandShapeOpInterface
     : public RuntimeVerifiableOpInterface::ExternalModel<ExpandShapeOpInterface,
                                                          ExpandShapeOp> {
   void generateRuntimeVerification(Operation *op, OpBuilder &builder,
                                    Location loc) const {
     auto expandShapeOp = cast<ExpandShapeOp>(op);
 
     // Verify that the expanded dim sizes are a product of the collapsed dim
for (auto it : llvm::enumerate(expandShapeOp.getReassociationIndices())) {
       // staticResultDimSz must divide srcDimSz evenly.
       Value mod =
           builder.create<arith::RemSIOp>(loc, srcDimSz, staticResultDimSz);
       Value isModZero = builder.create<arith::CmpIOp>(
           loc, arith::CmpIPredicate::eq, mod,
           builder.create<arith::ConstantIndexOp>(loc, 0));
       builder.create<cf::AssertOp>(
           loc, isModZero,
-          "static result dims in reassoc group do not divide src dim evenly");
+          "memref::ExpandShapeOp: static result dims in reassoc group do not "
+          "divide src dim evenly");
     }
   }
 };
 } // namespace
 } // namespace memref
 } // namespace mlir
 
 void mlir::memref::registerRuntimeVerifiableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+    CastOp::attachInterface<CastOpInterface>(*ctx);
     ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
 
     // Load additional dialects of which ops may get created.
     ctx->loadDialect<arith::ArithDialect>();
   });
 }

mlir/test/Dialect/MemRef/runtime-verification.mlir

 // RUN: mlir-opt %s -generate-runtime-verification -cse | FileCheck %s
 
 // CHECK-LABEL: func @expand_shape(
 //  CHECK-SAME:     %[[m:.*]]: memref<?xf32>
 //   CHECK-DAG:   %[[c0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[c5:.*]] = arith.constant 5 : index
 //   CHECK-DAG:   %[[dim:.*]] = memref.dim %[[m]], %[[c0]]
 //       CHECK:   %[[mod:.*]] = arith.remsi %[[dim]], %[[c5]]
 //       CHECK:   %[[cmpi:.*]] = arith.cmpi eq, %[[mod]], %[[c0]]
-//       CHECK:   cf.assert %[[cmpi]], "static result dims in reassoc group do not divide src dim evenly"
+//       CHECK:   cf.assert %[[cmpi]], "memref::ExpandShapeOp: static result dims in reassoc group do not divide src dim evenly"
 func.func @expand_shape(%m: memref<?xf32>) -> memref<?x5xf32> {
   %0 = memref.expand_shape %m [[0, 1]] : memref<?xf32> into memref<?x5xf32>
   return %0 : memref<?x5xf32>
 }

mlir/test/Integration/Dialect/Memref/cast-runtime-verification.mlir

+// RUN: mlir-opt %s -generate-runtime-verification -convert-memref-to-llvm \
+// RUN:     -test-cf-assert \
+// RUN:     -convert-func-to-llvm -reconcile-unrealized-casts | \
+// 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 @cast_to_static_dim(%m: memref<?xf32>) -> memref<10xf32> {
+  %0 = memref.cast %m : memref<?xf32> to memref<10xf32>
+  return %0 : memref<10xf32>
+}
+
+func.func @cast_to_ranked(%m: memref<*xf32>) -> memref<f32> {
+  %0 = memref.cast %m : memref<*xf32> to memref<f32>
+  return %0 : memref<f32>
+}
+
+func.func @cast_to_static_strides(%m: memref<?xf32, strided<[?], offset: ?>>)
+    -> memref<?xf32, strided<[9], offset: 5>> {
+  %0 = memref.cast %m : memref<?xf32, strided<[?], offset: ?>>
+                     to memref<?xf32, strided<[9], offset: 5>>
+  return %0 : memref<?xf32, strided<[9], offset: 5>>
+}
+
+func.func @main() {
+  // All casts inside the called functions are invalid at runtime.
+  %alloc = memref.alloc() : memref<5xf32>
+
+  // CHECK: memref::CastOp: size mismatch of dim 0
+  %1 = memref.cast %alloc : memref<5xf32> to memref<?xf32>
+  func.call @cast_to_static_dim(%1) : (memref<?xf32>) -> (memref<10xf32>)
+
+  // CHECK-NEXT: memref::CastOp: rank mismatch
+  %3 = memref.cast %alloc : memref<5xf32> to memref<*xf32>
+  func.call @cast_to_ranked(%3) : (memref<*xf32>) -> (memref<f32>)
+
+  // CHECK-NEXT: memref::CastOp: offset mismatch
+  // CHECK-NEXT: memref::CastOp: stride mismatch of dim 0
+  %4 = memref.cast %alloc
+      : memref<5xf32> to memref<?xf32, strided<[?], offset: ?>>
+  func.call @cast_to_static_strides(%4)
+      : (memref<?xf32, strided<[?], offset: ?>>)
+     -> (memref<?xf32, strided<[9], offset: 5>>)
+
+  return
+}