[Mlir-commits] [mlir] 4cbcc02 - [mlir][xegpu] Resolve layout conflict on scf.for init operands (#215911)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 28 19:43:04 PDT 2026
Author: Jianhui Li
Date: 2026-08-28T19:42:58-07:00
New Revision: 4cbcc02ed026f6f92c818656db4fca5a442870b5
URL: https://github.com/llvm/llvm-project/commit/4cbcc02ed026f6f92c818656db4fca5a442870b5
DIFF: https://github.com/llvm/llvm-project/commit/4cbcc02ed026f6f92c818656db4fca5a442870b5.diff
LOG: [mlir][xegpu] Resolve layout conflict on scf.for init operands (#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
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
Added:
Modified:
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index b11f7ecd3df06..ac4aea9d60d54 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2852,6 +2852,54 @@ xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
return nullptr;
}
+// For a loop terminator operand (scf.for's scf.yield, scf.while's
+// scf.condition), returns the layout of the region iter_arg it forwards into,
+// which is the authoritative loop-carried layout, or nullptr when that position
+// was never assigned a layout.
+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));
+ auto it = mapping.find(&operand);
+ if (it == mapping.end())
+ return nullptr;
+ xegpu::DistributeLayoutAttr iterArgLayout;
+ for (Value input : it->second) {
+ auto arg = dyn_cast<BlockArgument>(input);
+ if (!arg)
+ continue;
+ xegpu::DistributeLayoutAttr layout = xegpu::getDistributeLayoutAttr(arg);
+ assert((!iterArgLayout || !layout || iterArgLayout.isEqualTo(layout)) &&
+ "region inputs fed by one terminator operand disagree on layout");
+ if (!iterArgLayout)
+ iterArgLayout = layout;
+ }
+ return iterArgLayout;
+}
+
+// For the terminator of a region op that carries nothing back into its regions
+// (scf.if), returns the layout of the parent result the operand feeds.
+static xegpu::DistributeLayoutAttr getParentResultLayoutForYieldOperand(
+ RegionBranchTerminatorOpInterface terminator, OpOperand &operand) {
+ auto branch = dyn_cast<RegionBranchOpInterface>(terminator->getParentOp());
+ if (!branch)
+ return nullptr;
+ RegionBranchSuccessorMapping mapping;
+ branch.getSuccessorOperandInputMapping(mapping,
+ RegionBranchPoint(terminator));
+ auto it = mapping.find(&operand);
+ if (it == mapping.end())
+ return nullptr;
+ for (Value input : it->second)
+ if (auto result = dyn_cast<OpResult>(input))
+ return xegpu::getDistributeLayoutAttr(result);
+ 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.
@@ -2862,6 +2910,21 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// ResolveLayoutConflicts compares producer-vs-declared
if (isa<xegpu::AnchorLayoutInterface>(op))
return xegpu::getDistributeLayoutAttr(operand);
+ // Region ops with forwarded operands (scf.for's and scf.while's inits) carry
+ // the required operand layout as the layout_operand_N that
+ // propagateRegionArgsToInits back-propagated from the region argument.
+ // TODO: derive that layout from the region argument here instead, so this
+ // function is the only place an operand's required layout comes from.
+ if (isa<RegionBranchOpInterface>(op))
+ return xegpu::getDistributeLayoutAttr(operand);
+ // A region terminator requires the layout of the successor input its operand
+ // feeds: the region iter_arg for a loop, and the parent result for a region
+ // op with no loop-carried values (scf.if).
+ if (auto terminator = dyn_cast<RegionBranchTerminatorOpInterface>(op)) {
+ if (isa<LoopLikeOpInterface>(op->getParentOp()))
+ return getLoopCarriedLayoutForYieldOperand(terminator, operand);
+ return getParentResultLayoutForYieldOperand(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 19d02e02b02ba..895d7facad6df 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1658,18 +1658,20 @@ 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) {
+ // TODO: handle scf.while's "after" region arguments. They are tied to no
+ // init operand, so nothing records the layout they require, and the
+ // conflict on the scf.condition operand feeding them is left unresolved
+ // rather than converted.
+ 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))
@@ -1766,22 +1768,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 2b3ef1eea971a..8474ce1b3ac69 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,84 @@ 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
+}
+
+// The init operand already carries the loop-carried layout [8, 16], so it is passed
+// through unchanged. The body wants [16, 16], so the iter_arg is converted on the
+// way in and the result converted back to [8, 16] before the yield.
+// CHECK-LABEL: func.func @conflict_init_operand
+// CHECK: %[[INIT:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<16x16xf16>
+// CHECK-NEXT: scf.for {{.*}} iter_args(%[[ACC:.*]] = %[[INIT]]) -> (vector<16x16xf16>) {
+// CHECK-NEXT: %[[CVT_IN:.*]] = xegpu.convert_layout %[[ACC]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [16, 16]>}>
+// CHECK-NEXT: %[[EXP:.*]] = math.exp %[[CVT_IN]] {layout_result_0 = #xegpu.layout<inst_data = [16, 16]>} : vector<16x16xf16>
+// CHECK-NEXT: %[[CVT_OUT:.*]] = xegpu.convert_layout %[[EXP]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}>
+// CHECK-NEXT: scf.yield %[[CVT_OUT]] : vector<16x16xf16>
+// CHECK: } {layout_operand_3 = #xegpu.layout<inst_data = [8, 16]>, layout_result_0 = #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
+ %init = "some_op"() {layout_result_0 = #inst_data_8x16} : () -> vector<16x16xf16>
+ %0 = scf.for %i = %c0 to %c4 step %c1 iter_args(%acc = %init) -> vector<16x16xf16> {
+ %1 = math.exp %acc {layout_result_0 = #inst_data_16x16} : vector<16x16xf16>
+ scf.yield %1 : vector<16x16xf16>
+ } {layout_operand_3 = #inst_data_8x16, layout_result_0 = #inst_data_8x16}
+ return
+}
+
+// Nested loops carrying one value in two layouts: the outer loop carries [8, 16]
+// and the inner one [16, 16]. Both boundaries are region crossings, so the inner
+// init operand and the outer yield operand each need a convert_layout.
+// CHECK-LABEL: func.func @conflict_nested_loop_carried
+// CHECK: scf.for {{.*}} iter_args(%[[OUTER_ACC:.*]] = %{{.*}}) -> (vector<16x16xf16>) {
+// CHECK: %[[CVT_IN:.*]] = xegpu.convert_layout %[[OUTER_ACC]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [16, 16]>}>
+// CHECK: %[[INNER:.*]] = scf.for {{.*}} iter_args(%{{.*}} = %[[CVT_IN]]) -> (vector<16x16xf16>) {
+// CHECK: } {layout_operand_3 = #xegpu.layout<inst_data = [16, 16]>, layout_result_0 = #xegpu.layout<inst_data = [16, 16]>}
+// CHECK: %[[CVT_OUT:.*]] = xegpu.convert_layout %[[INNER]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}>
+// CHECK: scf.yield %[[CVT_OUT]] : vector<16x16xf16>
+// CHECK: } {layout_operand_3 = #xegpu.layout<inst_data = [8, 16]>, layout_result_0 = #xegpu.layout<inst_data = [8, 16]>}
+func.func @conflict_nested_loop_carried() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %cst = arith.constant {layout_result_0 = #inst_data_8x16} dense<0.0> : vector<16x16xf16>
+ %0 = scf.for %i = %c0 to %c4 step %c1 iter_args(%outer = %cst) -> vector<16x16xf16> {
+ %1 = scf.for %j = %c0 to %c4 step %c1 iter_args(%inner = %outer) -> vector<16x16xf16> {
+ %2 = "some_op"() {layout_result_0 = #inst_data_16x16} : () -> vector<16x16xf16>
+ %3 = arith.addf %inner, %2 {layout_result_0 = #inst_data_16x16} : vector<16x16xf16>
+ scf.yield %3 : vector<16x16xf16>
+ } {layout_operand_3 = #inst_data_16x16, layout_result_0 = #inst_data_16x16}
+ scf.yield %1 : vector<16x16xf16>
+ } {layout_operand_3 = #inst_data_8x16, layout_result_0 = #inst_data_8x16}
+ %4 = math.exp %0 {layout_result_0 = #inst_data_8x16} : vector<16x16xf16>
+ return
+}
+
+// TODO: this scf.condition conflict is not resolved yet. The "after" region
+// argument is tied to no init operand, so it carries no layout for this pass to
+// read, and the [16, 16] value is left unconverted. Recording current behavior so
+// a fix surfaces as a test change.
+// CHECK-LABEL: func.func @negative_while_condition_operand
+// CHECK: %[[V:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [16, 16]>} : () -> vector<16x16xf16>
+// CHECK-NEXT: scf.condition(%{{.*}}) %[[V]] : vector<16x16xf16>
+// CHECK-NOT: xegpu.convert_layout
+func.func @negative_while_condition_operand(%cond: i1) {
+ %cst = arith.constant {layout_result_0 = #inst_data_8x16} dense<0.0> : vector<16x16xf16>
+ %0 = scf.while (%before = %cst) : (vector<16x16xf16>) -> vector<16x16xf16> {
+ %1 = "some_op"() {layout_result_0 = #inst_data_16x16} : () -> vector<16x16xf16>
+ scf.condition(%cond) %1 : vector<16x16xf16>
+ } do {
+ ^bb0(%after: vector<16x16xf16>):
+ %2 = math.exp %after {layout_result_0 = #inst_data_8x16} : vector<16x16xf16>
+ scf.yield %2 : vector<16x16xf16>
+ } attributes {layout_operand_0 = #inst_data_8x16, layout_result_0 = #inst_data_8x16}
+ %3 = math.exp %0 {layout_result_0 = #inst_data_8x16} : vector<16x16xf16>
return
}
@@ -248,7 +325,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