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

[Thomas Preud'homme]

================
@@ -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) {
+    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");
+  }
+
+  // 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);
+
+  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) {
----------------
RoboTux wrote:

Because it allowed to use a single min/max for several gather/scatter. As you suggested I'll rework it to let this be done as a later pass.

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