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

[Luke Hutton]

================
@@ -0,0 +1,221 @@
+//===- 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 "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) {
----------------
lhutton1 wrote:

The op definition prevents use of dims with a value of 0, so this check is redundant

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