[Mlir-commits] [mlir] [mlir][xegpu] Resolve layout conflict on scf.for init operands (PR #215911)
Jianhui Li
llvmlistbot at llvm.org
Wed Aug 12 14:59:29 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/215911
Extends XeGPU layout conflict resolution to cover values carried across region boundaries. Previously
`ResolveLayoutConflicts::resolveVectorConsumer` skipped all region-branch operands, so a mismatch between a loop-carried
value's own layout and the layout its loop position requires was never reconciled and causing lowering issue.
assisted-by-claude
>From 08991e5887f0f73893dde9573e06f62d3dd7ab8b Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 12 Aug 2026 21:52:00 +0000
Subject: [PATCH] [mlir][xegpu] Resolve layout conflict on scf.for init
operands
ResolveLayoutConflicts::resolveVectorConsumer skipped all region-branch
operands, so a loop init value whose layout differs from its tied
iter_arg was never reconciled. This happens when a splat constant is
shared between a loop init and another consumer (e.g. a reduction acc)
that pins a different layout: propagation gives the value one layout but
pins the loop-carried position to another. XeGPUBlocking then blocks the
loop-carried value and the scf.if result differently and bridges them
with an unrealized_conversion_cast (vector<1xf32> -> vector<8xf32>),
which later crashes the lane propagate-layout run with "No consumer
layout found for vector operand."
Teach getConsumerLayoutAt to report the required layout for region-carried
operands: composite region ops (scf.for/scf.while/scf.if) expose it as the
layout_operand_N pinned by propagateRegionArgsToInits, and loop terminators
inherit it from the region iter_arg the operand feeds (not the parent
result, whose layout is a copy of the yielded value). resolveVectorConsumer
then reconciles the conflict (rematerializing the constant / inserting a
convert_layout). scf.if/scf.condition yields forward into a result rather
than an iter_arg, so they legitimately have no layout to reconcile and are
skipped.
Also drop a dead RegionBranchOpInterface early-return in
updateOpWithForwardFill (the TypeSwitch already routes those ops elsewhere)
and correct its stale doc comment.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 29 ++++++++++++++++++
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 30 +++++++------------
.../XeGPU/resolve-layout-conflicts.mlir | 26 ++++++++++++++--
3 files changed, 63 insertions(+), 22 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index eb3c3faf75572..b4d3a8866bd8a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2842,6 +2842,27 @@ xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
return nullptr;
}
+// For a yield operand, return the layout of the region iter_arg it forwards
+// into (the authoritative loop-carried layout), or nullptr if it feeds no block
+// argument (e.g. scf.if's yield, which only feeds results).
+static xegpu::DistributeLayoutAttr getLoopCarriedLayoutForYieldOperand(
+ RegionBranchTerminatorOpInterface terminator, OpOperand &operand) {
+ auto branch = dyn_cast<RegionBranchOpInterface>(terminator->getParentOp());
+ if (!branch)
+ return nullptr;
+ RegionBranchSuccessorMapping mapping;
+ branch.getSuccessorOperandInputMapping(mapping,
+ RegionBranchPoint(terminator));
+ for (const auto &[successorOperand, successorInputs] : mapping) {
+ if (successorOperand != &operand)
+ continue;
+ for (Value input : successorInputs)
+ if (auto arg = dyn_cast<BlockArgument>(input))
+ return xegpu::getDistributeLayoutAttr(arg);
+ }
+ return nullptr;
+}
+
/// Returns the layout required on `operand`: anchor ops report their declared
/// per-operand layout directly; non-anchor ops back-derive it from their result
/// layout via inferSourceLayoutFromResultForNonAnchorOp.
@@ -2852,6 +2873,14 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// ResolveLayoutConflicts compares producer-vs-declared
if (isa<xegpu::AnchorLayoutInterface>(op))
return xegpu::getDistributeLayoutAttr(operand);
+ // Composite region ops (scf.for/scf.while/scf.if) carry the required operand
+ // layout as the layout_operand_N pinned by propagateRegionArgsToInits.
+ if (isa<RegionBranchOpInterface>(op))
+ return xegpu::getDistributeLayoutAttr(operand);
+ // Region terminators (scf.yield/scf.condition) inherit the layout of the
+ // region iter_arg their operand feeds.
+ if (auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(op))
+ return getLoopCarriedLayoutForYieldOperand(terminator, operand);
// For non-anchor ops, derive the operand layout from the op's result
// layout via op-specific semantics.
xegpu::DistributeLayoutAttr resLayout;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index c4a103fad6c56..39455ccc0c758 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1641,18 +1641,17 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
return success(); // uniform non-tensor-data vector does not require
// layout
}
- // Region branch ops (e.g. scf.for) and their terminators (e.g. scf.yield)
- // forward their operands to successor region inputs / parent op results;
- // their consumer layout is resolved through that forwarding, not at this
- // use point.
- if (isa<RegionBranchOpInterface, RegionBranchTerminatorOpInterface>(
- consumerOp))
- return success();
-
+ // getConsumerLayoutAt also covers region-carried operands (loop init and
+ // yield operands), so a layout conflict there is reconciled below rather than
+ // silently trusted to region forwarding.
auto consumerLayout = xegpu::getConsumerLayoutAt(operand);
- if (!consumerLayout)
+ if (!consumerLayout) {
+ // A terminator (e.g.scf.if's yield) legitimately has no layout to reconcile
+ if (isa<RegionBranchTerminatorOpInterface>(consumerOp))
+ return success();
return consumerOp->emitError(
"No consumer layout found for vector operand.");
+ }
// If layouts are same, no conflict exists, return success.
if (consumerLayout.isEqualTo(producerLayout))
@@ -1749,22 +1748,15 @@ ResolveLayoutConflicts::resolveTensorDescConsumer(OpOperand &operand) {
using GetLayoutFnTy = function_ref<xegpu::DistributeLayoutAttr(Value)>;
-/// Update an operation with the layout of its results. If the result type is
-/// a vector type, a temporary layout attribute is added to the operation. If
-/// the result type is a tensor descriptor type, the type is updated with the
-/// layout attribute. The users of the result are also updated with the layout
-/// attribute.
+/// Update an operation with the layout of its results. For a vector result a
+/// temporary layout attribute is added to the op; for a tensor descriptor
+/// result the layout is written into its type.
///
/// If the global propagation left a result without a layout, forward-fill it
/// locally from the operand layouts.
static LogicalResult updateOpWithForwardFill(mlir::OpBuilder &builder,
mlir::Operation *op,
GetLayoutFnTy getLayoutOfValue) {
- // Region ops (like scf.for) are already handled by the
- // updateControlFlowOps.
- if (mlir::isa<mlir::RegionBranchOpInterface>(op))
- return success();
-
// Iterate over all the results.
for (OpResult result : op->getResults()) {
Type resultType = result.getType();
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 6a625ff515029..ee3679406e2f9 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -78,7 +78,7 @@ func.func @load_nd_with_conflicting_tensor_desc_in_loop(%arg0: memref<64x64xf16>
-> vector<16x16xf16>
%3 = arith.addf %acc, %2 {layout_result_0 = #inst_data_8x16} : vector<16x16xf16>
scf.yield %3, %tdesc : vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
- } {layout_result_0 = #inst_data_8x16}
+ } {layout_operand_3 = #inst_data_8x16, layout_result_0 = #inst_data_8x16}
xegpu.prefetch_nd %0 [%c0, %c0] {layout = #inst_data_16x16} : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
return
}
@@ -220,7 +220,27 @@ func.func @conflict_inside_loop() {
%1 = "some_op"() {layout_result_0 = #inst_data_16x16} : () -> vector<16x16xf16>
%2 = arith.addf %acc, %1 {layout_result_0 = #inst_data_8x16} : vector<16x16xf16>
scf.yield %2 : vector<16x16xf16>
- } {layout_result_0 = #inst_data_8x16}
+ } {layout_operand_3 = #inst_data_8x16, layout_result_0 = #inst_data_8x16}
+ return
+}
+
+// Conflict on the scf.for init operand: the init value %cst carries [16, 16],
+// but the loop-carried position is pinned to [8, 16] (layout_operand_3). The
+// splat constant is rematerialized with the loop-carried layout and the init
+// is repointed to the clone.
+// CHECK-LABEL: func.func @conflict_init_operand
+// CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16]>} dense<{{.*}}> : vector<16x16xf16>
+// CHECK-NEXT: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<{{.*}}> : vector<16x16xf16>
+// CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[CST]]) -> (vector<16x16xf16>)
+// CHECK: layout_operand_3 = #xegpu.layout<inst_data = [8, 16]>
+func.func @conflict_init_operand() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %cst = arith.constant {layout_result_0 = #inst_data_16x16} dense<0.0> : vector<16x16xf16>
+ %0 = scf.for %i = %c0 to %c4 step %c1 iter_args(%acc = %cst) -> vector<16x16xf16> {
+ scf.yield %acc : vector<16x16xf16>
+ } {layout_operand_3 = #inst_data_8x16, layout_result_0 = #inst_data_8x16}
return
}
@@ -248,7 +268,7 @@ func.func @conflict_postop() {
%1 = "some_op"() {layout_result_0 = #inst_data_16x16} : () -> vector<16x16xf16>
%2 = arith.addf %acc, %1 {layout_result_0 = #inst_data_16x16} : vector<16x16xf16>
scf.yield %2 : vector<16x16xf16>
- } {layout_result_0 = #inst_data_16x16}
+ } {layout_operand_3 = #inst_data_16x16, layout_result_0 = #inst_data_16x16}
%1 = math.exp %0 {layout_result_0 = #inst_data_8x16} : vector<16x16xf16>
return
}
More information about the Mlir-commits
mailing list