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

[Longsheng Mou]

CoTinker wrote:

> Hi, can you clarify the reason for the crash in #185366? Single block gpu function should not be a requirement for MoveFuncBodyToWarpOp. But multi block gpu function body may not be handled properly in vector distribution which this pass depends on. So issue could be in some place else.

The crash occurs in the `VectorShapeCastDistribution` pattern when calling `getTerminator()`, because `WarpExecuteOnLane0Op` ends up containing multiple blocks — which is invalid.
https://github.com/llvm/llvm-project/blob/fe990b9005260bcf4a5630b577483e954c6bb60e/mlir/lib/Dialect/GPU/Utils/DistributionUtils.cpp#L89-L93
The root cause is in `XeGPUSubgroupDistributePass`: after `MoveFuncBodyToWarpOp` wraps the function body, it can produce a malformed `WarpExecuteOnLane0Op` (one with multiple blocks). However, no verification fires between that step and the subsequent distribution, so the invalid IR passes through silently until the crash.
Adding an explicit verification between `MoveFuncBodyToWarpOp` and distribution feels heavy-handed, so instead I propose bailing out early in `MoveFuncBodyToWarpOp` if the function body contains multiple blocks — skipping the wrapping entirely in that case.
Open to other approaches if there's a cleaner way to handle this.

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