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

[llvmlistbot at llvm.org]

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

>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 1/2] [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]>

>From 28e49c5ec820405d5ba583e1c97b222d8100952f Mon Sep 17 00:00:00 2001
From: woruyu <1214539920 at qq.com>
Date: Thu, 11 Jun 2026 19:11:44 +0800
Subject: [PATCH 2/2] fix: review

---
 mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 37 +++++++++------------
 1 file changed, 15 insertions(+), 22 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 0fdd16287eec8..01d64505d3564 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -526,35 +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);
+    auto updateTensorResultTypes = [](ResultRange results, ValueRange values) {
+      for (auto [result, value] : llvm::zip_equal(results, values)) {
+        Type resultTy = result.getType();
+        Type valueTy = value.getType();
+        if (isa<RankedTensorType>(resultTy) && valueTy != resultTy)
+          result.setType(valueTy);
       }
+    };
+
+    // For scf.while, parent op results correspond to scf.condition operands,
+    // not to the after-region scf.yield operands.
+    op->walk([&](scf::ConditionOp conditionOp) {
+      auto whileOp = cast<scf::WhileOp>(conditionOp->getParentOp());
+      updateTensorResultTypes(whileOp->getOpResults(), conditionOp.getArgs());
     });
 
     // using yieldOp as anchor to update the result type of its ParentOp
-    op->walk([](scf::YieldOp yieldOp) {
+    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();
-        Type yieldTy = yieldOp.getResults()[idx].getType();
-        if (isa<RankedTensorType>(resultTy) && yieldTy != resultTy)
-          r.setType(yieldTy);
-      }
+      updateTensorResultTypes(parentOp->getOpResults(), yieldOp.getResults());
     });
   }