[Mlir-commits] [mlir] [mlir][GPU] Fix crash in WarpExecuteOnLane0Op::verify with wrong terminator (PR #183930)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/183930

WarpExecuteOnLane0Op::verify() called getTerminator() which performed an unconditional cast<gpu::YieldOp> on the block's last operation. When the op body was written with a different terminator (e.g. affine.yield), the cast asserted immediately instead of emitting a verifier diagnostic.

Fix by using dyn_cast in verify() before calling getTerminator(), and emitting a proper error message when the terminator is not gpu.yield.

Add a regression test to invalid.mlir.

Fixes #181450

>From fa1878cdae291b264618c366d744d1baed6c1174 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Feb 2026 08:55:06 -0800
Subject: [PATCH] [mlir][GPU] Fix crash in WarpExecuteOnLane0Op::verify with
 wrong terminator

WarpExecuteOnLane0Op::verify() called getTerminator() which performed an
unconditional cast<gpu::YieldOp> on the block's last operation. When the
op body was written with a different terminator (e.g. affine.yield), the
cast asserted immediately instead of emitting a verifier diagnostic.

Fix by using dyn_cast in verify() before calling getTerminator(), and
emitting a proper error message when the terminator is not gpu.yield.

Add a regression test to invalid.mlir.

Fixes #181450
---
 mlir/lib/Dialect/GPU/IR/GPUDialect.cpp |  4 +++-
 mlir/test/Dialect/GPU/invalid.mlir     | 17 +++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index d347cda64a9f0..2680c7311924f 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2535,7 +2535,9 @@ LogicalResult WarpExecuteOnLane0Op::verify() {
   if (getArgs().size() != getWarpRegion().getNumArguments())
     return emitOpError(
         "expected same number op arguments and block arguments.");
-  gpu::YieldOp yield = getTerminator();
+  auto yield = dyn_cast<gpu::YieldOp>(getBody()->getTerminator());
+  if (!yield)
+    return emitOpError("expected body to be terminated with 'gpu.yield'");
   if (yield.getNumOperands() != getNumResults())
     return emitOpError(
         "expected same number of yield operands and return values.");
diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index 20fe50469e0e4..f8e75cec9b7cb 100644
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/mlir/test/Dialect/GPU/invalid.mlir
@@ -1086,3 +1086,20 @@ func.func @warp_mismatch_rank(%laneid: index) {
   }
   return
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/181450:
+// gpu.warp_execute_on_lane_0 with a wrong terminator used to crash with an
+// unchecked cast in getTerminator(). The verifier should now emit a proper
+// error instead.
+func.func @warp_execute_wrong_terminator() {
+  %laneid = arith.constant 0 : index
+  %c0 = arith.constant 0 : index
+  %v = vector.create_mask %c0 : vector<4xi1>
+  // expected-error @+1 {{'gpu.warp_execute_on_lane_0' op expected body to be terminated with 'gpu.yield'}}
+  %out = gpu.warp_execute_on_lane_0(%laneid)[32] -> vector<4xi1> {
+    affine.yield %v : vector<4xi1>
+  }
+  return
+}