[Mlir-commits] [mlir] 7c0969c - [TOSA] Add pass to sanitize gather/scatter indices (#221200)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 10 10:56:01 PDT 2026
Author: Thomas Preud'homme
Date: 2026-09-10T18:55:56+01:00
New Revision: 7c0969c0842a1b19bcffd8bf3b67a76badb2c431
URL: https://github.com/llvm/llvm-project/commit/7c0969c0842a1b19bcffd8bf3b67a76badb2c431
DIFF: https://github.com/llvm/llvm-project/commit/7c0969c0842a1b19bcffd8bf3b67a76badb2c431.diff
LOG: [TOSA] Add pass to sanitize gather/scatter indices (#221200)
Add new TosaGatherScatterHardeningPass TOSA pass to clamp indices before
they are fed to a tosa.gather or tosa.scatter unless a suitable clamp
sequence already exists. This enables avoiding out-of-bounds accesses on
untrusted 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. The pass supports EXT-INT64
by using the same element type in the clamp min/max sequence as the
indices tensor. In case the indexed dimension is bigger than can be
represented in the element type bitwidth the maximum value in that
bitwidth is used.
Assisted-by: codex
Signed-off-by: Thomas Preud'homme <thomas.preudhomme at arm.com>
Added:
mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp
mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening-invalid.mlir
mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir
Modified:
mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
index 53c9ff57ca3a7..0ce6c4a059c60 100644
--- a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
@@ -199,6 +199,24 @@ def TosaDowngrade1p1To1p0Pass
}];
}
+def TosaGatherScatterHardeningPass
+ : Pass<"tosa-gather-scatter-hardening", "func::FuncOp"> {
+ let summary = "Clamp TOSA gather and scatter indices to valid bounds";
+ let description = [{
+ Bounds the indices used by each gather and scatter with TOSA minimum and
+ maximum operations so that all accesses are in bounds. This improves
+ security in the context of untrusted TOSA graph by preventing memory access
+ outside values (gather) and values_out (scatter) tensor.
+
+ Each operation is rewritten independently; a subsequent CSE pass is
+ required to share bounding ops between gather/scatter with the same input
+ and bounds. Other users continue to use the original indices.
+
+ Currently the pass only supports static index dimensions and fails
+ otherwise.
+ }];
+}
+
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..885ce0504deae
--- /dev/null
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaGatherScatterHardening.cpp
@@ -0,0 +1,191 @@
+//===- 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/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <type_traits>
+
+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 {
+
+/// Returns the effective upper bound when the indexed dimension is static.
+static FailureOr<int64_t> getIndexUpperBound(Operation *op) {
+ Value values = op->getOperand(0);
+ auto valuesType = dyn_cast<RankedTensorType>(values.getType());
+ if (!valuesType || valuesType.isDynamicDim(1)) {
+ op->emitOpError("requires a statically known indexed dimension for "
+ "gather/scatter hardening");
+ return failure();
+ }
+
+ auto indicesType = cast<ShapedType>(op->getOperand(1).getType());
+ auto elementType = cast<IntegerType>(indicesType.getElementType());
+
+ // The upper bound must be representable in the index element type.
+ int64_t maxRepresentable =
+ llvm::APInt::getSignedMaxValue(elementType.getWidth()).getSExtValue();
+ return std::min(valuesType.getDimSize(1) - 1, maxRepresentable);
+}
+
+/// Returns whether the indices already have sufficiently restrictive bounds.
+template <typename OuterOp, typename InnerOp>
+static bool isAlreadyHardened(Value indices, int64_t requiredUpperBound) {
+ static_assert(
+ (std::is_same_v<OuterOp, tosa::MinimumOp> &&
+ std::is_same_v<InnerOp, tosa::MaximumOp>) ||
+ (std::is_same_v<OuterOp, tosa::MaximumOp> &&
+ std::is_same_v<InnerOp, tosa::MinimumOp>),
+ "expected a tosa::MinimumOp/tosa::MaximumOp pair in either order");
+
+ auto outerOp = indices.getDefiningOp<OuterOp>();
+ if (!outerOp)
+ return false;
+
+ // Either operand can be the bound, including when both are constants. A
+ // constant match alone is not enough: try the other operand if it is unsafe.
+ for (unsigned boundOperand = 0; boundOperand < 2; ++boundOperand) {
+ llvm::APInt outerBound;
+ if (!matchPattern(outerOp->getOperand(boundOperand),
+ m_ConstantInt(&outerBound)))
+ continue;
+
+ llvm::APInt requiredUpper(outerBound.getBitWidth(),
+ static_cast<uint64_t>(requiredUpperBound));
+ // The outer bound must itself be in range, since it can override the inner
+ // bound, e.g. minimum(maximum(x, 0), -1) would produce -1. This guarantees
+ // the other operand is meeting the outer bound check (e.g. smaller or equal
+ // to the required upper bound if outer op is a minimum). The inner
+ // operation only needs to enforce the opposite bound.
+ if (outerBound.isNegative() || outerBound.sgt(requiredUpper))
+ continue;
+
+ auto matchesInnerBound = [&](Value value) {
+ llvm::APInt innerBound;
+ if (!matchPattern(value, m_ConstantInt(&innerBound)))
+ return false;
+ return isa<tosa::MinimumOp>(outerOp) ? !innerBound.isNegative()
+ : innerBound.sle(requiredUpper);
+ };
+
+ // Check whether other operand of the outer op is also a constant and is
+ // meeting the inner bound check. No inner op involved and the result is
+ // therefore completely within bound thanks to the earlier check.
+ Value innerResult = outerOp->getOperand(1 - boundOperand);
+ if (matchesInnerBound(innerResult))
+ return true;
+ // The other outer op operand is not a constant so check that the inner op
+ // enforces inner bound check.
+ if (auto innerOp = innerResult.getDefiningOp<InnerOp>())
+ if (llvm::any_of(innerOp->getOperands(), matchesInnerBound))
+ return true;
+ }
+ return false;
+}
+
+/// 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();
+}
+
+/// Independently hardens one gather or scatter operation's indices.
+template <typename OpTy>
+struct HardenIndexUsePattern final : OpRewritePattern<OpTy> {
+ HardenIndexUsePattern(MLIRContext *context, bool &hardeningFailed)
+ : OpRewritePattern<OpTy>(context), hardeningFailed(hardeningFailed) {}
+
+ LogicalResult matchAndRewrite(OpTy op,
+ PatternRewriter &rewriter) const override {
+ FailureOr<int64_t> upperBound = getIndexUpperBound(op.getOperation());
+ if (failed(upperBound)) {
+ hardeningFailed = true;
+ return rewriter.notifyMatchFailure(
+ op, "indexed dimension does not have a static upper bound");
+ }
+
+ Value indices = op->getOperand(1);
+ if (isAlreadyHardened<tosa::MinimumOp, tosa::MaximumOp>(indices,
+ *upperBound) ||
+ isAlreadyHardened<tosa::MaximumOp, tosa::MinimumOp>(indices,
+ *upperBound))
+ return rewriter.notifyMatchFailure(op, "indices are already hardened");
+
+ auto indicesType = cast<ShapedType>(indices.getType());
+ auto elementType = cast<IntegerType>(indicesType.getElementType());
+ Value lowerBound =
+ createIndexBoundConstant(rewriter, op.getLoc(), elementType, 0);
+ Value upperBoundValue = createIndexBoundConstant(rewriter, op.getLoc(),
+ elementType, *upperBound);
+ Value nonNegativeIndices =
+ tosa::MaximumOp::create(rewriter, op.getLoc(), indices.getType(),
+ indices, lowerBound)
+ .getResult();
+ Value clampedIndices =
+ tosa::MinimumOp::create(rewriter, op.getLoc(), indices.getType(),
+ nonNegativeIndices, upperBoundValue)
+ .getResult();
+
+ rewriter.modifyOpInPlace(
+ op, [&] { op->setOperand(/*indices=*/1, clampedIndices); });
+ return success();
+ }
+
+private:
+ bool &hardeningFailed;
+};
+
+struct TosaGatherScatterHardeningPass
+ : public tosa::impl::TosaGatherScatterHardeningPassBase<
+ TosaGatherScatterHardeningPass> {
+ using Base::Base;
+
+ void runOnOperation() override {
+ bool hardeningFailed = false;
+ RewritePatternSet patterns(&getContext());
+ patterns.add<HardenIndexUsePattern<tosa::GatherOp>,
+ HardenIndexUsePattern<tosa::ScatterOp>>(&getContext(),
+ hardeningFailed);
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))) ||
+ hardeningFailed)
+ signalPassFailure();
+ }
+};
+
+} // 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..4609f21962710
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening-invalid.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -tosa-gather-scatter-hardening -verify-diagnostics
+
+// A correctly hardened gather must not hide a scatter whose indexed dimension
+// is dynamic, even when they share the same bounded indices.
+func.func @dynamic_indexed_dimension(
+ %values: tensor<1x21x1xi8>, %dynamic_values: tensor<1x?x1xi8>,
+ %indices: tensor<1x2xi32>, %updates: tensor<1x2x1xi8>)
+ -> (tensor<1x2x1xi8>, tensor<1x?x1xi8>) {
+ %zero = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %upper = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %0 = tosa.maximum %indices, %zero : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %1 = tosa.minimum %0, %upper : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %gather = tosa.gather %values, %1 : (tensor<1x21x1xi8>, tensor<1x2xi32>) -> tensor<1x2x1xi8>
+ // expected-error at +1 {{'tosa.scatter' op requires a statically known indexed dimension for gather/scatter hardening}}
+ %scatter = tosa.scatter %dynamic_values, %1, %updates : (tensor<1x?x1xi8>, tensor<1x2xi32>, tensor<1x2x1xi8>) -> tensor<1x?x1xi8>
+ return %gather, %scatter : tensor<1x2x1xi8>, tensor<1x?x1xi8>
+}
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..5eabb07fece12
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/tosa-gather-scatter-hardening.mlir
@@ -0,0 +1,202 @@
+// RUN: mlir-opt %s -tosa-gather-scatter-hardening -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -tosa-gather-scatter-hardening -tosa-gather-scatter-hardening -split-input-file | FileCheck %s
+
+// Add a broadcastable clamp without changing other uses of the indices.
+// CHECK-LABEL: func.func @gather(
+// CHECK-DAG: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}>
+// CHECK-NEXT: %[[NONNEGATIVE:.*]] = tosa.maximum %arg1, %[[ZERO]]
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NEXT: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
+// CHECK-NEXT: return %[[GATHER]], %arg1
+func.func @gather(%values: tensor<3x21x5xi8>, %indices: tensor<3x6xi32>)
+ -> (tensor<3x6x5xi8>, tensor<3x6xi32>) {
+ %0 = tosa.gather %values, %indices : (tensor<3x21x5xi8>, tensor<3x6xi32>) -> tensor<3x6x5xi8>
+ return %0, %indices : tensor<3x6x5xi8>, tensor<3x6xi32>
+}
+
+// -----
+
+// Recognize an existing maximum followed by minimum, with commuted operands.
+// CHECK-LABEL: func.func @already_hardened_maximum_first(
+// CHECK: %[[NONNEGATIVE:.*]] = tosa.maximum %arg1,
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]],
+// CHECK-NEXT: %[[SCATTER:.*]] = tosa.scatter %arg0, %[[CLAMPED]], %arg2
+// CHECK-NEXT: return %[[SCATTER]]
+func.func @already_hardened_maximum_first(
+ %values: tensor<1x21x1xi8>, %indices: tensor<1x2xi32>,
+ %updates: tensor<1x2x1xi8>) -> tensor<1x21x1xi8> {
+ %zero = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %upper = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %0 = tosa.maximum %zero, %indices : (tensor<1x1xi32>, tensor<1x2xi32>) -> tensor<1x2xi32>
+ %1 = tosa.minimum %upper, %0 : (tensor<1x1xi32>, tensor<1x2xi32>) -> tensor<1x2xi32>
+ %2 = tosa.scatter %values, %1, %updates : (tensor<1x21x1xi8>, tensor<1x2xi32>, tensor<1x2x1xi8>) -> tensor<1x21x1xi8>
+ return %2 : tensor<1x21x1xi8>
+}
+
+// -----
+
+// Recognize minimum followed by maximum. Both inner operands are constants;
+// the data splat (25) must not hide the valid upper bound (20).
+// CHECK-LABEL: func.func @already_hardened_minimum_first(
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[DATA:.*]] = "tosa.const"() <{values = dense<25> : tensor<1x2xi32>}>
+// CHECK: %[[BOUNDED:.*]] = tosa.minimum %[[UPPER]], %[[DATA]]
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.maximum %[[BOUNDED]],
+// CHECK-NEXT: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
+// CHECK-NEXT: return %[[GATHER]]
+func.func @already_hardened_minimum_first(%values: tensor<1x21x1xi8>)
+ -> tensor<1x2x1xi8> {
+ %zero = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %upper = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %data = "tosa.const"() <{values = dense<25> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
+ %0 = tosa.minimum %upper, %data : (tensor<1x1xi32>, tensor<1x2xi32>) -> tensor<1x2xi32>
+ %1 = tosa.maximum %0, %zero : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %2 = tosa.gather %values, %1 : (tensor<1x21x1xi8>, tensor<1x2xi32>) -> tensor<1x2x1xi8>
+ return %2 : tensor<1x2x1xi8>
+}
+
+// -----
+
+// No inner op is needed: minimum(25, 20) produces 20, within [0, 20].
+// The first constant is out of bounds, so the second must be tried.
+// CHECK-LABEL: func.func @constant_indices(
+// CHECK-DAG: %[[DATA:.*]] = "tosa.const"() <{values = dense<25> : tensor<1x2xi32>}>
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}>
+// CHECK-NEXT: %[[BOUNDED:.*]] = tosa.minimum %[[DATA]], %[[UPPER]]
+// CHECK-NEXT: %[[GATHER:.*]] = tosa.gather %arg0, %[[BOUNDED]]
+// CHECK-NEXT: return %[[GATHER]]
+func.func @constant_indices(%values: tensor<1x21x1xi8>) -> tensor<1x2x1xi8> {
+ %data = "tosa.const"() <{values = dense<25> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
+ %upper = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %0 = tosa.minimum %data, %upper : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %1 = tosa.gather %values, %0 : (tensor<1x21x1xi8>, tensor<1x2xi32>) -> tensor<1x2x1xi8>
+ return %1 : tensor<1x2x1xi8>
+}
+
+// -----
+
+// Cap the upper bound at the largest representable index.
+// CHECK-LABEL: func.func @large_indexed_dimension(
+// CHECK-DAG: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<2147483647> : tensor<1x1xi32>}>
+// CHECK-NEXT: %[[NONNEGATIVE:.*]] = tosa.maximum %arg1, %[[ZERO]]
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NEXT: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
+// CHECK-NEXT: return %[[GATHER]]
+func.func @large_indexed_dimension(
+ %values: tensor<1x2147483649x1xi8>, %indices: tensor<1x1xi32>)
+ -> tensor<1x1x1xi8> {
+ %0 = tosa.gather %values, %indices : (tensor<1x2147483649x1xi8>, tensor<1x1xi32>) -> tensor<1x1x1xi8>
+ return %0 : tensor<1x1x1xi8>
+}
+
+// -----
+
+// An outer maximum can override a valid upper bound. Add a new clamp using
+// the i64 index type rather than accepting the existing [21, 21] range.
+// CHECK-LABEL: func.func @wrong_bound(
+// CHECK-DAG: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi64>}>
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi64>}>
+// CHECK-DAG: %[[BAD_BOUND:.*]] = "tosa.const"() <{values = dense<21> : tensor<1x1xi64>}>
+// CHECK-NEXT: %[[BOUNDED:.*]] = tosa.minimum %arg1, %[[UPPER]]
+// CHECK-NEXT: %[[ORIGINAL:.*]] = tosa.maximum %[[BOUNDED]], %[[BAD_BOUND]]
+// CHECK-NEXT: %[[NONNEGATIVE:.*]] = tosa.maximum %[[ORIGINAL]], %[[ZERO]]
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NEXT: %[[SCATTER:.*]] = tosa.scatter %arg0, %[[CLAMPED]], %arg2
+// CHECK-NEXT: return %[[SCATTER]]
+func.func @wrong_bound(
+ %values: tensor<1x21x1xi8>, %indices: tensor<1x2xi64>,
+ %updates: tensor<1x2x1xi8>) -> tensor<1x21x1xi8> {
+ %upper = "tosa.const"() <{values = dense<20> : tensor<1x1xi64>}> : () -> tensor<1x1xi64>
+ %bad_bound = "tosa.const"() <{values = dense<21> : tensor<1x1xi64>}> : () -> tensor<1x1xi64>
+ %0 = tosa.minimum %indices, %upper : (tensor<1x2xi64>, tensor<1x1xi64>) -> tensor<1x2xi64>
+ %1 = tosa.maximum %0, %bad_bound : (tensor<1x2xi64>, tensor<1x1xi64>) -> tensor<1x2xi64>
+ %2 = tosa.scatter %values, %1, %updates : (tensor<1x21x1xi8>, tensor<1x2xi64>, tensor<1x2x1xi8>) -> tensor<1x21x1xi8>
+ return %2 : tensor<1x21x1xi8>
+}
+
+// -----
+
+// A valid outer upper bound does not repair a negative inner lower bound.
+// CHECK-LABEL: func.func @invalid_inner_bound(
+// CHECK-DAG: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}>
+// CHECK: %[[INNER:.*]] = tosa.maximum %arg1,
+// CHECK-NEXT: %[[ORIGINAL:.*]] = tosa.minimum %[[INNER]], %[[UPPER]]
+// CHECK-NEXT: %[[NONNEGATIVE:.*]] = tosa.maximum %[[ORIGINAL]], %[[ZERO]]
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NEXT: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
+// CHECK-NEXT: return %[[GATHER]]
+func.func @invalid_inner_bound(%values: tensor<1x21x1xi8>, %indices: tensor<1x2xi32>)
+ -> tensor<1x2x1xi8> {
+ %lower = "tosa.const"() <{values = dense<-1> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %upper = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %0 = tosa.maximum %indices, %lower : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %1 = tosa.minimum %0, %upper : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %2 = tosa.gather %values, %1 : (tensor<1x21x1xi8>, tensor<1x2xi32>) -> tensor<1x2x1xi8>
+ return %2 : tensor<1x2x1xi8>
+}
+
+// -----
+
+// A negative outer minimum overrides the valid inner lower bound.
+// CHECK-LABEL: func.func @negative_outer_bound(
+// CHECK-DAG: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[NEGATIVE:.*]] = "tosa.const"() <{values = dense<-1> : tensor<1x1xi32>}>
+// CHECK-NEXT: %[[INNER:.*]] = tosa.maximum %arg1, %[[ZERO]]
+// CHECK-NEXT: %[[ORIGINAL:.*]] = tosa.minimum %[[INNER]], %[[NEGATIVE]]
+// CHECK-NEXT: %[[NONNEGATIVE:.*]] = tosa.maximum %[[ORIGINAL]], %[[ZERO]]
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NEXT: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
+// CHECK-NEXT: return %[[GATHER]]
+func.func @negative_outer_bound(%values: tensor<1x21x1xi8>, %indices: tensor<1x2xi32>)
+ -> tensor<1x2x1xi8> {
+ %zero = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %upper = "tosa.const"() <{values = dense<-1> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %0 = tosa.maximum %indices, %zero : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %1 = tosa.minimum %0, %upper : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %2 = tosa.gather %values, %1 : (tensor<1x21x1xi8>, tensor<1x2xi32>) -> tensor<1x2x1xi8>
+ return %2 : tensor<1x2x1xi8>
+}
+
+// -----
+
+// An upper bound alone does not prove that indices are nonnegative.
+// CHECK-LABEL: func.func @incomplete_clamp(
+// CHECK-DAG: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}>
+// CHECK-NEXT: %[[ORIGINAL:.*]] = tosa.minimum %arg1, %[[UPPER]]
+// CHECK-NEXT: %[[NONNEGATIVE:.*]] = tosa.maximum %[[ORIGINAL]], %[[ZERO]]
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NEXT: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
+// CHECK-NEXT: return %[[GATHER]]
+func.func @incomplete_clamp(%values: tensor<1x21x1xi8>, %indices: tensor<1x2xi32>)
+ -> tensor<1x2x1xi8> {
+ %upper = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %0 = tosa.minimum %indices, %upper : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %1 = tosa.gather %values, %0 : (tensor<1x21x1xi8>, tensor<1x2xi32>) -> tensor<1x2x1xi8>
+ return %1 : tensor<1x2x1xi8>
+}
+
+// -----
+
+// With two constants, one in-range operand is insufficient: minimum(-1, 20)
+// produces -1 and still needs hardening.
+// CHECK-LABEL: func.func @unsafe_constant_indices(
+// CHECK-DAG: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[UPPER:.*]] = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}>
+// CHECK-DAG: %[[DATA:.*]] = "tosa.const"() <{values = dense<-1> : tensor<1x2xi32>}>
+// CHECK-NEXT: %[[ORIGINAL:.*]] = tosa.minimum %[[DATA]], %[[UPPER]]
+// CHECK-NEXT: %[[NONNEGATIVE:.*]] = tosa.maximum %[[ORIGINAL]], %[[ZERO]]
+// CHECK-NEXT: %[[CLAMPED:.*]] = tosa.minimum %[[NONNEGATIVE]], %[[UPPER]]
+// CHECK-NEXT: %[[GATHER:.*]] = tosa.gather %arg0, %[[CLAMPED]]
+// CHECK-NEXT: return %[[GATHER]]
+func.func @unsafe_constant_indices(%values: tensor<1x21x1xi8>) -> tensor<1x2x1xi8> {
+ %data = "tosa.const"() <{values = dense<-1> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
+ %upper = "tosa.const"() <{values = dense<20> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+ %0 = tosa.minimum %data, %upper : (tensor<1x2xi32>, tensor<1x1xi32>) -> tensor<1x2xi32>
+ %1 = tosa.gather %values, %0 : (tensor<1x21x1xi8>, tensor<1x2xi32>) -> tensor<1x2x1xi8>
+ return %1 : tensor<1x2x1xi8>
+}
More information about the Mlir-commits
mailing list