[Mlir-commits] [mlir] [mlir][gpu] Fix crash in gpu-to-llvm with unranked memref and bare-ptr calling convention (PR #185062)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

When using `--gpu-to-llvm` with `use-bare-pointers-for-kernels=true` and a `gpu.launch_func` whose kernel has an `UnrankedMemRefType` argument, `LLVMTypeConverter::promoteOperands` would hit an `llvm_unreachable` because unranked memrefs are not supported with the bare-pointer calling convention.

Fix this by checking for unranked memref kernel arguments before calling `promoteOperands` and returning a proper conversion failure with a diagnostic instead of crashing.

Fixes #<!-- -->184939

Assisted-by: Claude Code

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


2 Files Affected:

- (modified) mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp (+10) 
- (added) mlir/test/Conversion/GPUCommon/lower-launch-func-bare-ptr-invalid.mlir (+23) 


``````````diff
diff --git a/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp b/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
index 0f72bf0c0d59e..d48a0db4d9de0 100644
--- a/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
+++ b/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
@@ -976,6 +976,16 @@ LogicalResult LegalizeLaunchFuncOpPattern::matchAndRewrite(
   // Note: If `useBarePtrCallConv` is set in the type converter's options,
   // the value of `kernelBarePtrCallConv` will be ignored.
   OperandRange origArguments = launchOp.getKernelOperands();
+  bool effectiveBarePtr = kernelBarePtrCallConv ||
+                          getTypeConverter()->getOptions().useBarePtrCallConv;
+  if (effectiveBarePtr) {
+    for (Value arg : origArguments) {
+      if (isa<UnrankedMemRefType>(arg.getType()))
+        return rewriter.notifyMatchFailure(
+            loc, "unranked memref kernel argument is not supported with "
+                 "the bare-pointer calling convention");
+    }
+  }
   SmallVector<Value, 8> llvmArguments = getTypeConverter()->promoteOperands(
       loc, origArguments, adaptor.getKernelOperands(), rewriter,
       /*useBarePtrCallConv=*/kernelBarePtrCallConv);
diff --git a/mlir/test/Conversion/GPUCommon/lower-launch-func-bare-ptr-invalid.mlir b/mlir/test/Conversion/GPUCommon/lower-launch-func-bare-ptr-invalid.mlir
new file mode 100644
index 0000000000000..19dbdc94b6d92
--- /dev/null
+++ b/mlir/test/Conversion/GPUCommon/lower-launch-func-bare-ptr-invalid.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s --gpu-to-llvm="use-bare-pointers-for-kernels=1" -split-input-file -verify-diagnostics
+
+// Test that gpu.launch_func with an unranked memref kernel argument fails
+// gracefully (no crash) when using the bare-pointer calling convention.
+// See issue #184939.
+
+module attributes {gpu.container_module} {
+  gpu.module @kernels {
+    gpu.func @kernel(%arg0: memref<*xi32>) kernel {
+      gpu.return
+    }
+  }
+  func.func @main(%arg: memref<*xi32>) {
+    %c1 = arith.constant 1 : index
+    %c0 = arith.constant 0 : i32
+    // expected-error at +1 {{failed to legalize operation 'gpu.launch_func'}}
+    gpu.launch_func @kernels::@kernel
+        blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1)
+        dynamic_shared_memory_size %c0
+        args(%arg : memref<*xi32>)
+    return
+  }
+}

``````````

</details>


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