[NFC] Move RaiseSpecialOps to GlobalOptimization (iree-org#15881)

This moves RaiseSpecialOps out of Flow and into GlobalOptimization. The reason for the pass being in Flow is mostly historical, and now that all of the relevant passes needed before or after are in GlobalOptimization, we can move this pass also.
ramiro050 · Dec 19, 2023 · f366140 · f366140
1 parent 25a729d
commit f366140
Show file tree

Hide file tree

Showing 16 changed files with 32 additions and 27 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
@@ -56,7 +56,6 @@ iree_compiler_cc_library(
         "OutlineDispatchRegions.cpp",
         "PassDetail.h",
         "Passes.cpp",
-        "RaiseSpecialOps.cpp",
         "RegionOpUtils.cpp",
         "SplitReduction.cpp",
         "TensorPadToTensorInsertSlice.cpp",

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -55,7 +55,6 @@ iree_cc_library(
     "OutlineDispatchRegions.cpp"
     "PassDetail.h"
     "Passes.cpp"
-    "RaiseSpecialOps.cpp"
     "RegionOpUtils.cpp"
     "SplitReduction.cpp"
     "TensorPadToTensorInsertSlice.cpp"

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -129,7 +129,6 @@ void buildFlowTransformPassPipeline(OpPassManager &passManager,
 
   FunctionLikeNest(passManager)
       // Preprocess the input to a form more amenable for fusion
-      .addPass(createRaiseSpecialOps)
       .addPass(createInterchangeGenericOpsPass)
       .addPass(memref::createResolveShapedTypeResultDimsPass)
       .addPass(mlir::createCanonicalizerPass)

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -190,10 +190,6 @@ createOutlineLargeConstantsPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>>
 createDeduplicateExecutablesPass();
 
-// Create a pass to raise sequence of ops to higher level linalg.ext
-// representation.
-std::unique_ptr<Pass> createRaiseSpecialOps();
-
 // Create a pass to split reduction dimension.
 std::unique_ptr<Pass> createSplitReductionPass();
 

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -241,12 +241,6 @@ def OutlineDispatchRegions :
   let constructor = "mlir::iree_compiler::IREE::Flow::createOutlineDispatchRegionsPass()";
 }
 
-def RaiseSpecialOps :
-    Pass<"iree-flow-raise-special-ops", ""> {
-  let summary = "raise special ops like softmax to the high level linalg.ext representation";
-  let constructor = "mlir::iree_compiler::IREE::Flow::createRaiseSpecialOps()";
-}
-
 def SplitReduction :
     Pass<"iree-flow-split-reduction-ops", ""> {
   let summary = "Split reduction dimension to increase parallelism.";

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
@@ -44,7 +44,6 @@ iree_lit_test_suite(
             "pad_fusion_with_consumer.mlir",
             "pad_fusion_with_producer.mlir",
             "pipeline_tests.mlir",
-            "raise_special_ops.mlir",
             "tensor_pad_to_tensor_insert_slice.mlir",
             "top_level_scf_to_cfg.mlir",
             "transform_dispatch_region_formation.mlir",

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -42,7 +42,6 @@ iree_lit_test_suite(
     "pad_fusion_with_consumer.mlir"
     "pad_fusion_with_producer.mlir"
     "pipeline_tests.mlir"
-    "raise_special_ops.mlir"
     "tensor_pad_to_tensor_insert_slice.mlir"
     "top_level_scf_to_cfg.mlir"
     "transform_dispatch_region_formation.mlir"

diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -59,6 +59,7 @@ iree_compiler_cc_library(
         "MaterializeHomogeneousEncodings.cpp",
         "OptimizeNumerics.cpp",
         "Passes.cpp",
+        "RaiseSpecialOps.cpp",
         "RemoveZeroExtentTensors.cpp",
         "SetEncoding.cpp",
         "Utils.cpp",
@@ -84,9 +85,11 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/Util/Transforms",
         "//compiler/src/iree/compiler/Pipelines:Options",
         "//compiler/src/iree/compiler/Utils",
+        "//llvm-external-projects/iree-dialects:IREEDialectsTransforms",
         "//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
         "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
         "//llvm-external-projects/iree-dialects:IREELinalgExtUtils",
+        "//llvm-external-projects/iree-dialects:IREELinalgTransformDialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:ArithDialect",

diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -55,15 +55,18 @@ iree_cc_library(
     "MaterializeHomogeneousEncodings.cpp"
     "OptimizeNumerics.cpp"
     "Passes.cpp"
+    "RaiseSpecialOps.cpp"
     "RemoveZeroExtentTensors.cpp"
     "SetEncoding.cpp"
     "Utils.cpp"
   DEPS
     ::PassHeaders
     ::PassesIncGen
+    IREEDialectsTransforms
     IREELinalgExtDialect
     IREELinalgExtTransforms
     IREELinalgExtUtils
+    IREELinalgTransformDialect
     LLVMSupport
     MLIRAffineDialect
     MLIRArithDialect

diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -89,7 +89,7 @@ void buildGlobalOptimizationPassPipeline(
       // RaiseSpecialOps, by virtue of implementing various peephole
       // optimizations, is sensitive to surrounding IR structure. Thus we run
       // this pass both before unit dim folding + consteval, as well as after.
-      .addPass(IREE::Flow::createRaiseSpecialOps)
+      .addPass(createRaiseSpecialOps)
       // 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.
@@ -163,12 +163,14 @@ void buildGlobalOptimizationPassPipeline(
   mainPassManager.addPass(
       IREE::Util::createFixedPointIteratorPass(std::move(pipeline)));
 
-  // Strip std.assert & co after we perform optimizations; prior to this we
-  // may use the assertions to derive information during analysis.
-  if (transformOptions.options.stripAssertions) {
-    FunctionLikeNest(mainPassManager)
-        .addPass(IREE::Util::createStripDebugOpsPass);
-  }
+  FunctionLikeNest(mainPassManager)
+      // After running const-eval to a fixed point and folding unit extent dims,
+      // try any new raising opportunities.
+      .addPass(createRaiseSpecialOps)
+      // Strip std.assert & co after we perform optimizations; prior to this we
+      // may use the assertions to derive information during analysis.
+      .addPredicatedPass(transformOptions.options.stripAssertions,
+                         IREE::Util::createStripDebugOpsPass);
 }
 
 namespace {

diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.h b/compiler/src/iree/compiler/GlobalOptimization/Passes.h
@@ -94,6 +94,10 @@ createMaterializeHomogeneousEncodingsPass();
 // iree-global-opt-infer-numeric-narrowing.
 std::unique_ptr<Pass> createOptimizeNumericsPass();
 
+// Create a pass to raise sequence of ops to higher level linalg.ext
+// representation.
+std::unique_ptr<Pass> createRaiseSpecialOps();
+
 // Removes tensors that have 0-extents.
 std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
 createRemoveZeroExtentTensorsPass();

diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -102,6 +102,12 @@ def OptimizeNumerics :
   let constructor = "mlir::iree_compiler::GlobalOptimization::createOptimizeNumericsPass()";
 }
 
+def RaiseSpecialOps :
+    Pass<"iree-global-opt-raise-special-ops", ""> {
+  let summary = "raise special ops like softmax to the high level linalg.ext representation";
+  let constructor = "mlir::iree_compiler::GlobalOptimization::createRaiseSpecialOps()";
+}
+
 def RemoveZeroExtentTensors :
     InterfacePass<"iree-global-opt-remove-zero-extent-tensors", "mlir::FunctionOpInterface"> {
   let summary = "Remove tensors that have 0-extents";

diff --git a/...alect/Flow/Transforms/RaiseSpecialOps.cpp → ...er/GlobalOptimization/RaiseSpecialOps.cpp b/...alect/Flow/Transforms/RaiseSpecialOps.cpp → ...er/GlobalOptimization/RaiseSpecialOps.cpp
@@ -9,8 +9,8 @@
 #include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree-dialects/Dialect/LinalgTransform/StructuredTransformOpsExt.h"
 #include "iree-dialects/Transforms/TransformMatchers.h"
-#include "iree/compiler/Dialect/Flow/Transforms/PassDetail.h"
-#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
+#include "iree/compiler/GlobalOptimization/PassDetail.h"
+#include "iree/compiler/GlobalOptimization/Passes.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -21,7 +21,7 @@
 using namespace mlir;
 using transform_ext::StructuredOpMatcher;
 
-namespace mlir::iree_compiler::IREE::Flow {
+namespace mlir::iree_compiler::GlobalOptimization {
 
 namespace {
 
@@ -806,4 +806,4 @@ std::unique_ptr<Pass> createRaiseSpecialOps() {
   return std::make_unique<RaiseSpecialOpsPass>();
 }
 
-} // namespace mlir::iree_compiler::IREE::Flow
+} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
@@ -29,6 +29,7 @@ iree_lit_test_suite(
             "lift_generic_to_transpose_batch_matmul.mlir",
             "materialize_homogeneous_encodings.mlir",
             "optimize_numerics.mlir",
+            "raise_special_ops.mlir",
             "remove_zero_extent_tensors.mlir",
             "set_encoding.mlir",
             "transformation_pipeline.mlir",

diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
@@ -27,6 +27,7 @@ iree_lit_test_suite(
     "lift_generic_to_transpose_batch_matmul.mlir"
     "materialize_homogeneous_encodings.mlir"
     "optimize_numerics.mlir"
+    "raise_special_ops.mlir"
     "remove_zero_extent_tensors.mlir"
     "set_encoding.mlir"
     "transformation_pipeline.mlir"

diff --git a/...ow/Transforms/test/raise_special_ops.mlir → ...lOptimization/test/raise_special_ops.mlir b/...ow/Transforms/test/raise_special_ops.mlir → ...lOptimization/test/raise_special_ops.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --iree-flow-raise-special-ops -canonicalize --split-input-file %s | FileCheck %s
+// RUN: iree-opt --iree-global-opt-raise-special-ops -canonicalize --split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: @softmax
 //  CHECK-SAME: %[[ARG:.+]]: tensor<?x?x?xf32>