Move multi-dimensional unrolling pass from Flow to VMLA.
Since now Linalg backend accepts multi-dimensional reductions, we can move the
pass to VMLA.
PiperOrigin-RevId: 300646590
diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD
index d4982ef..76c106a 100644
--- a/iree/compiler/Dialect/Flow/Transforms/BUILD
+++ b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -33,7 +33,6 @@
"Passes.cpp",
"PrePostPartitioningConversion.cpp",
"RematerializeDispatchConstants.cpp",
- "UnrollReductions.cpp",
],
hdrs = [
"Passes.h",
diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
index 5c47dd9..6ffa54b 100644
--- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -33,7 +33,6 @@
"Passes.cpp"
"PrePostPartitioningConversion.cpp"
"RematerializeDispatchConstants.cpp"
- "UnrollReductions.cpp"
DEPS
LLVMSupport
MLIRAnalysis
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
index 91a2e69..c7331ad 100644
--- a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -48,10 +48,6 @@
passManager.addNestedPass<FuncOp>(
IREE::Flow::createPrePartitioningConversionPass());
- // Unroll multi-dimensional reductions to one reduction per dimension.
- passManager.addNestedPass<FuncOp>(IREE::Flow::createUnrollReductionsPass());
- passManager.addNestedPass<FuncOp>(createCSEPass());
-
// First perform module-level analysis that following passes will use to query
// per-function dispatchability information. We run this first so that it only
// needs to run once and will be cached for all of the following passes.
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.h b/iree/compiler/Dialect/Flow/Transforms/Passes.h
index 42aaca9..1ac81b7 100644
--- a/iree/compiler/Dialect/Flow/Transforms/Passes.h
+++ b/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -57,10 +57,6 @@
// allowed to pass through successfully.
std::unique_ptr<OpPassBase<ModuleOp>> createLegalizeInputTypesPass();
-// Unrolls multi-dimensional reduction operations into reductions along each
-// dimension, from innermost to outermost.
-std::unique_ptr<OpPassBase<FuncOp>> createUnrollReductionsPass();
-
// Runs pre-partitioning conversion passes to convert to the flow dialect.
// This converts some input ops directly to flow ops when doing so has a
// benefit. Other ops are left unmodified and will be outlined later on.
diff --git a/iree/compiler/Dialect/VMLA/Transforms/BUILD b/iree/compiler/Dialect/VMLA/Transforms/BUILD
index 5c0ae3e..f2f6970 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/BUILD
+++ b/iree/compiler/Dialect/VMLA/Transforms/BUILD
@@ -22,6 +22,7 @@
srcs = [
"Conversion.cpp",
"Passes.cpp",
+ "UnrollReductions.cpp",
],
hdrs = [
"Passes.h",
diff --git a/iree/compiler/Dialect/VMLA/Transforms/CMakeLists.txt b/iree/compiler/Dialect/VMLA/Transforms/CMakeLists.txt
index e9a42ed..faac549 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/VMLA/Transforms/CMakeLists.txt
@@ -22,6 +22,7 @@
SRCS
"Conversion.cpp"
"Passes.cpp"
+ "UnrollReductions.cpp"
DEPS
LLVMSupport
MLIRIR
diff --git a/iree/compiler/Dialect/VMLA/Transforms/Passes.cpp b/iree/compiler/Dialect/VMLA/Transforms/Passes.cpp
index 0ab1c4f..f870d54 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/VMLA/Transforms/Passes.cpp
@@ -44,6 +44,10 @@
// TODO(benvanik): preserve these hints during conversion.
passManager.addNestedPass<FuncOp>(createDropCompilerHintsPass());
+ // Unroll multi-dimensional reductions to one reduction per dimension.
+ passManager.addNestedPass<FuncOp>(createUnrollReductionsPass());
+ passManager.addNestedPass<FuncOp>(createCSEPass());
+
// Convert from the various input dialects to the VMLA dialect.
passManager.addPass(createConversionPass());
diff --git a/iree/compiler/Dialect/VMLA/Transforms/Passes.h b/iree/compiler/Dialect/VMLA/Transforms/Passes.h
index dce3d1d..f71e596 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/Passes.h
+++ b/iree/compiler/Dialect/VMLA/Transforms/Passes.h
@@ -48,6 +48,10 @@
// Input canonicalization and legalization
//===----------------------------------------------------------------------===//
+// Unrolls multi-dimensional reduction operations into reductions along each
+// dimension, from innermost to outermost.
+std::unique_ptr<OpPassBase<FuncOp>> createUnrollReductionsPass();
+
//===----------------------------------------------------------------------===//
// Dialect conversion
//===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Flow/Transforms/UnrollReductions.cpp b/iree/compiler/Dialect/VMLA/Transforms/UnrollReductions.cpp
similarity index 97%
rename from iree/compiler/Dialect/Flow/Transforms/UnrollReductions.cpp
rename to iree/compiler/Dialect/VMLA/Transforms/UnrollReductions.cpp
index 2f626ed..5844397 100644
--- a/iree/compiler/Dialect/Flow/Transforms/UnrollReductions.cpp
+++ b/iree/compiler/Dialect/VMLA/Transforms/UnrollReductions.cpp
@@ -12,8 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
-#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h"
@@ -22,7 +20,7 @@
namespace mlir {
namespace iree_compiler {
namespace IREE {
-namespace Flow {
+namespace VMLA {
namespace {
@@ -167,10 +165,10 @@
}
static PassRegistration<UnrollReductionsPass> pass(
- "iree-flow-unroll-reductions",
+ "iree-vmla-unroll-reductions",
"Unrolls multi-dimensional reductions to one reduction per dimension.");
-} // namespace Flow
+} // namespace VMLA
} // namespace IREE
} // namespace iree_compiler
} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/unroll_reductions.mlir b/iree/compiler/Dialect/VMLA/Transforms/test/unroll_reductions.mlir
similarity index 95%
rename from iree/compiler/Dialect/Flow/Transforms/test/unroll_reductions.mlir
rename to iree/compiler/Dialect/VMLA/Transforms/test/unroll_reductions.mlir
index d41efed..5c3584f 100644
--- a/iree/compiler/Dialect/Flow/Transforms/test/unroll_reductions.mlir
+++ b/iree/compiler/Dialect/VMLA/Transforms/test/unroll_reductions.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -split-input-file -iree-flow-unroll-reductions -cse %s | IreeFileCheck %s
+// RUN: iree-opt -split-input-file -iree-vmla-unroll-reductions -cse %s | IreeFileCheck %s
// CHECK-LABEL: func @unrolled_reduction
func @unrolled_reduction(%arg0: tensor<4x2x8xf32>) -> tensor<4xf32> {
