[Mlir-commits] [mlir] [mlir][MemRefToLLVM] Cast dealloc pointer back to address space 0 (PR #237806)

[Cheng-Chieh Liu]

https://github.com/jay87406 created https://github.com/llvm/llvm-project/pull/237806

`DeallocOpLowering` passes the allocated pointer from the memref descriptor straight to `free`. For a memref in a non-default memory space, e.g. `memref<4xf32, 1>`, that pointer is `!llvm.ptr<1>`, because the alloc lowering casts the result of `malloc` into the memref's address space. `free` takes a plain `!llvm.ptr`, so the generated `llvm.call` fails the verifier with an operand type mismatch. Unranked memrefs such as `memref<*xf32, 1>` hit the same error.

Cast the pointer back to the default address space before the call, undoing the cast done on the alloc side, and add tests for the ranked and unranked cases.

Fixes #236924

>From e97fdd852cd4b9dcd52f37950b0edb5d3118bedd Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Sun, 1 Nov 2026 16:36:46 +0000
Subject: [PATCH] [mlir][MemRefToLLVM] Cast dealloc pointer back to address
 space 0

DeallocOpLowering passes the allocated pointer from the memref
descriptor straight to free. For a memref in a non-default memory
space, e.g. memref<4xf32, 1>, that pointer is !llvm.ptr<1>, because the
alloc lowering casts the result of malloc into the memref's address
space. free takes a plain !llvm.ptr, so the generated llvm.call fails
the verifier with an operand type mismatch. Unranked memrefs such as
memref<*xf32, 1> hit the same error.

Cast the pointer back to the default address space before the call,
undoing the cast done on the alloc side, and add tests for the ranked
and unranked cases.

Fixes #236924
---
 .../Conversion/MemRefToLLVM/MemRefToLLVM.cpp  |  6 +++++
 .../convert-dynamic-memref-ops.mlir           | 23 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 3ecaa49510d8f4..b1e533d108886c 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -559,6 +559,12 @@ class DeallocOpLowering : public ConvertOpToLLVMPattern<memref::DeallocOp> {
       allocatedPtr = MemRefDescriptor(adaptor.getMemref())
                          .allocatedPtr(rewriter, op.getLoc());
     }
+    // `free` takes a pointer in the default address space. Undo the cast that
+    // the alloc lowering inserts for memrefs in another address space.
+    Type freeArgTy = getPtrType();
+    if (allocatedPtr.getType() != freeArgTy)
+      allocatedPtr = LLVM::AddrSpaceCastOp::create(rewriter, op.getLoc(),
+                                                   freeArgTy, allocatedPtr);
     rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, freeFunc.value(),
                                               allocatedPtr);
     return success();
diff --git a/mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir b/mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir
index 49d67d8b0c23ba..c1d07036d77932 100644
--- a/mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir
@@ -53,6 +53,29 @@ func.func @unranked_dealloc(%arg0: memref<*xf32>) {
 
 // -----
 
+// CHECK-LABEL: func @dealloc_address_space
+func.func @dealloc_address_space(%arg0: memref<?xf32, 1>) {
+//      CHECK: %[[ptr:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<1>, ptr<1>, i64, array<1 x i64>, array<1 x i64>)>
+// CHECK-NEXT: %[[cast:.*]] = llvm.addrspacecast %[[ptr]] : !llvm.ptr<1> to !llvm.ptr
+// CHECK-NEXT: llvm.call @free(%[[cast]]) : (!llvm.ptr) -> ()
+  memref.dealloc %arg0 : memref<?xf32, 1>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @unranked_dealloc_address_space
+func.func @unranked_dealloc_address_space(%arg0: memref<*xf32, 1>) {
+//      CHECK: %[[memref:.*]] = llvm.extractvalue %{{.*}} : !llvm.struct<(i64, ptr)>
+//      CHECK: %[[ptr:.*]] = llvm.load %[[memref]] : !llvm.ptr -> !llvm.ptr<1>
+// CHECK-NEXT: %[[cast:.*]] = llvm.addrspacecast %[[ptr]] : !llvm.ptr<1> to !llvm.ptr
+// CHECK-NEXT: llvm.call @free(%[[cast]]) : (!llvm.ptr) -> ()
+  memref.dealloc %arg0 : memref<*xf32, 1>
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func @dynamic_alloc(
 //       CHECK:   %[[Marg:.*]]: index, %[[Narg:.*]]: index)
 func.func @dynamic_alloc(%arg0: index, %arg1: index) -> memref<?x?xf32> {