[flang-commits] [flang] 9089be3 - [flang][cuda] Add support for custom descriptor allocator in codegen (#216881)
via flang-commits
flang-commits at lists.llvm.org
Mon Aug 17 20:00:29 PDT 2026
Author: Valentin Clement (バレンタイン クレメン)
Date: 2026-08-17T20:00:24-07:00
New Revision: 9089be3caf952d1410287ba88b8debfd931a4d5c
URL: https://github.com/llvm/llvm-project/commit/9089be3caf952d1410287ba88b8debfd931a4d5c
DIFF: https://github.com/llvm/llvm-project/commit/9089be3caf952d1410287ba88b8debfd931a4d5c.diff
LOG: [flang][cuda] Add support for custom descriptor allocator in codegen (#216881)
Simlar to https://github.com/llvm/llvm-project/pull/216841 but for
CodeGen pass usage.
Added:
Modified:
flang/include/flang/Optimizer/CodeGen/CGPasses.td
flang/include/flang/Optimizer/CodeGen/CodeGen.h
flang/lib/Optimizer/CodeGen/CodeGen.cpp
flang/test/Fir/CUDA/cuda-code-gen.mlir
Removed:
################################################################################
diff --git a/flang/include/flang/Optimizer/CodeGen/CGPasses.td b/flang/include/flang/Optimizer/CodeGen/CGPasses.td
index b1ad48f7b9c78..2741e0206dfec 100644
--- a/flang/include/flang/Optimizer/CodeGen/CGPasses.td
+++ b/flang/include/flang/Optimizer/CodeGen/CGPasses.td
@@ -40,7 +40,12 @@ def FIRToLLVMLowering : Pass<"fir-to-llvm-ir", "mlir::ModuleOp"> {
Option<"typeDescriptorsRenamedForAssembly",
"type-descriptors-renamed-for-assembly", "bool", /*default=*/"false",
"Global variables created to describe derived types "
- "have been renamed to avoid special symbols in their names.">
+ "have been renamed to avoid special symbols in their names.">,
+ Option<"cudaDescriptorAllocFunction", "cuda-descriptor-alloc-function",
+ "std::string", /*default=*/"",
+ "Name of the function to call to allocate CUDA Fortran descriptors. "
+ "Must have the same signature as CUFAllocDescriptor. "
+ "Defaults to CUFAllocDescriptor.">
];
}
diff --git a/flang/include/flang/Optimizer/CodeGen/CodeGen.h b/flang/include/flang/Optimizer/CodeGen/CodeGen.h
index b7a9397edfe6d..948c240967c5b 100644
--- a/flang/include/flang/Optimizer/CodeGen/CodeGen.h
+++ b/flang/include/flang/Optimizer/CodeGen/CodeGen.h
@@ -16,6 +16,7 @@
#include "llvm/IR/Module.h"
#include "llvm/Support/raw_ostream.h"
#include <memory>
+#include <string>
namespace fir {
@@ -61,6 +62,11 @@ struct FIRToLLVMPassOptions {
// type's descriptor.
bool typeDescriptorsRenamedForAssembly = false;
+ // Name of the function to call when allocating CUDA Fortran descriptors
+ // during FIR-to-LLVM lowering. Must have the same signature as
+ // CUFAllocDescriptor. Empty defaults to CUFAllocDescriptor.
+ std::string cudaDescriptorAllocFunction;
+
// 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 4b7891bd46df1..4e33dc008e53a 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -1663,7 +1663,8 @@ static mlir::Value
genCUFAllocDescriptor(mlir::Location loc,
mlir::ConversionPatternRewriter &rewriter,
mlir::ModuleOp mod, fir::BaseBoxType boxTy,
- const fir::LLVMTypeConverter &typeConverter) {
+ const fir::LLVMTypeConverter &typeConverter,
+ llvm::StringRef cudaDescriptorAllocFunction) {
std::optional<mlir::DataLayout> dl =
fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true);
if (!dl)
@@ -1684,14 +1685,15 @@ genCUFAllocDescriptor(mlir::Location loc,
auto fctTy = mlir::LLVM::LLVMFunctionType::get(
llvmPointerType, {llvmIntPtrType, llvmPointerType, llvmInt32Type});
- auto llvmFunc = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(
- RTNAME_STRING(CUFAllocDescriptor));
- auto funcFunc =
- mod.lookupSymbol<mlir::func::FuncOp>(RTNAME_STRING(CUFAllocDescriptor));
+ llvm::StringRef funcName = cudaDescriptorAllocFunction.empty()
+ ? RTNAME_STRING(CUFAllocDescriptor)
+ : cudaDescriptorAllocFunction;
+
+ auto llvmFunc = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(funcName);
+ auto funcFunc = mod.lookupSymbol<mlir::func::FuncOp>(funcName);
if (!llvmFunc && !funcFunc) {
auto builder = mlir::OpBuilder::atBlockEnd(mod.getBody());
- mlir::LLVM::LLVMFuncOp::create(builder, loc,
- RTNAME_STRING(CUFAllocDescriptor), fctTy);
+ mlir::LLVM::LLVMFuncOp::create(builder, loc, funcName, fctTy);
}
mlir::Type structTy = typeConverter.convertBoxTypeAsStruct(boxTy);
@@ -1699,14 +1701,15 @@ genCUFAllocDescriptor(mlir::Location loc,
mlir::Value sizeInBytes =
fir::genConstantIndex(loc, llvmIntPtrType, rewriter, boxSize);
llvm::SmallVector args = {sizeInBytes, sourceFile, sourceLine};
- return mlir::LLVM::CallOp::create(rewriter, loc, fctTy,
- RTNAME_STRING(CUFAllocDescriptor), args)
+ return mlir::LLVM::CallOp::create(rewriter, loc, fctTy, funcName, args)
.getResult();
}
static bool isUsedByGPULaunchFunc(mlir::Value val);
-static bool isDeviceAllocation(mlir::Value val, mlir::Value adaptorVal);
+static bool
+isDeviceAllocation(mlir::Value val, mlir::Value adaptorVal,
+ llvm::StringRef cudaDescriptorAllocFunction = {});
/// Get the address of the type descriptor global variable that was created by
/// lowering for derived type \p recType.
@@ -2163,7 +2166,8 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
auto mod = boxValue.getDefiningOp()->getParentOfType<mlir::ModuleOp>();
auto baseBoxTy = mlir::dyn_cast<fir::BaseBoxType>(boxTy);
storage =
- genCUFAllocDescriptor(loc, rewriter, mod, baseBoxTy, this->lowerTy());
+ genCUFAllocDescriptor(loc, rewriter, mod, baseBoxTy, this->lowerTy(),
+ this->options.cudaDescriptorAllocFunction);
} else {
storage = this->genAllocaAndAddrCastWithType(loc, llvmBoxTy, defaultAlign,
rewriter);
@@ -2221,7 +2225,8 @@ struct EmboxOpConversion : public EmboxCommonConversion<fir::EmboxOp> {
return mlir::failure();
}
bool needsDeviceAlloc =
- isDeviceAllocation(embox.getMemref(), adaptor.getMemref()) ||
+ isDeviceAllocation(embox.getMemref(), adaptor.getMemref(),
+ this->options.cudaDescriptorAllocFunction) ||
isUsedByGPULaunchFunc(embox);
auto result = placeInMemoryIfNotGlobalInit(rewriter, embox.getLoc(), boxTy,
dest, needsDeviceAlloc);
@@ -2266,7 +2271,8 @@ static bool isUsedByOpenACCDataClause(mlir::Value val) {
return false;
}
-static bool isDeviceAllocation(mlir::Value val, mlir::Value adaptorVal) {
+static bool isDeviceAllocation(mlir::Value val, mlir::Value adaptorVal,
+ llvm::StringRef cudaDescriptorAllocFunction) {
if (val.getDefiningOp() &&
val.getDefiningOp()->getParentOfType<mlir::gpu::GPUModuleOp>())
return false;
@@ -2280,13 +2286,16 @@ static bool isDeviceAllocation(mlir::Value val, mlir::Value adaptorVal) {
return true;
if (auto loadOp = mlir::dyn_cast_or_null<fir::LoadOp>(val.getDefiningOp()))
- return isDeviceAllocation(loadOp.getMemref(), {});
+ return isDeviceAllocation(loadOp.getMemref(), {},
+ cudaDescriptorAllocFunction);
if (auto boxAddrOp =
mlir::dyn_cast_or_null<fir::BoxAddrOp>(val.getDefiningOp()))
- return isDeviceAllocation(boxAddrOp.getVal(), {});
+ return isDeviceAllocation(boxAddrOp.getVal(), {},
+ cudaDescriptorAllocFunction);
if (auto convertOp =
mlir::dyn_cast_or_null<fir::ConvertOp>(val.getDefiningOp()))
- return isDeviceAllocation(convertOp.getValue(), {});
+ return isDeviceAllocation(convertOp.getValue(), {},
+ cudaDescriptorAllocFunction);
if (!val.getDefiningOp() && adaptorVal) {
if (auto blockArg = llvm::cast<mlir::BlockArgument>(adaptorVal)) {
if (blockArg.getOwner() && blockArg.getOwner()->getParentOp() &&
@@ -2308,14 +2317,17 @@ static bool isDeviceAllocation(mlir::Value val, mlir::Value adaptorVal) {
}
}
if (auto callOp = mlir::dyn_cast_or_null<fir::CallOp>(val.getDefiningOp()))
- if (callOp.getCallee() &&
- (callOp.getCallee().value().getRootReference().getValue().starts_with(
- RTNAME_STRING(CUFMemAlloc)) ||
- callOp.getCallee().value().getRootReference().getValue().starts_with(
- RTNAME_STRING(CUFAllocDescriptor)) ||
- callOp.getCallee().value().getRootReference().getValue() ==
- "__tgt_acc_get_deviceptr"))
- return true;
+ if (callOp.getCallee()) {
+ llvm::StringRef calleeName =
+ callOp.getCallee().value().getRootReference().getValue();
+ llvm::StringRef allocDescName = cudaDescriptorAllocFunction.empty()
+ ? RTNAME_STRING(CUFAllocDescriptor)
+ : cudaDescriptorAllocFunction;
+ if (calleeName.starts_with(RTNAME_STRING(CUFMemAlloc)) ||
+ calleeName.starts_with(allocDescName) ||
+ calleeName == "__tgt_acc_get_deviceptr")
+ return true;
+ }
return false;
}
@@ -2353,7 +2365,8 @@ struct CreateBoxOpConversion : public EmboxCommonConversion<fir::CreateBoxOp> {
dest = insertBaseAddress(rewriter, loc, dest, adaptor.getMemref());
bool needsDeviceAlloc =
- isDeviceAllocation(createBox.getMemref(), adaptor.getMemref()) ||
+ isDeviceAllocation(createBox.getMemref(), adaptor.getMemref(),
+ this->options.cudaDescriptorAllocFunction) ||
isUsedByGPULaunchFunc(createBox);
mlir::Value result = placeInMemoryIfNotGlobalInit(rewriter, loc, boxTy,
dest, needsDeviceAlloc);
@@ -2552,7 +2565,8 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
TODO(loc, "fir.embox codegen of derived with length parameters");
bool needsDeviceAlloc =
isUsedByGPULaunchFunc(xbox) ||
- (isDeviceAllocation(xbox.getMemref(), adaptor.getMemref()) &&
+ (isDeviceAllocation(xbox.getMemref(), adaptor.getMemref(),
+ this->options.cudaDescriptorAllocFunction) &&
!isUsedByOpenACCDataClause(xbox));
mlir::Value result = placeInMemoryIfNotGlobalInit(rewriter, loc, boxTy,
dest, needsDeviceAlloc);
@@ -2673,7 +2687,8 @@ struct XReboxOpConversion : public EmboxCommonConversion<fir::cg::XReboxOp> {
dest = insertBaseAddress(rewriter, loc, dest, base);
bool needsDeviceAlloc =
isUsedByGPULaunchFunc(rebox) ||
- (isDeviceAllocation(rebox.getBox(), adaptor.getBox()) &&
+ (isDeviceAllocation(rebox.getBox(), adaptor.getBox(),
+ this->options.cudaDescriptorAllocFunction) &&
!isUsedByOpenACCDataClause(rebox));
mlir::Value result = placeInMemoryIfNotGlobalInit(
rewriter, rebox.getLoc(), destBoxTy, dest, needsDeviceAlloc);
@@ -3851,13 +3866,15 @@ struct LoadOpConversion : public fir::FIROpConversion<fir::LoadOp> {
// source is handled below if the load is used by a GPU launch.
auto mod = load->getParentOfType<mlir::ModuleOp>();
newBoxStorage =
- genCUFAllocDescriptor(loc, rewriter, mod, boxTy, lowerTy());
+ genCUFAllocDescriptor(loc, rewriter, mod, boxTy, lowerTy(),
+ this->options.cudaDescriptorAllocFunction);
}
}
if (!newBoxStorage && isUsedByGPULaunchFunc(load)) {
auto mod = load->getParentOfType<mlir::ModuleOp>();
newBoxStorage =
- genCUFAllocDescriptor(loc, rewriter, mod, boxTy, lowerTy());
+ genCUFAllocDescriptor(loc, rewriter, mod, boxTy, lowerTy(),
+ this->options.cudaDescriptorAllocFunction);
}
if (!newBoxStorage)
newBoxStorage = genAllocaAndAddrCastWithType(loc, llvmLoadTy,
@@ -4727,6 +4744,9 @@ class FIRToLLVMLowering
options.typeDescriptorsRenamedForAssembly =
typeDescriptorsRenamedForAssembly;
+ if (!cudaDescriptorAllocFunction.empty())
+ options.cudaDescriptorAllocFunction = cudaDescriptorAllocFunction;
+
// Run dynamic pass pipeline for converting Math dialect
// operations into other dialects (llvm, func, etc.).
// Some conversions of Math operations cannot be done
diff --git a/flang/test/Fir/CUDA/cuda-code-gen.mlir b/flang/test/Fir/CUDA/cuda-code-gen.mlir
index 43623b4a00c5d..1ac72f68eb294 100644
--- a/flang/test/Fir/CUDA/cuda-code-gen.mlir
+++ b/flang/test/Fir/CUDA/cuda-code-gen.mlir
@@ -1,4 +1,5 @@
// 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" %s | FileCheck %s --check-prefix=CUSTOM
module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<f80, dense<128> : vector<2xi64>>, #dlti.dl_entry<i128, dense<128> : vector<2xi64>>, #dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<270>, dense<32> : vector<4xi64>>, #dlti.dl_entry<f128, dense<128> : vector<2xi64>>, #dlti.dl_entry<f64, dense<64> : vector<2xi64>>, #dlti.dl_entry<f16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<i16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>} {
func.func @_QQmain() attributes {fir.bindc_name = "cufkernel_global"} {
@@ -187,6 +188,8 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<!llvm.ptr<270> = dense<32> : vec
// CHECK-LABEL: llvm.func @_QPouter
// CHECK: _FortranACUFAllocDescriptor
+// CUSTOM-LABEL: llvm.func @_QPouter
+// CUSTOM: custom_alloc_desc
// -----
@@ -198,6 +201,8 @@ func.func @_QMm1Psub1(%arg0: !fir.box<!fir.array<?xi32>> {cuf.data_attr = #cuf.c
// CHECK-LABEL: llvm.func @_QMm1Psub1
// CHECK-COUNT-2: _FortranACUFAllocDescriptor
+// CUSTOM-LABEL: llvm.func @_QMm1Psub1
+// CUSTOM-COUNT-2: custom_alloc_desc
// -----
More information about the flang-commits
mailing list