huge refactor for high RewritePatterns usage and less ad-hoc cpp code

remove Spatial many ops in favor of tensor ops like in pim

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -8,21 +8,17 @@
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
-#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
-#include "llvm/Support/raw_os_ostream.h"
-
-#include <fstream>
-#include <iterator>
-#include <utility>
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostLegality.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PrePatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Compiler/CompilerOptions.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -33,8 +29,6 @@ namespace onnx_mlir {
 
 namespace {
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
-
 struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ONNXToSpatialPass)
   StringRef getArgument() const override { return "convert-onnx-to-spatial"; }
@@ -44,71 +38,64 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
   ONNXToSpatialPass(const ONNXToSpatialPass& pass) {}
 
   void runOnOperation() override;
-
-private:
-  void annotateWeightsConstants(func::FuncOp funcOp) const;
-  LogicalResult encapsulateGlobalInstruction(func::FuncOp funcOp);
-  LogicalResult promoteConstantInputsToWeights(func::FuncOp funcOp);
-  void populateEmptyFunction(func::FuncOp funcOp);
 };
 
 } // namespace
 
-static void foldSingleLaneComputeBatches(func::FuncOp funcOp) {
+static void populateEmptyFunction(func::FuncOp funcOp) {
   IRRewriter rewriter(funcOp.getContext());
-  SmallVector<spatial::SpatComputeBatch> batchOps;
-  funcOp.walk([&](spatial::SpatComputeBatch batchOp) { batchOps.push_back(batchOp); });
+  IRMapping mapper;
+  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
+  if (!computes.empty())
+    return;
 
-  for (auto batchOp : batchOps) {
-    if (batchOp.getLaneCount() != 1)
-      continue;
+  auto returnOp = cast<func::ReturnOp>(funcOp.getFunctionBody().front().getTerminator());
+  rewriter.setInsertionPoint(returnOp);
 
-    auto loc = batchOp.getLoc();
-    rewriter.setInsertionPoint(batchOp);
-    auto computeOp =
-      spatial::SpatCompute::create(rewriter, loc, batchOp.getResultTypes(), batchOp.getWeights(), batchOp.getInputs());
-    computeOp.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(batchOp.getWeights().size()), static_cast<int>(batchOp.getInputs().size())});
+  SmallVector<Type> sourceTypes;
+  SmallVector<Location> sourceLocs;
+  sourceTypes.reserve(funcOp.getNumArguments());
+  sourceLocs.reserve(funcOp.getNumArguments());
+  for (Value source : funcOp.getArguments()) {
+    sourceTypes.push_back(source.getType());
+    sourceLocs.push_back(source.getLoc());
+  }
 
-    Block& templateBlock = batchOp.getBody().front();
-    SmallVector<Type> blockArgTypes;
-    SmallVector<Location> blockArgLocs;
-    for (BlockArgument arg : templateBlock.getArguments()) {
-      blockArgTypes.push_back(arg.getType());
-      blockArgLocs.push_back(loc);
-    }
-    auto* newBlock =
-      rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
+  auto newCompute = spatial::SpatCompute::create(
+    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
+  auto* newBlock = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLocs);
+  for (auto [blockArg, computeArg] : llvm::zip(newBlock->getArguments(), newCompute.getOperands()))
+    mapper.map(computeArg, blockArg);
+  newCompute.getProperties().setOperandSegmentSizes({0, static_cast<int>(sourceTypes.size())});
 
-    IRMapping mapper;
-    for (auto [oldArg, newArg] : llvm::zip(templateBlock.getArguments(), newBlock->getArguments()))
-      mapper.map(oldArg, newArg);
-    rewriter.setInsertionPointToEnd(newBlock);
-    for (Operation& op : templateBlock)
+  rewriter.setInsertionPointToEnd(newBlock);
+  for (Operation& op : funcOp.getOps())
+    if (!isa<spatial::SpatCompute, func::ReturnOp>(&op))
       rewriter.clone(op, mapper);
 
-    batchOp.replaceAllUsesWith(computeOp.getResults());
-    rewriter.eraseOp(batchOp);
-  }
+  auto yield = spatial::SpatYieldOp::create(rewriter, funcOp.getLoc(), returnOp.getOperands());
+  for (size_t i = 0; i < yield.getNumOperands(); ++i)
+    yield.setOperand(i, mapper.lookupOrDefault(yield.getOperand(i)));
+
+  for (Operation& op : llvm::make_early_inc_range(funcOp.getOps()))
+    if (!isa<spatial::SpatCompute, func::ReturnOp>(&op)) {
+      op.dropAllUses();
+      rewriter.eraseOp(&op);
+    }
+
+  for (auto [index, computeResult] : llvm::enumerate(newCompute.getResults()))
+    returnOp.setOperand(index, computeResult);
 }
 
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   MLIRContext* ctx = &getContext();
 
-  RewritePatternSet mergeActivationPatterns(ctx);
-  mergeActivationPatterns.add<onnxToArithConstant>(ctx);
-  mergeActivationPatterns.add<convAddToConvWithBiasLeft>(ctx);
-  mergeActivationPatterns.add<convAddToConvWithBiasRight>(ctx);
-  mergeActivationPatterns.add<matMulAddToGemm>(ctx);
-  mergeActivationPatterns.add<matMulToGemm>(ctx);
-  mergeActivationPatterns.add<removeFlattenSameShape>(ctx);
-  populateMatMulRewritePatterns(mergeActivationPatterns, ctx);
+  RewritePatternSet prePatterns(ctx);
+  populatePrePatterns(prePatterns, ctx);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(prePatterns))))
+    llvm::dbgs() << "Failed to apply pre-patterns, continuing...\n";
 
-  if (failed(applyPatternsGreedily(moduleOp, std::move(mergeActivationPatterns))))
-    llvm::dbgs() << "Failed to merge activation patterns, continuing...\n";
-
-  IRRewriter rewriter(moduleOp);
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
     signalPassFailure();
@@ -140,34 +127,23 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXReduceMeanV13Op>();
   target.addIllegalOp<ONNXSplitOp>();
 
-  RewritePatternSet patterns(ctx);
-  patterns.add<removeLRN>(ctx);
-
-  populateElementwisePatterns(patterns, ctx);
-  populateGemmPatterns(patterns, ctx);
-  populateConvPatterns(patterns, ctx);
-  populatePoolPatterns(patterns, ctx);
-  populateReduceMeanPatterns(patterns, ctx);
-  populateReluPatterns(patterns, ctx);
-  populateSigmoidPatterns(patterns, ctx);
-  populateSoftmaxPatterns(patterns, ctx);
-  populateConcatPatterns(patterns, ctx);
-  populateGatherPatterns(patterns, ctx);
-  populateResizePatterns(patterns, ctx);
-  populateReshapePatterns(patterns, ctx);
-  populateSplitPatterns(patterns, ctx);
-
-  if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
+  RewritePatternSet conversionPatterns(ctx);
+  populateConversionPatterns(conversionPatterns, ctx);
+  if (failed(applyPartialConversion(moduleOp, target, std::move(conversionPatterns)))) {
     signalPassFailure();
     return;
   }
 
-  foldSingleLaneComputeBatches(*entryFunc);
+  RewritePatternSet earlyPostPatterns(ctx);
+  populateEarlyPostPatterns(earlyPostPatterns, ctx);
+  if (failed(applyPatternsGreedily(*entryFunc, std::move(earlyPostPatterns)))) {
+    signalPassFailure();
+    return;
+  }
 
-  // Count the number of compute ops and check they do not exceed the core count
   if (coresCount != -1) {
     int computeOpsCount = 0;
-    for (auto& op : entryFunc->getFunctionBody().front().getOperations())
+    for (Operation& op : entryFunc->getFunctionBody().front().getOperations())
       if (isa<spatial::SpatCompute>(op))
         computeOpsCount++;
 
@@ -185,355 +161,23 @@ void ONNXToSpatialPass::runOnOperation() {
 
   annotateWeightsConstants(*entryFunc);
 
+  RewritePatternSet postPatterns(ctx);
+  populatePostPatterns(postPatterns, ctx);
+  if (failed(applyPatternsGreedily(*entryFunc, std::move(postPatterns)))) {
+    signalPassFailure();
+    return;
+  }
+
+  if (failed(verifyONNXToSpatialHostLegality(*entryFunc))) {
+    signalPassFailure();
+    return;
+  }
+
   populateEmptyFunction(*entryFunc);
 
-  if (failed(encapsulateGlobalInstruction(*entryFunc))) {
-    signalPassFailure();
-    return;
-  }
-
-  if (failed(promoteConstantInputsToWeights(*entryFunc))) {
-    signalPassFailure();
-    return;
-  }
-
-  // Dump to file for debug
   dumpModule(moduleOp, "spatial0");
 }
 
-template <typename T>
-bool encapsulator(IRRewriter& rewriter, Location loc, Operation* inst, std::function<Value(T)> funcSource) {
-  if (T toRemoveOp = llvm::dyn_cast_if_present<T>(inst)) {
-    Value source = funcSource(toRemoveOp);
-    rewriter.setInsertionPointAfter(toRemoveOp);
-    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), source);
-    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
-    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
-    rewriter.setInsertionPointToEnd(BB);
-    IRMapping mapper;
-    mapper.map(source, BB->getArgument(0));
-    auto newInst = rewriter.clone(*inst, mapper);
-    spatial::SpatYieldOp::create(rewriter, loc, newInst->getResults());
-    inst->replaceAllUsesWith(newCompute->getResults());
-    inst->erase();
-    return true;
-  }
-  return false;
-}
-
-bool encapsulateSlice(IRRewriter& rewriter, Location loc, Operation* inst) {
-  if (tensor::ExtractSliceOp toRemoveOp = llvm::dyn_cast_if_present<tensor::ExtractSliceOp>(inst)) {
-    auto source = toRemoveOp.getSource();
-    rewriter.setInsertionPointAfter(toRemoveOp);
-    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), source);
-    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
-    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
-    rewriter.setInsertionPointToEnd(BB);
-    IRMapping mapper;
-    mapper.map(source, BB->getArgument(0));
-    auto newInst = rewriter.clone(*inst, mapper);
-    spatial::SpatYieldOp::create(rewriter, loc, newInst->getResults());
-    inst->replaceAllUsesWith(newCompute->getResults());
-    inst->erase();
-    return true;
-  }
-  return false;
-}
-
-bool encapsulateConcat(IRRewriter& rewriter, Location loc, Operation* inst) {
-  if (auto toRemoveOp = llvm::dyn_cast_if_present<tensor::ConcatOp>(inst)) {
-    auto sources = toRemoveOp.getInputs();
-    rewriter.setInsertionPointAfter(toRemoveOp);
-    if (llvm::any_of(sources,
-                     [](auto source) { return isa_and_present<spatial::SpatCompute>(source.getDefiningOp()); })) {
-      auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), sources);
-      SmallVector<Type> sourceTypes;
-      SmallVector<Location> sourceLoc;
-      for (auto source : sources) {
-        sourceTypes.push_back(source.getType());
-        sourceLoc.push_back(loc);
-      }
-      auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
-      newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sources.size()});
-      rewriter.setInsertionPointToEnd(BB);
-      IRMapping mapper;
-      for (auto [source, bbArg] : llvm::zip(sources, BB->getArguments()))
-        mapper.map(source, bbArg);
-      auto newConcat = spatial::SpatConcatOp::create(rewriter,
-                                                     loc,
-                                                     toRemoveOp.getType(),
-                                                     rewriter.getI64IntegerAttr(toRemoveOp.getDim()),
-                                                     ValueRange(BB->getArguments()));
-      spatial::SpatYieldOp::create(rewriter, loc, newConcat.getOutput());
-      inst->replaceAllUsesWith(newCompute->getResults());
-      inst->erase();
-      return true;
-    }
-    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), sources);
-    SmallVector<Type> sourceTypes;
-    SmallVector<Location> sourceLoc;
-    for (auto source : sources) {
-      sourceTypes.push_back(source.getType());
-      sourceLoc.push_back(loc);
-    }
-    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
-    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sources.size()});
-    rewriter.setInsertionPointToEnd(BB);
-    IRMapping mapper;
-    for (auto [source, bbArg] : llvm::zip(sources, BB->getArguments()))
-      mapper.map(source, bbArg);
-    auto newConcat = rewriter.clone(*inst, mapper);
-    spatial::SpatYieldOp::create(rewriter, loc, newConcat->getResults());
-    inst->replaceAllUsesWith(newCompute->getResults());
-    inst->erase();
-    return true;
-  }
-  return false;
-}
-
-static FailureOr<bool> sourceOperandHasWeightAlways(Operation* op) {
-  if (op == nullptr)
-    return false;
-
-  Operation* source = nullptr;
-  do {
-
-    if (isa<spatial::SpatCompute, spatial::SpatComputeBatch>(*op)) {
-      return false;
-    }
-    else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(*op)) {
-      auto tmpSource = extractSliceOp.getSource();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto extractRowsOp = dyn_cast<spatial::SpatExtractRowsOp>(*op)) {
-      auto tmpSource = extractRowsOp.getInput();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(*op)) {
-      auto tmpSource = expandShapeOp.getSrc();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto transposeOp = dyn_cast<ONNXTransposeOp>(*op)) {
-      auto tmpSource = transposeOp.getData();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(*op)) {
-      auto tmpSource = collapseShapeOp.getSrc();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto constantOp = dyn_cast<arith::ConstantOp>(*op)) {
-      source = constantOp;
-    }
-    else if (auto concatOp = dyn_cast<tensor::ConcatOp>(*op)) {
-      bool res = false;
-      for (auto operand : concatOp.getOperands()) {
-        res |= hasWeightAlways(operand.getDefiningOp());
-        if (res)
-          return res;
-      }
-      return res;
-    }
-    else if (auto concatOp = dyn_cast<spatial::SpatConcatOp>(*op)) {
-      bool res = false;
-      for (auto operand : concatOp.getOperands()) {
-        res |= hasWeightAlways(operand.getDefiningOp());
-        if (res)
-          return res;
-      }
-      return res;
-    }
-    else {
-      op->emitOpError("unsupported global instruction while promoting weight-backed operands into Spatial computes");
-      return failure();
-    }
-  }
-  while (source == nullptr);
-
-  return hasWeightAlways(source);
-}
-
-// TODO what we want to keep in global?
-LogicalResult ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
-  Location loc = funcOp.getLoc();
-  IRRewriter rewriter(&getContext());
-  bool keep = true;
-  while (keep) {
-    keep = false;
-    for (auto& instruction : llvm::make_early_inc_range(funcOp.getOps())) {
-      if (isa<spatial::SpatCompute, spatial::SpatComputeBatch, spatial::SpatExtractRowsOp>(instruction)
-          || isa<func::ReturnOp>(instruction))
-        continue;
-
-      auto weightBacked = sourceOperandHasWeightAlways(&instruction);
-      if (failed(weightBacked))
-        return failure();
-      if (*weightBacked)
-        continue;
-
-      keep |= encapsulateSlice(rewriter, loc, &instruction);
-
-      keep |= encapsulator<tensor::ExpandShapeOp>(
-        rewriter, loc, &instruction, [](tensor::ExpandShapeOp expand) { return expand.getSrc(); });
-
-      keep |= encapsulator<ONNXTransposeOp>(
-        rewriter, loc, &instruction, [](ONNXTransposeOp transpose) { return transpose.getData(); });
-
-      keep |= encapsulator<tensor::CollapseShapeOp>(
-        rewriter, loc, &instruction, [](tensor::CollapseShapeOp collapse) { return collapse.getSrc(); });
-
-      keep |= encapsulateConcat(rewriter, loc, &instruction);
-    }
-  }
-  return success();
-}
-
-void ONNXToSpatialPass::annotateWeightsConstants(func::FuncOp funcOp) const {
-  funcOp.walk([&](arith::ConstantOp constantOp) {
-    if (hasOnlySpatialMvmVmmWeightUses(constantOp.getResult()))
-      markWeightAlways(constantOp);
-  });
-}
-
-LogicalResult ONNXToSpatialPass::promoteConstantInputsToWeights(func::FuncOp funcOp) {
-  IRRewriter rewriter(&getContext());
-  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
-
-  for (auto compute : computes) {
-    SmallVector<bool> promoteInput(compute.getInputs().size(), false);
-    bool needsRewrite = false;
-    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      if (!isWeightLikeComputeOperand(input))
-        continue;
-      promoteInput[inputIdx] = true;
-      needsRewrite = true;
-    }
-    if (!needsRewrite)
-      continue;
-
-    rewriter.setInsertionPointAfter(compute);
-
-    SmallVector<Value> newWeights(compute.getWeights().begin(), compute.getWeights().end());
-    SmallVector<Value> newInputs;
-    SmallVector<Type> newInputTypes;
-    SmallVector<Location> newInputLocs;
-    newWeights.reserve(compute.getWeights().size() + compute.getInputs().size());
-    newInputs.reserve(compute.getInputs().size());
-    newInputTypes.reserve(compute.getInputs().size());
-    newInputLocs.reserve(compute.getInputs().size());
-
-    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      if (promoteInput[inputIdx]) {
-        newWeights.push_back(input);
-        continue;
-      }
-      newInputs.push_back(input);
-      newInputTypes.push_back(input.getType());
-      newInputLocs.push_back(input.getLoc());
-    }
-
-    auto newCompute =
-      spatial::SpatCompute::create(rewriter, compute.getLoc(), compute.getResultTypes(), newWeights, newInputs);
-    auto* newBlock =
-      rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), newInputTypes, newInputLocs);
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(newWeights.size()), static_cast<int>(newInputs.size())});
-    rewriter.setInsertionPointToStart(newBlock);
-
-    IRMapping mapper;
-    auto& oldBlock = compute.getBody().front();
-    size_t newInputIdx = 0;
-    for (auto [oldInputIdx, oldArg] : llvm::enumerate(oldBlock.getArguments())) {
-      if (!promoteInput[oldInputIdx]) {
-        mapper.map(oldArg, newBlock->getArgument(newInputIdx++));
-        continue;
-      }
-
-      auto clonedValue = materializeWeightLikeValueInBlock(compute.getInputs()[oldInputIdx], rewriter, mapper);
-      if (failed(clonedValue))
-        return compute.emitError("failed to materialize promoted weight-like operand inside compute body");
-      mapper.map(oldArg, *clonedValue);
-    }
-
-    for (auto& op : oldBlock.without_terminator())
-      rewriter.clone(op, mapper);
-
-    auto oldYield = cast<spatial::SpatYieldOp>(oldBlock.getTerminator());
-    SmallVector<Value> newYieldOperands;
-    newYieldOperands.reserve(oldYield.getOutputs().size());
-    for (Value operand : oldYield.getOutputs()) {
-      auto mapped = mapper.lookupOrNull(operand);
-      newYieldOperands.push_back(mapped ? cast<Value>(mapped) : operand);
-    }
-    spatial::SpatYieldOp::create(rewriter, oldYield.getLoc(), newYieldOperands);
-
-    compute.replaceAllUsesWith(newCompute);
-    compute.erase();
-  }
-
-  return success();
-}
-
-void ONNXToSpatialPass::populateEmptyFunction(func::FuncOp funcOp) {
-  IRRewriter rewriter(&getContext());
-  IRMapping mapper;
-  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
-  if (!computes.empty())
-    return;
-  auto returnOp = llvm::cast<func::ReturnOp>(funcOp.getRegion().front().getTerminator());
-  rewriter.setInsertionPoint(returnOp);
-
-  SmallVector<Type> sourceTypes;
-  SmallVector<Location> sourceLoc;
-  for (auto source : funcOp.getArguments()) {
-    sourceTypes.push_back(source.getType());
-    sourceLoc.push_back(source.getLoc());
-  }
-
-  auto newCompute = spatial::SpatCompute::create(
-    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
-  auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
-  for (auto [bbArg, computeArg] : llvm::zip(BB->getArguments(), newCompute.getOperands()))
-    mapper.map(computeArg, bbArg);
-  newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sourceTypes.size()});
-  rewriter.setInsertionPointToEnd(BB);
-  for (Operation& inst : funcOp.getOps())
-    if (!isa<spatial::SpatCompute, func::ReturnOp>(&inst))
-      rewriter.clone(inst, mapper);
-
-  auto yield = spatial::SpatYieldOp::create(rewriter, funcOp.getLoc(), returnOp.getOperands());
-  for (size_t i = 0; i < yield.getNumOperands(); ++i)
-    yield.setOperand(i, mapper.lookupOrDefault(yield.getOperand(i)));
-
-  for (Operation& inst : llvm::make_early_inc_range(funcOp.getOps()))
-    if (!isa<spatial::SpatCompute, func::ReturnOp>(&inst)){
-      inst.dropAllUses();
-      rewriter.eraseOp(&inst);
-  }
-
-  for (auto [index, computeResult] : llvm::enumerate(newCompute.getResults()))
-    returnOp.setOperand(index, computeResult);
-}
-
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
 } // namespace onnx_mlir