From 352dcb01fcdd1b8a7ce04c4b77ed400eff91c8ea Mon Sep 17 00:00:00 2001
From: Max Dawkins <max.dawkins@gmail.com>
Date: Thu, 19 Sep 2024 11:51:33 -0400
Subject: [PATCH] remove createPass wrapper

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
---
 .../Codegen/Common/DecomposePackUnPackOps.cpp     |  3 ++-
 .../src/iree/compiler/Codegen/Common/Passes.h     | 11 -----------
 .../src/iree/compiler/Codegen/Common/Passes.td    | 10 +++++++++-
 .../src/iree/compiler/Codegen/LLVMGPU/Passes.cpp  | 15 ++++++++++-----
 4 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp b/compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp
index e8b18370a2c37..6caf0b5a1ed7c 100644
--- a/compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp
@@ -19,7 +19,6 @@
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/Visitors.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #define DEBUG_TYPE "iree-codegen-decompose-pack-unpack-ops"
@@ -31,6 +30,8 @@ namespace mlir::iree_compiler {
 
 namespace {
 
+using PackUnPackControlFn = std::function<LogicalResult(Operation *)>;
+
 /// A wrapper pattern that calls linalg::lowerPack on tensor::PackOp. It lowers
 /// a tensor.pack op to tensor.pad + tensor.expand_shape + linalg.transpose ops.
 struct LowerPackPattern : public OpRewritePattern<tensor::PackOp> {
diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.h b/compiler/src/iree/compiler/Codegen/Common/Passes.h
index 0abb8e2abb33e..94192d52a76b6 100644
--- a/compiler/src/iree/compiler/Codegen/Common/Passes.h
+++ b/compiler/src/iree/compiler/Codegen/Common/Passes.h
@@ -67,17 +67,6 @@ using ConfigFn =
 std::unique_ptr<InterfacePass<FunctionOpInterface>>
 createConvolutionToIGEMMPass(ConfigFn configFn);
 
-using PackUnPackControlFn = std::function<LogicalResult(Operation *)>;
-// Pass to decompose pack and unpack ops into pad/extract_slice and reshape ops.
-// If specified, `controlFn` controls which ops get decomposed. The `controlFn`
-// should be used with `useOnlyReshapes` set to true.
-// TODO(Max191): Add a controlFn upstream for `GeneralizeOuterUnitDim*` patterns
-// and remove the need to have `useOnlyReshapes = true` when using `controlFn`.
-std::unique_ptr<InterfacePass<FunctionOpInterface>>
-createDecomposePackUnPackOpsPass(
-    bool tileOuterToOne, bool useOnlyReshapes = false,
-    std::optional<PackUnPackControlFn> controlFn = std::nullopt);
-
 std::unique_ptr<Pass> createDecomposeSoftmaxPass(bool useFusion);
 
 /// Pass to perform linalg on tensor bufferization. The function passed into
diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.td b/compiler/src/iree/compiler/Codegen/Common/Passes.td
index 7d912f8a445c3..f221df122e044 100644
--- a/compiler/src/iree/compiler/Codegen/Common/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Common/Passes.td
@@ -155,7 +155,15 @@ def DecomposePackUnPackOpsPass :
     Option<"tileOuterToOne", "tile-outer-to-one", "bool", "false",
            "Always apply tiling to make outer dimension be ones">,
     Option<"useOnlyReshapes", "use-only-reshapes", "bool", "false",
-           "Use decomposition into reshape ops, even when packing unit dimensions.">
+           "Use decomposition into reshape ops, even when packing unit dimensions.">,
+    // If specified, `controlFn` controls which ops get decomposed. 
+    // TODO(Max191): The `controlFn` should be used with `useOnlyReshapes` set
+    // to true. We should add a controlFn upstream for `GeneralizeOuterUnitDim*`
+    // patterns and remove the need to have `useOnlyReshapes = true` when using
+    // `controlFn`.
+    Option<"controlFn", "control-fn",
+           "std::optional<std::function<LogicalResult(Operation *)>>", "std::nullopt",
+           "Controls which pack and unpack ops get decomposed.">
   ];
 }
 
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index 1f4ef92066b6a..a392a8c1d9d18 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -347,9 +347,11 @@ void addGPUTileAndFusePassPipeline(OpPassManager &funcPassManager,
   }
 
   // Step 3. Decompose pack and unpack ops and propagate the resulting reshapes.
-  funcPassManager.addPass(
-      createDecomposePackUnPackOpsPass(/*tileOuterToOne=*/false,
-                                       /*useOnlyReshapes=*/true));
+  {
+    DecomposePackUnPackOpsPassOptions options;
+    options.useOnlyReshapes = true;
+    funcPassManager.addPass(createDecomposePackUnPackOpsPass(options));
+  }
 
   // Step 3.5. Expand the inner dimensions of MultiMma ops in preparation for
   // distribution to lanes.
@@ -945,8 +947,11 @@ void addGPUPackUnPackPasses(OpPassManager &funcPassManager) {
   funcPassManager.addPass(createCanonicalizerPass());
   funcPassManager.addPass(createCSEPass());
 
-  funcPassManager.addPass(
-      createDecomposePackUnPackOpsPass(/*tileOuterToOne=*/true));
+  {
+    DecomposePackUnPackOpsPassOptions options;
+    options.tileOuterToOne = true;
+    funcPassManager.addPass(createDecomposePackUnPackOpsPass(options));
+  }
   funcPassManager.addPass(createCanonicalizerPass());
   funcPassManager.addPass(createCSEPass());
   addGPUVectorizationPasses(funcPassManager);