better MaterializeMergeSchedule.cpp (something still broken downstream)

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -14,7 +14,6 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/FoldUtils.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/WalkPatternRewriteDriver.h"
 
 #include "llvm/ADT/StringRef.h"
@@ -28,7 +27,6 @@
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.hpp"
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/Cleanup.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.hpp"
@@ -67,6 +65,7 @@ private:
   LogicalResult allocateAndInitializeCoreLocalVariables(func::FuncOp funcOp, IRRewriter& rewriter);
 
   void markOpToRemove(Operation* op);
+  void eraseOpsToRemove();
 
   void enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter);
 };
@@ -268,13 +267,7 @@ void SpatialToPimPass::runOnOperation() {
 
   enlargeVMMOutTensorsToCrossbarSize(funcOp, rewriter);
   replaceReturnWithOutputBuffers(returnOp, rewriter, returnPathState);
-
-  SmallVector<Operation*> pendingRemovals(operationsToRemove.begin(), operationsToRemove.end());
-  if (failed(erasePendingOps(pendingRemovals, rewriter))) {
-    funcOp.emitOpError("failed to erase obsolete Spatial ops after lowering to PIM");
-    signalPassFailure();
-    return;
-  }
+  eraseOpsToRemove();
 
   RewritePatternSet finalTensorPackingPatterns(ctx);
   populateTensorPackingPatterns(finalTensorPackingPatterns);
@@ -399,6 +392,13 @@ void SpatialToPimPass::markOpToRemove(Operation* op) {
     operationsToRemove.push_back(op);
 }
 
+void SpatialToPimPass::eraseOpsToRemove() {
+  for (Operation* op : operationsToRemove) {
+    op->dropAllUses();
+    op->erase();
+  }
+}
+
 std::unique_ptr<Pass> createSpatialToPimPass() { return std::make_unique<SpatialToPimPass>(); }
 
 } // namespace onnx_mlir