[Mlir-commits] [mlir] [mlir][xegpu] Explain why a region op's operand layout is not derived locally (PR #219620)

[Jianhui Li]

https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/219620

getConsumerLayoutAt reads a region op's forwarded-operand layout from the layout_operand_N attribute that propagateRegionArgsToInits writes. A TODO suggested deriving it from the region argument instead, so that this function would be the only source of an operand's required layout. Replace the TODO comment since it is not practical: the pass mutates the very uses such a derivation would read, so the answer would depend on walk progress.

>From 3e7816e2204d24214bae58dce36991a820b13e13 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 29 Aug 2026 03:54:52 +0000
Subject: [PATCH] [mlir][xegpu] Explain why a region op's operand layout is not
 derived locally

getConsumerLayoutAt reads a region op's forwarded-operand layout from the
layout_operand_N attribute that propagateRegionArgsToInits writes. A TODO
suggested deriving it from the region argument instead, so that this function
would be the only source of an operand's required layout. Two ways of doing that
were tried and both are wrong:

- Deriving from the region argument's use points fails on scf.while, whose init
  operand feeds a "before" argument used only by scf.condition, which cannot
  answer either because its own target is an untied "after" argument. It also
  only appeared to work for scf.for: conflict resolution walks post-order, so by
  the time the loop's init operand is visited the body's operands have already
  been rewritten to freshly inserted convert_layout ops, and the "derived"
  layout is really that convert's input_layout read back. Visiting the operand
  first yields the body's demand instead of the loop-carried layout.

- Deriving from the tied loop result fails on scf.while as well, since
  getTiedLoopResult does not answer for its init operands.

Replace the TODO with the reason: the pass mutates the very uses such a
derivation would read, so the answer would depend on walk progress.

Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
 mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index ac4aea9d60d54..c1b71f65d4054 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2912,9 +2912,10 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
     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.
+  // propagateRegionArgsToInits back-propagated from the region argument. Do not
+  // re-derive it from that argument here: conflict resolution inserts
+  // convert_layout ops as it walks, rewriting the argument's uses, so what
+  // those uses require depends on how far the walk has progressed.
   if (isa<RegionBranchOpInterface>(op))
     return xegpu::getDistributeLayoutAttr(operand);
   // A region terminator requires the layout of the successor input its operand