Basic shape calculation hoisting pass.
This should be enough to unblock basic dispatch region formation.
PiperOrigin-RevId: 298924505
diff --git a/iree/compiler/Dialect/Shape/Transforms/BUILD b/iree/compiler/Dialect/Shape/Transforms/BUILD
index 70e2747..a32cad3 100644
--- a/iree/compiler/Dialect/Shape/Transforms/BUILD
+++ b/iree/compiler/Dialect/Shape/Transforms/BUILD
@@ -23,6 +23,7 @@
"CleanupPlaceholders.cpp",
"ConvertHLOToShapeDialect.cpp",
"ExpandFunctionDynamicDims.cpp",
+ "HoistShapeCalculations.cpp",
"MaterializeShapeCalculations.cpp",
"Passes.cpp",
"TieDynamicShapes.cpp",
@@ -33,6 +34,7 @@
deps = [
"//iree/compiler/Dialect/Shape/IR",
"//iree/compiler/Dialect/Shape/Plugins/XLA:XlaHloShapeBuilder",
+ "@llvm-project//mlir:Analysis",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:StandardOps",
diff --git a/iree/compiler/Dialect/Shape/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Shape/Transforms/CMakeLists.txt
index e2bd871..accadbd 100644
--- a/iree/compiler/Dialect/Shape/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/Shape/Transforms/CMakeLists.txt
@@ -23,10 +23,12 @@
"CleanupPlaceholders.cpp"
"ConvertHLOToShapeDialect.cpp"
"ExpandFunctionDynamicDims.cpp"
+ "HoistShapeCalculations.cpp"
"MaterializeShapeCalculations.cpp"
"Passes.cpp"
"TieDynamicShapes.cpp"
DEPS
+ MLIRAnalysis
MLIRIR
MLIRPass
MLIRStandardOps
diff --git a/iree/compiler/Dialect/Shape/Transforms/HoistShapeCalculations.cpp b/iree/compiler/Dialect/Shape/Transforms/HoistShapeCalculations.cpp
new file mode 100644
index 0000000..2f78ee2
--- /dev/null
+++ b/iree/compiler/Dialect/Shape/Transforms/HoistShapeCalculations.cpp
@@ -0,0 +1,159 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <algorithm>
+#include <iterator>
+
+#include "iree/compiler/Dialect/Shape/IR/ShapeDialect.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeTypes.h"
+#include "mlir/Analysis/Dominance.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace {
+
+bool isSimpleShapeCalculationOp(Operation *op) {
+ // The op must have no side effects.
+ if (!op->hasNoSideEffect()) {
+ return false;
+ }
+ // The op should operate on types that are likely shape calculations.
+ // The exact predicate used here isn't too important. The main important thing
+ // is that we don't want to include ops on tensors.
+ for (Type type : op->getOperandTypes()) {
+ if (type.isa<IntegerType>() || type.isa<Shape::RankedShapeType>()) {
+ return false;
+ }
+ }
+ return true;
+}
+
+// Return an operation in `block` that defines `v`, if one exists.
+Operation *getDefiningOpInBlock(Value v, Block &block) {
+ if (OpResult opResult = v.dyn_cast<OpResult>()) {
+ if (opResult.getOwner()->getBlock() == &block) {
+ return opResult.getOwner();
+ }
+ }
+ return nullptr;
+}
+
+DenseSet<Operation *> calculateOpsToHoist(Block &block) {
+ // Strategy:
+ // Backward DFS from shapex.tie_shape shape operands (second
+ // operands), staying within the block and not incorporating ops that don't
+ // satisfy `isSimpleShapeCalculationOp`.
+
+ SmallVector<Operation *, 16> worklist;
+ // The return value, and also used as a "visited" set for our DFS below.
+ DenseSet<Operation *> opsToHoistSet;
+
+ // The worklist is initially populated with shape-producing ops defined in
+ // this block.
+ //
+ // We are only hoisting ops within the block, so block arguments and any
+ // values defined outside the block (which already dominate the entire block)
+ // don't matter.
+ for (Operation &op : block) {
+ if (auto tieShape = dyn_cast<Shape::TieShapeOp>(op)) {
+ if (Operation *op = getDefiningOpInBlock(tieShape.shape(), block)) {
+ worklist.push_back(op);
+ }
+ }
+ }
+ while (!worklist.empty()) {
+ Operation *op = worklist.pop_back_val();
+ if (!isSimpleShapeCalculationOp(op)) {
+ continue;
+ }
+ if (opsToHoistSet.insert(op).second) {
+ for (Value v : op->getOperands()) {
+ if (Operation *op = getDefiningOpInBlock(v, block)) {
+ worklist.push_back(op);
+ }
+ }
+ }
+ }
+ return opsToHoistSet;
+}
+
+void hoistOps(DenseSet<Operation *> opsToHoistSet, Block &block,
+ DominanceInfo &domInfo) {
+ auto opsToHoist = llvm::to_vector<16>(opsToHoistSet);
+ llvm::sort(opsToHoist, [&](Operation *lhs, Operation *rhs) {
+ return domInfo.properlyDominates(lhs, rhs);
+ });
+
+ for (Operation *op : opsToHoist) {
+ Operation *insertAfter = nullptr;
+ for (Value operand : op->getOperands()) {
+ if (Operation *definingOp = getDefiningOpInBlock(operand, block)) {
+ if (insertAfter == nullptr ||
+ domInfo.properlyDominates(definingOp, insertAfter)) {
+ insertAfter = definingOp;
+ }
+ }
+ }
+ if (insertAfter != nullptr) {
+ op->moveBefore(&*std::next(insertAfter->getIterator()));
+ } else {
+ op->moveBefore(&block, block.begin());
+ }
+ }
+}
+
+// Best-effort pass for hoisting shape calculations earlier in the program.
+// We currently don't provide any hard guarantees about exactly what invariants
+// are established by this pass.
+//
+// The goal of this pass is to unblock further progress on dynamic shape
+// support. One pragmatic thing we observe is that for IREE, dispatch region
+// formation requires that when there is a `shapex.tie_shape %tensor, %shape`
+// op, to even properly form the dispatch region, IREE needs `%shape` to
+// dominate `%tensor` since the dispatch region's "workload" is derived from the
+// shape.
+//
+// This pass doesn't have a cost model, so it shouldn't be considered a generic
+// "hoist stuff to make things faster" type of pass. It's strictly a
+// best-effort pass to make certain lowerings work, albeit on somewhat shaky
+// ground. Longer-term, IREE's dispatch region formation will use a more
+// sophisticated algorithm and the analysis/hoisting done here will be a
+// byproduct of the dispatch region formation legality analysis/preparation.
+class HoistShapeCalculations : public FunctionPass<HoistShapeCalculations> {
+ public:
+ void runOnFunction() override {
+ auto func = getFunction();
+ DominanceInfo domInfo(func);
+ for (Block &block : func) {
+ DenseSet<Operation *> opsToHoist = calculateOpsToHoist(block);
+ hoistOps(opsToHoist, block, domInfo);
+ }
+ }
+};
+} // namespace
+
+std::unique_ptr<OpPassBase<FuncOp>> createHoistShapeCalculationsPass() {
+ return std::make_unique<HoistShapeCalculations>(); // NOLINT
+}
+
+static PassRegistration<HoistShapeCalculations> pass(
+ "iree-shape-hoist-shape-calculations",
+ "Best-effort shape calculation hoisting.");
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Shape/Transforms/Passes.h b/iree/compiler/Dialect/Shape/Transforms/Passes.h
index 1ed4997..32fe813 100644
--- a/iree/compiler/Dialect/Shape/Transforms/Passes.h
+++ b/iree/compiler/Dialect/Shape/Transforms/Passes.h
@@ -49,6 +49,10 @@
// dialect.
std::unique_ptr<OpPassBase<FuncOp>> createConvertHLOToShapePass();
+// Best-effort hoisting of shape calculations to attempt to establish the
+// invariant that shape.tie_shape second operand dominates the first operand.
+std::unique_ptr<OpPassBase<FuncOp>> createHoistShapeCalculationsPass();
+
} // namespace iree_compiler
} // namespace mlir
diff --git a/iree/compiler/Dialect/Shape/Transforms/test/hoist_shape_calculations.mlir b/iree/compiler/Dialect/Shape/Transforms/test/hoist_shape_calculations.mlir
new file mode 100644
index 0000000..8a2cdad
--- /dev/null
+++ b/iree/compiler/Dialect/Shape/Transforms/test/hoist_shape_calculations.mlir
@@ -0,0 +1,55 @@
+// RUN: iree-opt -split-input-file -verify-diagnostics -iree-shape-hoist-shape-calculations %s | IreeFileCheck %s
+
+// CHECK-LABEL: func @f
+func @f(%arg0: tensor<?xf32>, %arg1: index) {
+ // CHECK: shapex.make_ranked_shape
+ // CHECK: addf
+ %t = addf %arg0, %arg0 : tensor<?xf32>
+ %shape = shapex.make_ranked_shape %arg1 -> !shapex.ranked_shape<[?]>
+ shapex.tie_shape %t, %shape : tensor<?xf32>, !shapex.ranked_shape<[?]>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @f
+func @f(%arg0: tensor<?xf32>) {
+ // CHECK: addf
+ // CHECK: dim
+ // CHECK: shapex.make_ranked_shape
+ %t = addf %arg0, %arg0 : tensor<?xf32>
+ %dim = dim %t, 0 : tensor<?xf32>
+ %shape = shapex.make_ranked_shape %dim -> !shapex.ranked_shape<[?]>
+ shapex.tie_shape %t, %shape : tensor<?xf32>, !shapex.ranked_shape<[?]>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @f
+func @f(%arg0: tensor<?xf32>, %arg1: index) {
+ // CHECK: addi
+ // CHECK: muli
+ // CHECK: shapex.make_ranked_shape
+ // CHECK: some_dialect.some_op
+ "some_dialect.some_op"() : () -> ()
+ %addi = addi %arg1, %arg1 : index
+ %dim = muli %addi, %addi : index
+ %shape = shapex.make_ranked_shape %dim -> !shapex.ranked_shape<[?]>
+ shapex.tie_shape %arg0, %shape : tensor<?xf32>, !shapex.ranked_shape<[?]>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @f
+func @f(%arg0: tensor<?xf32>, %arg1: index) {
+ // CHECK: some_dialect.some_op
+ // CHECK: some_dialect.side_effecting_muli
+ // CHECK: shapex.make_ranked_shape
+ "some_dialect.some_op"() : () -> ()
+ %dim = "some_dialect.side_effecting_muli"(%arg1, %arg1) : (index, index) -> index
+ %shape = shapex.make_ranked_shape %dim -> !shapex.ranked_shape<[?]>
+ shapex.tie_shape %arg0, %shape : tensor<?xf32>, !shapex.ranked_shape<[?]>
+ return
+}
