[Mlir-commits] [mlir] [mlir][gpu] Fix address space casts in gpu.alloc/gpu.dealloc lowering (PR #234394)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Jared Hoberock (jaredhoberock)

<details>
<summary>Changes</summary>

In `GPUToLLVMConversion`, fix mismatches in address space when lowering `gpu.alloc` and `gpu.dealloc`.

Memory spaces aren't simple integers (they can be attributes) so the `gpu.alloc` lowering must take the address space from the type converter instead of simply reading the memory space off the memref as an integer.

Additionally, the `gpu.dealloc` lowering passed the descriptor's pointer directly to `mgpuMemFree`. Now it gets cast to the default address space expected by the runtime.

Adds unit tests for `gpu.alloc` on an attribute memory space and for `gpu.dealloc` on a non-default address space.

Assisted-by: Claude

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


2 Files Affected:

- (modified) mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp (+6-1) 
- (modified) mlir/test/Conversion/GPUCommon/lower-alloc-to-gpu-runtime-calls.mlir (+19) 


``````````diff
diff --git a/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp b/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
index da807d732beece..7744d3b21e885c 100644
--- a/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
+++ b/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
@@ -797,7 +797,8 @@ LogicalResult ConvertAllocOpToGpuRuntimeCallPattern::matchAndRewrite(
 
   // Cast the runtime-returned pointer to the memref's address space if it
   // isn't there already, so it fits the descriptor's pointer slots.
-  unsigned dstAddrSpace = memRefType.getMemorySpaceAsInt();
+  unsigned dstAddrSpace =
+      *getTypeConverter()->getMemRefAddressSpace(memRefType);
   unsigned srcAddrSpace =
       cast<LLVM::LLVMPointerType>(allocatedPtr.getType()).getAddressSpace();
   if (dstAddrSpace != srcAddrSpace) {
@@ -837,6 +838,10 @@ LogicalResult ConvertDeallocOpToGpuRuntimeCallPattern::matchAndRewrite(
 
   Value pointer =
       MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc);
+  // The runtime frees a pointer in the default address space.
+  if (pointer.getType() != llvmPointerType)
+    pointer =
+        LLVM::AddrSpaceCastOp::create(rewriter, loc, llvmPointerType, pointer);
   auto nullPtr = mlir::LLVM::ZeroOp::create(rewriter, loc, llvmPointerType);
   Value stream = adaptor.getAsyncDependencies().empty()
                      ? nullPtr
diff --git a/mlir/test/Conversion/GPUCommon/lower-alloc-to-gpu-runtime-calls.mlir b/mlir/test/Conversion/GPUCommon/lower-alloc-to-gpu-runtime-calls.mlir
index 5d5c0093edf719..84072c8a3ea208 100644
--- a/mlir/test/Conversion/GPUCommon/lower-alloc-to-gpu-runtime-calls.mlir
+++ b/mlir/test/Conversion/GPUCommon/lower-alloc-to-gpu-runtime-calls.mlir
@@ -36,6 +36,25 @@ module attributes {gpu.container_module} {
     return
   }
 
+  // CHECK-LABEL: llvm.func @dealloc_address_space
+  func.func @dealloc_address_space(%arg0 : memref<10xf32, 1>) {
+    // CHECK: %[[pointer:.*]] = llvm.extractvalue {{.*}}[0]
+    // CHECK: %[[generic:.*]] = llvm.addrspacecast %[[pointer]] : !llvm.ptr<1> to !llvm.ptr
+    // CHECK: llvm.call @mgpuMemFree(%[[generic]], {{.*}})
+    gpu.dealloc %arg0 : memref<10xf32, 1>
+    return
+  }
+
+  // CHECK-LABEL: llvm.func @alloc_dealloc_generic_space
+  func.func @alloc_dealloc_generic_space() {
+    // CHECK: llvm.call @mgpuMemAlloc
+    // CHECK-NOT: llvm.addrspacecast
+    // CHECK: llvm.call @mgpuMemFree
+    %0 = gpu.alloc () : memref<10xf32, #ptr.generic_space>
+    gpu.dealloc %0 : memref<10xf32, #ptr.generic_space>
+    return
+  }
+
   // CHECK-LABEL: llvm.func @alloc_sync
   // CHECK-SAME: %[[size:.*]]: i64
   func.func @alloc_sync(%size : index) {

``````````

</details>


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