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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: woruyu (woruyu)

<details>
<summary>Changes</summary>

  ### 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.

---
Full diff: https://github.com/llvm/llvm-project/pull/202931.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp (+19) 
- (modified) mlir/test/Dialect/XeGPU/xegpu-blocking.mlir (+16) 


``````````diff
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]>

``````````

</details>


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