Fixed pass pipeline and some minor edits

plaidml · Mar 1, 2024 · a6134d6 · a6134d6
1 parent eb7360b
commit a6134d6
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 18 deletions.
diff --git a/include/TPP/Passes.td b/include/TPP/Passes.td
@@ -204,7 +204,10 @@ def DefaultTppPasses : Pass<"default-tpp-passes", "ModuleOp"> {
   let options= [
     Option<"linalgToLoops", "linalg-to-loops",
            "bool", /*default=*/"false",
-           "Skip all TPP transformations. Lower linalg directly to loops.">
+           "Skip all TPP transformations. Lower linalg directly to loops.">,
+    ListOption<"parallelTaskGrid", "parallel-task-grid",
+           "unsigned", "Grid-sizes for parallel tasks.">
+
   ];
 }
 
@@ -289,6 +292,12 @@ def LocalDialectsLowering : Pass<"lower-local-dialects", "ModuleOp"> {
                            "tensor::TensorDialect",
                            "xsmm::XsmmDialect",
                            "LLVM::LLVMDialect"];
+  let options = [
+    ListOption<"parallelTaskGrid", "parallel-task-grid",
+           "unsigned", "Grid-sizes for parallel tasks.">
+
+  ];
+
 }
 
 def Postprocessing : Pass<"postprocess", "func::FuncOp"> {
@@ -474,7 +483,7 @@ def FoldXsmmFlags : Pass<"fold-xsmm-flags", "func::FuncOp"> {
 def SCFParallelLoopTiling : Pass<"scf-parallel-loop-tiling-pass"> {
   let summary = "Tile parallel loops";
   let options = [
-    ListOption<"tileSizes", "parallel-loop-tile-sizes", "int64_t",
+    ListOption<"tileSizes", "parallel-loop-tile-sizes", "unsigned",
                "Factors to tile parallel loops by">,
     Option<"noMinMaxBounds", "no-min-max-bounds", "bool",
            /*default=*/"false",

diff --git a/lib/TPP/DefaultPipeline.cpp b/lib/TPP/DefaultPipeline.cpp
@@ -47,10 +47,10 @@ llvm::cl::opt<bool>
                 llvm::cl::init(false));
 
 // Control grid parallelism sizes.
-llvm::cl::list<int64_t>
+llvm::cl::list<unsigned>
     parallelTaskGrid("parallel-task-grid",
                      llvm::cl::desc("Grid-sizes for parallel tasks"),
-                     llvm::cl::list_init<int64_t>(SmallVector<int64_t>{2, 8}),
+                     llvm::cl::list_init<unsigned>(SmallVector<unsigned>{2, 8}),
                      llvm::cl::CommaSeparated);
 
 namespace mlir {
@@ -127,7 +127,8 @@ struct DefaultPipeline : public tpp::impl::DefaultPipelineBase<DefaultPipeline>,
       pm.addPass(createGpuPipeline(GpuPipelineOptions{gpuBackend}));
     } else {
       // Apply the default preprocessing pass
-      DefaultTppPassesOptions tppDefaultOptions{linalgToLoops};
+      DefaultTppPassesOptions tppDefaultOptions{linalgToLoops,
+                                                parallelTaskGrid};
       pm.addPass(createDefaultTppPasses(tppDefaultOptions));
     }
 
@@ -140,15 +141,8 @@ struct DefaultPipeline : public tpp::impl::DefaultPipelineBase<DefaultPipeline>,
     pm.addPass(tpp::createConvertPerfToFunc());
     pm.addPass(createConvertTensorToLinalgPass());
     pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());
-    if (defParallel) {
-      mlir::tpp::SCFParallelLoopTilingOptions tilingOptions;
-      tilingOptions.tileSizes = parallelTaskGrid;
-      pm.addPass(createSCFParallelLoopTiling(tilingOptions));
-      pm.addNestedPass<func::FuncOp>(createIntelAMXTileConfigInsertionPass());
-      pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
-      pm.addNestedPass<func::FuncOp>(createIntelAMXTileConfigHoistingPass());
+    if (defParallel)
       pm.addPass(createConvertSCFToOpenMPPass());
-    }
     pm.addPass(createConvertVectorToSCFPass());
     pm.addPass(arith::createArithExpandOpsPass());
     pm.addPass(createLowerAffinePass());

diff --git a/lib/TPP/DefaultTppPasses.cpp b/lib/TPP/DefaultTppPasses.cpp
@@ -77,6 +77,10 @@ struct LocalDialectsLowering
     : public tpp::impl::LocalDialectsLoweringBase<LocalDialectsLowering>,
       UtilityPassBase<ModuleOp> {
 
+  LocalDialectsLowering() {}
+  LocalDialectsLowering(const LocalDialectsLoweringOptions &options) {
+    parallelTaskGrid = options.parallelTaskGrid;
+  }
   void runOnOperation() override {
     auto module = getOperation();
 
@@ -106,6 +110,16 @@ struct LocalDialectsLowering
     // that they are hoisted out of loops.
     pm.addNestedPass<func::FuncOp>(createCleanup());
 
+    mlir::tpp::SCFParallelLoopTilingOptions tilingOptions;
+    tilingOptions.tileSizes = parallelTaskGrid;
+    pm.addPass(createSCFParallelLoopTiling(tilingOptions));
+
+    pm.addNestedPass<func::FuncOp>(createIntelAMXTileConfigInsertionPass());
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+    pm.addNestedPass<func::FuncOp>(createLoopInvariantCodeMotionPass());
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+    pm.addNestedPass<func::FuncOp>(createIntelAMXTileConfigHoistingPass());
+
     pm.addPass(createConvertXsmmToFunc());
     pm.addPass(createConvertPerfToFunc());
   }
@@ -310,7 +324,8 @@ struct DefaultTppPasses
     pm.addPass(createConvertForAllToParallelOp());
 
     // Covert all local TPP-related dialects.
-    pm.addPass(createLocalDialectsLowering());
+    LocalDialectsLoweringOptions localDialectsLowering{parallelTaskGrid};
+    pm.addPass(createLocalDialectsLowering(localDialectsLowering));
 
     // Clean up after the default pipeline.
     pm.addNestedPass<func::FuncOp>(createPostprocessing());

diff --git a/lib/TPP/Transforms/IntelAMXTileConfig.cpp b/lib/TPP/Transforms/IntelAMXTileConfig.cpp
@@ -1,4 +1,4 @@
-//===- IntelAMXTileConfig.cpp ---------------------------------------------===//
+//===- IntelAMXTileConfig.cpp ------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

diff --git a/lib/TPP/Transforms/IntelAMXTileConfigHoisting.cpp b/lib/TPP/Transforms/IntelAMXTileConfigHoisting.cpp
@@ -1,4 +1,4 @@
-//===- IntelAMXTileConfigHoisting.cpp -------------------------------------===//
+//===- IntelAMXTileConfigHoisting.cpp ----------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

diff --git a/lib/TPP/Transforms/SCFParallelLoopTiling.cpp b/lib/TPP/Transforms/SCFParallelLoopTiling.cpp
@@ -54,7 +54,7 @@ using namespace mlir::scf;
 /// %i0 + j0 and %i1 + %j1.
 ///
 /// The old loop is replaced with the new one.
-void tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes,
+void tileParallelLoop(ParallelOp op, ArrayRef<unsigned> tileSizes,
                       bool noMinMaxBounds) {
   bool useParallelOp = false;
   /* TODO, need to implement this case */
@@ -244,7 +244,7 @@ namespace {
 struct SCFParallelLoopTiling
     : public tpp::impl::SCFParallelLoopTilingBase<SCFParallelLoopTiling> {
   SCFParallelLoopTiling(){};
-  SCFParallelLoopTiling(ArrayRef<int64_t> tileSizes,
+  SCFParallelLoopTiling(ArrayRef<unsigned> tileSizes,
                         bool noMinMaxBounds = false) {
     this->tileSizes = tileSizes;
     this->noMinMaxBounds = noMinMaxBounds;