[Mlir-commits] [mlir] [mlir][xegpu] Resolve scf.while "after" region argument layouts (PR #220094)

[Jianhui Li]

================
@@ -273,6 +273,22 @@ LogicalResult xegpu::propagateYieldOperandsToRegionResults(
       // Assign the yield operand's layout to the region op result it feeds.
       if (auto result = dyn_cast<OpResult>(successorInput))
         xegpu::setDistributeLayoutAttr(result, successorOperandLayout);
+      // A block argument has no attribute of its own, and scf.while's "after"
+      // arguments are tied to no init operand to borrow one from. Their layout
+      // is only recoverable when the terminator forwards a region argument
+      // unchanged, so that getDistributeLayoutAttr can read it from there. If a
+      // layout is required and the forwarded value is anything else, say so
+      // rather than lowering the region with an unknown layout.
+      if (auto arg = dyn_cast<BlockArgument>(successorInput)) {
+        auto loop =
+            dyn_cast<LoopLikeOpInterface>(arg.getOwner()->getParentOp());
+        bool tiedToInit = loop && loop.getTiedLoopInit(arg);
+        if (!tiedToInit && getLayoutOfValue(arg) &&
+            !isa<BlockArgument>(successorOperand->get()))
+          return terminator->emitError(
+              "region argument requires a layout, but the forwarded value is "
----------------
Jianhui-Li wrote:

Thanks! Comment updated to make it clear: unsupported region structure: the successor argument it feeds is not tied to an init operand, so its value must be passed through from predecessor region argument."


https://github.com/llvm/llvm-project/pull/220094