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

[Jared Hoberock]

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

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 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

>From 133967398bb2308f3ffa361d953560654c2937ac Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Tue, 20 Oct 2026 12:49:50 -0500
Subject: [PATCH] [mlir][gpu] Fix gpu.alloc and gpu.dealloc lowering for
 address spaces

The `gpu.alloc` lowering casts the pointer returned by `mgpuMemAlloc` to
the address space given by `memRefType.getMemorySpaceAsInt()`. The memref
descriptor's pointers are in the address space the type converter gives
the memref's memory space. For a memory space attribute the converter
maps to an address space, such as `#ptr.generic_space`,
`getMemorySpaceAsInt` asserts. The cast now uses the address space the
type converter computes.

The `gpu.dealloc` lowering passes the descriptor's allocated pointer to
`mgpuMemFree`, whose parameter is in the default address space. For a
memref in another address space, the call fails to verify ("operand type
mismatch"). The pointer is now cast to the default address space first.

Assisted-by: Claude
---
 .../GPUCommon/GPUToLLVMConversion.cpp         |  7 ++++++-
 .../lower-alloc-to-gpu-runtime-calls.mlir     | 19 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

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) {