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

[Longsheng Mou]

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

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

>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] [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:        }