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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Cheng-Chieh Liu (jay87406)

<details>
<summary>Changes</summary>

`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

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


2 Files Affected:

- (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (+6) 
- (modified) mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir (+23) 


``````````diff
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> {

``````````

</details>


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