Cold-Air-Networks/mirror-llvm-project
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 7

[flang][cuda] Add option for allocator registration function to CUFAddConstructor (#211098)

Browse source 
This allows to register different allocators when needed.
This commit is contained in:
Valentin Clement (バレンタイン クレメン) 2026-07-21 13:41:14 -07:00 committed by GitHub
parent ded4902fe1
commit 009e841dbd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 19 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
flang/include/flang/Optimizer/Transforms/Passes.td
View file

@ -546,6 +546,10 @@ def CUFAddConstructor : Pass<"cuf-add-constructor", "mlir::ModuleOp"> {
"cuf::CUFDialect", "mlir::func::FuncDialect", "mlir::DLTIDialect"
];
let options = [
Option<"allocatorRegistrationFunction",
"allocator-registration-function", "std::string", /*default=*/"",
"Name of the function to call to register CUDA Fortran allocators. "
"Defaults to CUFRegisterAllocator.">,
Option<"cudaUnified", "cuda-unified", "bool", /*default=*/"false",
"Treat host module globals as unified memory (-gpu=mem:unified): "
"register every plain host module global that is mirrored in the "

10
flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
View file

@ -302,13 +302,17 @@ struct CUFAddConstructor
builder.setInsertionPointToStart(entryBlock);
if (needAllocatorRegistration) {
// Symbol reference to CUFRegisterAllocator.
llvm::StringRef allocatorRegistrationFunctionName =
RTNAME_STRING(CUFRegisterAllocator);
if (!allocatorRegistrationFunction.empty())
allocatorRegistrationFunctionName = allocatorRegistrationFunction;
// Symbol reference to the allocator registration function.
builder.setInsertionPointToEnd(mod.getBody());
auto registerFuncOp = mlir::LLVM::LLVMFuncOp::create(
builder, loc, RTNAME_STRING(CUFRegisterAllocator), funcTy);
builder, loc, allocatorRegistrationFunctionName, funcTy);
registerFuncOp.setVisibility(mlir::SymbolTable::Visibility::Private);
auto cufRegisterAllocatorRef = mlir::SymbolRefAttr::get(
mod.getContext(), RTNAME_STRING(CUFRegisterAllocator));
mod.getContext(), allocatorRegistrationFunctionName);
builder.setInsertionPointToStart(entryBlock);
mlir::LLVM::CallOp::create(builder, loc, funcTy, cufRegisterAllocatorRef);
}

8
flang/test/Fir/CUDA/cuda-constructor.f90
View file

@ -1,4 +1,7 @@
! RUN: bbc -fcuda -emit-fir %s -o - | fir-opt --cuf-add-constructor | FileCheck %s
! RUN: bbc -fcuda -emit-fir %s -o - | fir-opt \
! RUN: --cuf-add-constructor="allocator-registration-function=custom_register" \
! RUN: | FileCheck %s --check-prefix=CUSTOM
program main
real, allocatable, device :: ahost(:)
@ -10,3 +13,8 @@ end
! CHECK: }
! CHECK: llvm.func @_FortranACUFRegisterAllocator() attributes {sym_visibility = "private"}
! CHECK: llvm.mlir.global_ctors ctors = [@__cudaFortranConstructor], priorities = [0 : i32], data = [#llvm.zero]
! CUSTOM-LABEL: llvm.func internal @__cudaFortranConstructor() {
! CUSTOM: llvm.call @custom_register() : () -> ()
! CUSTOM: }
! CUSTOM: llvm.func @custom_register() attributes {sym_visibility = "private"}