Fix compile error when hoisting constants. (#24672)
The constant-expression hoisting pass (`HoistIntoGlobalsPass`) could
hoist a const-expr defined *inside* a `flow.dispatch.region`, placing
the resulting `util.global.load` within the region. When the region is
later outlined into an isolated executable module the global is no
longer visible, producing an "undefined global" verification error.
Skip any const-expr whose defining op is strictly nested inside a
`flow.dispatch.region` / `flow.dispatch.workgroups`. Hoisting the
dispatch op itself remains valid and is unaffected.
Fixes issue https://github.com/iree-org/iree/issues/24671
---------
Signed-off-by: Paul Stark <paul.stark@cdprojektred.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp b/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
index d4360c8..12e7afc 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
@@ -37,6 +37,16 @@
// Maps an original value in the program to the symbol name of a global.
using HoistedValueMap = llvm::DenseMap<Value, GlobalOp>;
+// Returns true if |op| is a flow dispatch region/workgroups op, i.e. a region
+// that is later outlined into an isolated executable module.
+// The dialect/op names are checked by string since Flow depends on Util
+// and we don't want to create a circular dependency.
+static bool isDispatchContainer(Operation *op) {
+ StringRef opName = op->getName().getStringRef();
+ return opName == "flow.dispatch.region" ||
+ opName == "flow.dispatch.workgroups";
+}
+
static std::string getHoistedName(Type type) {
std::string str;
llvm::raw_string_ostream os(str);
@@ -119,26 +129,33 @@
auto walkRes = funcOp.walk<WalkOrder::PreOrder>([&](Operation *iterOp) {
// We only want to look at const-expr ops (non roots) since they may
- // have interesting escapes. Early exit here for efficiency.
+ // have interesting escapes. Non const-expr ops are ignored here, but
+ // we still fall through to the dispatch check below so that we stop
+ // descending into dispatch bodies even when the dispatch op itself is
+ // not a const-expr.
auto *iterInfo = constExprs.lookup(iterOp);
- if (!iterInfo) {
- return WalkResult::advance();
+ if (iterInfo) {
+ // Record operation order for deterministic sorting. Since we walk in
+ // PreOrder, producers are visited before their users.
+ opOrder[iterOp] = orderIdx++;
+ for (Value constExprResult : iterOp->getResults()) {
+ auto *resultInfo = constExprs.lookup(constExprResult);
+ assert(resultInfo && "must have const-expr info");
+ if (policy.getDecision(resultInfo)->getOutcome() !=
+ ConstExprHoistingPolicy::ENABLE_HOIST) {
+ continue;
+ }
+ if (failed(hoistConstExpr(constExprResult, hoistedMap,
+ moduleSymbols, constExprs, opOrder))) {
+ return WalkResult::interrupt();
+ }
+ }
}
- // Record operation order for deterministic sorting. Since we walk in
- // PreOrder, producers are visited before their users.
- opOrder[iterOp] = orderIdx++;
- for (Value constExprResult : iterOp->getResults()) {
- auto *resultInfo = constExprs.lookup(constExprResult);
- assert(resultInfo && "must have const-expr info");
- if (policy.getDecision(resultInfo)->getOutcome() !=
- ConstExprHoistingPolicy::ENABLE_HOIST) {
- continue;
- }
- if (failed(hoistConstExpr(constExprResult, hoistedMap, moduleSymbols,
- constExprs, opOrder))) {
- return WalkResult::interrupt();
- }
+ // FIXME: this is a temporary workaround for issue #24671
+ // Never hoist values defined inside a dispatch region
+ if (isDispatchContainer(iterOp)) {
+ return WalkResult::skip();
}
return WalkResult::advance();
});
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/hoist_into_globals.mlir b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/hoist_into_globals.mlir
index e2502b1..2fd972e 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/hoist_into_globals.mlir
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/hoist_into_globals.mlir
@@ -391,6 +391,93 @@
// -----
+// A const-expr that is nested *inside* a dispatch region must not be hoisted:
+// doing so would place a `util.global.load` inside the dispatch region, which
+// later gets outlined into an isolated executable where the global is not
+// visible (producing an "undefined global" error). Here the constants are
+// defined inside the dispatch and the region result is captured into a second
+// dispatch, so the inner const-expr (not the whole region) would otherwise be
+// the hoisting candidate.
+
+// CHECK-LABEL: @do_not_hoist_const_expr_nested_in_dispatch
+module @do_not_hoist_const_expr_nested_in_dispatch {
+ // CHECK-NOT: util.global
+ // CHECK-NOT: util.initializer
+ util.func public @main(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+ // CHECK: flow.dispatch.region
+ // The inner const-expr stays inside the dispatch; no global load is placed
+ // inside the region.
+ // CHECK-NOT: util.global.load
+ // CHECK: linalg.generic
+ // CHECK: arith.mulf
+ %0 = flow.dispatch.region -> (tensor<2x2xf32>) {
+ %cst0 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
+ %cst1 = arith.constant dense<[[5.0, 6.0], [7.0, 8.0]]> : tensor<2x2xf32>
+ %empty = tensor.empty() : tensor<2x2xf32>
+ %g = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%cst0, %cst1 : tensor<2x2xf32>, tensor<2x2xf32>) outs(%empty : tensor<2x2xf32>) {
+ ^bb0(%a: f32, %b: f32, %o: f32):
+ %m = arith.mulf %a, %b : f32
+ linalg.yield %m : f32
+ } -> tensor<2x2xf32>
+ flow.return %g : tensor<2x2xf32>
+ }
+ %1 = flow.dispatch.region -> (tensor<2x2xf32>) {
+ %empty = tensor.empty() : tensor<2x2xf32>
+ %r = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%0, %arg0 : tensor<2x2xf32>, tensor<2x2xf32>) outs(%empty : tensor<2x2xf32>) {
+ ^bb0(%a: f32, %b: f32, %o: f32):
+ %s = arith.addf %a, %b : f32
+ linalg.yield %s : f32
+ } -> tensor<2x2xf32>
+ flow.return %r : tensor<2x2xf32>
+ }
+ util.return %1 : tensor<2x2xf32>
+ }
+}
+
+// -----
+
+// A whole flow.dispatch.region that is itself a const-expr (all of its captured
+// operands are constant) is still hoisted as a unit: the entire region moves
+// into the initializer and a single util.global.load replaces it in the
+// function body. The hoisting walk performs this hoist *before* it stops
+// descending into the region body, so whole-region hoisting keeps working; and
+// because the body is then skipped, no extra util.global.load is stranded
+// inside the region for the nested const-expr (which would become an undefined
+// global once the region is outlined). This complements
+// @do_not_hoist_const_expr_nested_in_dispatch above.
+
+// CHECK-LABEL: @hoist_whole_dispatch_region_skips_body
+module @hoist_whole_dispatch_region_skips_body {
+ // CHECK: util.global private @[[SYM:.+]] : tensor<2x2xi32>
+ // CHECK: util.initializer {
+ // CHECK: %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<2x2xi32>) {
+ // CHECK: linalg.generic
+ // CHECK-NOT: util.global.load
+ // CHECK: flow.return
+ // CHECK: }
+ // CHECK: util.global.store %[[DISPATCH]], @[[SYM]] : tensor<2x2xi32>
+ // CHECK: util.return
+ // CHECK: }
+ // CHECK: util.func public @main
+ util.func public @main() -> tensor<2x2xi32> {
+ %cst = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+ %empty = tensor.empty() : tensor<2x2xi32>
+ // CHECK: %[[VAL:.+]] = util.global.load immutable @[[SYM]] : tensor<2x2xi32>
+ // CHECK: util.return %[[VAL]]
+ %0 = flow.dispatch.region -> (tensor<2x2xi32>) {
+ %g = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%cst : tensor<2x2xi32>) outs(%empty : tensor<2x2xi32>) {
+ ^bb0(%a: i32, %o: i32):
+ %m = arith.addi %a, %a : i32
+ linalg.yield %m : i32
+ } -> tensor<2x2xi32>
+ flow.return %g : tensor<2x2xi32>
+ }
+ util.return %0 : tensor<2x2xi32>
+ }
+}
+
+// -----
+
// The --iree-util-const-expr-max-size-increase-threshold flag controls the
// maximum size increase (vs sum of size of it's roots) allowed for hoisting a
// constant expression. The threshold is set to 64 bytes in this test suite.
