Move SCFToCFG into Util dialect and later in the pipeline (#14583)

It is difficult to investigate buffer / dispatch behavior across CFG,
while SCF is much simpler to trace their usage. Pushed SCF to CFG later
in the pipeline.

compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
index 91d5562..6b19c08 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
@@ -69,6 +69,7 @@
         "SplitReduction.cpp",
         "StripSignedness.cpp",
         "TensorPadToTensorInsertSlice.cpp",
+        "TopLevelSCFToCFG.cpp",
         "VerifyInputLegality.cpp",
     ],
     hdrs = [
@@ -117,6 +118,7 @@
         "@llvm-project//mlir:Parser",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:SCFToControlFlow",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:TensorTransforms",

compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
index 296d646..873cdbd 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -68,6 +68,7 @@
     "SplitReduction.cpp"
     "StripSignedness.cpp"
     "TensorPadToTensorInsertSlice.cpp"
+    "TopLevelSCFToCFG.cpp"
     "VerifyInputLegality.cpp"
   DEPS
     ::PassesIncGen
@@ -98,6 +99,7 @@
     MLIRParser
     MLIRPass
     MLIRSCFDialect
+    MLIRSCFToControlFlow
     MLIRSupport
     MLIRTensorDialect
     MLIRTensorTransforms

compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
index 4915afd..dd4ac8b 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -227,6 +227,7 @@
   // Expand tensor shapes into SSA values and optimize the whole program.
   // The more we are able to equate shape dimensions at this level the better
   // our fusions will be.
+  FunctionLikeNest(passManager).addPass(createTopLevelSCFToCFGPass);
   passManager.addPass(IREE::Flow::createExpandTensorShapesPass());
   buildGlobalOptimizationPassPipeline(passManager, transformOptions);
 

compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
index f8c9b01..c8de4f8 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -141,6 +141,10 @@
 std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
 createRemoveZeroExtentTensorsPass();
 
+// Decomposes top-level SCF operations to CFG.
+std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
+createTopLevelSCFToCFGPass();
+
 // Verifies that the input to the Flow transformation pipeline is legal.
 // This includes checking for operations from dialects that are expected
 // to be legalized before this pass.

compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
index 6953102..b20e1ac 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -317,6 +317,12 @@
   let constructor = "mlir::iree_compiler::IREE::Flow::createStripSignednessPass()";
 }
 
+def TopLevelSCFToCFG :
+    InterfacePass<"iree-top-level-scf-to-cfg", "mlir::FunctionOpInterface"> {
+  let summary = "Converts non-nested SCF constructs to CFG (not traversing into opaque operations).";
+  let constructor = "mlir::iree_compiler::IREE::Flow::createTopLevelSCFToCFGPass()";
+}
+
 def VerifyInputLegality: Pass<"iree-verify-input-legality", ""> {
   let summary = "Checks the legality of the IR at the start of IREE flow transformation pipeline.";
   let constructor = "mlir::iree_compiler::IREE::Flow::createVerifyInputLegalityPass()";

compiler/src/iree/compiler/Dialect/Flow/Transforms/TopLevelSCFToCFG.cpp

diff --git a/compiler/src/iree/compiler/InputConversion/Common/TopLevelSCFToCFG.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/TopLevelSCFToCFG.cpp
similarity index 83%
rename from compiler/src/iree/compiler/InputConversion/Common/TopLevelSCFToCFG.cpp
rename to compiler/src/iree/compiler/Dialect/Flow/Transforms/TopLevelSCFToCFG.cpp
index 264382d..b3e0d6b 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/TopLevelSCFToCFG.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/TopLevelSCFToCFG.cpp
@@ -4,17 +4,20 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree/compiler/InputConversion/Common/PassDetail.h"
-#include "iree/compiler/InputConversion/Common/Passes.h"
+#include "iree/compiler/Dialect/Flow/Transforms/PassDetail.h"
+#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 namespace iree_compiler {
+namespace IREE {
+namespace Flow {
 
 namespace {
 
@@ -46,9 +49,12 @@
     signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<func::FuncOp>> createTopLevelSCFToCFGPass() {
+std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
+createTopLevelSCFToCFGPass() {
   return std::make_unique<TopLevelSCFToCFGPass>();
 }
 
+} // namespace Flow
+} // namespace IREE
 } // namespace iree_compiler
 } // namespace mlir

compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
index e11e6c7..3c06e64 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
@@ -54,6 +54,7 @@
             "set_encoding.mlir",
             "strip_signedness.mlir",
             "tensor_pad_to_tensor_insert_slice.mlir",
+            "top_level_scf_to_cfg.mlir",
             "transform_dispatch_region_formation.mlir",
             "transformation_pipeline.mlir",
             "verify_input_ir.mlir",

compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
index 4ae1f62..34fd688 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -52,6 +52,7 @@
     "set_encoding.mlir"
     "strip_signedness.mlir"
     "tensor_pad_to_tensor_insert_slice.mlir"
+    "top_level_scf_to_cfg.mlir"
     "transform_dispatch_region_formation.mlir"
     "transformation_pipeline.mlir"
     "verify_input_ir.mlir"

compiler/src/iree/compiler/Dialect/Flow/Transforms/test/top_level_scf_to_cfg.mlir

diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/top_level_scf_to_cfg.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/top_level_scf_to_cfg.mlir
similarity index 94%
rename from compiler/src/iree/compiler/InputConversion/Common/test/top_level_scf_to_cfg.mlir
rename to compiler/src/iree/compiler/Dialect/Flow/Transforms/test/top_level_scf_to_cfg.mlir
index 7bf1c0f..6b4987e 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/test/top_level_scf_to_cfg.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/top_level_scf_to_cfg.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --split-input-file --iree-top-level-scf-to-cfg %s | FileCheck %s
+// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(iree-top-level-scf-to-cfg))" %s | FileCheck %s
 
 // CHECK-LABEL: @generic_nested_for
 // While not super recommended, we do have cases of SCF constructs embedded

compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel

diff --git a/compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel b/compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel
index 919a2c6..462d313 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel
+++ b/compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel
@@ -50,7 +50,6 @@
         "QuantizedConvToConv.cpp",
         "QuantizedMatmulToMatmul.cpp",
         "SanitizeModuleNames.cpp",
-        "TopLevelSCFToCFG.cpp",
         "Utils.cpp",
     ],
     hdrs = [

compiler/src/iree/compiler/InputConversion/Common/CMakeLists.txt

diff --git a/compiler/src/iree/compiler/InputConversion/Common/CMakeLists.txt b/compiler/src/iree/compiler/InputConversion/Common/CMakeLists.txt
index 0ee94e7..a4a0c56 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/CMakeLists.txt
+++ b/compiler/src/iree/compiler/InputConversion/Common/CMakeLists.txt
@@ -54,7 +54,6 @@
     "QuantizedConvToConv.cpp"
     "QuantizedMatmulToMatmul.cpp"
     "SanitizeModuleNames.cpp"
-    "TopLevelSCFToCFG.cpp"
     "Utils.cpp"
   DEPS
     ::PassHeaders

compiler/src/iree/compiler/InputConversion/Common/Passes.cpp

diff --git a/compiler/src/iree/compiler/InputConversion/Common/Passes.cpp b/compiler/src/iree/compiler/InputConversion/Common/Passes.cpp
index f7342a1..e03379d 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/Passes.cpp
+++ b/compiler/src/iree/compiler/InputConversion/Common/Passes.cpp
@@ -21,7 +21,6 @@
 
 void buildCommonInputConversionPassPipeline(OpPassManager &passManager) {
   // Currently we don't handle SCF ops well and have to convert them all to CFG.
-  passManager.addNestedPass<func::FuncOp>(createTopLevelSCFToCFGPass());
   passManager.addPass(createIREEImportPublicPass());
   passManager.addPass(createImportMLProgramPass());
   passManager.addPass(createSanitizeModuleNamesPass());

compiler/src/iree/compiler/InputConversion/Common/Passes.h

diff --git a/compiler/src/iree/compiler/InputConversion/Common/Passes.h b/compiler/src/iree/compiler/InputConversion/Common/Passes.h
index edbfcb7..7810bd2 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/Passes.h
+++ b/compiler/src/iree/compiler/InputConversion/Common/Passes.h
@@ -41,7 +41,6 @@
 std::unique_ptr<OperationPass<func::FuncOp>>
 createLinalgQuantizedMatmulToMatmulPass();
 std::unique_ptr<OperationPass<ModuleOp>> createSanitizeModuleNamesPass();
-std::unique_ptr<OperationPass<func::FuncOp>> createTopLevelSCFToCFGPass();
 
 //===----------------------------------------------------------------------===//
 // Register all Passes

compiler/src/iree/compiler/InputConversion/Common/Passes.td

diff --git a/compiler/src/iree/compiler/InputConversion/Common/Passes.td b/compiler/src/iree/compiler/InputConversion/Common/Passes.td
index 950cce6..6f4d3c6 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/Passes.td
+++ b/compiler/src/iree/compiler/InputConversion/Common/Passes.td
@@ -39,12 +39,6 @@
   let constructor = "mlir::iree_compiler::createSanitizeModuleNamesPass()";
 }
 
-def TopLevelSCFToCFG :
-    Pass<"iree-top-level-scf-to-cfg", "func::FuncOp"> {
-  let summary = "Converts non-nested SCF constructs to CFG (not traversing into opaque operations).";
-  let constructor = "mlir::iree_compiler::createTopLevelSCFToCFGPass()";
-}
-
 def AutoInputConversionPipeline :
     Pass<"iree-auto-input-conversion", "ModuleOp"> {
   let summary = "Analyzes and runs appropriate input pipeline.";

compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel

diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel b/compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel
index a99b55e..89a24b8 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel
@@ -24,7 +24,6 @@
             "linalg_quantized_conv_to_conv.mlir",
             "linalg_quantized_matmul_to_matmul.mlir",
             "sanitize_module_names.mlir",
-            "top_level_scf_to_cfg.mlir",
         ],
         include = ["*.mlir"],
     ),

compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt

diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt
index 6c5ee1f..8b967d3 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt
@@ -20,7 +20,6 @@
     "linalg_quantized_conv_to_conv.mlir"
     "linalg_quantized_matmul_to_matmul.mlir"
     "sanitize_module_names.mlir"
-    "top_level_scf_to_cfg.mlir"
   TOOLS
     FileCheck
     iree-opt

compiler/src/iree/compiler/InputConversion/TOSA/Passes.cpp

diff --git a/compiler/src/iree/compiler/InputConversion/TOSA/Passes.cpp b/compiler/src/iree/compiler/InputConversion/TOSA/Passes.cpp
index 0d6da19..de6ff8e 100644
--- a/compiler/src/iree/compiler/InputConversion/TOSA/Passes.cpp
+++ b/compiler/src/iree/compiler/InputConversion/TOSA/Passes.cpp
@@ -36,7 +36,6 @@
   // In the future it would be nice if we could have all of flow be both scf
   // and cfg compatible.
   passManager.addNestedPass<func::FuncOp>(tosa::createTosaToSCF());
-  passManager.addNestedPass<func::FuncOp>(createTopLevelSCFToCFGPass());
 
   // We also don't handle calls well on the old codepath; until we remove the
   // use of the CFG we can continue inlining.