[flang-commits] [flang] [flang][cuda] Avoid leaking managed descriptor created during code generation phase (PR #234851)
via flang-commits
flang-commits at lists.llvm.org
Thu Oct 22 22:54:39 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-flang-codegen
Author: Valentin Clement (バレンタイン クレメン) (clementval)
<details>
<summary>Changes</summary>
During the codegen phase, descriptor that are created from emboss/rebox operation and that are passed to a cuda kernel need to be created in managed memory through a cudaMallocManaged call. The corresponding cudaFree was never generated. This patch adds a free after last use when it is safe to do so.
---
Full diff: https://github.com/llvm/llvm-project/pull/234851.diff
4 Files Affected:
- (modified) flang/include/flang/Optimizer/CodeGen/CGPasses.td (+6)
- (modified) flang/include/flang/Optimizer/CodeGen/CodeGen.h (+6)
- (modified) flang/lib/Optimizer/CodeGen/CodeGen.cpp (+88-8)
- (added) flang/test/Fir/CUDA/cuda-descriptor-free.mlir (+39)
``````````diff
diff --git a/flang/include/flang/Optimizer/CodeGen/CGPasses.td b/flang/include/flang/Optimizer/CodeGen/CGPasses.td
index 6f0ae6588c2de..54460b453a9c5 100644
--- a/flang/include/flang/Optimizer/CodeGen/CGPasses.td
+++ b/flang/include/flang/Optimizer/CodeGen/CGPasses.td
@@ -46,6 +46,12 @@ def FIRToLLVMLowering : Pass<"fir-to-llvm-ir", "mlir::ModuleOp"> {
"Name of the function to call to allocate CUDA Fortran descriptors. "
"Must have the same signature as CUFAllocDescriptor. "
"Defaults to CUFAllocDescriptor.">,
+ Option<"cudaDescriptorFreeFunction", "cuda-descriptor-free-function",
+ "std::string", /*default=*/"",
+ "Name of the function to call to free CUDA Fortran descriptors "
+ "allocated by the alloc function. "
+ "Must have the same signature as CUFFreeDescriptor. "
+ "Defaults to CUFFreeDescriptor.">,
Option<"unifiedHeapAllocSuffix", "unified-heap-alloc-suffix", "std::string",
/*default=*/"", "Suffix of the allocator entry points used for "
"allocations marked with the unified heap allocation mode.">,
diff --git a/flang/include/flang/Optimizer/CodeGen/CodeGen.h b/flang/include/flang/Optimizer/CodeGen/CodeGen.h
index d034cc3f91947..c1eafc4f3cf30 100644
--- a/flang/include/flang/Optimizer/CodeGen/CodeGen.h
+++ b/flang/include/flang/Optimizer/CodeGen/CodeGen.h
@@ -71,6 +71,12 @@ struct FIRToLLVMPassOptions {
// CUFAllocDescriptor. Empty defaults to CUFAllocDescriptor.
std::string cudaDescriptorAllocFunction;
+ // Name of the function to call when freeing CUDA Fortran descriptors
+ // allocated with cudaDescriptorAllocFunction during FIR-to-LLVM lowering.
+ // Must have the same signature as CUFFreeDescriptor. Empty defaults to
+ // CUFFreeDescriptor.
+ std::string cudaDescriptorFreeFunction;
+
// Specify the calculation method for complex number division used by the
// Conversion pass of the MLIR complex dialect.
Fortran::frontend::CodeGenOptions::ComplexRangeKind ComplexRange =
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index d5097aa6f7603..d32091df86d0d 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -1872,6 +1872,73 @@ genCUFAllocDescriptor(mlir::Location loc,
.getResult();
}
+/// Free the managed descriptor \p desc after the last use of the FIR box value
+/// \p box it was allocated for. The kernel launch may be asynchronous, so the
+/// device is synchronized first. Nothing is emitted, and the descriptor is
+/// left allocated, unless every user of \p box (looking through conversions)
+/// is a call or a kernel launch in the block defining \p box, so that the end
+/// of the lifetime is known.
+static void genCUFFreeDescriptorAfterLastUse(
+ mlir::Location loc, mlir::ConversionPatternRewriter &rewriter,
+ mlir::ModuleOp mod, mlir::Value desc, mlir::Value box,
+ llvm::StringRef cudaDescriptorFreeFunction,
+ mlir::SymbolTableCollection *symbolTables) {
+ mlir::Operation *boxOp = box.getDefiningOp();
+ if (!boxOp || boxOp->getParentOfType<mlir::gpu::GPUModuleOp>())
+ return;
+ mlir::Block *block = boxOp->getBlock();
+ mlir::Operation *lastUser = nullptr;
+ llvm::SmallVector<mlir::Value, 4> worklist{box};
+ while (!worklist.empty()) {
+ mlir::Value current = worklist.pop_back_val();
+ for (mlir::Operation *user : current.getUsers()) {
+ if (user->getBlock() != block)
+ return;
+ if (auto convert = mlir::dyn_cast<fir::ConvertOp>(user))
+ worklist.push_back(convert.getResult());
+ else if (!mlir::isa<fir::CallOp, mlir::gpu::LaunchFuncOp>(user))
+ return;
+ if (!lastUser || lastUser->isBeforeInBlock(user))
+ lastUser = user;
+ }
+ }
+ if (!lastUser || lastUser->hasTrait<mlir::OpTrait::IsTerminator>())
+ return;
+
+ mlir::MLIRContext *ctx = mod.getContext();
+ mlir::LLVM::LLVMPointerType llvmPointerType =
+ mlir::LLVM::LLVMPointerType::get(ctx);
+ mlir::Type llvmInt32Type = mlir::IntegerType::get(ctx, 32);
+ auto declareIfMissing = [&](llvm::StringRef name,
+ mlir::LLVM::LLVMFunctionType fctTy) {
+ if (fir::lookupSymbolIn<mlir::LLVM::LLVMFuncOp>(mod, name, symbolTables) ||
+ fir::lookupSymbolIn<mlir::func::FuncOp>(mod, name, symbolTables))
+ return;
+ auto builder = mlir::OpBuilder::atBlockEnd(mod.getBody());
+ auto decl = mlir::LLVM::LLVMFuncOp::create(builder, loc, name, fctTy);
+ fir::registerSymbol(symbolTables, mod, decl);
+ };
+ auto syncTy = mlir::LLVM::LLVMFunctionType::get(llvmInt32Type, {});
+ auto freeTy = mlir::LLVM::LLVMFunctionType::get(
+ mlir::LLVM::LLVMVoidType::get(ctx),
+ {llvmPointerType, llvmPointerType, llvmInt32Type});
+ llvm::StringRef syncName = "cudadevicesynchronize_";
+ llvm::StringRef freeName = cudaDescriptorFreeFunction.empty()
+ ? RTNAME_STRING(CUFFreeDescriptor)
+ : cudaDescriptorFreeFunction;
+ declareIfMissing(syncName, syncTy);
+ declareIfMissing(freeName, freeTy);
+
+ mlir::OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointAfter(lastUser);
+ mlir::LLVM::CallOp::create(rewriter, loc, syncTy, syncName,
+ mlir::ValueRange{});
+ mlir::Value sourceFile = genSourceFile(loc, mod, rewriter, symbolTables);
+ mlir::Value sourceLine = genSourceLine(loc, rewriter);
+ llvm::SmallVector args = {desc, sourceFile, sourceLine};
+ mlir::LLVM::CallOp::create(rewriter, loc, freeTy, freeName, args);
+}
+
static bool isUsedByGPULaunchFunc(mlir::Value val);
/// Get the address of the type descriptor global variable that was created by
@@ -2329,7 +2396,8 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
placeInMemoryIfNotGlobalInit(mlir::ConversionPatternRewriter &rewriter,
mlir::Location loc, mlir::Type boxTy,
mlir::Value boxValue,
- bool needDeviceAllocation = false) const {
+ bool needDeviceAllocation = false,
+ mlir::Value firBox = {}) const {
if (isInGlobalOp(rewriter))
return boxValue;
mlir::Type llvmBoxTy = boxValue.getType();
@@ -2340,6 +2408,10 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
storage = genCUFAllocDescriptor(
loc, rewriter, mod, baseBoxTy, this->lowerTy(),
this->options.cudaDescriptorAllocFunction, this->symbolTables);
+ if (firBox)
+ genCUFFreeDescriptorAfterLastUse(
+ loc, rewriter, mod, storage, firBox,
+ this->options.cudaDescriptorFreeFunction, this->symbolTables);
} else {
storage = this->genAllocaAndAddrCastWithType(loc, llvmBoxTy, defaultAlign,
rewriter);
@@ -2397,8 +2469,9 @@ struct EmboxOpConversion : public EmboxCommonConversion<fir::EmboxOp> {
return mlir::failure();
}
bool needsDeviceAlloc = isUsedByGPULaunchFunc(embox);
- auto result = placeInMemoryIfNotGlobalInit(rewriter, embox.getLoc(), boxTy,
- dest, needsDeviceAlloc);
+ auto result = placeInMemoryIfNotGlobalInit(
+ rewriter, embox.getLoc(), boxTy, dest, needsDeviceAlloc,
+ embox.getResult());
rewriter.replaceOp(embox, result);
return mlir::success();
}
@@ -2486,8 +2559,8 @@ struct CreateBoxOpConversion : public EmboxCommonConversion<fir::CreateBoxOp> {
dest = insertBaseAddress(rewriter, loc, dest, adaptor.getMemref());
bool needsDeviceAlloc = isUsedByGPULaunchFunc(createBox);
- mlir::Value result = placeInMemoryIfNotGlobalInit(rewriter, loc, boxTy,
- dest, needsDeviceAlloc);
+ mlir::Value result = placeInMemoryIfNotGlobalInit(
+ rewriter, loc, boxTy, dest, needsDeviceAlloc, createBox.getResult());
rewriter.replaceOp(createBox, result);
return mlir::success();
}
@@ -2682,8 +2755,8 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
if (fir::isDerivedTypeWithLenParams(boxTy))
TODO(loc, "fir.embox codegen of derived with length parameters");
bool needsDeviceAlloc = isUsedByGPULaunchFunc(xbox);
- mlir::Value result = placeInMemoryIfNotGlobalInit(rewriter, loc, boxTy,
- dest, needsDeviceAlloc);
+ mlir::Value result = placeInMemoryIfNotGlobalInit(
+ rewriter, loc, boxTy, dest, needsDeviceAlloc, xbox.getResult());
rewriter.replaceOp(xbox, result);
return mlir::success();
}
@@ -2801,7 +2874,8 @@ struct XReboxOpConversion : public EmboxCommonConversion<fir::cg::XReboxOp> {
dest = insertBaseAddress(rewriter, loc, dest, base);
bool needsDeviceAlloc = isUsedByGPULaunchFunc(rebox);
mlir::Value result = placeInMemoryIfNotGlobalInit(
- rewriter, rebox.getLoc(), destBoxTy, dest, needsDeviceAlloc);
+ rewriter, rebox.getLoc(), destBoxTy, dest, needsDeviceAlloc,
+ rebox.getResult());
rewriter.replaceOp(rebox, result);
return mlir::success();
}
@@ -4066,6 +4140,9 @@ struct LoadOpConversion : public fir::FIROpConversion<fir::LoadOp> {
newBoxStorage = genCUFAllocDescriptor(
loc, rewriter, mod, boxTy, lowerTy(),
this->options.cudaDescriptorAllocFunction, this->symbolTables);
+ genCUFFreeDescriptorAfterLastUse(
+ loc, rewriter, mod, newBoxStorage, load.getResult(),
+ this->options.cudaDescriptorFreeFunction, this->symbolTables);
}
if (!newBoxStorage)
newBoxStorage = genAllocaAndAddrCastWithType(loc, llvmLoadTy,
@@ -5097,6 +5174,9 @@ class FIRToLLVMLowering
if (!cudaDescriptorAllocFunction.empty())
options.cudaDescriptorAllocFunction = cudaDescriptorAllocFunction;
+ if (!cudaDescriptorFreeFunction.empty())
+ options.cudaDescriptorFreeFunction = cudaDescriptorFreeFunction;
+
if (!unifiedHeapAllocSuffix.empty())
options.unifiedHeapAllocSuffix = unifiedHeapAllocSuffix;
if (!managedHeapAllocSuffix.empty())
diff --git a/flang/test/Fir/CUDA/cuda-descriptor-free.mlir b/flang/test/Fir/CUDA/cuda-descriptor-free.mlir
new file mode 100644
index 0000000000000..d287ab2719f32
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-descriptor-free.mlir
@@ -0,0 +1,39 @@
+// RUN: fir-opt --split-input-file --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu" %s | FileCheck %s
+// RUN: fir-opt --split-input-file --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu cuda-descriptor-alloc-function=custom_alloc_desc cuda-descriptor-free-function=custom_free_desc" %s | FileCheck %s --check-prefix=CUSTOM
+
+// A managed descriptor passed to a BIND(C) device dummy is synchronized and
+// freed after the call.
+
+func.func @bindc_call(%arg0: !fir.ref<!fir.array<10xf32>>) {
+ %box = fir.embox %arg0 : (!fir.ref<!fir.array<10xf32>>) -> !fir.box<!fir.array<10xf32>>
+ fir.call @c_sub(%box) proc_attrs<bind_c> : (!fir.box<!fir.array<10xf32>> {cuf.data_attr = #cuf.cuda<device>}) -> ()
+ return
+}
+func.func private @c_sub(!fir.box<!fir.array<10xf32>> {cuf.data_attr = #cuf.cuda<device>}) attributes {fir.proc_attrs = #fir.proc_attrs<bind_c>}
+
+// CHECK-LABEL: llvm.func @bindc_call
+// CHECK: %[[DESC:.*]] = llvm.call @_FortranACUFAllocDescriptor
+// CHECK: llvm.call @c_sub(%[[DESC]])
+// CHECK-NEXT: llvm.call @cudadevicesynchronize_() : () -> i32
+// CHECK: llvm.call @_FortranACUFFreeDescriptor(%[[DESC]], %{{.*}}, %{{.*}}) : (!llvm.ptr, !llvm.ptr, i32) -> ()
+// CUSTOM-LABEL: llvm.func @bindc_call
+// CUSTOM: %[[DESC:.*]] = llvm.call @custom_alloc_desc
+// CUSTOM: llvm.call @cudadevicesynchronize_() : () -> i32
+// CUSTOM: llvm.call @custom_free_desc(%[[DESC]], %{{.*}}, %{{.*}})
+
+// -----
+
+// The descriptor escapes through a store, so its lifetime is unknown and it
+// is not freed.
+
+func.func @escapes(%arg0: !fir.ref<!fir.array<10xf32>>, %arg1: !fir.ref<!fir.box<!fir.array<10xf32>>>) {
+ %box = fir.embox %arg0 : (!fir.ref<!fir.array<10xf32>>) -> !fir.box<!fir.array<10xf32>>
+ fir.call @c_sub(%box) proc_attrs<bind_c> : (!fir.box<!fir.array<10xf32>> {cuf.data_attr = #cuf.cuda<device>}) -> ()
+ fir.store %box to %arg1 : !fir.ref<!fir.box<!fir.array<10xf32>>>
+ return
+}
+func.func private @c_sub(!fir.box<!fir.array<10xf32>> {cuf.data_attr = #cuf.cuda<device>}) attributes {fir.proc_attrs = #fir.proc_attrs<bind_c>}
+
+// CHECK-LABEL: llvm.func @escapes
+// CHECK: llvm.call @_FortranACUFAllocDescriptor
+// CHECK-NOT: FreeDescriptor
``````````
</details>
https://github.com/llvm/llvm-project/pull/234851
More information about the flang-commits
mailing list