[Mlir-commits] [mlir] [TOSA] Add pass to sanitize gather/scatter indices (PR #221200)

[Luke Hutton]

================
@@ -0,0 +1,161 @@
+//===- 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 <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 {
+
+/// 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.
+static bool isAlreadyHardened(Value indices, int64_t requiredUpperBound) {
+  auto minimumOp = indices.getDefiningOp<tosa::MinimumOp>();
+  if (!minimumOp)
+    return false;
+
+  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>();
----------------
lhutton1 wrote:

We might want to consider the case where the user has already attempted some form of input sanitisation but used `min(max())` rather than `max(min())`. I appreciate this is somewhat a consequence of not being able to use `tosa.clamp`

https://github.com/llvm/llvm-project/pull/221200