[Mlir-commits] [mlir] [mlir][Tensor] Check for out-of-bounds extraction in `extract_slice` verifier (PR #130487)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/130487

Also fix a test case that had an invalid op.


>From 106cc327b13893c1bd2ee7b43ba4b6117ea42dcc Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Sun, 9 Mar 2025 13:52:49 +0100
Subject: [PATCH] [mlir][Tensor] Check for out-of-bounds extraction in
 `extract_slice` verifier

---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp      | 37 ++++++++++++++++++-
 .../Tensor/bubble-up-extract-slice-op.mlir    |  2 +-
 mlir/test/Dialect/Tensor/invalid.mlir         | 16 ++++++++
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index e3c8899cd44d9..6dfc05eab08e6 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2354,11 +2354,44 @@ static LogicalResult produceSliceErrorMsg(SliceVerificationResult result,
 
 /// Verifier for ExtractSliceOp.
 LogicalResult ExtractSliceOp::verify() {
+  RankedTensorType sourceType = getSourceType();
+  SmallVector<OpFoldResult> mixedOffsets = getMixedOffsets();
+  SmallVector<OpFoldResult> mixedSizes = getMixedSizes();
+  SmallVector<OpFoldResult> mixedStrides = getMixedStrides();
+
   // Verify result type against inferred type.
   RankedTensorType expectedType = ExtractSliceOp::inferResultType(
-      getSourceType(), getMixedOffsets(), getMixedSizes(), getMixedStrides());
+      sourceType, mixedOffsets, mixedSizes, mixedStrides);
   SliceVerificationResult result = isRankReducedType(expectedType, getType());
-  return produceSliceErrorMsg(result, *this, expectedType);
+  if (result != SliceVerificationResult::Success)
+    return produceSliceErrorMsg(result, *this, expectedType);
+
+  // Verify that offsets, sizes, strides do not run out-of-bounds with respect
+  // to the source tensor.
+  for (int64_t i = 0, e = sourceType.getRank(); i < e; ++i) {
+    // Nothing to verify for dynamic source dims.
+    if (sourceType.isDynamicDim(i))
+      continue;
+    auto offsetOfr = dyn_cast<Attribute>(mixedOffsets[i]);
+    // Nothing to verify if the offset is dynamic.
+    if (!offsetOfr)
+      continue;
+    int64_t staticOffset = *getConstantIntValue(offsetOfr);
+    if (staticOffset >= sourceType.getDimSize(i))
+      return emitOpError("offset ") << i << " is out-of-bounds";
+    auto sizeOfr = dyn_cast<Attribute>(mixedSizes[i]);
+    auto strideOfr = dyn_cast<Attribute>(mixedStrides[i]);
+    if (!sizeOfr || !strideOfr)
+      continue;
+    int64_t staticSize = *getConstantIntValue(sizeOfr);
+    int64_t staticStride = *getConstantIntValue(strideOfr);
+    if (staticOffset + (staticSize - 1) * staticStride >=
+        sourceType.getDimSize(i))
+      return emitOpError("extraction along source dimension ")
+             << i << " runs out-of-bounds";
+  }
+
+  return success();
 }
 
 llvm::SmallBitVector ExtractSliceOp::getDroppedDims() {
diff --git a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
index 252e7494bff79..3900bc56f433d 100644
--- a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
+++ b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
@@ -31,7 +31,7 @@ func.func @no_bubble_up_extract_slice_on_non_contiguous(%src: tensor<60xf32>) ->
 
 func.func @no_bubble_up_extract_slice_on_stride(%src: tensor<60xf32>) -> tensor<1x1x5xf32> {
   %expand = tensor.expand_shape %src [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>
-  %extract = tensor.extract_slice %expand[0, 0, 5][1, 1, 5][1, 1, 2] : tensor<2x3x10xf32> to tensor<1x1x5xf32>
+  %extract = tensor.extract_slice %expand[0, 0, 0][1, 1, 5][1, 1, 2] : tensor<2x3x10xf32> to tensor<1x1x5xf32>
   return %extract : tensor<1x1x5xf32>
 }
 
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index 654169841c1c1..64a8585738913 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -258,6 +258,22 @@ func.func @illegal_num_offsets(%arg0 : tensor<?x?x?xf32>, %arg1 : index, %arg2 :
 
 // -----
 
+func.func @extract_slice_offset_out_of_bounds(%arg0: tensor<10xf32>) {
+  // expected-error at +1 {{offset 0 is out-of-bounds}}
+  %0 = tensor.extract_slice %arg0 [10][1][1] : tensor<10xf32> to tensor<1xf32>
+  return
+}
+
+// -----
+
+func.func @extract_slice_runs_out_of_bounds(%arg0: tensor<9xf32>) {
+  // expected-error at +1 {{extraction along source dimension 0 runs out-of-bounds}}
+  %0 = tensor.extract_slice %arg0 [3][4][2] : tensor<9xf32> to tensor<4xf32>
+  return
+}
+
+// -----
+
 func.func @insert_slice_wrong_result_rank(%t1: tensor<?xf32>, %t2: tensor<?x?xf32>, %idx : index) {
   // expected-error @+1 {{expected rank to be smaller or equal to the other rank.}}
   %0 = tensor.insert_slice %t2 into %t1[0][4][1] : tensor<?x?xf32> into tensor<?xf32>