Create function declaration in the proper module (#161281)

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.

Author: Renaud-K

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;

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