[Mlir-commits] [mlir] [TOSA] Prevent OoB accesses in gather/scatter (PR #213242)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 31 03:51:49 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Thomas Preud'homme (RoboTux)
<details>
<summary>Changes</summary>
Clamp indices in tosa.gather and tosa.scatter lowerings to prevent out
of bound accesses of the input tensor.
---
Full diff: https://github.com/llvm/llvm-project/pull/213242.diff
4 Files Affected:
- (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp (+7)
- (modified) mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp (+9-2)
- (modified) mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir (+17-4)
- (modified) mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir (+4-1)
``````````diff
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index e3f40c57eb312..425dea8c0fbd6 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -2551,6 +2551,7 @@ class GatherConverter : public OpConversionPattern<tosa::GatherOp> {
rewriter.getContext()),
rewriter.getMultiDimIdentityMap(resultTy.getRank())};
+ Value kSzVal = rewriter.createOrFold<tensor::DimOp>(loc, input, 1);
auto genericOp = linalg::GenericOp::create(
rewriter, loc, ArrayRef<Type>({resultTy}), ValueRange{indices},
ValueRange{emptyTensor}, affineMaps,
@@ -2560,6 +2561,12 @@ class GatherConverter : public OpConversionPattern<tosa::GatherOp> {
auto index0 = linalg::IndexOp::create(rewriter, loc, 0);
Value index1 = arith::IndexCastOp::create(
rewriter, loc, rewriter.getIndexType(), indexValue);
+ auto outOfBound = arith::CmpIOp::create(
+ rewriter, loc, rewriter.getI1Type(), arith::CmpIPredicate::uge,
+ index1, kSzVal);
+ index1 =
+ arith::SelectOp::create(rewriter, loc, rewriter.getIndexType(),
+ outOfBound, kSzVal, index1);
auto index2 = linalg::IndexOp::create(rewriter, loc, 2);
Value extract = tensor::ExtractOp::create(
rewriter, loc, input, ValueRange{index0, index1, index2});
diff --git a/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp b/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
index 7e9c9090c51df..c8c1897faab3d 100644
--- a/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
+++ b/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
@@ -120,15 +120,22 @@ class ScatterOpConverter : public OpRewritePattern<tosa::ScatterOp> {
auto lbs = Repeated<Value>(2, zero);
auto steps = Repeated<Value>(2, one);
auto ubs = llvm::SmallVector<Value>{{dimN, dimW}};
+ Value kSzVal = rewriter.createOrFold<tensor::DimOp>(loc, valuesIn, 1);
auto buildBody = [&](OpBuilder &builder, Location loc, ValueRange ivs,
ValueRange args) -> scf::ValueVector {
auto n = ivs[0];
- // Read the index and cast it to index type
+ // Read the index, cast it to index type and clamp it
auto index = tensor::ExtractOp::create(builder, loc, indices, ivs);
auto castIndex = arith::IndexCastOp::create(
builder, loc, builder.getIndexType(), index);
+ auto outOfBound = arith::CmpIOp::create(
+ builder, loc, builder.getI1Type(), arith::CmpIPredicate::uge,
+ castIndex, kSzVal);
+ auto clampedIndex =
+ arith::SelectOp::create(builder, loc, builder.getIndexType(),
+ outOfBound, kSzVal, castIndex);
// Offset, sizes, and strides for the input tensor
auto inputOffset = llvm::to_vector(ivs);
@@ -141,7 +148,7 @@ class ScatterOpConverter : public OpRewritePattern<tosa::ScatterOp> {
inputOffset, sizes, strides);
// Insert the slice into the output accumulator tensor.
- llvm::SmallVector<Value> outputOffset = {n, castIndex, zero};
+ llvm::SmallVector<Value> outputOffset = {n, clampedIndex, zero};
auto updated = tensor::InsertSliceOp::create(
builder, loc, slice, args[0], outputOffset, sizes, strides);
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index a803ee7d99153..ac3273f1148e4 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -1884,12 +1884,15 @@ func.func @argmax_dyn_axis(%arg0 : tensor<3x?xi32>) -> () {
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
func.func @gather_float(%arg0: tensor<2x3x2xf32>, %arg1: tensor<2x3xi32>) -> () {
// CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[C3:.+]] = arith.constant 3
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xf32>)
// CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: f32)
// CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]]
+ // CHECK: %[[COND:.+]] = arith.cmpi uge, %[[CAST]], %[[C3]]
+ // CHECK: %[[SELECT:.+]] = arith.select %[[COND]], %[[C3]], %[[CAST]]
// CHECK: %[[IDX2:.+]] = linalg.index 2
- // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<2x3x2xf32>
+ // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[SELECT]], %[[IDX2]]] : tensor<2x3x2xf32>
// CHECK: linalg.yield %[[EXTRACT]]
%0 = tosa.gather %arg0, %arg1 : (tensor<2x3x2xf32>, tensor<2x3xi32>) -> tensor<2x3x2xf32>
return
@@ -1904,12 +1907,15 @@ func.func @gather_float_dyn(%arg0: tensor<?x3x2xf32>, %arg1: tensor<?x3xi32>) ->
// CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]])
+ // CHECK: %[[C3:.+]] = arith.constant 3
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<?x3xi32>) outs(%[[INIT]] : tensor<?x3x2xf32>)
// CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: f32)
// CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]]
+ // CHECK: %[[COND:.+]] = arith.cmpi uge, %[[CAST]], %[[C3]]
+ // CHECK: %[[SELECT:.+]] = arith.select %[[COND]], %[[C3]], %[[CAST]]
// CHECK: %[[IDX2:.+]] = linalg.index 2
- // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<?x3x2xf32>
+ // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[SELECT]], %[[IDX2]]] : tensor<?x3x2xf32>
// CHECK: linalg.yield %[[EXTRACT]]
%0 = tosa.gather %arg0, %arg1 : (tensor<?x3x2xf32>, tensor<?x3xi32>) -> tensor<?x3x2xf32>
return
@@ -1928,12 +1934,16 @@ func.func @gather_float_all_dynamic(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?x
// CHECK: %[[C2:.+]] = arith.constant 2
// CHECK: %[[CHANNEL:.+]] = tensor.dim %[[ARG0]], %[[C2]]
// CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]], %[[INDEX]], %[[CHANNEL]])
+ // CHECK: %[[C1_2:.+]] = arith.constant 1
+ // CHECK: %[[RANGE:.+]] = tensor.dim %[[ARG0]], %[[C1_2]]
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<?x?xi32>) outs(%[[INIT]] : tensor<?x?x?xf32>)
// CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: f32)
// CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]]
+ // CHECK: %[[COND:.+]] = arith.cmpi uge, %[[CAST]], %[[RANGE]]
+ // CHECK: %[[SELECT:.+]] = arith.select %[[COND]], %[[RANGE]], %[[CAST]]
// CHECK: %[[IDX2:.+]] = linalg.index 2
- // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<?x?x?xf32>
+ // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[SELECT]], %[[IDX2]]] : tensor<?x?x?xf32>
// CHECK: linalg.yield %[[EXTRACT]]
%0 = tosa.gather %arg0, %arg1 : (tensor<?x?x?xf32>, tensor<?x?xi32>) -> tensor<?x?x?xf32>
return
@@ -1946,12 +1956,15 @@ func.func @gather_float_all_dynamic(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?x
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
func.func @gather_int(%arg0: tensor<2x3x2xi32>, %arg1: tensor<2x3xi32>) -> () {
// CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[C3:.+]] = arith.constant 3
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xi32>)
// CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: i32)
// CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]]
+ // CHECK: %[[COND:.+]] = arith.cmpi uge, %[[CAST]], %[[C3]]
+ // CHECK: %[[SELECT:.+]] = arith.select %[[COND]], %[[C3]], %[[CAST]]
// CHECK: %[[IDX2:.+]] = linalg.index 2
- // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<2x3x2xi32>
+ // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[SELECT]], %[[IDX2]]] : tensor<2x3x2xi32>
// CHECK: linalg.yield %[[EXTRACT]]
%0 = tosa.gather %arg0, %arg1 : (tensor<2x3x2xi32>, tensor<2x3xi32>) -> tensor<2x3x2xi32>
return
diff --git a/mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir b/mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir
index b6f2383ac81fc..3fdcb2ad37536 100644
--- a/mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir
+++ b/mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir
@@ -66,14 +66,17 @@ func.func @scatter_test(%values_in: tensor<3x7x5xi32>, %indices : tensor<3x6xi32
// CHECK-DAG: [[C_3:%.+]] = arith.constant 3 : index
// CHECK-DAG: [[C_5:%.+]] = arith.constant 5 : index
// CHECK-DAG: [[C_6:%.+]] = arith.constant 6 : index
+ // CHECK-DAG: [[C_7:%.+]] = arith.constant 7 : index
// CHECK-DAG: [[C_0_0:%.+]] = arith.constant 0 : index
// CHECK-DAG: [[C_1_0:%.+]] = arith.constant 1 : index
// CHECK: [[RESULT_0:%.+]] = scf.for [[ITER_VAR_0:%.+]] = [[C_0_0]] to [[C_3]] step [[C_1_0]] iter_args([[ITER_ARG_0:%.+]] = [[VALUES_IN]]) -> (tensor<3x7x5xi32>) {
// CHECK: [[RESULT_1:%.+]] = scf.for [[ITER_VAR_1:%.+]] = [[C_0_0]] to [[C_6]] step [[C_1_0]] iter_args([[ITER_ARG_1:%.+]] = [[ITER_ARG_0]]) -> (tensor<3x7x5xi32>) {
// CHECK-DAG: [[EXTRACTED:%.+]] = tensor.extract [[INDICES]][[[ITER_VAR_0]], [[ITER_VAR_1]]] : tensor<3x6xi32>
// CHECK-DAG: [[EXTRACTED_CAST:%.+]] = arith.index_cast [[EXTRACTED]] : i32 to index
+ // CHECK-DAG: [[EXTRACTED_COND:%.+]] = arith.cmpi uge, [[EXTRACTED_CAST]], [[C_7]]
+ // CHECK-DAG: [[EXTRACTED_CLAMP:%.+]] = arith.select [[EXTRACTED_COND]], [[C_7]], [[EXTRACTED_CAST]]
// CHECK-DAG: [[EXTRACTED_SLICE:%.+]] = tensor.extract_slice [[INPUT]][[[ITER_VAR_0]], [[ITER_VAR_1]], [[C_0_0]]] [[[C_1_0]], [[C_1_0]], [[C_5]]] [[[C_1_0]], [[C_1_0]], [[C_1_0]]] : tensor<3x6x5xi32> to tensor<?x?x?xi32>
- // CHECK-DAG: [[INSERTED_SLICE:%.+]] = tensor.insert_slice [[EXTRACTED_SLICE]] into [[ITER_ARG_1]][[[ITER_VAR_0]], [[EXTRACTED_CAST]], [[C_0_0]]] [[[C_1_0]], [[C_1_0]], [[C_5]]] [[[C_1_0]], [[C_1_0]], [[C_1_0]]] : tensor<?x?x?xi32> into tensor<3x7x5xi32>
+ // CHECK-DAG: [[INSERTED_SLICE:%.+]] = tensor.insert_slice [[EXTRACTED_SLICE]] into [[ITER_ARG_1]][[[ITER_VAR_0]], [[EXTRACTED_CLAMP]], [[C_0_0]]] [[[C_1_0]], [[C_1_0]], [[C_5]]] [[[C_1_0]], [[C_1_0]], [[C_1_0]]] : tensor<?x?x?xi32> into tensor<3x7x5xi32>
// CHECK: scf.yield [[INSERTED_SLICE]] : tensor<3x7x5xi32>
// CHECK: }
// CHECK: scf.yield [[RESULT_1]] : tensor<3x7x5xi32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/213242
More information about the Mlir-commits
mailing list