[Mlir-commits] [mlir] [mlir][XeGPU] Validate single-block requirement in MoveFuncBodyToWarpOp (PR #188471)

[Longsheng Mou]

https://github.com/CoTinker updated https://github.com/llvm/llvm-project/pull/188471

>From f80b9b49b2dbf70d07330b8c61913fd863c0f11b Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Wed, 25 Mar 2026 17:13:27 +0800
Subject: [PATCH 1/3] [mlir][XeGPU] Validate single-block requirement in
 MoveFuncBodyToWarpOp

Add validation to ensure GPU function body has a single block before applying the MoveFuncBodyToWarpOp transformation in XeGPU subgroup distribution. This check prevents incorrect handling of multi-block GPU functions.
---
 .../Transforms/XeGPUSubgroupDistribute.cpp    |  4 ++++
 .../XeGPU/move-gpu-func-to-warp-op.mlir       | 23 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index 69c2fb7493086..b16e7c360bb88 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
@@ -152,6 +152,10 @@ struct MoveFuncBodyToWarpOp : public OpRewritePattern<gpu::GPUFuncOp> {
       return rewriter.notifyMatchFailure(
           gpuFuncOp, "Subgroup distribution requires target attribute attached "
                      "to set the warp size");
+    if (!gpuFuncOp.getBody().hasOneBlock())
+        return rewriter.notifyMatchFailure(
+            gpuFuncOp, "expected gpu.func to have a single block");
+
     // If the function only contains a single void return, skip.
     if (llvm::all_of(gpuFuncOp.getBody().getOps(), [](Operation &op) {
           return isa<gpu::ReturnOp>(op) && !op.getNumOperands();
diff --git a/mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir b/mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir
index 57a26fa4d7a7a..f95353d357452 100644
--- a/mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir
+++ b/mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir
@@ -70,3 +70,26 @@ gpu.module @test {
 // Regression test for MoveFuncBodyToWarpOp on malformed generic gpu.func.
 // CHECK-LABEL: gpu.func @missing_return_terminator
 // CHECK-NEXT:    "test.unknown"() : () -> ()
+
+// -----
+
+gpu.module @test {
+  gpu.func @multiple_blocks(%cond: i1) {
+    cf.cond_br %cond, ^bb1, ^bb2
+  ^bb1:  // pred: ^bb0
+    "test.unknown"() : () -> ()
+    cf.br ^bb2
+  ^bb2:  // 2 preds: ^bb0, ^bb1
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  gpu.func @multiple_blocks(
+// CHECK-SAME:                             %[[VAL_0:.*]]: i1) {
+// CHECK:          cf.cond_br %[[VAL_0]], ^bb1, ^bb2
+// CHECK:        ^bb1:
+// CHECK:          "test.unknown"() : () -> ()
+// CHECK:          cf.br ^bb2
+// CHECK:        ^bb2:
+// CHECK:          gpu.return
+// CHECK:        }

>From daf6afe39998539dd93160be049e4b23978d6e35 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Wed, 25 Mar 2026 20:13:12 +0800
Subject: [PATCH 2/3] format

---
 mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index b16e7c360bb88..51637b1402fb0 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
@@ -153,8 +153,8 @@ struct MoveFuncBodyToWarpOp : public OpRewritePattern<gpu::GPUFuncOp> {
           gpuFuncOp, "Subgroup distribution requires target attribute attached "
                      "to set the warp size");
     if (!gpuFuncOp.getBody().hasOneBlock())
-        return rewriter.notifyMatchFailure(
-            gpuFuncOp, "expected gpu.func to have a single block");
+      return rewriter.notifyMatchFailure(
+          gpuFuncOp, "expected gpu.func to have a single block");
 
     // If the function only contains a single void return, skip.
     if (llvm::all_of(gpuFuncOp.getBody().getOps(), [](Operation &op) {

>From 8f70dba67cc78a40220cb7781f44ce340e55070f Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Fri, 27 Mar 2026 14:47:35 +0800
Subject: [PATCH 3/3] adjust the check style

---
 .../Dialect/XeGPU/move-gpu-func-to-warp-op.mlir | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir b/mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir
index f95353d357452..3c2d987039840 100644
--- a/mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir
+++ b/mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir
@@ -84,12 +84,11 @@ gpu.module @test {
   }
 }
 
-// CHECK-LABEL:  gpu.func @multiple_blocks(
-// CHECK-SAME:                             %[[VAL_0:.*]]: i1) {
-// CHECK:          cf.cond_br %[[VAL_0]], ^bb1, ^bb2
-// CHECK:        ^bb1:
-// CHECK:          "test.unknown"() : () -> ()
-// CHECK:          cf.br ^bb2
-// CHECK:        ^bb2:
-// CHECK:          gpu.return
-// CHECK:        }
+// CHECK-LABEL: gpu.func @multiple_blocks
+// CHECK-SAME:  %[[COND:.*]]: i1
+// CHECK-NEXT:  cf.cond_br %[[COND]], ^bb1, ^bb2
+// CHECK:       ^bb1:
+// CHECK-NEXT:    "test.unknown"() : () -> ()
+// CHECK-NEXT:    cf.br ^bb2
+// CHECK:       ^bb2:
+// CHECK-NEXT:    gpu.return