[mlir][tensor] Add runtime verification for cast/dim/extract/insert/extract_slice

mlir/include/mlir/Dialect/Tensor/Transforms/RuntimeOpVerification.h

@@ -0,0 +1,21 @@
+//===- RuntimeOpVerification.h - Op Verification ----------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TENSOR_RUNTIMEOPVERIFICATION_H
+#define MLIR_DIALECT_TENSOR_RUNTIMEOPVERIFICATION_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace tensor {
+void registerRuntimeVerifiableOpInterfaceExternalModels(
+    DialectRegistry &registry);
+} // namespace tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TENSOR_RUNTIMEOPVERIFICATION_H

mlir/include/mlir/InitAllDialects.h

@@ -84,6 +84,7 @@
 #include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/Transforms/RuntimeOpVerification.h"
 #include "mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/ShardingInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
@@ -186,6 +187,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerFindPayloadReplacementOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
+  tensor::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   tensor::registerSubsetOpInterfaceExternalModels(registry);
   tensor::registerTilingInterfaceExternalModels(registry);
   tensor::registerValueBoundsOpInterfaceExternalModels(registry);

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

@@ -30,6 +30,7 @@ add_mlir_dialect_library(MLIRTensorDialect
   MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRParallelCombiningOpInterface
+  MLIRRuntimeVerifiableOpInterface
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
   MLIRSupport

mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp

@@ -12,6 +12,7 @@
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
+#include "mlir/Interfaces/RuntimeVerifiableOpInterface.h"
 #include "mlir/Interfaces/SubsetOpInterface.h"
 #include "mlir/Transforms/InliningUtils.h"
 
@@ -58,6 +59,8 @@ void TensorDialect::initialize() {
                             InsertSliceOp, ReshapeOp>();
   declarePromisedInterfaces<ReifyRankedShapedTypeOpInterface, ExpandShapeOp,
                             CollapseShapeOp, PadOp>();
+  declarePromisedInterfaces<RuntimeVerifiableOpInterface, CastOp, DimOp,
+                            ExtractOp, InsertOp, ExtractSliceOp>();
   declarePromisedInterfaces<SubsetOpInterface, ExtractSliceOp, InsertSliceOp,
                             ParallelInsertSliceOp>();
   declarePromisedInterfaces<SubsetInsertionOpInterface, InsertSliceOp,

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

@@ -8,6 +8,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
   MergeConsecutiveInsertExtractSlicePatterns.cpp
   ReshapePatterns.cpp
   RewriteAsConstant.cpp
+  RuntimeOpVerification.cpp
   SwapExtractSliceWithProducerPatterns.cpp
   SubsetInsertionOpInterfaceImpl.cpp
 
@@ -28,6 +29,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
   MLIRLinalgDialect
   MLIRMemRefDialect
   MLIRPass
+  MLIRRuntimeVerifiableOpInterface
   MLIRSCFDialect
   MLIRSubsetOpInterface
   MLIRTensorDialect

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

@@ -0,0 +1,208 @@
+//===- RuntimeOpVerification.cpp - Op Verification ------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tensor/Transforms/RuntimeOpVerification.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Interfaces/RuntimeVerifiableOpInterface.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace tensor {
+namespace {
+/// Generate a runtime check for lb <= value < ub.
+Value generateInBoundsCheck(OpBuilder &builder, Location loc, Value value,
+                            Value lb, Value ub) {
+  Value inBounds1 = builder.createOrFold<arith::CmpIOp>(
+      loc, arith::CmpIPredicate::sge, value, lb);
+  Value inBounds2 = builder.createOrFold<arith::CmpIOp>(
+      loc, arith::CmpIPredicate::slt, value, ub);
+  Value inBounds =
+      builder.createOrFold<arith::AndIOp>(loc, inBounds1, inBounds2);
+  return inBounds;
+}
+
+struct CastOpInterface
+    : public RuntimeVerifiableOpInterface::ExternalModel<CastOpInterface,
+                                                         CastOp> {
+  void generateRuntimeVerification(Operation *op, OpBuilder &builder,
+                                   Location loc) const {
+    auto castOp = cast<CastOp>(op);
+    auto srcType = cast<TensorType>(castOp.getSource().getType());
+
+    // Nothing to check if the result is an unranked tensor.
+    auto resultType = dyn_cast<RankedTensorType>(castOp.getType());
+    if (!resultType)
+      return;
+
+    if (isa<UnrankedTensorType>(srcType)) {
+      // 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,
+          RuntimeVerifiableOpInterface::generateErrorMessage(op,
+                                                             "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 = dyn_cast<RankedTensorType>(srcType))
+        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,
+          RuntimeVerifiableOpInterface::generateErrorMessage(
+              op, "size mismatch of dim " + std::to_string(it.index())));
+    }
+  }
+};
+
+struct DimOpInterface
+    : public RuntimeVerifiableOpInterface::ExternalModel<DimOpInterface,
+                                                         DimOp> {
+  void generateRuntimeVerification(Operation *op, OpBuilder &builder,
+                                   Location loc) const {
+    auto dimOp = cast<DimOp>(op);
+    Value rank = builder.create<RankOp>(loc, dimOp.getSource());
+    Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
+    builder.create<cf::AssertOp>(
+        loc, generateInBoundsCheck(builder, loc, dimOp.getIndex(), zero, rank),
+        RuntimeVerifiableOpInterface::generateErrorMessage(
+            op, "index is out of bounds"));
+  }
+};
+
+/// Verifies that the indices on extract/insert ops are in-bounds of the
+/// tensor's index space: 0 <= index#i < dim#i
+template <typename OpTy>
+struct ExtractInsertOpInterface
+    : public RuntimeVerifiableOpInterface::ExternalModel<
+          ExtractInsertOpInterface<OpTy>, OpTy> {
+  void generateRuntimeVerification(Operation *op, OpBuilder &builder,
+                                   Location loc) const {
+    auto extractInsertOp = cast<OpTy>(op);
+
+    Value tensor;
+    if constexpr (std::is_same_v<OpTy, ExtractOp>) {
+      tensor = extractInsertOp.getTensor();
+    } else if constexpr (std::is_same_v<OpTy, InsertOp>) {
+      tensor = extractInsertOp.getDest();
+    } else {
+      llvm_unreachable("invalid op");
+    }
+    auto tensorType = cast<RankedTensorType>(tensor.getType());
+    auto rank = tensorType.getRank();
+    if (rank == 0) {
+      // Nothing to check for 0-d tensors.
+      return;
+    }
+
+    auto indices = extractInsertOp.getIndices();
+    auto zero = builder.create<arith::ConstantIndexOp>(loc, 0);
+    Value assertCond;
+    for (auto i : llvm::seq<int64_t>(0, rank)) {
+      Value dimOp = builder.createOrFold<tensor::DimOp>(loc, tensor, i);
+      Value inBounds =
+          generateInBoundsCheck(builder, loc, indices[i], zero, dimOp);
+      assertCond =
+          i > 0 ? builder.createOrFold<arith::AndIOp>(loc, assertCond, inBounds)
+                : inBounds;
+    }
+    builder.create<cf::AssertOp>(
+        loc, assertCond,
+        RuntimeVerifiableOpInterface::generateErrorMessage(
+            op, "out-of-bounds access"));
+  }
+};
+
+struct ExtractSliceOpInterface
+    : public RuntimeVerifiableOpInterface::ExternalModel<
+          ExtractSliceOpInterface, ExtractSliceOp> {
+  void generateRuntimeVerification(Operation *op, OpBuilder &builder,
+                                   Location loc) const {
+    auto extractSliceOp = cast<ExtractSliceOp>(op);
+    RankedTensorType sourceType = extractSliceOp.getSource().getType();
+
+    // For each dimension, assert that:
+    // 0 <= offset < dim_size
+    // 0 <= offset + (size - 1) * stride < dim_size
+    Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
+    Value one = builder.create<arith::ConstantIndexOp>(loc, 1);
+    for (int64_t i = 0, e = sourceType.getRank(); i < e; ++i) {
+      Value offset = getValueOrCreateConstantIndexOp(
+          builder, loc, extractSliceOp.getMixedOffsets()[i]);
+      Value size = getValueOrCreateConstantIndexOp(
+          builder, loc, extractSliceOp.getMixedSizes()[i]);
+      Value stride = getValueOrCreateConstantIndexOp(
+          builder, loc, extractSliceOp.getMixedStrides()[i]);
+
+      // Verify that offset is in-bounds.
+      Value dimSize = builder.createOrFold<tensor::DimOp>(
+          loc, extractSliceOp.getSource(), i);
+      Value offsetInBounds =
+          generateInBoundsCheck(builder, loc, offset, zero, dimSize);
+      builder.create<cf::AssertOp>(
+          loc, offsetInBounds,
+          RuntimeVerifiableOpInterface::generateErrorMessage(
+              op, "offset " + std::to_string(i) + " is out-of-bounds"));
+
+      // Verify that slice does not run out-of-bounds.
+      Value sizeMinusOne = builder.create<arith::SubIOp>(loc, size, one);
+      Value sizeMinusOneTimesStride =
+          builder.create<arith::MulIOp>(loc, sizeMinusOne, stride);
+      Value lastPos =
+          builder.create<arith::AddIOp>(loc, offset, sizeMinusOneTimesStride);
+      Value lastPosInBounds =
+          generateInBoundsCheck(builder, loc, lastPos, zero, dimSize);
+      builder.create<cf::AssertOp>(
+          loc, lastPosInBounds,
+          RuntimeVerifiableOpInterface::generateErrorMessage(
+              op, "extract_slice runs out-of-bounds along dimension " +
+                      std::to_string(i)));
+    }
+  }
+};
+} // namespace
+} // namespace tensor
+} // namespace mlir
+
+void mlir::tensor::registerRuntimeVerifiableOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+    CastOp::attachInterface<CastOpInterface>(*ctx);
+    DimOp::attachInterface<DimOpInterface>(*ctx);
+    ExtractOp::attachInterface<ExtractInsertOpInterface<ExtractOp>>(*ctx);
+    ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
+    InsertOp::attachInterface<ExtractInsertOpInterface<InsertOp>>(*ctx);
+
+    // Load additional dialects of which ops may get created.
+    ctx->loadDialect<arith::ArithDialect, cf::ControlFlowDialect>();
+  });
+}

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

@@ -0,0 +1,50 @@
+// RUN: mlir-opt %s -generate-runtime-verification \
+// RUN:     -one-shot-bufferize="bufferize-function-boundaries" \
+// RUN:     -buffer-deallocation-pipeline=private-function-dynamic-ownership \
+// RUN:     -test-cf-assert \
+// RUN:     -convert-scf-to-cf \
+// RUN:     -convert-to-llvm | \
+// RUN: mlir-runner -e main -entry-point-result=void \
+// RUN:     -shared-libs=%tlir_runner_utils 2>&1 | \
+// RUN: FileCheck %s
+
+func.func private @cast_to_static_dim(%t: tensor<?xf32>) -> tensor<10xf32> {
+  %0 = tensor.cast %t : tensor<?xf32> to tensor<10xf32>
+  return %0 : tensor<10xf32>
+}
+
+func.func private @cast_to_ranked(%t: tensor<*xf32>) -> tensor<f32> {
+  %0 = tensor.cast %t : tensor<*xf32> to tensor<f32>
+  return %0 : tensor<f32>
+}
+
+func.func private @valid_cast(%t: tensor<*xf32>) -> tensor<?xf32> {
+  %0 = tensor.cast %t : tensor<*xf32> to tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+func.func @main() {
+  // All casts inside the called functions are invalid at runtime, except for
+  // the last one.
+  %alloc = tensor.empty() : tensor<5xf32>
+
+  //      CHECK: ERROR: Runtime op verification failed
+  // CHECK-NEXT: "tensor.cast"(%{{.*}}) : (tensor<?xf32>) -> tensor<10xf32>
+  // CHECK-NEXT: ^ size mismatch of dim 0
+  // CHECK-NEXT: Location: loc({{.*}})
+  %1 = tensor.cast %alloc : tensor<5xf32> to tensor<?xf32>
+  func.call @cast_to_static_dim(%1) : (tensor<?xf32>) -> (tensor<10xf32>)
+
+  // CHECK-NEXT: ERROR: Runtime op verification failed
+  // CHECK-NEXT: "tensor.cast"(%{{.*}}) : (tensor<*xf32>) -> tensor<f32>
+  // CHECK-NEXT: ^ rank mismatch
+  // CHECK-NEXT: Location: loc({{.*}})
+  %3 = tensor.cast %alloc : tensor<5xf32> to tensor<*xf32>
+  func.call @cast_to_ranked(%3) : (tensor<*xf32>) -> (tensor<f32>)
+
+  // A last cast that actually succeeds.
+  // CHECK-NOT: ERROR: Runtime op verification failed
+  func.call @valid_cast(%3) : (tensor<*xf32>) -> (tensor<?xf32>)
+
+  return
+}

mlir/test/Integration/Dialect/Tensor/dim-runtime-verification.mlir

@@ -0,0 +1,21 @@
+// RUN: mlir-opt %s -generate-runtime-verification \
+// RUN:     -one-shot-bufferize \
+// RUN:     -buffer-deallocation-pipeline \
+// RUN:     -test-cf-assert \
+// RUN:     -convert-to-llvm | \
+// RUN: mlir-runner -e main -entry-point-result=void \
+// RUN:     -shared-libs=%mlir_runner_utils 2>&1 | \
+// RUN: FileCheck %s
+
+func.func @main() {
+  %c4 = arith.constant 4 : index
+  %tensor = tensor.empty() : tensor<1xf32>
+
+  //      CHECK: ERROR: Runtime op verification failed
+  // CHECK-NEXT: "tensor.dim"(%{{.*}}, %{{.*}}) : (tensor<1xf32>, index) -> index
+  // CHECK-NEXT: ^ index is out of bounds
+  // CHECK-NEXT: Location: loc({{.*}})
+  %dim = tensor.dim %tensor, %c4 : tensor<1xf32>
+
+  return
+}