replace all-in-one pass with real pipeline (intel#174)

Author: xurui1995

include/gc-c/Passes.h

@@ -26,6 +26,9 @@ extern "C" {
 
 #include "gc/Dialect/CPURuntime/Transforms/CPURuntimePasses.capi.h.inc"
 #include "gc/Transforms/Passes.capi.h.inc"
+
+MLIR_CAPI_EXPORTED void mlirRegisterAllGCPassesAndPipelines(void);
+
 #ifdef __cplusplus
 }
 #endif

include/gc/Transforms/Passes.h

@@ -40,12 +40,12 @@ namespace xegpu {
 class XeGPUDialect;
 }
 
-class PassManager;
+class OpPassManager;
 
 namespace gc {
 
-void populateFrontendPasses(mlir::PassManager &);
-void populateCPUPipeline(mlir::PassManager &);
+void populateFrontendPasses(mlir::OpPassManager &);
+void populateCPUPipeline(mlir::OpPassManager &);
 
 #define GEN_PASS_DECL
 #include "gc/Transforms/Passes.h.inc"

include/gc/Transforms/Passes.td

@@ -32,19 +32,6 @@ def ConvertOneDNNGraphToLinalg : Pass<"convert-onednn-graph-to-linalg"> {
   ];
 }
 
-def GCCPUPipeline: Pass<"gc-cpu-pipeline"> {
-  let summary = "All-in-one pipeline for GC for CPU";
-  let dependentDialects = ["onednn_graph::OneDNNGraphDialect",
-      "tensor::TensorDialect",
-      "memref::MemRefDialect",
-      "linalg::LinalgDialect",
-      "linalgx::LinalgxDialect",
-      "LLVM::LLVMDialect",
-      "scf::SCFDialect",
-      "bufferization::BufferizationDialect",
-      "omp::OpenMPDialect",
-      "vector::VectorDialect"];
-}
 
 def LinalgToXeGPU : Pass<"linalg-to-xegpu", "func::FuncOp"> {
   let summary = "Convert linalg dialect to XeGPU dialect.";

lib/gc/CAPI/Passes.cpp

@@ -16,13 +16,22 @@
 using namespace mlir::gc;
 using namespace mlir::cpuruntime;
 
+namespace mlir::gc {
+void registerCPUPipeline();
+} // namespace mlir::gc
+
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 #include "gc/Dialect/CPURuntime/Transforms/CPURuntimePasses.capi.cpp.inc"
 #include "gc/Transforms/Passes.capi.cpp.inc"
 
+MLIR_CAPI_EXPORTED void mlirRegisterAllGCPassesAndPipelines() {
+  registerCPUPipeline();
+  mlirRegisterCPURuntimePasses();
+  mlirRegisterGraphCompilerPasses();
+}
 #ifdef __cplusplus
 }
 #endif

lib/gc/Transforms/Pipeline.cpp

@@ -19,6 +19,7 @@
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/DialectRegistry.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/LogicalResult.h"
@@ -32,12 +33,12 @@
 namespace mlir::gc {
 
 // linalg + linalgX + tensor
-void populateFrontendPasses(mlir::PassManager &pm) {
+void populateFrontendPasses(mlir::OpPassManager &pm) {
   pm.addPass(createConvertOneDNNGraphToLinalg());
 }
 
 // scf + arith + math + vector + tensor + linalg.brgemm + tensor.pack/unpack
-void populateTensorPasses(mlir::PassManager &pm) {
+void populateTensorPasses(mlir::OpPassManager &pm) {
   // todo: padding propagation pass
   // todo: layout propagation pass
   // todo: tensor constant propagation pass
@@ -51,7 +52,7 @@ void populateTensorPasses(mlir::PassManager &pm) {
 }
 
 // scf + arith + math + vector + tensor + linalg.brgemm
-void populateVectorPasses(mlir::PassManager &pm) {
+void populateVectorPasses(mlir::OpPassManager &pm) {
   // Do promotion for math / arith ops
   pm.addNestedPass<func::FuncOp>(math::createMathLegalizeToF32());
   // sourceTypeStrs can be extended
@@ -69,7 +70,7 @@ void populateVectorPasses(mlir::PassManager &pm) {
 }
 
 // scf + arith + math + vector + memref + linalg.brgemm
-void populateBufferizationPasses(mlir::PassManager &pm) {
+void populateBufferizationPasses(mlir::OpPassManager &pm) {
   bufferization::OneShotBufferizationOptions options;
   options.bufferizeFunctionBoundaries = true;
   options.setFunctionBoundaryTypeConversion(
@@ -88,7 +89,7 @@ void populateBufferizationPasses(mlir::PassManager &pm) {
 }
 
 // scf + arith + math + vector + memref + func/microkernel
-void populateMicroKernelPasses(mlir::PassManager &pm) {
+void populateMicroKernelPasses(mlir::OpPassManager &pm) {
   // todo: ConvertLinalgToMicrokernel pass
   // todo: CleanupInvalidMicrokernel pass
   // todo: InvariantMicrokernelMotion pass
@@ -98,13 +99,13 @@ void populateMicroKernelPasses(mlir::PassManager &pm) {
   // todo: DispatchMicrokernel
 }
 
-void populateCPURuntimePasses(mlir::PassManager &pm) {
+void populateCPURuntimePasses(mlir::OpPassManager &pm) {
   // todo: flatten nested parallel pass to support coarse-grain usion
   // remove this pass after we add FlattenNestedParallel
   pm.addPass(createConvertSCFToOpenMPPass());
 }
 
-void populateLoweringToLLVMPasses(mlir::PassManager &pm) {
+void populateLoweringToLLVMPasses(mlir::OpPassManager &pm) {
   pm.addPass(createFinalizeMemRefToLLVMConversionPass());
   pm.addPass(createConvertSCFToCFPass());
   pm.addPass(cpuruntime::createCPURuntimeToLLVM());
@@ -120,13 +121,13 @@ void populateLoweringToLLVMPasses(mlir::PassManager &pm) {
   pm.addPass(createSymbolDCEPass());
 }
 
-void populateLLVMPasses(mlir::PassManager &pm) {
+void populateLLVMPasses(mlir::OpPassManager &pm) {
   pm.addPass(memref::createExpandOpsPass());
   pm.addPass(memref::createExpandStridedMetadataPass());
   populateLoweringToLLVMPasses(pm);
 }
 
-void populateCPUPipeline(mlir::PassManager &pm) {
+void populateCPUPipeline(mlir::OpPassManager &pm) {
   // front-end, oneDNN graph dialect
   populateFrontendPasses(pm);
   // middle-end, LinalgX/Linalg/tensor dialects
@@ -144,24 +145,10 @@ void populateCPUPipeline(mlir::PassManager &pm) {
   populateLLVMPasses(pm);
 }
 
-#define GEN_PASS_DEF_GCCPUPIPELINE
-#include "gc/Transforms/Passes.h.inc"
-namespace {
-
-class GCCPUPipeline : public impl::GCCPUPipelineBase<GCCPUPipeline> {
-public:
-  friend struct PassHelper;
-  using impl::GCCPUPipelineBase<GCCPUPipeline>::GCCPUPipelineBase;
-  void runOnOperation() final {
-    auto op = getOperation();
-    PassManager pm{op->getContext()};
-    populateCPUPipeline(pm);
-    // TODO(longsheng): add a option to
-    // disable threading and enable pm.enableIRPrinting();
-    if (failed(pm.run(op)))
-      signalPassFailure();
-  }
-};
+void registerCPUPipeline() {
+  PassPipelineRegistration<>("gc-cpu-pipeline",
+                             "The CPU pipeline for Graph Compiler",
+                             populateCPUPipeline);
+}
 
-} // namespace
 } // namespace mlir::gc

python/MainModule.cpp

@@ -24,7 +24,8 @@
 PYBIND11_MODULE(_gc_mlir, m) {
   m.doc() = "Graph-compiler MLIR Python binding";
 
-  mlirRegisterGraphCompilerPasses();
+  mlirRegisterAllGCPassesAndPipelines();
+
   //===----------------------------------------------------------------------===//
   // OneDNNGraph
   //===----------------------------------------------------------------------===//
@@ -44,7 +45,6 @@ PYBIND11_MODULE(_gc_mlir, m) {
   //===----------------------------------------------------------------------===//
   // CPURuntime
   //===----------------------------------------------------------------------===//
-  mlirRegisterCPURuntimePasses();
   auto cpuruntimeM = m.def_submodule("cpuruntime");
   cpuruntimeM.def(
       "register_dialect",

src/gc-opt/gc-opt.cpp

@@ -30,6 +30,10 @@
 #include <imex/InitIMEXPasses.h>
 #endif
 
+namespace mlir::gc {
+void registerCPUPipeline();
+} // namespace mlir::gc
+
 int main(int argc, char *argv[]) {
 #ifdef GC_USE_GPU
   imex::registerTransformsPasses();
@@ -41,6 +45,7 @@ int main(int argc, char *argv[]) {
   imex::registerConvertXeTileToXeGPU();
 #endif
   mlir::registerAllPasses();
+  mlir::gc::registerCPUPipeline();
   mlir::gc::registerGraphCompilerPasses();
   mlir::cpuruntime::registerCPURuntimePasses();
   mlir::DialectRegistry registry;