Skip to content

Conversation

@Renaud-K
Copy link
Contributor

Using memref.dealloc in the gpu module would add a function definition for @free in the the top level module instead of the gpu module. The fix is to do what is already done for memref.alloc which is to use op->getParentWithTrait<OpTrait::SymbolTable>() instead of op->getParentOfType<ModuleOp>() to create the call in the proper module.

@llvmbot
Copy link
Member

llvmbot commented Sep 29, 2025

@llvm/pr-subscribers-mlir-gpu

Author: Renaud Kauffmann (Renaud-K)

Changes

Using memref.dealloc in the gpu module would add a function definition for @<!-- -->free in the the top level module instead of the gpu module. The fix is to do what is already done for memref.alloc which is to use op-&gt;getParentWithTrait&lt;OpTrait::SymbolTable&gt;() instead of op-&gt;getParentOfType&lt;ModuleOp&gt;() to create the call in the proper module.


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

2 Files Affected:

  • (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (+4-4)
  • (added) mlir/test/Dialect/GPU/memref-to-llvm.mlir (+33)
diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 262e0e7a30c63..cc6314cbd1ffe 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -48,8 +48,8 @@ static bool isStaticStrideOrOffset(int64_t strideOrOffset) {
 }
 
 static FailureOr<LLVM::LLVMFuncOp>
-getFreeFn(OpBuilder &b, const LLVMTypeConverter *typeConverter, ModuleOp module,
-          SymbolTableCollection *symbolTables) {
+getFreeFn(OpBuilder &b, const LLVMTypeConverter *typeConverter,
+          Operation *module, SymbolTableCollection *symbolTables) {
   bool useGenericFn = typeConverter->getOptions().useGenericFunctions;
 
   if (useGenericFn)
@@ -483,8 +483,8 @@ class DeallocOpLowering : public ConvertOpToLLVMPattern<memref::DeallocOp> {
                   ConversionPatternRewriter &rewriter) const override {
     // Insert the `free` declaration if it is not already present.
     FailureOr<LLVM::LLVMFuncOp> freeFunc =
-        getFreeFn(rewriter, getTypeConverter(), op->getParentOfType<ModuleOp>(),
-                  symbolTables);
+        getFreeFn(rewriter, getTypeConverter(),
+                  op->getParentWithTrait<OpTrait::SymbolTable>(), symbolTables);
     if (failed(freeFunc))
       return failure();
     Value allocatedPtr;
diff --git a/mlir/test/Dialect/GPU/memref-to-llvm.mlir b/mlir/test/Dialect/GPU/memref-to-llvm.mlir
new file mode 100644
index 0000000000000..81a96bf29e84f
--- /dev/null
+++ b/mlir/test/Dialect/GPU/memref-to-llvm.mlir
@@ -0,0 +1,33 @@
+// RUN: mlir-opt --convert-to-llvm %s | FileCheck %s
+
+// Checking that malloc and free are declared in the proper module.
+
+// CHECK: module attributes {gpu.container_module} {
+// CHECK:   llvm.func @free(!llvm.ptr)
+// CHECK:   llvm.func @malloc(i64) -> !llvm.ptr
+// CHECK:   gpu.module @kernels {
+// CHECK:     llvm.func @free(!llvm.ptr)
+// CHECK:     llvm.func @malloc(i64) -> !llvm.ptr
+// CHECK:     gpu.func @kernel_1
+// CHECK:       llvm.call @malloc({{.*}}) : (i64) -> !llvm.ptr
+// CHECK:       llvm.call @free({{.*}}) : (!llvm.ptr) -> ()
+// CHECK:       gpu.return
+// CHECK:     }
+// CHECK:   }
+// CHECK: }
+module attributes {gpu.container_module} {
+
+  gpu.module @kernels {
+    gpu.func @kernel_1() kernel {
+      %memref_a = memref.alloc() : memref<8x16xf32>
+      memref.dealloc %memref_a : memref<8x16xf32>
+      gpu.return
+    }
+  }
+
+  func.func @main() {
+    %memref_a = memref.alloc() : memref<8x16xf32>
+    memref.dealloc %memref_a : memref<8x16xf32>
+    return
+  }
+}

@llvmbot
Copy link
Member

llvmbot commented Sep 29, 2025

@llvm/pr-subscribers-mlir

Author: Renaud Kauffmann (Renaud-K)

Changes

Using memref.dealloc in the gpu module would add a function definition for @<!-- -->free in the the top level module instead of the gpu module. The fix is to do what is already done for memref.alloc which is to use op-&gt;getParentWithTrait&lt;OpTrait::SymbolTable&gt;() instead of op-&gt;getParentOfType&lt;ModuleOp&gt;() to create the call in the proper module.


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

2 Files Affected:

  • (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (+4-4)
  • (added) mlir/test/Dialect/GPU/memref-to-llvm.mlir (+33)
diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 262e0e7a30c63..cc6314cbd1ffe 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -48,8 +48,8 @@ static bool isStaticStrideOrOffset(int64_t strideOrOffset) {
 }
 
 static FailureOr<LLVM::LLVMFuncOp>
-getFreeFn(OpBuilder &b, const LLVMTypeConverter *typeConverter, ModuleOp module,
-          SymbolTableCollection *symbolTables) {
+getFreeFn(OpBuilder &b, const LLVMTypeConverter *typeConverter,
+          Operation *module, SymbolTableCollection *symbolTables) {
   bool useGenericFn = typeConverter->getOptions().useGenericFunctions;
 
   if (useGenericFn)
@@ -483,8 +483,8 @@ class DeallocOpLowering : public ConvertOpToLLVMPattern<memref::DeallocOp> {
                   ConversionPatternRewriter &rewriter) const override {
     // Insert the `free` declaration if it is not already present.
     FailureOr<LLVM::LLVMFuncOp> freeFunc =
-        getFreeFn(rewriter, getTypeConverter(), op->getParentOfType<ModuleOp>(),
-                  symbolTables);
+        getFreeFn(rewriter, getTypeConverter(),
+                  op->getParentWithTrait<OpTrait::SymbolTable>(), symbolTables);
     if (failed(freeFunc))
       return failure();
     Value allocatedPtr;
diff --git a/mlir/test/Dialect/GPU/memref-to-llvm.mlir b/mlir/test/Dialect/GPU/memref-to-llvm.mlir
new file mode 100644
index 0000000000000..81a96bf29e84f
--- /dev/null
+++ b/mlir/test/Dialect/GPU/memref-to-llvm.mlir
@@ -0,0 +1,33 @@
+// RUN: mlir-opt --convert-to-llvm %s | FileCheck %s
+
+// Checking that malloc and free are declared in the proper module.
+
+// CHECK: module attributes {gpu.container_module} {
+// CHECK:   llvm.func @free(!llvm.ptr)
+// CHECK:   llvm.func @malloc(i64) -> !llvm.ptr
+// CHECK:   gpu.module @kernels {
+// CHECK:     llvm.func @free(!llvm.ptr)
+// CHECK:     llvm.func @malloc(i64) -> !llvm.ptr
+// CHECK:     gpu.func @kernel_1
+// CHECK:       llvm.call @malloc({{.*}}) : (i64) -> !llvm.ptr
+// CHECK:       llvm.call @free({{.*}}) : (!llvm.ptr) -> ()
+// CHECK:       gpu.return
+// CHECK:     }
+// CHECK:   }
+// CHECK: }
+module attributes {gpu.container_module} {
+
+  gpu.module @kernels {
+    gpu.func @kernel_1() kernel {
+      %memref_a = memref.alloc() : memref<8x16xf32>
+      memref.dealloc %memref_a : memref<8x16xf32>
+      gpu.return
+    }
+  }
+
+  func.func @main() {
+    %memref_a = memref.alloc() : memref<8x16xf32>
+    memref.dealloc %memref_a : memref<8x16xf32>
+    return
+  }
+}

Copy link
Contributor

@clementval clementval left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@Renaud-K Renaud-K merged commit 5d739cf into llvm:main Sep 29, 2025
12 checks passed
mahesh-attarde pushed a commit to mahesh-attarde/llvm-project that referenced this pull request Oct 3, 2025
Using `memref.dealloc` in the gpu module would add a function definition
for `@free` in the the top level module instead of the gpu module. The
fix is to do what is already done for memref.alloc which is to use
`op->getParentWithTrait<OpTrait::SymbolTable>()` instead of
`op->getParentOfType<ModuleOp>()` to create the call in the proper
module.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants