[Mlir-commits] [mlir] [XeGPU] Handle scf.while result types from scf.condition (PR #202931)

[llvmlistbot at llvm.org]

https://github.com/woruyu created https://github.com/llvm/llvm-project/pull/202931

  ### Summary
This PR resolves https://github.com/llvm/llvm-project/issues/202797. XeGPU blocking updates SCF structural op result types after converting vectors through tensor types with layout encodings. The existing code used scf.yield operands as the source for parent op result types, which is not valid for scf.while.

For scf.while, parent results correspond to the operands of scf.condition in the before region. The after-region scf.yield operands are loop-carried values for the next iteration and may have a different count, including zero. Skip scf.while in the scf.yield walk and update its result types from scf.condition instead.

This avoids crashing on while ops with results but no yield operands.

>From 61da5adfa1b86e8ddb2bd226b87878468330dca9 Mon Sep 17 00:00:00 2001
From: woruyu <1214539920 at qq.com>
Date: Wed, 10 Jun 2026 19:05:02 +0800
Subject: [PATCH] [XeGPU] Handle scf.while result types from scf.condition

  XeGPU blocking updates SCF structural op result types after converting
  vectors through tensor types with layout encodings. The existing code used
  scf.yield operands as the source for parent op result types, which is not
  valid for scf.while.

  For scf.while, parent results correspond to the operands of scf.condition
  in the before region. The after-region scf.yield operands are loop-carried
  values for the next iteration and may have a different count, including
  zero. Skip scf.while in the scf.yield walk and update its result types from
  scf.condition instead.

  This avoids crashing on while ops with results but no yield operands.
---
 mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 19 +++++++++++++++++++
 mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 16 ++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 0fb0ac6e3416d..0fdd16287eec8 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -526,9 +526,28 @@ void xegpu::doSCFStructuralTypeConversionWithTensorType(
       return WalkResult::advance();
     });
 
+    // For scf.while, the parent op results correspond to scf.condition
+    // operands, not to the after-region scf.yield operands.
+    op->walk([](scf::ConditionOp conditionOp) {
+      auto whileOp = dyn_cast<scf::WhileOp>(conditionOp->getParentOp());
+      if (!whileOp)
+        return;
+
+      ValueRange args = conditionOp.getArgs();
+      for (OpResult r : whileOp->getOpResults()) {
+        unsigned idx = r.getResultNumber();
+        Type resultTy = r.getType();
+        Type conditionTy = args[idx].getType();
+        if (isa<RankedTensorType>(resultTy) && conditionTy != resultTy)
+          r.setType(conditionTy);
+      }
+    });
+
     // using yieldOp as anchor to update the result type of its ParentOp
     op->walk([](scf::YieldOp yieldOp) {
       Operation *parentOp = yieldOp->getParentOp();
+      if (isa<scf::WhileOp>(parentOp))
+        return;
       for (OpResult r : parentOp->getOpResults()) {
         unsigned idx = r.getResultNumber();
         Type resultTy = r.getType();
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index c0ea112edc818..18c0751c036e1 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -1,5 +1,21 @@
 // RUN: mlir-opt --xegpu-blocking -split-input-file %s | FileCheck %s
 
+// CHECK-LABEL: func.func @while_results_without_iter_args
+func.func @while_results_without_iter_args() -> i32 {
+  %c0_i32 = arith.constant 0 : i32
+  %c0_i64 = arith.constant 0 : i64
+  %true = arith.constant true
+  %0:2 = scf.while : () -> (i32, i64) {
+    scf.condition(%true) %c0_i32, %c0_i64 : i32, i64
+  } do {
+  ^bb0(%arg0: i32, %arg1: i64):
+    scf.yield
+  }
+  return %0#0 : i32
+}
+
+// -----
+
 #a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
 #b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [16, 1]>
 #c = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>