[Mlir-commits] [mlir] [TOSA] Add pass to sanitize gather/scatter indices (PR #221200)
Thomas Preud'homme
llvmlistbot at llvm.org
Fri Sep 4 08:38:50 PDT 2026
https://github.com/RoboTux updated https://github.com/llvm/llvm-project/pull/221200
>From 62738c2a40281881c8888415945d63e0c1b0d586 Mon Sep 17 00:00:00 2001
From: Thomas Preud'homme <thomas.preudhomme at arm.com>
Date: Fri, 4 Sep 2026 11:30:26 +0100
Subject: [PATCH 1/2] [TOSA] Add pass to sanitize gather/scatter indices
Add new TosaGatherScatterHardeningPass TOSA pass to clamp indices before
they are fed to a tosa.gather or tosa.scatter unless a suitable clamp
already exists. This enables avoiding out of bound accesses on
unstrusted indices in tosa.gather and tosa.scatter lowering.
The pass fails if the dimension indexed by the indices tensor of any of
the tosa.gather or tosa.scatter is dynamic since clamp has min/max as
attributes. The pass supports EXT-INT64 by using the same bitwidth in
clamp min/max attributes as the indices tensor element type. In case the
indexed dimension is bigger than can be represented in the element type
bitwidth the maximum bitwidth is used.
To keep IR minimal, several tosa.gather and tosa.scatter sharing the same
dimension sizes and indices tensor will share a single clamp.
Code was generated interactively with codex until I was satisfied with
the resulting patch. This description is my own.
Signed-off-by: Thomas Preud'homme <thomas.preudhomme at arm.com>
---
.../mlir/Dialect/Tosa/Transforms/Passes.td | 17 ++
.../Dialect/Tosa/Transforms/CMakeLists.txt | 1 +
.../Transforms/TosaGatherScatterHardening.cpp | 189 ++++++++++++++++++
...tosa-gather-scatter-hardening-invalid.mlir | 9 +
.../Tosa/tosa-gather-scatter-hardening.mlir | 107 ++++++++++
5 files changed, 323 insertions(+)
create mode 100644 mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp
create mode 100644 mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening-invalid.mlir
create mode 100644 mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir
diff --git a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
index 96a0b080d6a6f..9e4ce659a84dd 100644
--- a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
@@ -199,6 +199,23 @@ def TosaDowngrade1p1To1p0Pass
}];
}
+def TosaGatherScatterHardeningPass
+ : Pass<"tosa-gather-scatter-hardening", "func::FuncOp"> {
+ let summary = "Clamp TOSA gather and scatter indices to valid bounds";
+ let description = [{
+ Inserts a TOSA clamp on the indices used by each gather and scatter so that
+ all accesses are in bounds. When the same indices value is used by several
+ gather or scatter operations that require identical clamp bounds, one clamp
+ is shared by all of them. Other users continue to use the original indices.
+
+ The indexed dimension must be statically known and non-empty.
+ }];
+
+ let dependentDialects = [
+ "tosa::TosaDialect"
+ ];
+}
+
def TosaNarrowI64ToI32Pass : Pass<"tosa-narrow-i64-to-i32", "func::FuncOp"> {
let summary = "Narrow I64 TOSA operations to I32";
let description = [{
diff --git a/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
index 1fd18bb5a395b..da23374ed65e1 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_dialect_library(MLIRTosaTransforms
TosaDecomposeTransposeConv.cpp
TosaDecomposeDepthwise.cpp
TosaFolders.cpp
+ TosaGatherScatterHardening.cpp
TosaInferShapes.cpp
TosaLayerwiseConstantFoldPass.cpp
TosaMakeBroadcastable.cpp
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp
new file mode 100644
index 0000000000000..96ecb92954116
--- /dev/null
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp
@@ -0,0 +1,189 @@
+//===- TosaGatherScatterHardening.cpp -------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass that clamps gather and scatter indices to the
+// statically known bounds of their indexed tensors.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+#include <algorithm>
+#include <cstdint>
+
+namespace mlir {
+namespace tosa {
+#define GEN_PASS_DEF_TOSAGATHERSCATTERHARDENINGPASS
+#include "mlir/Dialect/Tosa/Transforms/Passes.h.inc"
+} // namespace tosa
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::tosa;
+
+namespace {
+
+struct IndexUseGroup {
+ IndexUseGroup(Value indices, int64_t upperBound)
+ : indices(indices), upperBound(upperBound) {}
+
+ Value indices;
+ SmallVector<Operation *> users;
+ int64_t upperBound;
+};
+
+using GroupsForIndices = DenseMap<Value, SmallVector<unsigned>>;
+
+/// Validates one gather/scatter index use and adds it to a compatible group.
+static LogicalResult collectIndexUse(Operation *op, Value indices, Value values,
+ GroupsForIndices &groupsForIndices,
+ SmallVectorImpl<IndexUseGroup> &groups) {
+ auto valuesType = dyn_cast<RankedTensorType>(values.getType());
+ if (!valuesType || valuesType.getRank() <= 1 || valuesType.isDynamicDim(1)) {
+ return op->emitOpError(
+ "requires a statically known indexed dimension for gather/scatter "
+ "hardening");
+ }
+
+ int64_t indexedSize = valuesType.getDimSize(1);
+ if (indexedSize <= 0) {
+ return op->emitOpError(
+ "requires a non-empty indexed dimension for gather/scatter hardening");
+ }
+
+ auto indicesType = dyn_cast<ShapedType>(indices.getType());
+ auto elementType = indicesType
+ ? dyn_cast<IntegerType>(indicesType.getElementType())
+ : IntegerType();
+ if (!elementType || elementType.getWidth() > 64) {
+ return op->emitOpError(
+ "requires indices with an integer element type of at most 64 bits for "
+ "gather/scatter hardening");
+ }
+
+ // Clamp bounds must be representable in the index element type. Group only
+ // users whose resulting clamp bounds are identical.
+ int64_t maxRepresentable =
+ llvm::APInt::getSignedMaxValue(elementType.getWidth()).getSExtValue();
+ int64_t upperBound = std::min(indexedSize - 1, maxRepresentable);
+
+ unsigned groupIndex = groups.size();
+ SmallVector<unsigned> &candidateGroups = groupsForIndices[indices];
+ for (unsigned candidate : candidateGroups) {
+ if (groups[candidate].upperBound == upperBound) {
+ groupIndex = candidate;
+ break;
+ }
+ }
+ if (groupIndex == groups.size()) {
+ groups.emplace_back(indices, upperBound);
+ candidateGroups.push_back(groupIndex);
+ }
+
+ groups[groupIndex].users.push_back(op);
+ return success();
+}
+
+/// Collects and validates all gather/scatter index uses in the function.
+static LogicalResult
+collectIndexUseGroups(func::FuncOp funcOp,
+ SmallVectorImpl<IndexUseGroup> &groups) {
+ GroupsForIndices groupsForIndices;
+ bool analysisFailed = false;
+
+ funcOp.walk([&](Operation *op) {
+ llvm::TypeSwitch<Operation *>(op)
+ .Case<tosa::GatherOp>([&](tosa::GatherOp gatherOp) {
+ analysisFailed |= failed(collectIndexUse(op, gatherOp.getIndices(),
+ gatherOp.getValues(),
+ groupsForIndices, groups));
+ })
+ .Case<tosa::ScatterOp>([&](tosa::ScatterOp scatterOp) {
+ analysisFailed |= failed(collectIndexUse(op, scatterOp.getIndices(),
+ scatterOp.getValuesIn(),
+ groupsForIndices, groups));
+ });
+ });
+
+ return analysisFailed ? failure() : success();
+}
+
+/// Returns whether the indices already have sufficiently restrictive bounds.
+static bool isAlreadyHardened(Value indices, IntegerAttr maxAttr) {
+ auto clampOp = indices.getDefiningOp<tosa::ClampOp>();
+ if (!clampOp)
+ return false;
+
+ auto existingMin = dyn_cast<IntegerAttr>(clampOp.getMinValAttr());
+ auto existingMax = dyn_cast<IntegerAttr>(clampOp.getMaxValAttr());
+ return existingMin && existingMax && !existingMin.getValue().isNegative() &&
+ existingMax.getValue().sle(maxAttr.getValue());
+}
+
+/// Creates a clamp for the group and rewires its gather/scatter users.
+static void hardenIndexUseGroup(IndexUseGroup &group, OpBuilder &builder) {
+ Value indices = group.indices;
+ auto indicesType = cast<ShapedType>(indices.getType());
+ auto elementType = cast<IntegerType>(indicesType.getElementType());
+ unsigned bitWidth = elementType.getWidth();
+
+ IntegerAttr minAttr =
+ IntegerAttr::get(elementType, llvm::APInt::getZero(bitWidth));
+ IntegerAttr maxAttr = IntegerAttr::get(
+ elementType,
+ llvm::APInt(bitWidth, static_cast<uint64_t>(group.upperBound)));
+
+ // Keep the pass idempotent and avoid nesting a second clamp around an
+ // existing clamp that is at least as restrictive as the required one.
+ if (isAlreadyHardened(indices, maxAttr))
+ return;
+
+ if (Operation *definingOp = indices.getDefiningOp())
+ builder.setInsertionPointAfter(definingOp);
+ else
+ builder.setInsertionPointToStart(cast<BlockArgument>(indices).getOwner());
+
+ Value clampedIndices =
+ tosa::ClampOp::create(builder, group.users.front()->getLoc(),
+ indices.getType(), indices, minAttr, maxAttr)
+ .getResult();
+
+ for (Operation *user : group.users)
+ user->setOperand(/*indices=*/1, clampedIndices);
+}
+
+struct TosaGatherScatterHardeningPass
+ : public tosa::impl::TosaGatherScatterHardeningPassBase<
+ TosaGatherScatterHardeningPass> {
+ using Base::Base;
+
+ void runOnOperation() override {
+ SmallVector<IndexUseGroup> groups;
+ // Do not partially harden the function when one operation cannot be made
+ // safe using a statically bounded tosa.clamp.
+ if (failed(collectIndexUseGroups(getOperation(), groups))) {
+ signalPassFailure();
+ return;
+ }
+
+ OpBuilder builder(&getContext());
+ for (IndexUseGroup &group : groups)
+ hardenIndexUseGroup(group, builder);
+ }
+};
+
+} // namespace
diff --git a/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening-invalid.mlir b/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening-invalid.mlir
new file mode 100644
index 0000000000000..9f3c4937f20e0
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening-invalid.mlir
@@ -0,0 +1,9 @@
+// RUN: not mlir-opt %s -tosa-gather-scatter-hardening 2>&1 | FileCheck %s
+
+func.func @dynamic_indexed_dimension(%arg0: tensor<3x?x5xi8>,
+ %arg1: tensor<3x6xi32>)
+ -> tensor<3x6x5xi8> {
+ // CHECK: error: 'tosa.gather' op requires a statically known indexed dimension for gather/scatter hardening
+ %0 = tosa.gather %arg0, %arg1 : (tensor<3x?x5xi8>, tensor<3x6xi32>) -> tensor<3x6x5xi8>
+ return %0 : tensor<3x6x5xi8>
+}
diff --git a/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir b/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir
new file mode 100644
index 0000000000000..1964a415e6180
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir
@@ -0,0 +1,107 @@
+// RUN: mlir-opt %s -tosa-gather-scatter-hardening -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @gather(
+// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg1 {max_val = 20 : i32, min_val = 0 : i32} : (tensor<3x6xi32>) -> tensor<3x6xi32>
+// CHECK: %[[OTHER_USE:.*]] = tosa.add %arg1, %arg2
+// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMP]]
+// CHECK: return %[[GATHER]], %[[OTHER_USE]]
+func.func @gather(%arg0: tensor<3x21x5xi8>, %arg1: tensor<3x6xi32>,
+ %arg2: tensor<3x6xi32>)
+ -> (tensor<3x6x5xi8>, tensor<3x6xi32>) {
+ %0 = tosa.add %arg1, %arg2 : (tensor<3x6xi32>, tensor<3x6xi32>) -> tensor<3x6xi32>
+ %1 = tosa.gather %arg0, %arg1 : (tensor<3x21x5xi8>, tensor<3x6xi32>) -> tensor<3x6x5xi8>
+ return %1, %0 : tensor<3x6x5xi8>, tensor<3x6xi32>
+}
+
+// -----
+
+// One clamp is shared by gather and scatter when their bounds match.
+
+// CHECK-LABEL: func.func @shared_indices(
+// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg2 {max_val = 20 : i32, min_val = 0 : i32} : (tensor<2x4xi32>) -> tensor<2x4xi32>
+// CHECK-NOT: tosa.clamp
+// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMP]]
+// CHECK: %[[SCATTER:.*]] = tosa.scatter %arg1, %[[CLAMP]], %arg3
+// CHECK-NOT: tosa.clamp
+// CHECK: return %[[GATHER]], %[[SCATTER]]
+func.func @shared_indices(%arg0: tensor<2x21x3xf32>,
+ %arg1: tensor<2x21x3xf32>,
+ %arg2: tensor<2x4xi32>,
+ %arg3: tensor<2x4x3xf32>)
+ -> (tensor<2x4x3xf32>, tensor<2x21x3xf32>) {
+ %0 = tosa.gather %arg0, %arg2 : (tensor<2x21x3xf32>, tensor<2x4xi32>) -> tensor<2x4x3xf32>
+ %1 = tosa.scatter %arg1, %arg2, %arg3 : (tensor<2x21x3xf32>, tensor<2x4xi32>, tensor<2x4x3xf32>) -> tensor<2x21x3xf32>
+ return %0, %1 : tensor<2x4x3xf32>, tensor<2x21x3xf32>
+}
+
+// -----
+
+// Gather and scatter get separate clamps when their bounds differ.
+
+// CHECK-LABEL: func.func @different_bounds(
+// CHECK-DAG: %[[GATHER_CLAMP:.*]] = tosa.clamp %arg2 {max_val = 20 : i32, min_val = 0 : i32} : (tensor<2x4xi32>) -> tensor<2x4xi32>
+// CHECK-DAG: %[[SCATTER_CLAMP:.*]] = tosa.clamp %arg2 {max_val = 51 : i32, min_val = 0 : i32} : (tensor<2x4xi32>) -> tensor<2x4xi32>
+// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[GATHER_CLAMP]]
+// CHECK: %[[SCATTER:.*]] = tosa.scatter %arg1, %[[SCATTER_CLAMP]], %arg3
+// CHECK: return %[[GATHER]], %[[SCATTER]]
+func.func @different_bounds(%arg0: tensor<2x21x3xf32>,
+ %arg1: tensor<2x52x3xf32>,
+ %arg2: tensor<2x4xi32>,
+ %arg3: tensor<2x4x3xf32>)
+ -> (tensor<2x4x3xf32>, tensor<2x52x3xf32>) {
+ %0 = tosa.gather %arg0, %arg2 : (tensor<2x21x3xf32>, tensor<2x4xi32>) -> tensor<2x4x3xf32>
+ %1 = tosa.scatter %arg1, %arg2, %arg3 : (tensor<2x52x3xf32>, tensor<2x4xi32>, tensor<2x4x3xf32>) -> tensor<2x52x3xf32>
+ return %0, %1 : tensor<2x4x3xf32>, tensor<2x52x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @i64_indices(
+// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg1 {max_val = 26 : i64, min_val = 0 : i64} : (tensor<13x4xi64>) -> tensor<13x4xi64>
+// CHECK: tosa.scatter %arg0, %[[CLAMP]], %arg2
+func.func @i64_indices(%arg0: tensor<13x27x3xi16>,
+ %arg1: tensor<13x4xi64>,
+ %arg2: tensor<13x4x3xi16>) -> tensor<13x27x3xi16> {
+ %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x27x3xi16>, tensor<13x4xi64>, tensor<13x4x3xi16>) -> tensor<13x27x3xi16>
+ return %0 : tensor<13x27x3xi16>
+}
+
+// -----
+
+// An i32 index cannot represent K - 1 here, so its signed maximum is a safe
+// upper clamp bound.
+
+// CHECK-LABEL: func.func @large_indexed_dimension(
+// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg1 {max_val = 2147483647 : i32, min_val = 0 : i32}
+// CHECK: tosa.gather %arg0, %[[CLAMP]]
+func.func @large_indexed_dimension(
+ %arg0: tensor<1x2147483649x1xi8>, %arg1: tensor<1x1xi32>)
+ -> tensor<1x1x1xi8> {
+ %0 = tosa.gather %arg0, %arg1 : (tensor<1x2147483649x1xi8>, tensor<1x1xi32>) -> tensor<1x1x1xi8>
+ return %0 : tensor<1x1x1xi8>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @already_hardened(
+// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg1 {max_val = 20 : i32, min_val = 0 : i32}
+// CHECK-NOT: tosa.clamp
+// CHECK: tosa.gather %arg0, %[[CLAMP]]
+func.func @already_hardened(%arg0: tensor<3x21x5xi8>,
+ %arg1: tensor<3x6xi32>) -> tensor<3x6x5xi8> {
+ %0 = tosa.clamp %arg1 {max_val = 20 : i32, min_val = 0 : i32} : (tensor<3x6xi32>) -> tensor<3x6xi32>
+ %1 = tosa.gather %arg0, %0 : (tensor<3x21x5xi8>, tensor<3x6xi32>) -> tensor<3x6x5xi8>
+ return %1 : tensor<3x6x5xi8>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @no_gather_or_scatter(
+// CHECK-NOT: tosa.clamp
+// CHECK: %[[ADD:.*]] = tosa.add %arg0, %arg1
+// CHECK: return %[[ADD]]
+func.func @no_gather_or_scatter(%arg0: tensor<2x4xi32>,
+ %arg1: tensor<2x4xi32>) -> tensor<2x4xi32> {
+ %0 = tosa.add %arg0, %arg1 : (tensor<2x4xi32>, tensor<2x4xi32>) -> tensor<2x4xi32>
+ return %0 : tensor<2x4xi32>
+}
>From f9899508cf56938f7551e154e526e98704fbb9f3 Mon Sep 17 00:00:00 2001
From: Thomas Preud'homme <thomas.preudhomme at arm.com>
Date: Fri, 4 Sep 2026 16:19:12 +0100
Subject: [PATCH 2/2] Use maximum+minimum instead of clamp
Clamp does not support int32 or int64 so use const(x2)+maximum+minimum
sequences instead.
Signed-off-by: Thomas Preud'homme <thomas.preudhomme at arm.com>
---
.../mlir/Dialect/Tosa/Transforms/Passes.td | 9 ++-
.../Transforms/TosaGatherScatterHardening.cpp | 78 ++++++++++++------
.../Tosa/tosa-gather-scatter-hardening.mlir | 80 +++++++++++++------
3 files changed, 115 insertions(+), 52 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
index 9e4ce659a84dd..174f138a66312 100644
--- a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
@@ -203,10 +203,11 @@ def TosaGatherScatterHardeningPass
: Pass<"tosa-gather-scatter-hardening", "func::FuncOp"> {
let summary = "Clamp TOSA gather and scatter indices to valid bounds";
let description = [{
- Inserts a TOSA clamp on the indices used by each gather and scatter so that
- all accesses are in bounds. When the same indices value is used by several
- gather or scatter operations that require identical clamp bounds, one clamp
- is shared by all of them. Other users continue to use the original indices.
+ Bounds the indices used by each gather and scatter with TOSA minimum and
+ maximum operations so that all accesses are in bounds. When the same
+ indices value is used by several gather or scatter operations that require
+ identical bounds, one bounding sequence is shared by all of them. Other
+ users continue to use the original indices.
The indexed dimension must be statically known and non-empty.
}];
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp
index 96ecb92954116..0b0676663785c 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp
@@ -16,7 +16,9 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
@@ -75,8 +77,8 @@ static LogicalResult collectIndexUse(Operation *op, Value indices, Value values,
"gather/scatter hardening");
}
- // Clamp bounds must be representable in the index element type. Group only
- // users whose resulting clamp bounds are identical.
+ // Bounds must be representable in the index element type. Group only users
+ // whose resulting bounds are identical.
int64_t maxRepresentable =
llvm::APInt::getSignedMaxValue(elementType.getWidth()).getSExtValue();
int64_t upperBound = std::min(indexedSize - 1, maxRepresentable);
@@ -123,33 +125,55 @@ collectIndexUseGroups(func::FuncOp funcOp,
}
/// Returns whether the indices already have sufficiently restrictive bounds.
-static bool isAlreadyHardened(Value indices, IntegerAttr maxAttr) {
- auto clampOp = indices.getDefiningOp<tosa::ClampOp>();
- if (!clampOp)
+static bool isAlreadyHardened(Value indices, int64_t requiredUpperBound) {
+ auto minimumOp = indices.getDefiningOp<tosa::MinimumOp>();
+ if (!minimumOp)
return false;
- auto existingMin = dyn_cast<IntegerAttr>(clampOp.getMinValAttr());
- auto existingMax = dyn_cast<IntegerAttr>(clampOp.getMaxValAttr());
- return existingMin && existingMax && !existingMin.getValue().isNegative() &&
- existingMax.getValue().sle(maxAttr.getValue());
+ Value maximumResult = minimumOp.getInput1();
+ llvm::APInt upperBound;
+ if (!matchPattern(minimumOp.getInput2(), m_ConstantInt(&upperBound))) {
+ maximumResult = minimumOp.getInput2();
+ if (!matchPattern(minimumOp.getInput1(), m_ConstantInt(&upperBound)))
+ return false;
+ }
+
+ auto maximumOp = maximumResult.getDefiningOp<tosa::MaximumOp>();
+ if (!maximumOp)
+ return false;
+
+ llvm::APInt lowerBound;
+ if (!matchPattern(maximumOp.getInput2(), m_ConstantInt(&lowerBound)) &&
+ !matchPattern(maximumOp.getInput1(), m_ConstantInt(&lowerBound)))
+ return false;
+
+ unsigned bitWidth = upperBound.getBitWidth();
+ llvm::APInt requiredUpper(bitWidth,
+ static_cast<uint64_t>(requiredUpperBound));
+ return lowerBound.getBitWidth() == bitWidth && !lowerBound.isNegative() &&
+ !upperBound.isNegative() && upperBound.sle(requiredUpper);
+}
+
+/// Creates a rank-two splat constant suitable for index broadcasting.
+static Value createIndexBoundConstant(OpBuilder &builder, Location loc,
+ IntegerType elementType, int64_t value) {
+ auto type = RankedTensorType::get({1, 1}, elementType);
+ auto valueAttr =
+ IntegerAttr::get(elementType, llvm::APInt(elementType.getWidth(),
+ static_cast<uint64_t>(value)));
+ auto values = DenseElementsAttr::get(type, valueAttr);
+ return tosa::ConstOp::create(builder, loc, type, values).getResult();
}
-/// Creates a clamp for the group and rewires its gather/scatter users.
+/// Creates minimum/maximum bounds and rewires the group's index users.
static void hardenIndexUseGroup(IndexUseGroup &group, OpBuilder &builder) {
Value indices = group.indices;
auto indicesType = cast<ShapedType>(indices.getType());
auto elementType = cast<IntegerType>(indicesType.getElementType());
- unsigned bitWidth = elementType.getWidth();
- IntegerAttr minAttr =
- IntegerAttr::get(elementType, llvm::APInt::getZero(bitWidth));
- IntegerAttr maxAttr = IntegerAttr::get(
- elementType,
- llvm::APInt(bitWidth, static_cast<uint64_t>(group.upperBound)));
-
- // Keep the pass idempotent and avoid nesting a second clamp around an
- // existing clamp that is at least as restrictive as the required one.
- if (isAlreadyHardened(indices, maxAttr))
+ // Keep the pass idempotent and avoid nesting a second min/max pair around
+ // indices that are already at least as restricted as required.
+ if (isAlreadyHardened(indices, group.upperBound))
return;
if (Operation *definingOp = indices.getDefiningOp())
@@ -157,9 +181,17 @@ static void hardenIndexUseGroup(IndexUseGroup &group, OpBuilder &builder) {
else
builder.setInsertionPointToStart(cast<BlockArgument>(indices).getOwner());
+ Location loc = group.users.front()->getLoc();
+ Value lowerBound = createIndexBoundConstant(builder, loc, elementType, 0);
+ Value upperBound =
+ createIndexBoundConstant(builder, loc, elementType, group.upperBound);
+ Value nonNegativeIndices =
+ tosa::MaximumOp::create(builder, loc, indices.getType(), indices,
+ lowerBound)
+ .getResult();
Value clampedIndices =
- tosa::ClampOp::create(builder, group.users.front()->getLoc(),
- indices.getType(), indices, minAttr, maxAttr)
+ tosa::MinimumOp::create(builder, loc, indices.getType(),
+ nonNegativeIndices, upperBound)
.getResult();
for (Operation *user : group.users)
@@ -174,7 +206,7 @@ struct TosaGatherScatterHardeningPass
void runOnOperation() override {
SmallVector<IndexUseGroup> groups;
// Do not partially harden the function when one operation cannot be made
- // safe using a statically bounded tosa.clamp.
+ // safe using statically bounded minimum and maximum operations.
if (failed(collectIndexUseGroups(getOperation(), groups))) {
signalPassFailure();
return;
diff --git a/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir b/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir
index 1964a415e6180..fa044308468f2 100644
--- a/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir
@@ -1,9 +1,12 @@
// RUN: mlir-opt %s -tosa-gather-scatter-hardening -split-input-file | FileCheck %s
// CHECK-LABEL: func.func @gather(
-// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg1 {max_val = 20 : i32, min_val = 0 : i32} : (tensor<3x6xi32>) -> tensor<3x6xi32>
+// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[NONNEGATIVE:.*]] = tosa.maximum %arg1, %[[ZERO]]
+// CHECK: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
// CHECK: %[[OTHER_USE:.*]] = tosa.add %arg1, %arg2
-// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMP]]
+// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
// CHECK: return %[[GATHER]], %[[OTHER_USE]]
func.func @gather(%arg0: tensor<3x21x5xi8>, %arg1: tensor<3x6xi32>,
%arg2: tensor<3x6xi32>)
@@ -15,14 +18,20 @@ func.func @gather(%arg0: tensor<3x21x5xi8>, %arg1: tensor<3x6xi32>,
// -----
-// One clamp is shared by gather and scatter when their bounds match.
+// One bounding sequence is shared by gather and scatter when their bounds
+// match.
// CHECK-LABEL: func.func @shared_indices(
-// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg2 {max_val = 20 : i32, min_val = 0 : i32} : (tensor<2x4xi32>) -> tensor<2x4xi32>
-// CHECK-NOT: tosa.clamp
-// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMP]]
-// CHECK: %[[SCATTER:.*]] = tosa.scatter %arg1, %[[CLAMP]], %arg3
-// CHECK-NOT: tosa.clamp
+// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[NONNEGATIVE:.*]] = tosa.maximum %arg2, %[[ZERO]]
+// CHECK: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NOT: tosa.maximum
+// CHECK-NOT: tosa.minimum
+// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
+// CHECK: %[[SCATTER:.*]] = tosa.scatter %arg1, %[[CLAMPED]], %arg3
+// CHECK-NOT: tosa.maximum
+// CHECK-NOT: tosa.minimum
// CHECK: return %[[GATHER]], %[[SCATTER]]
func.func @shared_indices(%arg0: tensor<2x21x3xf32>,
%arg1: tensor<2x21x3xf32>,
@@ -36,13 +45,19 @@ func.func @shared_indices(%arg0: tensor<2x21x3xf32>,
// -----
-// Gather and scatter get separate clamps when their bounds differ.
+// Gather and scatter get separate bounding sequences when their bounds differ.
// CHECK-LABEL: func.func @different_bounds(
-// CHECK-DAG: %[[GATHER_CLAMP:.*]] = tosa.clamp %arg2 {max_val = 20 : i32, min_val = 0 : i32} : (tensor<2x4xi32>) -> tensor<2x4xi32>
-// CHECK-DAG: %[[SCATTER_CLAMP:.*]] = tosa.clamp %arg2 {max_val = 51 : i32, min_val = 0 : i32} : (tensor<2x4xi32>) -> tensor<2x4xi32>
-// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[GATHER_CLAMP]]
-// CHECK: %[[SCATTER:.*]] = tosa.scatter %arg1, %[[SCATTER_CLAMP]], %arg3
+// CHECK: %[[SCATTER_ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[SCATTER_UPPER:.*]] = "tosa.const"() <{values = dense<51> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[SCATTER_NONNEGATIVE:.*]] = tosa.maximum %arg2, %[[SCATTER_ZERO]]
+// CHECK: %[[SCATTER_CLAMPED:.*]] = tosa.minimum %[[SCATTER_NONNEGATIVE]], %[[SCATTER_UPPER]]
+// CHECK: %[[GATHER_ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[GATHER_UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[GATHER_NONNEGATIVE:.*]] = tosa.maximum %arg2, %[[GATHER_ZERO]]
+// CHECK: %[[GATHER_CLAMPED:.*]] = tosa.minimum %[[GATHER_NONNEGATIVE]], %[[GATHER_UPPER]]
+// CHECK: %[[GATHER:.*]] = tosa.gather %arg0, %[[GATHER_CLAMPED]]
+// CHECK: %[[SCATTER:.*]] = tosa.scatter %arg1, %[[SCATTER_CLAMPED]], %arg3
// CHECK: return %[[GATHER]], %[[SCATTER]]
func.func @different_bounds(%arg0: tensor<2x21x3xf32>,
%arg1: tensor<2x52x3xf32>,
@@ -57,8 +72,11 @@ func.func @different_bounds(%arg0: tensor<2x21x3xf32>,
// -----
// CHECK-LABEL: func.func @i64_indices(
-// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg1 {max_val = 26 : i64, min_val = 0 : i64} : (tensor<13x4xi64>) -> tensor<13x4xi64>
-// CHECK: tosa.scatter %arg0, %[[CLAMP]], %arg2
+// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi64>}> : () -> tensor<1x1xi64>
+// CHECK: %[[UPPER:.*]] = "tosa.const"() <{values = dense<26> : tensor<1x1xi64>}> : () -> tensor<1x1xi64>
+// CHECK: %[[NONNEGATIVE:.*]] = tosa.maximum %arg1, %[[ZERO]]
+// CHECK: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK: tosa.scatter %arg0, %[[CLAMPED]], %arg2
func.func @i64_indices(%arg0: tensor<13x27x3xi16>,
%arg1: tensor<13x4xi64>,
%arg2: tensor<13x4x3xi16>) -> tensor<13x27x3xi16> {
@@ -69,11 +87,14 @@ func.func @i64_indices(%arg0: tensor<13x27x3xi16>,
// -----
// An i32 index cannot represent K - 1 here, so its signed maximum is a safe
-// upper clamp bound.
+// upper bound.
// CHECK-LABEL: func.func @large_indexed_dimension(
-// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg1 {max_val = 2147483647 : i32, min_val = 0 : i32}
-// CHECK: tosa.gather %arg0, %[[CLAMP]]
+// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[UPPER:.*]] = "tosa.const"() <{values = dense<2147483647> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[NONNEGATIVE:.*]] = tosa.maximum %arg1, %[[ZERO]]
+// CHECK: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK: tosa.gather %arg0, %[[CLAMPED]]
func.func @large_indexed_dimension(
%arg0: tensor<1x2147483649x1xi8>, %arg1: tensor<1x1xi32>)
-> tensor<1x1x1xi8> {
@@ -84,20 +105,29 @@ func.func @large_indexed_dimension(
// -----
// CHECK-LABEL: func.func @already_hardened(
-// CHECK: %[[CLAMP:.*]] = tosa.clamp %arg1 {max_val = 20 : i32, min_val = 0 : i32}
-// CHECK-NOT: tosa.clamp
-// CHECK: tosa.gather %arg0, %[[CLAMP]]
+// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+// CHECK: %[[NONNEGATIVE:.*]] = tosa.maximum %arg1, %[[ZERO]]
+// CHECK: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NOT: tosa.maximum
+// CHECK-NOT: tosa.minimum
+// CHECK: tosa.gather %arg0, %[[CLAMPED]]
func.func @already_hardened(%arg0: tensor<3x21x5xi8>,
%arg1: tensor<3x6xi32>) -> tensor<3x6x5xi8> {
- %0 = tosa.clamp %arg1 {max_val = 20 : i32, min_val = 0 : i32} : (tensor<3x6xi32>) -> tensor<3x6xi32>
- %1 = tosa.gather %arg0, %0 : (tensor<3x21x5xi8>, tensor<3x6xi32>) -> tensor<3x6x5xi8>
- return %1 : tensor<3x6x5xi8>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %1 = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %2 = tosa.maximum %arg1, %0 : (tensor<3x6xi32>, tensor<1x1xi32>) -> tensor<3x6xi32>
+ %3 = tosa.minimum %2, %1 : (tensor<3x6xi32>, tensor<1x1xi32>) -> tensor<3x6xi32>
+ %4 = tosa.gather %arg0, %3 : (tensor<3x21x5xi8>, tensor<3x6xi32>) -> tensor<3x6x5xi8>
+ return %4 : tensor<3x6x5xi8>
}
// -----
// CHECK-LABEL: func.func @no_gather_or_scatter(
-// CHECK-NOT: tosa.clamp
+// CHECK-NOT: tosa.const
+// CHECK-NOT: tosa.maximum
+// CHECK-NOT: tosa.minimum
// CHECK: %[[ADD:.*]] = tosa.add %arg0, %arg1
// CHECK: return %[[ADD]]
func.func @no_gather_or_scatter(%arg0: tensor<2x4xi32>,
More information about the Mlir-commits
mailing list