Rework global hoisting to be policy based. (#7950)
* Adds a policy/decision layer on top of the analysis.
* Reworks the transformation to be based purely on the policy with no inline decisions.
* Tightens up the policy a bit to better handle:
* Never hoisting a standalone init_tensor
* Working properly with a chain of non-hoistable leaves (i.e. several broadcast or metadata ops which should not hoist)
diff --git a/iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.cpp b/iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.cpp
index d681ef6..9de6344 100644
--- a/iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.cpp
+++ b/iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.cpp
@@ -21,6 +21,31 @@
namespace IREE {
namespace Util {
+//===----------------------------------------------------------------------===//
+// ConstExprAnalysis
+//===----------------------------------------------------------------------===//
+
+namespace {
+OpOperand *findOperandFor(Operation *op, Value input) {
+ for (OpOperand &operand : op->getOpOperands()) {
+ if (operand.get() == input) return &operand;
+ }
+ return nullptr;
+}
+
+} // namespace
+
+bool ConstExprAnalysis::ConstValueInfo::hasNonAnalyzedConsumer() const {
+ // The analysis cannot represent zero-result operations, so detect that
+ // and return.
+ for (Operation *user : getOperation()->getUsers()) {
+ if (user->getNumResults() == 0) {
+ return true;
+ }
+ }
+ return false;
+}
+
ConstExprAnalysis::ConstExprAnalysis(Operation *rootOp) {
Explorer explorer(rootOp, TraversalAction::SHALLOW);
explorer.initialize();
@@ -116,6 +141,14 @@
}
}
}
+
+ // Go through and populate all consumer sets now that producers are known.
+ for (auto it : constInfoMap) {
+ ConstValueInfo *consumer = it.second;
+ for (ConstValueInfo *producer : consumer->producers) {
+ producer->consumers.insert(consumer);
+ }
+ }
}
ConstExprAnalysis::ConstValueInfo *ConstExprAnalysis::addInfo(
@@ -182,6 +215,136 @@
void ConstExprAnalysis::dump() const { print(llvm::errs()); }
+//===----------------------------------------------------------------------===//
+// ConstExprHoistingPolicy
+//===----------------------------------------------------------------------===//
+
+ConstExprHoistingPolicy::ConstExprHoistingPolicy(
+ const ConstExprAnalysis &analysis)
+ : analysis(analysis), decisions(analysis.allocedConstInfos.size()) {
+ for (auto &it : analysis.allocedConstInfos) {
+ decisions[it.get()] = {};
+ }
+}
+
+void ConstExprHoistingPolicy::initialize() {
+ // Bootstrap the worklist in analysis order, which is topological (def, use)
+ // order.
+ // TODO: Do a secondary sort?
+ Worklist worklist;
+ worklist.reserve(analysis.allocedConstInfos.size());
+ for (auto &it : analysis.allocedConstInfos) {
+ worklist.push_back(it.get());
+ }
+
+ // Since just initializing invariants, which are local, iteration order
+ // doesn't matter.
+ for (auto *info : worklist) {
+ Decision *decision = getDecision(info);
+ makeInvariantDecision(info, decision);
+ Outcome outcome = decision->getOutcome();
+ if (outcome != UNDECIDED) {
+ LLVM_DEBUG(dbgs() << "ConstExprHoistPolicy(INVARIANT, ");
+ if (outcome == ENABLE_HOIST) {
+ LLVM_DEBUG(dbgs() << "ENABLE_HOIST");
+ } else if (outcome == DISABLE_HOIST) {
+ LLVM_DEBUG(dbgs() << "DISABLE_HOIST");
+ }
+ LLVM_DEBUG(dbgs() << "): " << info->constValue << "\n");
+ }
+ }
+
+ // Work iteratively until converged.
+ for (int i = 0;; ++i) {
+ bool madeChange = false;
+ for (auto *info : worklist) {
+ Decision *decision = getDecision(info);
+ if (decision->getOutcome() != UNDECIDED) continue;
+ makeDecision(info, decision);
+
+ if (decision->getOutcome() != UNDECIDED) {
+ madeChange = true;
+ LLVM_DEBUG(dbgs() << "ConstExprHoistPolicy(" << i << ", ");
+ if (decision->getOutcome() == ENABLE_HOIST) {
+ LLVM_DEBUG(dbgs() << "ENABLE_HOIST");
+ } else if (decision->getOutcome() == DISABLE_HOIST) {
+ LLVM_DEBUG(dbgs() << "DISABLE_HOIST");
+ }
+ LLVM_DEBUG(dbgs() << "): " << info->constValue << "\n");
+ }
+ }
+
+ if (!madeChange) {
+ LLVM_DEBUG(dbgs() << "ConstExprHoistPolicy(" << i << ", CONVERGED)\n");
+ break;
+ }
+ }
+
+ for (auto *info : worklist) {
+ Decision *decision = getDecision(info);
+ if (decision->getOutcome() == UNDECIDED) {
+ LLVM_DEBUG(dbgs() << "ConstExprHoistPolicy: Value did not converge: "
+ << info->constValue << "\n");
+ }
+ }
+}
+
+void ConstExprHoistingPolicy::makeInvariantDecision(
+ const ConstExprAnalysis::ConstValueInfo *info, Decision *decision) {
+ // Check 1: Is it not const-expr.
+ if (!info->isConstExpr()) {
+ return decision->disableHoist();
+ }
+
+ // Check 2: Is it a root (these are already hoisted).
+ if (info->isRoot) {
+ decision->disableHoist();
+ }
+
+ // Check 3: Is the op itself a valid "leaf" that can become a global.
+ if (!isHoistableConstExprLeaf(info)) {
+ return decision->disableHoist();
+ }
+}
+
+void ConstExprHoistingPolicy::makeDecision(
+ const ConstExprAnalysis::ConstValueInfo *info, Decision *decision) {
+ // A const-expr value has a legal escape if:
+ // - Has a non analyzed consumer
+ // - It has an anlyzed consumer that:
+ // - Has been marked as DISABLE_HOIST (must feed into something that is
+ // not being hoisted).
+ // - Is consumed by a hoistable operand or no operand (signals implicit
+ // capture).
+ bool hasLegalEscape = info->hasNonAnalyzedConsumer();
+ if (!hasLegalEscape) {
+ for (auto *consumerInfo : info->consumers) {
+ Decision *consumerDecision = getDecision(consumerInfo);
+ if (consumerDecision->getOutcome() != DISABLE_HOIST) continue;
+
+ Operation *consumerOp = consumerInfo->getOperation();
+ OpOperand *consumerOperand = findOperandFor(consumerOp, info->constValue);
+ if (!consumerOperand) {
+ // Must be an implicit capture.
+ hasLegalEscape = true;
+ break;
+ } else if (isHoistableConstExprConsumingOperand(consumerOperand)) {
+ hasLegalEscape = true;
+ }
+ }
+ }
+
+ // If there is no legal escape, we can concretely disable.
+ if (!hasLegalEscape) {
+ decision->disableHoist();
+ return;
+ }
+
+ // Otherwise, we can conditionally enable hoisting (based on cost model, etc).
+ // TODO: Implement further conditions.
+ decision->enableHoist();
+}
+
} // namespace Util
} // namespace IREE
} // namespace iree_compiler
diff --git a/iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.h b/iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.h
index c59bed0..91ff33b 100644
--- a/iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.h
+++ b/iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.h
@@ -7,6 +7,8 @@
#ifndef IREE_COMPILER_DIALECT_IREE_UTIL_ANALYSIS_CONSTANT_CONST_EXPR_H_
#define IREE_COMPILER_DIALECT_IREE_UTIL_ANALYSIS_CONSTANT_CONST_EXPR_H_
+#include <vector>
+
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
@@ -101,12 +103,20 @@
// Direct producers that feed into this constant value.
SmallPtrSet<ConstValueInfo *, 8> producers;
+ // Direct consumers (constant and non-constant) of this value.
+ SmallPtrSet<ConstValueInfo *, 8> consumers;
+
// Whether this is a root.
bool isRoot = false;
// Whether this is a const-expr value.
bool isConstExpr() const { return state == CONSTANT; }
+ // If the value is consumed by an operation that was not analyzed, returns
+ // true. This can be considered a non-constexpr escape.
+ bool hasNonAnalyzedConsumer() const;
+
+ // Gets the defining operation.
Operation *getOperation() const {
Operation *ret = constValue.getDefiningOp();
assert(ret && "const-expr must have a defining op");
@@ -139,6 +149,66 @@
// Worklist of const value info structs which need more resolution.
using ConstValueWorklist = llvm::SmallVector<ConstValueInfo *>;
ConstValueWorklist worklist;
+ friend class ConstExprHoistingPolicy;
+};
+
+// Mutable base class for implementing policies that make decisions on
+// which expressions to hoist. This wraps a read-only ConstExprAnalysis,
+// overlaying it with cost and decisions about which specific expressions to
+// hoist.
+//
+// The default base class will hoist everything that is eligible.
+class ConstExprHoistingPolicy {
+ public:
+ using Worklist = llvm::SmallVector<const ConstExprAnalysis::ConstValueInfo *>;
+ enum Outcome {
+ UNDECIDED = 0,
+ ENABLE_HOIST = 1,
+ DISABLE_HOIST = 2,
+ };
+ class Decision {
+ public:
+ void disableHoist() {
+ assert(outcome == UNDECIDED &&
+ "can only disable hoisting of an undecided decision");
+ outcome = DISABLE_HOIST;
+ }
+ void enableHoist() {
+ assert(outcome == UNDECIDED &&
+ "can only disable hoisting of an undecided decision");
+ outcome = ENABLE_HOIST;
+ }
+
+ Outcome getOutcome() const { return outcome; }
+
+ private:
+ Outcome outcome = UNDECIDED;
+ };
+
+ ConstExprHoistingPolicy(const ConstExprAnalysis &analysis);
+ void initialize();
+ Decision *getDecision(const ConstExprAnalysis::ConstValueInfo *info) {
+ return &decisions[info];
+ }
+
+ private:
+ // At initialization time, makes any fixed decisions. This hook can only
+ // make decisions that do not depend on any const-exprs outside of what is
+ // passed.
+ void makeInvariantDecision(const ConstExprAnalysis::ConstValueInfo *info,
+ Decision *decision);
+ // Makes a decision that depends on producers and consumers of a value. This
+ // may be called repeatedly until convergence. The implementation should call
+ // decision.disableHoist() or decision.enableHoist() if it can reach a
+ // decision.
+ void makeDecision(const ConstExprAnalysis::ConstValueInfo *info,
+ Decision *decision);
+
+ const ConstExprAnalysis &analysis;
+
+ // Map of ConstValueInfo * to decision structs. All are allocated at
+ // initialization and then the structure is not changed.
+ llvm::DenseMap<const ConstExprAnalysis::ConstValueInfo *, Decision> decisions;
};
inline raw_ostream &operator<<(raw_ostream &os,
diff --git a/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp b/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp
index 15bbc9f..6ea7f8e 100644
--- a/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp
+++ b/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp
@@ -95,19 +95,25 @@
return {};
}
+ // Forbid if part of a parent that should be treated atomically.
+ if (op->getParentOfType<linalg::LinalgOp>()) {
+ return {};
+ }
+
return getInfoForDefaultConstExprOp(op);
}
bool isHoistableConstExprLeaf(const ConstExprAnalysis::ConstValueInfo *info) {
- if (!info->getOperation()->isRegistered()) {
- if (info->getOperation()->getName().getStringRef() ==
+ Operation *op = info->getOperation();
+ if (!op->isRegistered()) {
+ if (op->getName().getStringRef() ==
"iree_unregistered.non_leaf_const_expr") {
return false;
}
}
// Generally, we prefer to not hoist broadcasts.
- if (auto genericOp = dyn_cast<linalg::GenericOp>(info->getOperation())) {
+ if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
// Detect op that only broadcast input as fusing them makes the new
// op cheaper.
if (genericOp.getNumParallelLoops() == genericOp.getNumLoops() &&
@@ -122,6 +128,12 @@
}
}
+ // Never hoist init_tensor. These are sometimes used for pure shape metadata
+ // and must not be separated from their consumers.
+ if (isa<linalg::InitTensorOp>(op)) {
+ return false;
+ }
+
return true;
}
diff --git a/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp b/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
index 5b1ac52..21f08ca 100644
--- a/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
+++ b/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
@@ -53,6 +53,8 @@
const auto &constExprs = getAnalysis<ConstExprAnalysis>();
LLVM_DEBUG(dbgs() << constExprs);
LLVM_DEBUG(dbgs() << "\n\n");
+ ConstExprHoistingPolicy policy(constExprs);
+ policy.initialize();
// Maps original values to newly materialized values.
HoistedValueMap hoistedMap;
@@ -64,66 +66,43 @@
// yet.
getOperation().walk<WalkOrder::PreOrder>([&](Operation *iterOp) {
// We only want to look at const-expr ops (non roots) since they may
- // have interesting escapes.
+ // have interesting escapes. Early exit here for efficiency.
auto *iterInfo = constExprs.lookup(iterOp);
- if (!iterInfo || iterInfo->isRoot || !iterInfo->isConstExpr()) {
+ if (!iterInfo) {
return WalkResult::advance();
}
- // Skip if our policy prohibits this being treated as a hoistable leaf.
- if (!isHoistableConstExprLeaf(iterInfo)) return WalkResult::advance();
-
- // This op is hoistable - for each result find eligible escapes and
- // hoist them.
- LLVM_DEBUG(dbgs() << "PROCESSING CONST-EXPR OP: " << *iterOp << "\n");
for (Value constExprResult : iterOp->getResults()) {
- // Need to snapshot the uses since we modify them during iteration.
- SmallVector<OpOperand *> uses;
- for (OpOperand &use : constExprResult.getUses()) {
- uses.push_back(&use);
+ auto *resultInfo = constExprs.lookup(constExprResult);
+ assert(resultInfo && "must have const-expr info");
+
+ if (policy.getDecision(resultInfo)->getOutcome() !=
+ ConstExprHoistingPolicy::ENABLE_HOIST) {
+ continue;
}
- // Iterate over uses snapshots.
- for (OpOperand *operand : uses) {
- auto *targetInfo = constExprs.lookup(operand->getOwner());
-
- // We do not treat is as a const escape if the target is:
- // - Not a valid operand to convert to a constant.
- // - Const-expr and is an allowed leaf by policy
- // Note that we never touch an operand that is part of a const-expr.
- if (targetInfo && targetInfo->isConstExpr() &&
- isHoistableConstExprLeaf(targetInfo)) {
- LLVM_DEBUG(dbgs() << " - SKIP (CONST-EXPR): "
- << *operand->getOwner() << "\n");
- continue;
- }
- if (!isHoistableConstExprConsumingOperand(operand)) {
- LLVM_DEBUG(dbgs() << " - SKIP (INVALID OPERAND): "
- << *operand->getOwner() << "\n");
- continue;
- }
-
- // Bingo.
- LLVM_DEBUG(dbgs() << " + HOIST CONST-EXPR:\n");
- LLVM_DEBUG(dbgs() << " : Operand #" << operand->getOperandNumber()
- << " of " << *operand->getOwner() << "\n");
- LLVM_DEBUG(dbgs() << " : From " << operand->get() << "\n\n");
-
- hoistConstExpr(operand, hoistedMap, moduleSymbols, constExprs);
- }
+ hoistConstExpr(constExprResult, hoistedMap, moduleSymbols, constExprs);
}
-
return WalkResult::advance();
});
+ // Apply any remaining RAUW cleanups. We have to do these at the cleanup
+ // phase since modifying the source program can invalidate the analysis.
+ // Up to this point, we have only been cloning.
+ OpBuilder builder(&getContext());
+ for (auto it : hoistedMap) {
+ Value originalValue = it.first;
+ GlobalOp globalOp = it.second;
+ builder.setInsertionPointAfterValue(originalValue);
+ auto load = builder.create<GlobalLoadOp>(globalOp->getLoc(), globalOp);
+ originalValue.replaceAllUsesWith(load);
+ }
cleanupDeadOps(constExprs);
}
- void hoistConstExpr(OpOperand *operand, HoistedValueMap &hoistedMap,
- SymbolTable &moduleSymbols,
- const ConstExprAnalysis &constExprs) {
- Operation *targetOp = operand->getOwner();
- Value originalValue = operand->get();
+ GlobalOp hoistConstExpr(Value originalValue, HoistedValueMap &hoistedMap,
+ SymbolTable &moduleSymbols,
+ const ConstExprAnalysis &constExprs) {
GlobalOp existingGlobal = hoistedMap.lookup(originalValue);
if (!existingGlobal) {
@@ -142,12 +121,7 @@
assert(existingGlobal &&
"hoisting const-expr should have mapped a global for the requested "
"value");
-
- // Already hoisted - just convert to a load.
- OpBuilder builder(targetOp);
- auto load =
- builder.create<GlobalLoadOp>(targetOp->getLoc(), existingGlobal);
- operand->set(load);
+ return existingGlobal;
}
void cloneProducerTreeInto(
diff --git a/iree/compiler/Dialect/Util/Transforms/test/hoist_into_globals.mlir b/iree/compiler/Dialect/Util/Transforms/test/hoist_into_globals.mlir
index df59eb2..2c1d4d9 100644
--- a/iree/compiler/Dialect/Util/Transforms/test/hoist_into_globals.mlir
+++ b/iree/compiler/Dialect/Util/Transforms/test/hoist_into_globals.mlir
@@ -64,10 +64,9 @@
// CHECK: func @main
builtin.func @main() -> (i32, i32, i32) {
- // CHECK: %[[LOAD_HOISTED_1:.*]] = util.global.load @[[HOISTED_1]] : i32
- // CHECK: %[[RESULT:.*]] = "iree_unregistered.var_expr"(%[[LOAD_HOISTED_1]])
// CHECK-DAG: %[[LOAD_HOISTED_0:.*]] = util.global.load @[[HOISTED_0]] : i32
// CHECK-DAG: %[[LOAD_HOISTED_1:.*]] = util.global.load @[[HOISTED_1]] : i32
+ // CHECK-DAG: %[[RESULT:.*]] = "iree_unregistered.var_expr"(%[[LOAD_HOISTED_1]])
// CHECK: return %[[LOAD_HOISTED_0]], %[[LOAD_HOISTED_1]], %[[RESULT]]
%0 = arith.constant 0 : i32
%1 = arith.constant 1 : i32
