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

[llvmlistbot at llvm.org]

Author: Longsheng Mou
Date: 2026-03-27T15:12:47+08:00
New Revision: f4b524c69e854712c7a1a8e0bf0695e7f3adad1b

URL: https://github.com/llvm/llvm-project/commit/f4b524c69e854712c7a1a8e0bf0695e7f3adad1b
DIFF: https://github.com/llvm/llvm-project/commit/f4b524c69e854712c7a1a8e0bf0695e7f3adad1b.diff

LOG: [mlir][XeGPU] Validate single-block requirement in MoveFuncBodyToWarpOp (#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.

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
    mlir/test/Dialect/XeGPU/move-gpu-func-to-warp-op.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index 69c2fb7493086..51637b1402fb0 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..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
@@ -70,3 +70,25 @@ 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:  %[[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