[mlir][Tensor] Check for out-of-bounds slice in insert/extract_slice verifier

[matthias-springer]

Also fix test cases that had invalid ops.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tensor

Author: Matthias Springer (matthias-springer)

Changes

Also fix a test case that had an invalid op.

---

Full diff: https://github.com/llvm/llvm-project/pull/130487.diff

3 Files Affected:

• (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+35-2)
• (modified) mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir (+1-1)
• (modified) mlir/test/Dialect/Tensor/invalid.mlir (+16)

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@+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@+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>

[matthias-springer]

Note: This verification is not part of the SliceVerificationResult infrastructure because memref.subview is currently allowed to extract out-of-bounds (according to the op documentation).

[MaheshRavishankar]

Making it all zero might be losing some test coverage. Maybe increase the size of the tensor instead?

[matthias-springer]

This must be [0, 0], otherwise the two insert_slice cannot be merged.

[MaheshRavishankar]

I dont see why. The existing test is taking an offset of [1, 2] and composing it with [3, %pos] to get [4, %pos + 2] . That seems correct to me. This change is changing the test behavior and potentially dropping coverage.

[matthias-springer]

The previous test was incorrect:

// This op is invalid because it inserts into %r0 at out-of-bounds offsets.
%0 = tensor.insert_slice %t into %r0[1, 2] [1, 1] [1, 1]
    : tensor<f32> into tensor<1x1xf32>

To fix it, let's try to make the tensor larger but keep the offsets, as you suggested (actually, we have to make both tensors larger now).

%0 = tensor.insert_slice %t into %r0[1, 2] [1, 1] [1, 1]
    : tensor<f32> into tensor<2x3xf32>
%1 = tensor.insert_slice %0 into %r1[0, %pos] [2, 3] [1, 1] 
    : tensor<2x3xf32> into tensor<2x14xf32>

These two tensor.insert_slice cannot be folded into a single tensor.insert_slice. The reason is because we need to insert both %t and some data from %r0 into %r1.

So the only way to fix the test (so that it still exercises the folding pattern) is to set the offsets to 0.

The existing test is taking an offset of [1, 2] and composing it with [3, %pos] to get [4, %pos + 2] . That seems correct to me. This change is changing the test behavior and potentially dropping coverage.

I think there's no way trigger this composition without the first insert_slice being out-of-bounds. The folding pattern code can probably be simplified.

[MaheshRavishankar]

oh, that makes sense. Thanks for the explanation.

[MaheshRavishankar]

I dont see why. The existing test is taking an offset of [1, 2] and composing it with [3, %pos] to get [4, %pos + 2] . That seems correct to me. This change is changing the test behavior and potentially dropping coverage.

[MaheshRavishankar]

oh, that makes sense. Thanks for the explanation.