add PIM accelerator

src/PIM/Conversion/ONNXToSpatial/Utils/SpatialReducer.hpp

@@ -0,0 +1,83 @@
+#pragma once
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialCommon.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/Casting.h"
+
+namespace onnx_mlir {
+
+using ResNum = unsigned int;
+
+using ComputeAndResNum = std::pair<spatial::SpatWeightedCompute, ResNum>;
+
+struct SpatialReducerChange {
+  Operation *fromOp;
+  unsigned int fromOpResNum;
+  Operation *toOp;
+  unsigned int toOpOperandNum;
+};
+
+using OpAndResNum = std::pair<Operation *, ResNum>;
+
+class SpatialReducer {
+
+public:
+  SpatialReducer(ConversionPatternRewriter &rewriter) : rewriter(rewriter) {}
+
+  OpAndResNum applyReducePattern(
+      SmallVector<ComputeAndResNum> &computeOpsAndResNum,
+      std::function<Value(const Value &, const Value &)> reduce,
+      std::function<Value(const Value &)> preprocess,
+      std::function<Value(const Value &)> postprocess);
+
+  OpAndResNum applyAddMapReduction(SmallVector<ComputeAndResNum> &computeOps,
+      ConversionPatternRewriter &rewriter, Value biasTile, MapOperations mapOp);
+
+  void finalizeReduceUpdates();
+
+  ~SpatialReducer() {
+    if (!reducesFinalized) {
+      finalizeReduceUpdates();
+    }
+  }
+
+  Value createImgConcatOp(
+      llvm::SmallVector<llvm::SmallVector<llvm::SmallVector<OpAndResNum>>>
+          &outputTiles,
+      Location &loc, Type outputType);
+
+  Value resolveValueFromOpAndResNum(OpAndResNum &opAndResNum);
+
+private:
+  [[nodiscard("computeOp result number gets updated")]] ResNum
+  applyResultProcessing(ComputeAndResNum computeOpAndResNum,
+      std::function<Value(const Value &)> processFun,
+      ConversionPatternRewriter &rewriter);
+
+  /**
+   * @brief Update the results of a ComputeOp.
+   *
+   * This function updates the results of a ComputeOp by taking a look at the
+   operands of its yieldOp.
+   * If the ComputeOp was replaced, it updates `opToReplacedCompute` with the
+   replaced ComputeOp.
+   *
+   * @param computeOp The ComputeOp to update the results of.
+   */
+  void updateResultsOfCompute(Operation *computeOp);
+
+  ConversionPatternRewriter &rewriter;
+  bool reducesFinalized = false;
+
+  // List of changes to be applied after the reduction is finalized
+  SmallVector<SpatialReducerChange, 4> reducerChanges;
+  // List of computeOps that need to be replaced with new results
+  SmallVector<spatial::SpatWeightedCompute> computeOpNeedingResUpdate;
+
+  std::unordered_map<Operation *, spatial::SpatWeightedCompute> opToReplacedCompute;
+
+  static llvm::SmallPtrSet<Operation *, 16> oldComputeOpsReplaced;
+};
+
+} // namespace onnx_mlir