[Mlir-commits] [mlir] [mlir][tosa] Add slice bounds validation (PR #238385)

[Michael Platings]

https://github.com/mplatings created https://github.com/llvm/llvm-project/pull/238385

Assisted-by: Codex

>From 8877205d21dab0e87a09f4470f96df8fb4fb2c66 Mon Sep 17 00:00:00 2001
From: Michael Platings <michael.platings at arm.com>
Date: Mon, 2 Nov 2026 18:09:02 +0000
Subject: [PATCH] [mlir][tosa] Add slice bounds validation

Assisted-by: Codex
---
 .../Tosa/Transforms/TosaValidation.cpp        | 22 ++++++++
 .../Dialect/Tosa/tosa-validation-invalid.mlir | 53 +++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 69e1682ffd075b..f8a2c114c5bc89 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -1450,6 +1450,28 @@ LogicalResult checkErrorIfSlice(Operation *op) {
     }
   }
 
+  // ERROR_IF(any(start + size > shape1))
+  if (hasStartValues && hasSizeValues) {
+    if (const auto inputType =
+            dyn_cast<RankedTensorType>(sliceOp.getInput1().getType())) {
+      for (const auto &[index, vals] : llvm::enumerate(llvm::zip_equal(
+               startValues, sizeValues, inputType.getShape()))) {
+        const auto &[start, size, inputDim] = vals;
+        if (ShapedType::isDynamic(inputDim))
+          continue;
+        // Earlier checks already ensured that start and slice are non-negative.
+        // Avoid overflowing start + size.
+        if (start > inputDim || size > inputDim - start)
+          return op->emitOpError(
+                     "start + size must be less than or equal to input "
+                     "dimension size, got start=")
+                 << start << ", size=" << size
+                 << " vs input dim size=" << inputDim << " at dimension "
+                 << index;
+      }
+    }
+  }
+
   // ERROR_IF(is_block_scale<in_out_t>() && start[rank(shape1) - 1] %
   // get_innermost_block_size<in_out_t>() != 0)
   if (hasStartValues && !startValues.empty()) {
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-invalid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-invalid.mlir
index 9a73d52eb812b3..f5b4395a1cf393 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-invalid.mlir
@@ -437,3 +437,56 @@ func.func @test_slice_block_scaled_rank0_start(%arg0: tensor<*x!tosa.block_scale
   %0 = tosa.slice %arg0, %start, %size : (tensor<*x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, !tosa.shape<0>, !tosa.shape<0>) -> tensor<*x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
   return %0 : tensor<*x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
 }
+
+// -----
+
+// Reject the CTS slice that begins at the last input element.
+func.func @test_slice_start_size_outside_bounds(%arg0: tensor<11x11x7xi32>) -> tensor<2x4x3xi32> {
+  %start = tosa.const_shape values(dense<[10, 10, 6]> : tensor<3xindex>) : () -> !tosa.shape<3>
+  %size = tosa.const_shape values(dense<[2, 4, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
+  // expected-error at +1 {{'tosa.slice' op start + size must be less than or equal to input dimension size, got start=10, size=2 vs input dim size=11 at dimension 0}}
+  %0 = tosa.slice %arg0, %start, %size : (tensor<11x11x7xi32>, !tosa.shape<3>, !tosa.shape<3>) -> tensor<2x4x3xi32>
+  return %0 : tensor<2x4x3xi32>
+}
+
+// -----
+
+// Check known dimensions even when an earlier input dimension is dynamic.
+func.func @test_slice_static_dim_outside_bounds(%arg0: tensor<?x11x7xi32>) -> tensor<1x2x1xi32> {
+  %start = tosa.const_shape values(dense<[0, 10, 0]> : tensor<3xindex>) : () -> !tosa.shape<3>
+  %size = tosa.const_shape values(dense<[1, 2, 1]> : tensor<3xindex>) : () -> !tosa.shape<3>
+  // expected-error at +1 {{'tosa.slice' op start + size must be less than or equal to input dimension size, got start=10, size=2 vs input dim size=11 at dimension 1}}
+  %0 = tosa.slice %arg0, %start, %size : (tensor<?x11x7xi32>, !tosa.shape<3>, !tosa.shape<3>) -> tensor<1x2x1xi32>
+  return %0 : tensor<1x2x1xi32>
+}
+
+// -----
+
+// Reject an oversized start without signed overflow in start + size.
+func.func @test_slice_start_size_overflow(%arg0: tensor<11xi32>) -> tensor<2xi32> {
+  %start = tosa.const_shape values(dense<[9223372036854775807]> : tensor<1xindex>) : () -> !tosa.shape<1>
+  %size = tosa.const_shape values(dense<[2]> : tensor<1xindex>) : () -> !tosa.shape<1>
+  // expected-error at +1 {{'tosa.slice' op start + size must be less than or equal to input dimension size, got start=9223372036854775807, size=2 vs input dim size=11 at dimension 0}}
+  %0 = tosa.slice %arg0, %start, %size : (tensor<11xi32>, !tosa.shape<1>, !tosa.shape<1>) -> tensor<2xi32>
+  return %0 : tensor<2xi32>
+}
+
+// -----
+
+// Accept a slice whose end exactly matches every input dimension.
+func.func @test_slice_exact_boundary(%arg0: tensor<11x11x7xi32>) -> tensor<2x4x3xi32> {
+  %start = tosa.const_shape values(dense<[9, 7, 4]> : tensor<3xindex>) : () -> !tosa.shape<3>
+  %size = tosa.const_shape values(dense<[2, 4, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
+  %0 = tosa.slice %arg0, %start, %size : (tensor<11x11x7xi32>, !tosa.shape<3>, !tosa.shape<3>) -> tensor<2x4x3xi32>
+  return %0 : tensor<2x4x3xi32>
+}
+
+// -----
+
+// Do not reject bounds that depend on unknown input dimensions.
+func.func @test_slice_dynamic_input_dimension(%arg0: tensor<?x11x7xi32>) -> tensor<2x4x3xi32> {
+  %start = tosa.const_shape values(dense<[10, 7, 4]> : tensor<3xindex>) : () -> !tosa.shape<3>
+  %size = tosa.const_shape values(dense<[2, 4, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
+  %0 = tosa.slice %arg0, %start, %size : (tensor<?x11x7xi32>, !tosa.shape<3>, !tosa.shape<3>) -> tensor<2x4x3xi32>
+  return %0 : tensor<2x4x3xi32>
+}