NFC: Merge ConvertToFlow passes into dedicated before/after passes. (#6850)
This refactoring prepares for a similar before/after split within `populateTensorToFlowPatterns` so I can move the `tensor.extract` to `flow.tensor.load` pattern from `PromoteTensorLoads` into that common location.
This also follows work in https://github.com/google/iree/pull/6651 to organize conversions from `tensor` and `linalg` to `flow`.
<sub>[Related Discord discussion here](https://discord.com/channels/689900678990135345/689906000043573354/879860687747289099)</sub>
diff --git a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/cast.mlir b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/cast.mlir
index 5167626..ef90610 100644
--- a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/cast.mlir
+++ b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/cast.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -allow-unregistered-dialect -split-input-file -iree-flow-convert-tensor-ops-pass %s | IreeFileCheck %s
+// RUN: iree-opt -allow-unregistered-dialect -split-input-file -iree-flow-convert-to-flow-before-dispatch-formation %s | IreeFileCheck %s
func @static_tensor_cast_to_dynamic(%arg0: tensor<4x4xf32>) -> tensor<?x?xf32> {
// CHECK-DAG: %[[C4:.*]] = constant 4 : index
diff --git a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/extract_slice.mlir b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/extract_slice.mlir
index ce1bff7..dc5b0d2 100644
--- a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/extract_slice.mlir
+++ b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/extract_slice.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -allow-unregistered-dialect -split-input-file -iree-flow-convert-tensor-ops-pass %s | IreeFileCheck %s
+// RUN: iree-opt -allow-unregistered-dialect -split-input-file -iree-flow-convert-to-flow-before-dispatch-formation %s | IreeFileCheck %s
func @extract_slice1(%arg0 : tensor<5x24x48xf32>) -> tensor<4xf32> {
%0 = tensor.extract_slice %arg0[2, 3, 4] [1, 1, 4] [1, 1, 1]
diff --git a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/from_elements.mlir b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/from_elements.mlir
index d380abe..a2953b8 100644
--- a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/from_elements.mlir
+++ b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/from_elements.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -allow-unregistered-dialect -split-input-file -iree-flow-convert-tensor-ops-pass %s | IreeFileCheck %s
+// RUN: iree-opt -allow-unregistered-dialect -split-input-file -iree-flow-convert-to-flow-before-dispatch-formation %s | IreeFileCheck %s
// CHECK: func @tensor.from_elements__to__flow.tensor.splat(%[[arg0:.*]]: i8)
func @tensor.from_elements__to__flow.tensor.splat(%arg0: i8) -> (i8) {
diff --git a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/insert_slice.mlir b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/insert_slice.mlir
index 1c29bf4..70f0fff 100644
--- a/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/insert_slice.mlir
+++ b/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/test/insert_slice.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -allow-unregistered-dialect -split-input-file -iree-flow-convert-tensor-ops-pass %s | IreeFileCheck %s
+// RUN: iree-opt -allow-unregistered-dialect -split-input-file -iree-flow-convert-to-flow-before-dispatch-formation %s | IreeFileCheck %s
func @insert_slice_convert
(%arg0 : tensor<?x24x48xf32>, %arg1 : tensor<1x4x48xf32>) ->
diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD
index 40bb2ed..41c6aae 100644
--- a/iree/compiler/Dialect/Flow/Transforms/BUILD
+++ b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -33,8 +33,7 @@
srcs = [
"ConvertConv2D1x1ToMatmulPass.cpp",
"ConvertConv2DToImg2ColPass.cpp",
- "ConvertLinalgTensorOps.cpp",
- "ConvertTensorOps.cpp",
+ "ConvertToFlow.cpp",
"DeduplicateExecutables.cpp",
"DestructiveUpdateUtils.cpp",
"DispatchLinalgOnTensors.cpp",
diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
index 167839f..4152bbf 100644
--- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -30,8 +30,7 @@
SRCS
"ConvertConv2D1x1ToMatmulPass.cpp"
"ConvertConv2DToImg2ColPass.cpp"
- "ConvertLinalgTensorOps.cpp"
- "ConvertTensorOps.cpp"
+ "ConvertToFlow.cpp"
"DeduplicateExecutables.cpp"
"DestructiveUpdateUtils.cpp"
"DispatchLinalgOnTensors.cpp"
diff --git a/iree/compiler/Dialect/Flow/Transforms/ConvertTensorOps.cpp b/iree/compiler/Dialect/Flow/Transforms/ConvertTensorOps.cpp
deleted file mode 100644
index a4ba75b..0000000
--- a/iree/compiler/Dialect/Flow/Transforms/ConvertTensorOps.cpp
+++ /dev/null
@@ -1,55 +0,0 @@
-// Copyright 2021 The IREE Authors
-//
-// Licensed under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-#include "iree/compiler/Dialect/Flow/Conversion/TensorToFlow/ConvertTensorToFlow.h"
-#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
-#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
-#include "iree/compiler/Dialect/Flow/Transforms/PassDetail.h"
-#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
-#include "mlir/Dialect/StandardOps/IR/Ops.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Pass/PassManager.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-
-#define DEBUG_TYPE "iree-flow-convert-tensor-ops"
-
-namespace mlir {
-namespace iree_compiler {
-namespace IREE {
-namespace Flow {
-
-namespace {
-
-/// Converts operations that can map to flow.tensor.* operations.
-struct ConvertTensorOpsPass
- : public ConvertTensorOpsBase<ConvertTensorOpsPass> {
- void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<IREE::Flow::FlowDialect, mlir::StandardOpsDialect,
- tensor::TensorDialect>();
- }
- void runOnOperation() override {
- auto funcOp = getOperation();
- MLIRContext *context = funcOp->getContext();
- context->allowUnregisteredDialects(true);
- RewritePatternSet patterns(&getContext());
- populateTensorToFlowPatterns(&getContext(), patterns);
- if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns)))) {
- return signalPassFailure();
- }
- }
-};
-} // namespace
-
-std::unique_ptr<OperationPass<mlir::FuncOp>> createConvertTensorOpsPass() {
- return std::make_unique<ConvertTensorOpsPass>();
-}
-
-} // namespace Flow
-} // namespace IREE
-} // namespace iree_compiler
-} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/ConvertLinalgTensorOps.cpp b/iree/compiler/Dialect/Flow/Transforms/ConvertToFlow.cpp
similarity index 71%
rename from iree/compiler/Dialect/Flow/Transforms/ConvertLinalgTensorOps.cpp
rename to iree/compiler/Dialect/Flow/Transforms/ConvertToFlow.cpp
index 455a7ef..c8da8d0 100644
--- a/iree/compiler/Dialect/Flow/Transforms/ConvertLinalgTensorOps.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/ConvertToFlow.cpp
@@ -4,6 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#include "iree/compiler/Dialect/Flow/Conversion/TensorToFlow/ConvertTensorToFlow.h"
#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/Flow/IR/FlowTypes.h"
@@ -18,7 +19,7 @@
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#define DEBUG_TYPE "iree-flow-convert-linalg-tensor-ops"
+#define DEBUG_TYPE "iree-flow-convert-to-flow"
namespace mlir {
namespace iree_compiler {
@@ -90,16 +91,9 @@
}
};
-/// Converts linalg operations that can map to flow.tensor.* operations.
-struct ConvertLinalgTensorOpsPass
- : public ConvertLinalgTensorOpsBase<ConvertLinalgTensorOpsPass> {
- ConvertLinalgTensorOpsPass(bool runBefore) {
- runBeforeDispatchRegionFormation = runBefore;
- }
- ConvertLinalgTensorOpsPass(const ConvertLinalgTensorOpsPass &that) {
- runBeforeDispatchRegionFormation = that.runBeforeDispatchRegionFormation;
- }
-
+struct ConvertToFlowBeforeDispatchFormation
+ : public ConvertToFlowBeforeDispatchFormationBase<
+ ConvertToFlowBeforeDispatchFormation> {
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<IREE::Flow::FlowDialect, tensor::TensorDialect,
linalg::LinalgDialect, mlir::StandardOpsDialect>();
@@ -109,26 +103,52 @@
MLIRContext *context = funcOp->getContext();
context->allowUnregisteredDialects(true);
RewritePatternSet patterns(&getContext());
- if (runBeforeDispatchRegionFormation) {
- patterns.insert<
- LinalgTensorReshapeToFlowTensorReshape<linalg::TensorCollapseShapeOp>,
- LinalgTensorReshapeToFlowTensorReshape<linalg::TensorExpandShapeOp>>(
- context);
- } else {
- patterns.insert<LinalgFillToFlowTensorSplat>(context);
- }
+
+ patterns.insert<
+ LinalgTensorReshapeToFlowTensorReshape<linalg::TensorCollapseShapeOp>,
+ LinalgTensorReshapeToFlowTensorReshape<linalg::TensorExpandShapeOp>>(
+ context);
+ populateTensorToFlowPatterns(context, patterns);
IREE::Flow::TensorReshapeOp::getCanonicalizationPatterns(patterns, context);
+
if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns)))) {
return signalPassFailure();
}
}
};
+
+struct ConvertToFlowAfterDispatchFormation
+ : public ConvertToFlowAfterDispatchFormationBase<
+ ConvertToFlowAfterDispatchFormation> {
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::Flow::FlowDialect, tensor::TensorDialect,
+ linalg::LinalgDialect, mlir::StandardOpsDialect>();
+ }
+ void runOnOperation() override {
+ auto funcOp = getOperation();
+ MLIRContext *context = funcOp->getContext();
+ context->allowUnregisteredDialects(true);
+ RewritePatternSet patterns(&getContext());
+
+ patterns.insert<LinalgFillToFlowTensorSplat>(context);
+ IREE::Flow::TensorReshapeOp::getCanonicalizationPatterns(patterns, context);
+
+ if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns)))) {
+ return signalPassFailure();
+ }
+ }
+};
+
} // namespace
-std::unique_ptr<OperationPass<mlir::FuncOp>> createConvertLinalgTensorOpsPass(
- bool runBeforeDispatchRegionFormation) {
- return std::make_unique<ConvertLinalgTensorOpsPass>(
- runBeforeDispatchRegionFormation);
+std::unique_ptr<OperationPass<mlir::FuncOp>>
+createConvertToFlowBeforeDispatchFormation() {
+ return std::make_unique<ConvertToFlowBeforeDispatchFormation>();
+}
+
+std::unique_ptr<OperationPass<mlir::FuncOp>>
+createConvertToFlowAfterDispatchFormation() {
+ return std::make_unique<ConvertToFlowAfterDispatchFormation>();
}
} // namespace Flow
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
index 47685da..415f0f3 100644
--- a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -143,17 +143,13 @@
}
passManager.addPass(memref::createResolveShapedTypeResultDimsPass());
passManager.addNestedPass<mlir::FuncOp>(
- IREE::Flow::createConvertTensorOpsPass());
- passManager.addNestedPass<mlir::FuncOp>(
- IREE::Flow::createConvertLinalgTensorOpsPass(
- /*runBeforeDispatchRegionFormation=*/true));
+ IREE::Flow::createConvertToFlowBeforeDispatchFormation());
passManager.addNestedPass<mlir::FuncOp>(mlir::createCanonicalizerPass());
passManager.addNestedPass<mlir::FuncOp>(
IREE::Flow::createDispatchLinalgOnTensorsPass());
passManager.addPass(memref::createResolveShapedTypeResultDimsPass());
passManager.addNestedPass<mlir::FuncOp>(
- IREE::Flow::createConvertLinalgTensorOpsPass(
- /*runBeforeDispatchRegionFormation=*/false));
+ IREE::Flow::createConvertToFlowAfterDispatchFormation());
// NOTE: required because the current dispatch-linalg-on-tensors pass
// creates a lot of dead IR that needs to be cleaned up.
passManager.addNestedPass<mlir::FuncOp>(mlir::createCanonicalizerPass());
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.h b/iree/compiler/Dialect/Flow/Transforms/Passes.h
index e35210c..91e65ba 100644
--- a/iree/compiler/Dialect/Flow/Transforms/Passes.h
+++ b/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -68,18 +68,14 @@
/// the most inner loops.
std::unique_ptr<OperationPass<mlir::FuncOp>> createInterchangeGenericOpsPass();
-// Convert tensor operations to equivalent flow.tensor.* operations
-// `runBeforeDispatchRegionFormation` controls whether to run before dispatch
-// region creation. If run after, it will catch operations that were left
-// outside of dispatch regions and could be represented as flow.tensor.* ops.
-std::unique_ptr<OperationPass<mlir::FuncOp>> createConvertTensorOpsPass();
+// Convert operations to equivalent flow ops before dispatch region creation.
+std::unique_ptr<OperationPass<mlir::FuncOp>>
+createConvertToFlowBeforeDispatchFormation();
-// Convert linalg.tensor operations to equivalent flow.tensor.* ops.
-// `runBeforeDispatchRegionFormation` controls whether to run before dispatch
-// region creation. If run after, it will catch operations that were left
-// outside of dispatch regions and could be represented as flow.tensor.* ops.
-std::unique_ptr<OperationPass<mlir::FuncOp>> createConvertLinalgTensorOpsPass(
- bool runBeforeDispatchRegionFormation = true);
+// Convert remaining operations that were left outside of dispatch regions to
+// equivalent flow ops.
+std::unique_ptr<OperationPass<mlir::FuncOp>>
+createConvertToFlowAfterDispatchFormation();
// Promote I1 tensor constants to I8 tensors to match later operations.
std::unique_ptr<OperationPass<mlir::FuncOp>> createPromoteI1ToI8Pass();
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.td b/iree/compiler/Dialect/Flow/Transforms/Passes.td
index 8a349e4..e8b6956 100644
--- a/iree/compiler/Dialect/Flow/Transforms/Passes.td
+++ b/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -21,20 +21,16 @@
let constructor = "mlir::iree_compiler::IREE::Flow::createConvertConv2DToImg2ColPass()";
}
-def ConvertTensorOps :
- Pass<"iree-flow-convert-tensor-ops-pass", "mlir::FuncOp"> {
- let summary = "Convert tensor operations to equivalent flow.tensor.* operations";
- let constructor = "mlir::iree_compiler::IREE::Flow::createConvertTensorOpsPass()";
+def ConvertToFlowBeforeDispatchFormation :
+ Pass<"iree-flow-convert-to-flow-before-dispatch-formation", "mlir::FuncOp"> {
+ let summary = "Convert operations to flow before dispatch formation";
+ let constructor = "mlir::iree_compiler::IREE::Flow::createConvertToFlowBeforeDispatchFormation()";
}
-def ConvertLinalgTensorOps :
- Pass<"iree-flow-convert-linalg-tensor-ops-pass", "mlir::FuncOp"> {
- let summary = "Convert linalg operations to equivalent flow.tensor.* operations";
- let constructor = "mlir::iree_compiler::IREE::Flow::createConvertLinalgTensorOpsPass()";
- let options = [
- Option<"runBeforeDispatchRegionFormation", "run-before-dispatch-region-formation",
- "bool", /*default=*/"true", "Run the pass before dispatch region formation">
- ];
+def ConvertToFlowAfterDispatchFormation :
+ Pass<"iree-flow-convert-to-flow-after-dispatch-formation", "mlir::FuncOp"> {
+ let summary = "Convert operations to flow after dispatch formation";
+ let constructor = "mlir::iree_compiler::IREE::Flow::createConvertToFlowAfterDispatchFormation()";
}
def DeduplicateExecutables :
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/convert_linalg_tensor_ops_after.mlir b/iree/compiler/Dialect/Flow/Transforms/test/convert_linalg_tensor_ops_after.mlir
index a7716d2..12aa304 100644
--- a/iree/compiler/Dialect/Flow/Transforms/test/convert_linalg_tensor_ops_after.mlir
+++ b/iree/compiler/Dialect/Flow/Transforms/test/convert_linalg_tensor_ops_after.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -iree-flow-convert-linalg-tensor-ops-pass='run-before-dispatch-region-formation=false' -canonicalize -cse -split-input-file %s | IreeFileCheck %s
+// RUN: iree-opt -iree-flow-convert-to-flow-after-dispatch-formation -canonicalize -cse -split-input-file %s | IreeFileCheck %s
func @turn_fill_into_splat(%arg0: tensor<?x?xf32>, %arg1: tensor<f32>, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> tensor<?x?xf32> {
%c0 = constant 0 : index
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/convert_linalg_tensor_ops_before.mlir b/iree/compiler/Dialect/Flow/Transforms/test/convert_linalg_tensor_ops_before.mlir
index 621912c..1fe2725 100644
--- a/iree/compiler/Dialect/Flow/Transforms/test/convert_linalg_tensor_ops_before.mlir
+++ b/iree/compiler/Dialect/Flow/Transforms/test/convert_linalg_tensor_ops_before.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt -iree-flow-convert-linalg-tensor-ops-pass -canonicalize -cse -split-input-file %s | IreeFileCheck %s
+// RUN: iree-opt -iree-flow-convert-to-flow-before-dispatch-formation -canonicalize -cse -split-input-file %s | IreeFileCheck %s
func @tensor_reshape(%arg0 : tensor<?x4x?x5x?x6xf32>, %arg1 : tensor<20x?x40xf32>)
-> (tensor<?x5x?xf32>, tensor<5x4x?x4x2x4x5xf32>)
