[Mlir-commits] [mlir] [mlir][async] Fix crash when lowering async values of non-LLVM types (PR #216487)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hamza Qureshi (hamzaqureshi5)

<details>
<summary>Changes</summary>

Fixes #<!-- -->216213

## The problem

`mlir-opt -convert-async-to-llvm` segfaults on valid input:

```mlir
func.func @<!-- -->test() -> !async.value<tensor<4xf32>> {
  %0 = async.runtime.create : !async.value<tensor<4xf32>>
  return %0 : !async.value<tensor<4xf32>>
}
```

To create an async value the lowering needs to know how much storage the payload takes, which it computes by emitting a GEP over the converted payload type:

```cpp
auto storedType = converter->convertType(valueType.getValueType());
...
LLVM::GEPOp::create(rewriter, loc, storagePtrType, storedType, ...);
```

`tensor<4xf32>` has no LLVM counterpart, so `convertType` returns a null `Type`. That null then reaches `TypeAttr::get`, which calls `Type::getContext()` on it, and the pass crashes.

## The fix

Check the converted type and bail out if the payload cannot be represented:

```cpp
Type storedType = converter->convertType(value.getValueType());
if (!storedType)
  return rewriter.notifyMatchFailure(
      op, "failed to convert async value type to LLVM type");
```

The `RuntimeStoreOp` and `RuntimeLoadOp` patterns in this same file already perform this check; `RuntimeCreateOpLowering` was missing it.

With the fix, the op is left unconverted and the conversion framework reports it normally:

```
error: failed to legalize operation 'async.runtime.create' that was explicitly marked illegal
```

## Testing

Added `mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm-invalid.mlir` covering two non-convertible payload types.

Verified with an assertions-enabled build:
- the reproducer segfaults before the fix and reports a clean error after
- the new test fails without the fix and passes with it
- async values with convertible payloads (`!async.value<f32>`) still lower to `mlirAsyncRuntimeCreateValue`
- `mlir/test/Conversion/AsyncToLLVM` and `mlir/test/Dialect/Async` are 18/18, no regressions

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


2 Files Affected:

- (modified) mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp (+9-3) 
- (added) mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm-invalid.mlir (+19) 


``````````diff
diff --git a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
index 46e53e71d35f5..a143876c1f357 100644
--- a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
+++ b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
@@ -592,12 +592,18 @@ class RuntimeCreateOpLowering : public ConvertOpToLLVMPattern<RuntimeCreateOp> {
 
     // To create a value we need to compute the storage requirement.
     if (auto value = dyn_cast<ValueType>(resultType)) {
+      // Computing the storage requirement emits a GEP over the stored type, so
+      // the payload type must have an LLVM counterpart.
+      Type storedType = converter->convertType(value.getValueType());
+      if (!storedType)
+        return rewriter.notifyMatchFailure(
+            op, "failed to convert async value type to LLVM type");
+
       // Returns the size requirements for the async value storage.
-      auto sizeOf = [&](ValueType valueType) -> Value {
+      auto sizeOf = [&]() -> Value {
         auto loc = op->getLoc();
         auto i64 = rewriter.getI64Type();
 
-        auto storedType = converter->convertType(valueType.getValueType());
         auto storagePtrType =
             AsyncAPI::opaquePointerType(rewriter.getContext());
 
@@ -611,7 +617,7 @@ class RuntimeCreateOpLowering : public ConvertOpToLLVMPattern<RuntimeCreateOp> {
       };
 
       rewriter.replaceOpWithNewOp<func::CallOp>(op, kCreateValue, resultType,
-                                                sizeOf(value));
+                                                sizeOf());
 
       return success();
     }
diff --git a/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm-invalid.mlir b/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm-invalid.mlir
new file mode 100644
index 0000000000000..f04ec00268831
--- /dev/null
+++ b/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm-invalid.mlir
@@ -0,0 +1,19 @@
+// RUN: mlir-opt %s -convert-async-to-llvm -split-input-file -verify-diagnostics
+
+// Computing the storage size of an async value emits a GEP over the converted
+// payload type. Payload types without an LLVM counterpart must be rejected
+// instead of building the GEP with a null type.
+
+func.func @create_value_unsupported_payload() {
+  // expected-error @below {{failed to legalize operation 'async.runtime.create'}}
+  %0 = async.runtime.create : !async.value<tensor<4xf32>>
+  return
+}
+
+// -----
+
+func.func @create_value_unsupported_payload_index() {
+  // expected-error @below {{failed to legalize operation 'async.runtime.create'}}
+  %0 = async.runtime.create : !async.value<tensor<?xindex>>
+  return
+}

``````````

</details>


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