Func Input converted to symbol

2026-04-27 13:48:03 +02:00
parent 951baca106
commit 5c839e62c1
8 changed files with 151 additions and 75 deletions
--- a/src/PIM/Conversion/SpatialToPim/Patterns.cpp
+++ b/src/PIM/Conversion/SpatialToPim/Patterns.cpp
@@ -0,0 +1,88 @@
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+
+#include "llvm/ADT/STLExtras.h"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::func::FuncOp funcOp, PatternRewriter& rewriter) const override {
+
+    if (funcOp.getArguments().empty())
+      return failure();
+
+    if (llvm::all_of(funcOp.getArguments(),
+                     [](mlir::BlockArgument blockArgument) { return blockArgument.use_empty(); }))
+      return failure();
+
+    Location loc = funcOp.getLoc();
+
+    for (auto [index, arg] : llvm::enumerate(funcOp.getArguments())) {
+      if (arg.getUses().empty())
+        continue;
+
+      rewriter.setInsertionPoint(funcOp.getOperation());
+
+      assert(isa<mlir::RankedTensorType>(arg.getType()));
+
+      auto argRankedTensorType = llvm::dyn_cast<mlir::RankedTensorType>(arg.getType());
+      mlir::MemRefType memRefType =
+        mlir::MemRefType::get(argRankedTensorType.getShape(), argRankedTensorType.getElementType());
+
+      std::string argName = "arg_" + std::to_string(index);
+
+      memref::GlobalOp::create(rewriter,
+                                             loc,
+                                             rewriter.getStringAttr(argName),
+                                             rewriter.getStringAttr("private"),
+                                             TypeAttr::get(memRefType),
+                                             {},
+                                             {},
+                                             {});
+
+      for (auto& argUses : llvm::make_early_inc_range(arg.getUses())) {
+        auto argUser = argUses.getOwner();
+        if (auto spatCompute = dyn_cast<spatial::SpatCompute>(argUser)) {
+          auto BBArgIndex = argUses.getOperandNumber() - spatCompute.getInputs().getBeginOperandIndex();
+          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+          auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+          auto toTensor = bufferization::ToTensorOp::create(rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+
+
+          rewriter.startOpModification(spatCompute.getOperation());
+          BBArgValue.replaceAllUsesWith(toTensor);
+          spatCompute.getInputsMutable().erase(BBArgIndex);
+          spatCompute.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatCompute.getOperation());
+        }
+        else {
+          rewriter.setInsertionPoint(argUser);
+          auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+          rewriter.startOpModification(argUser);
+          argUses.set(getGlobalOp);
+          rewriter.finalizeOpModification(argUser);
+        }
+      }
+    }
+
+    return success();
+  }
+};
+
+} // namespace
+void populateGlobalTensorToMemrefPatterns(RewritePatternSet& patterns) {
+  patterns.add<FuncOpArgToGlobalMemoryPattern>(patterns.getContext());
+}
+
+} // namespace onnx_mlir