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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-gpu

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/183930.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/GPU/IR/GPUDialect.cpp (+3-1) 
- (modified) mlir/test/Dialect/GPU/invalid.mlir (+17) 


``````````diff
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
+}

``````````

</details>


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