Move Decompose concat to Dispatch creation pipeline.

MaheshRavishankar · MaheshRavishankar · commit 1ba85a9b4b11 · 2024-11-13T20:19:00.000-08:00
Signed-off-by: MaheshRavishankar &lt;mahesh.ravishankar@gmail.com&gt;
diff --git a/compiler/src/iree/compiler/DispatchCreation/BUILD.bazel b/compiler/src/iree/compiler/DispatchCreation/BUILD.bazel
@@ -22,6 +22,7 @@ iree_compiler_cc_library(
         "CollapseReductionDimensions.cpp",
         "ConvertDispatchRegionsToWorkgroups.cpp",
         "ConvertTensorToFlow.cpp",
+        "DecomposeConcat.cpp",
         "DispatchWithTransformDialect.cpp",
         "ElementwiseOpFusion.cpp",
         "FoldUnitExtentDims.cpp",
diff --git a/compiler/src/iree/compiler/DispatchCreation/CMakeLists.txt b/compiler/src/iree/compiler/DispatchCreation/CMakeLists.txt
@@ -24,6 +24,7 @@ iree_cc_library(
     "CollapseReductionDimensions.cpp"
     "ConvertDispatchRegionsToWorkgroups.cpp"
     "ConvertTensorToFlow.cpp"
+    "DecomposeConcat.cpp"
     "DispatchWithTransformDialect.cpp"
     "ElementwiseOpFusion.cpp"
     "FoldUnitExtentDims.cpp"
diff --git a/compiler/src/iree/compiler/DispatchCreation/DecomposeConcat.cpp b/compiler/src/iree/compiler/DispatchCreation/DecomposeConcat.cpp
@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "iree/compiler/DispatchCreation/Passes.h"
 #include "llvm/ADT/STLExtras.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -13,10 +13,10 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
-namespace mlir::iree_compiler::GlobalOptimization {
+namespace mlir::iree_compiler::DispatchCreation {
 
 #define GEN_PASS_DEF_DECOMPOSECONCATPASS
-#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+#include "iree/compiler/DispatchCreation/Passes.h.inc"
 
 namespace {
 
diff --git a/compiler/src/iree/compiler/DispatchCreation/Passes.cpp b/compiler/src/iree/compiler/DispatchCreation/Passes.cpp
@@ -204,9 +204,20 @@ void addDispatchRegionCreationPreprocessingPasses(OpPassManager &passManager) {
 
 // Pipeline to first create `flow.dispatch.region` ops and then lower to
 // `flow.dispatch.workgroup` ops.
-static void addDispatchRegionCreationPasses(OpPassManager &passManager) {
+static void
+addDispatchRegionCreationPasses(OpPassManager &passManager,
+                                const TransformOptions &transformOptions) {
   FunctionLikeNest(passManager)
-      // Only want use the transform dialect for some dispatch regions and let
+      // We decompose and transpose concatenations immediately before folding
+      // unit extent dims because this allows decoupling unit dims in the
+      // concatenation from the transposes that are introduced.
+      .addPass([&]() {
+        DecomposeConcatPassOptions options;
+        options.enableConcatTransposition =
+            transformOptions.options.outerDimConcat;
+        return createDecomposeConcatPass(options);
+      }) // Only want use the transform dialect for some dispatch regions and
+         // let
       // the FormDispatchRegions handle the rest. This only moves the root
       // compute op into the dispatch region, so that we can run additional
       // transformations afterwards with a simple region and without bothering
@@ -303,7 +314,7 @@ void buildDispatchCreationPassPipeline(
       .addPass(mlir::createCSEPass);
 
   addDispatchRegionCreationPreprocessingPasses(passManager);
-  addDispatchRegionCreationPasses(passManager);
+  addDispatchRegionCreationPasses(passManager, transformOptions);
 
   FunctionLikeNest(passManager)
       .addPass(DispatchCreation::createConvertDispatchRegionsToWorkgroupsPass)
diff --git a/compiler/src/iree/compiler/DispatchCreation/Passes.h b/compiler/src/iree/compiler/DispatchCreation/Passes.h
@@ -23,7 +23,9 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 /// This is a placeholder for future. We should pass all the options through the
 /// struct.
-struct TransformOptions : public PassPipelineOptions<TransformOptions> {};
+struct TransformOptions : public PassPipelineOptions<TransformOptions> {
+  DispatchCreationOptions options;
+};
 
 void buildDispatchCreationPassPipeline(
     OpPassManager &passManager, const TransformOptions &transformOptions);
diff --git a/compiler/src/iree/compiler/DispatchCreation/Passes.td b/compiler/src/iree/compiler/DispatchCreation/Passes.td
@@ -39,6 +39,17 @@ def CollapseReductionDimensionsPass :
   ];
 }
 
+def DecomposeConcatPass :
+    Pass<"iree-dispatch-creation-decompose-concat", ""> {
+  let summary = "Decomposes concatenations into a destination and a sequence of slice inserts.";
+  let options = [
+    Option<"enableConcatTransposition", "enable-concat-transposition", "bool",
+           /*default=*/"false", "Allows transposing concatenations such that "
+                                "they occur on the inner most dims.">,
+  ];
+}
+
+
 def ElementwiseOpFusionPass :
     Pass<"iree-dispatch-creation-elementwise-op-fusion", ""> {
   let summary = "Fuse elementwise operations.";
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/BUILD.bazel b/compiler/src/iree/compiler/DispatchCreation/test/BUILD.bazel
@@ -46,6 +46,7 @@ iree_lit_test_suite(
             "split_reduction.mlir",
             "tensor_pad_to_tensor_insert_slice.mlir",
             "transform_dispatch_region_formation.mlir",
+            "transpose_and_decompose_concat.mlir",
             "transpose_generic_ops.mlir",
         ],
         include = ["*.mlir"],
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/CMakeLists.txt b/compiler/src/iree/compiler/DispatchCreation/test/CMakeLists.txt
@@ -44,6 +44,7 @@ iree_lit_test_suite(
     "split_reduction.mlir"
     "tensor_pad_to_tensor_insert_slice.mlir"
     "transform_dispatch_region_formation.mlir"
+    "transpose_and_decompose_concat.mlir"
     "transpose_generic_ops.mlir"
   TOOLS
     FileCheck
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/transpose_and_decompose_concat.mlir b/compiler/src/iree/compiler/DispatchCreation/test/transpose_and_decompose_concat.mlir
diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -47,7 +47,6 @@ iree_compiler_cc_library(
         "CleanupNumericNarrowing.cpp",
         "Convert1X1FilterConv2DToMatmul.cpp",
         "DataLayoutPropagation.cpp",
-        "DecomposeConcat.cpp",
         "DemoteContractionInputsToBF16.cpp",
         "DetachElementwiseFromNamedOps.cpp",
         "EraseUnusedLinalgOperands.cpp",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -43,7 +43,6 @@ iree_cc_library(
     "CleanupNumericNarrowing.cpp"
     "Convert1X1FilterConv2DToMatmul.cpp"
     "DataLayoutPropagation.cpp"
-    "DecomposeConcat.cpp"
     "DemoteContractionInputsToBF16.cpp"
     "DetachElementwiseFromNamedOps.cpp"
     "EraseUnusedLinalgOperands.cpp"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -120,13 +120,6 @@ void buildGlobalOptimizationPassPipeline(
       // optimizations, is sensitive to surrounding IR structure. Thus we run
       // this pass both before unit dim folding + consteval, as well as after.
       .addPass(createRaiseSpecialOpsPass)
-      // We decompose and transpose concatenations immediately before folding
-      // unit extent dims because this allows decoupling unit dims in the
-      // concatenation from the transposes that are introduced.
-      .addPass([&]() {
-        return createDecomposeConcatPass(
-            transformOptions.options.outerDimConcat);
-      })
       // We generalize certain named ops immediately before folding unit extent
       // dims as the unit dim folding pass updates indexing maps and is better
       // at working with generics. By this point we have already done any
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -19,16 +19,6 @@ def Convert1X1FilterConv2DToMatmulPass:
   let summary = "Convert linalg convolution ops with 1x1 kernels into linalg matrix multiplication ops.";
 }
 
-def DecomposeConcatPass :
-    Pass<"iree-global-opt-decompose-concat", ""> {
-  let summary = "Decomposes concatenations into a destination and a sequence of slice inserts.";
-  let options = [
-    Option<"enableConcatTransposition", "enable-concat-transposition", "bool",
-           /*default=*/"false", "Allows transposing concatenations such that "
-                                "they occur on the inner most dims.">,
-  ];
-}
-
 def DemoteContractionInputsToBF16Pass
     : Pass<"iree-global-opt-demote-contraction-inputs-to-bf16", ""> {
   let summary =
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
@@ -35,7 +35,6 @@ iree_lit_test_suite(
             "raise_special_ops.mlir",
             "remove_zero_extent_tensors.mlir",
             "transformation_pipeline.mlir",
-            "transpose_and_decompose_concat.mlir",
         ],
         include = ["*.mlir"],
     ),
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
@@ -33,7 +33,6 @@ iree_lit_test_suite(
     "raise_special_ops.mlir"
     "remove_zero_extent_tensors.mlir"
     "transformation_pipeline.mlir"
-    "transpose_and_decompose_concat.mlir"
   TOOLS
     FileCheck
     iree-opt
diff --git a/compiler/src/iree/compiler/Pipelines/Options.cpp b/compiler/src/iree/compiler/Pipelines/Options.cpp
@@ -129,10 +129,6 @@ void GlobalOptimizationOptions::bindOptions(OptionsBinder &binder) {
           "Propagates transposes to named ops even when the resulting op will "
           "be a linalg.generic"),
       llvm::cl::cat(category));
-  binder.opt<bool>("iree-opt-outer-dim-concat", outerDimConcat,
-                   llvm::cl::desc("Transposes all concatenations to happen"
-                                  "along the outer most dimension."),
-                   llvm::cl::cat(category));
   binder.opt<bool>("iree-opt-data-tiling", dataTiling,
                    llvm::cl::desc("Enables data tiling path."),
                    llvm::cl::cat(category));
@@ -197,6 +193,16 @@ void GlobalOptimizationOptions::bindOptions(OptionsBinder &binder) {
       llvm::cl::cat(category));
 }
 
+void DispatchCreationOptions::bindOptions(OptionsBinder &binder) {
+  static llvm::cl::OptionCategory category(
+      "IREE options for controlling dispatch creation.");
+
+  binder.opt<bool>("iree-opt-outer-dim-concat", outerDimConcat,
+                   llvm::cl::desc("Transposes all concatenations to happen"
+                                  "along the outer most dimension."),
+                   llvm::cl::cat(category));
+}
+
 void SchedulingOptions::bindOptions(OptionsBinder &binder) {
   static llvm::cl::OptionCategory category(
       "IREE options for controlling host/device scheduling.");
diff --git a/compiler/src/iree/compiler/Pipelines/Options.h b/compiler/src/iree/compiler/Pipelines/Options.h
@@ -109,9 +109,6 @@ struct GlobalOptimizationOptions {
   // rewriting named ops as fused generics.
   bool aggressiveTransposePropagation = false;
 
-  // Enables transposing all concatenations to the outer most dimension.
-  bool outerDimConcat = false;
-
   // Enables data tiling.
   bool dataTiling = true;
 
@@ -132,6 +129,15 @@ struct GlobalOptimizationOptions {
   using FromFlags = OptionsFromFlags<GlobalOptimizationOptions>;
 };
 
+// Options controlling dispatch creation.
+struct DispatchCreationOptions {
+  // Enables transposing all concatenations to the outer most dimension.
+  bool outerDimConcat = false;
+
+  void bindOptions(OptionsBinder &binder);
+  using FromFlags = OptionsFromFlags<GlobalOptimizationOptions>;
+};
+
 // Options controlling scheduling across host/device.
 struct SchedulingOptions {
   // Defines the execution model used for scheduling work.
