Refactor + ReduceMean batched

src/PIM/Conversion/ONNXToSpatial/Common/IndexingUtils.cpp

@@ -0,0 +1,104 @@
+#include "IndexingUtils.hpp"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+
+#include "llvm/ADT/APInt.h"
+
+#include <algorithm>
+
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+int64_t normalizeAxis(int64_t axis, int64_t rank) { return axis >= 0 ? axis : rank + axis; }
+
+FailureOr<int64_t> normalizeAxisChecked(int64_t axis, int64_t rank) {
+  int64_t normalizedAxis = normalizeAxis(axis, rank);
+  if (normalizedAxis < 0 || normalizedAxis >= rank)
+    return failure();
+  return normalizedAxis;
+}
+
+int64_t normalizeIndex(int64_t index, int64_t dimSize) { return index >= 0 ? index : dimSize + index; }
+
+static SmallVector<int64_t> normalizeAxesImpl(std::optional<ArrayAttr> axesAttr, int64_t rank) {
+  SmallVector<int64_t> normalizedAxes;
+  if (!axesAttr) {
+    normalizedAxes.reserve(rank);
+    for (int64_t axis = 0; axis < rank; ++axis)
+      normalizedAxes.push_back(axis);
+  }
+  else {
+    normalizedAxes.reserve(axesAttr->size());
+    for (Attribute attr : *axesAttr)
+      normalizedAxes.push_back(normalizeAxis(cast<IntegerAttr>(attr).getInt(), rank));
+    llvm::sort(normalizedAxes);
+    normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end());
+  }
+  return normalizedAxes;
+}
+
+SmallVector<int64_t> normalizeAxes(ArrayAttr axesAttr, int64_t rank) {
+  return normalizeAxesImpl(std::optional<ArrayAttr>(axesAttr), rank);
+}
+
+SmallVector<int64_t> normalizeAxes(std::optional<ArrayAttr> axesAttr, int64_t rank) {
+  return normalizeAxesImpl(axesAttr, rank);
+}
+
+FailureOr<SmallVector<int64_t>> normalizeAxesChecked(std::optional<ArrayAttr> axesAttr, int64_t rank) {
+  SmallVector<int64_t> normalizedAxes = normalizeAxesImpl(axesAttr, rank);
+  for (int64_t axis : normalizedAxes)
+    if (axis < 0 || axis >= rank)
+      return failure();
+  return normalizedAxes;
+}
+
+FailureOr<SmallVector<int64_t>> normalizeAxesChecked(ArrayAttr axesAttr, int64_t rank) {
+  return normalizeAxesChecked(std::optional<ArrayAttr>(axesAttr), rank);
+}
+
+Value createAffineApplyOrConstant(PatternRewriter& rewriter, Location loc, AffineExpr expr, ValueRange operands) {
+  AffineMap map = AffineMap::get(/*dimCount=*/operands.size(), /*symbolCount=*/0, expr);
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  return createAffineApplyOrFoldedConstant(rewriter, loc, map, operands, anchorOp);
+}
+
+Value multiplyIndexByConstant(PatternRewriter& rewriter, Operation* anchorOp, Value value, int64_t multiplier) {
+  if (multiplier == 0)
+    return getOrCreateHostIndexConstant(rewriter, anchorOp, 0);
+  if (multiplier == 1)
+    return value;
+
+  MLIRContext* context = rewriter.getContext();
+  AffineExpr d0 = getAffineDimExpr(0, context);
+  return createAffineApplyOrConstant(rewriter, anchorOp->getLoc(), d0 * multiplier, ValueRange {value});
+}
+
+Value modIndexByConstant(PatternRewriter& rewriter, Location loc, Value value, int64_t divisor) {
+  if (divisor == 1)
+    return getOrCreateHostIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+
+  MLIRContext* context = rewriter.getContext();
+  AffineExpr d0 = getAffineDimExpr(0, context);
+  return createAffineApplyOrConstant(rewriter, loc, d0 % divisor, ValueRange {value});
+}
+
+Value floorDivIndexByConstant(PatternRewriter& rewriter, Location loc, Value value, int64_t divisor) {
+  if (divisor == 1)
+    return value;
+
+  MLIRContext* context = rewriter.getContext();
+  AffineExpr d0 = getAffineDimExpr(0, context);
+  return createAffineApplyOrConstant(rewriter, loc, d0.floorDiv(divisor), ValueRange {value});
+}
+
+Value getOrMaterializeIndexValue(PatternRewriter& rewriter, Location loc, OpFoldResult value) {
+  if (auto attr = dyn_cast<Attribute>(value))
+    return getOrCreateHostIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), cast<IntegerAttr>(attr).getInt());
+  return cast<Value>(value);
+}
+
+} // namespace onnx_mlir