blazingly faster

src/PIM/Compiler/CMakeLists.txt

@@ -26,6 +26,7 @@ add_pim_library(OMPimCompilerUtils
   ${PIM_COMPILER_INCLUDE_DIRS}
 
   LINK_LIBS PUBLIC
+  MLIRAffineToStandard
   OMPimCompilerOptions
   OMPimCommon
   OMPimBufferization

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -15,13 +15,6 @@ llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget(
   llvm::cl::init(EmitPimCodegen),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<PimMergeSchedulerType>
-  pimMergeScheduler("pim-merge-scheduler",
-                    llvm::cl::desc("Scheduler used by the Spatial merge-compute-nodes pass"),
-                    llvm::cl::values(clEnumValN(MergeSchedulerPeft, "peft", "Use PEFT scheduling")),
-                    llvm::cl::init(MergeSchedulerPeft),
-                    llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport(
   "pim-memory-report",
   llvm::cl::desc("Emit a human-readable PIM memory planning report"),
@@ -64,9 +57,11 @@ llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow(
   llvm::cl::values(
     clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit spatial1 graph dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 scheduled dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 trivially merged graph dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 realized dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 scheduled dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial4, "spatial4", "Emit spatial4 realized dataflow CSV reports")),
   llvm::cl::values(clEnumValN(SpatialDataflowExportAll, "all", "Emit all Spatial dataflow CSV reports")),
   llvm::cl::init(SpatialDataflowExportNone),
   llvm::cl::cat(OnnxMlirOptions));

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -20,10 +20,6 @@ typedef enum {
   EmitPimCodegen = 3
 } PimEmissionTargetType;
 
-typedef enum {
-  MergeSchedulerPeft = 0,
-} PimMergeSchedulerType;
-
 typedef enum {
   PimMemoryReportNone = 0,
   PimMemoryReportSummary = 1,
@@ -47,12 +43,12 @@ typedef enum {
   SpatialDataflowExportSpatial1 = 1,
   SpatialDataflowExportSpatial2 = 2,
   SpatialDataflowExportSpatial3 = 3,
-  SpatialDataflowExportAll = 4,
+  SpatialDataflowExportSpatial4 = 4,
+  SpatialDataflowExportAll = 5,
 } PimSpatialDataflowExportType;
 
 extern llvm::cl::OptionCategory OnnxMlirOptions;
 extern llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget;
-extern llvm::cl::opt<PimMergeSchedulerType> pimMergeScheduler;
 extern llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport;
 extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
 extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -31,6 +32,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(createONNXToSpatialPass());
     pm.addPass(createSpatialLayoutPlanningPass());
     pm.addPass(createLowerSpatialPlansPass());
+    pm.addPass(createTrivialGraphComputeMergePass());
     pm.addPass(createMergeComputeNodesPass());
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
@@ -46,6 +48,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
   }
 
   if (pimEmissionTarget >= EmitPimCodegen) {
+    pm.addPass(mlir::createLowerAffinePass());
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
     if (!pimDisableMemoryCoalescing)

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -154,7 +154,10 @@ static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) {
 }
 
 static bool isSupportedAliasOp(Operation* op) {
-  return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
+  return isa<memref::SubViewOp,
+             memref::CastOp,
+             memref::CollapseShapeOp,
+             memref::ExpandShapeOp>(op);
 }
 
 static bool isRuntimeMemoryTouchOp(Operation* op) {