[flang-commits] [flang] [flang][cuda] Allocate descriptors in managed memory only when the device reads them (PR #234579)
via flang-commits
flang-commits at lists.llvm.org
Thu Oct 22 04:11:31 PDT 2026
https://github.com/khaki3 created https://github.com/llvm/llvm-project/pull/234579
Example:
```fortran
subroutine s(x)
real, device :: x(:)
end subroutine
real, device, allocatable :: a(:)
allocate(a(8))
do i = 1, n
call s(a)
end do
```
In this code, each `call s(a)` builds a temporary descriptor for `a`. Codegen stores a temporary descriptor either in an `alloca` or, when `isDeviceAllocation` says the data is device-resident, in managed memory from `CUFAllocDescriptor`. A managed descriptor may be read by an asynchronous kernel, so codegen emits no free for it. This one takes the managed path even though no kernel reads it, so every call leaks one.
Fix: allocate a descriptor with `CUFAllocDescriptor` only when the device reads it, meaning it reaches a `gpu.launch_func` or a device dummy of a BIND(C) procedure. Otherwise it uses an `alloca`.
- `isDeviceAllocation` is removed. Where the data lives doesn't matter, and a kernel launched on a dummy gets its own managed copy from the callee's rebox.
- The OpenACC data-clause check is removed, since a box used only by `acc` data clauses now stays on the stack anyway.
- Lowering records the CUDA data attribute of BIND(C) box arguments on the `fir.call`, so codegen sees it for direct and indirect calls.
- TargetRewrite now keeps a rewritten call's procedure attributes and the attributes of arguments passed unchanged. It used to drop them, which lost the BIND(C) marking when a call also passed, for example, a complex by value.
>From 7ddc412bfbbfb4fbb1331bbe12e26292c730749f Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 21 Oct 2026 10:38:03 -0700
Subject: [PATCH 1/2] [flang][cuda] Allocate descriptors in managed memory only
when the device reads them
Codegen placed the descriptor of any fir.embox, fir.rebox or
fir.create_box built over device-resident data in managed memory. Such a
descriptor is never freed, so every call passing a device array to an
assumed-shape dummy, including calls to host procedures, allocated one
that was never released.
Decide the placement from how the descriptor is used instead. It is
placed in device-accessible memory only when it may reach a
gpu.launch_func operand, looking through fir.convert, arith.select and
branch block arguments (as for the fir.if around an OPTIONAL argument),
or when it is passed to a device dummy of a BIND(C) procedure, whose
foreign code may launch a kernel on it. A Fortran callee that launches a
kernel with its dummy reboxes it first, so the launch check covers it.
Lowering now records the CUDA data attribute of BIND(C) box arguments
on the fir.call, so codegen also sees it for indirect calls through
procedure pointers.
---
flang/include/flang/Lower/CallInterface.h | 5 +
flang/lib/Lower/ConvertCall.cpp | 36 +++-
flang/lib/Optimizer/CodeGen/CodeGen.cpp | 154 +++++-------------
flang/test/Fir/CUDA/cuda-code-gen.mlir | 127 ++++++++++++---
.../Lower/CUDA/cuda-bindc-call-data-attr.cuf | 45 +++++
5 files changed, 226 insertions(+), 141 deletions(-)
create mode 100644 flang/test/Lower/CUDA/cuda-bindc-call-data-attr.cuf
diff --git a/flang/include/flang/Lower/CallInterface.h b/flang/include/flang/Lower/CallInterface.h
index 2673e61b1a915e..060eeaa3657cf8 100644
--- a/flang/include/flang/Lower/CallInterface.h
+++ b/flang/include/flang/Lower/CallInterface.h
@@ -202,6 +202,11 @@ class CallInterface {
/// Number of MLIR inputs/outputs of the created FuncOp.
std::size_t getNumFIRArguments() const { return inputs.size(); }
std::size_t getNumFIRResults() const { return outputs.size(); }
+ /// Return the MLIR attributes of the MLIR input at \p index.
+ llvm::ArrayRef<mlir::NamedAttribute>
+ getFIRArgumentAttributes(std::size_t index) const {
+ return inputs[index].attributes;
+ }
/// Return the MLIR output types.
llvm::SmallVector<mlir::Type> getResultType() const;
diff --git a/flang/lib/Lower/ConvertCall.cpp b/flang/lib/Lower/ConvertCall.cpp
index 2f91d59c54311e..027a8351f6a096 100644
--- a/flang/lib/Lower/ConvertCall.cpp
+++ b/flang/lib/Lower/ConvertCall.cpp
@@ -363,6 +363,31 @@ static bool mustDestroyOrFinalizeFunctionResult(
hlfir::mayHaveAllocatableComponent(callSiteType.getResult(0));
}
+/// Return the CUDA data attribute of each descriptor in \p args as call
+/// argument attributes, or null if none has one. Codegen reads them to place
+/// descriptors passed to BIND(C) device dummies in device-accessible memory,
+/// since foreign code may launch a kernel on them.
+static mlir::ArrayAttr
+genCUDADataArgAttrs(fir::FirOpBuilder &builder,
+ const Fortran::lower::CallerInterface &caller,
+ mlir::ValueRange args) {
+ llvm::SmallVector<mlir::Attribute> argAttrs(args.size(),
+ builder.getDictionaryAttr({}));
+ bool hasDataAttr = false;
+ for (unsigned i = 0, e = std::min<std::size_t>(args.size(),
+ caller.getNumFIRArguments());
+ i < e; ++i) {
+ if (!mlir::isa<fir::BaseBoxType>(args[i].getType()))
+ continue;
+ for (const mlir::NamedAttribute &attr : caller.getFIRArgumentAttributes(i))
+ if (attr.getName() == cuf::getDataAttrName()) {
+ argAttrs[i] = builder.getDictionaryAttr(attr);
+ hasDataAttr = true;
+ }
+ }
+ return hasDataAttr ? builder.getArrayAttr(argAttrs) : mlir::ArrayAttr{};
+}
+
std::tuple<Fortran::lower::LoweredResult, bool, mlir::Operation *>
Fortran::lower::genCallOpAndResult(
mlir::Location loc, Fortran::lower::AbstractConverter &converter,
@@ -802,9 +827,16 @@ Fortran::lower::genCallOpAndResult(
else if (caller.getCallDescription().hasAlwaysInline())
inlineAttr = fir::FortranInlineEnumAttr::get(
builder.getContext(), fir::FortranInlineEnum::always_inline);
+ mlir::ArrayAttr argAttrs;
+ if (procAttrs &&
+ fir::bitEnumContainsAny(procAttrs.getValue(),
+ fir::FortranProcedureFlagsEnum::bind_c))
+ argAttrs = genCUDADataArgAttrs(
+ builder, caller,
+ mlir::ValueRange(operands).drop_front(funcPointer ? 1 : 0));
auto call = fir::CallOp::create(
- builder, loc, funcType.getResults(), funcSymbolAttr, operands,
- /*arg_attrs=*/nullptr, /*res_attrs=*/nullptr, procAttrs, inlineAttr,
+ builder, loc, funcType.getResults(), funcSymbolAttr, operands, argAttrs,
+ /*res_attrs=*/nullptr, procAttrs, inlineAttr,
/*accessGroups=*/mlir::ArrayAttr{});
callOp = call;
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 79605d0c6eaca7..d7b35a272293a7 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -64,6 +64,7 @@
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Target/LLVMIR/Import.h"
@@ -1845,10 +1846,6 @@ genCUFAllocDescriptor(mlir::Location loc,
static bool isUsedByGPULaunchFunc(mlir::Value val);
-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.
template <typename ModOpTy>
@@ -2367,10 +2364,7 @@ struct EmboxOpConversion : public EmboxCommonConversion<fir::EmboxOp> {
"fir.embox codegen of derived with length parameters");
return mlir::failure();
}
- bool needsDeviceAlloc =
- isDeviceAllocation(embox.getMemref(), adaptor.getMemref(),
- this->options.cudaDescriptorAllocFunction) ||
- isUsedByGPULaunchFunc(embox);
+ bool needsDeviceAlloc = isUsedByGPULaunchFunc(embox);
auto result = placeInMemoryIfNotGlobalInit(rewriter, embox.getLoc(), boxTy,
dest, needsDeviceAlloc);
rewriter.replaceOp(embox, result);
@@ -2378,8 +2372,30 @@ struct EmboxOpConversion : public EmboxCommonConversion<fir::EmboxOp> {
}
};
-/// Return true if any user of \p val is a gpu.launch_func operation,
-/// indicating the descriptor must be in device-accessible memory.
+/// Return true if \p use passes a descriptor to a device dummy of a BIND(C)
+/// procedure. Foreign code may launch a kernel on the descriptor as is.
+static bool isPassedToBindCDeviceDummy(mlir::OpOperand &use) {
+ auto call = mlir::dyn_cast<fir::CallOp>(use.getOwner());
+ if (!call || !call.getArgAttrsAttr() || !call.getProcedureAttrs() ||
+ !fir::bitEnumContainsAny(*call.getProcedureAttrs(),
+ fir::FortranProcedureFlagsEnum::bind_c))
+ return false;
+ // An indirect call's first operand is the callee, not an argument.
+ unsigned firstArg = call.getCallee() ? 0 : 1;
+ unsigned argIdx = use.getOperandNumber() - firstArg;
+ mlir::ArrayAttr argAttrs = call.getArgAttrsAttr();
+ if (use.getOperandNumber() < firstArg || argIdx >= argAttrs.size())
+ return false;
+ auto dataAttr = mlir::cast<mlir::DictionaryAttr>(argAttrs[argIdx])
+ .getAs<cuf::DataAttributeAttr>(cuf::getDataAttrName());
+ return dataAttr && dataAttr.getValue() != cuf::DataAttribute::Pinned &&
+ dataAttr.getValue() != cuf::DataAttribute::Unified;
+}
+
+/// Return true if \p val may reach a gpu.launch_func operand or a BIND(C)
+/// device dummy, indicating the descriptor must be in device-accessible memory.
+/// Values are followed through conversions, selects and branches (e.g. the
+/// fir.if around an OPTIONAL argument).
static bool isUsedByGPULaunchFunc(mlir::Value val) {
llvm::SmallPtrSet<mlir::Value, 4> visited;
llvm::SmallVector<mlir::Value, 4> worklist{val};
@@ -2388,103 +2404,22 @@ static bool isUsedByGPULaunchFunc(mlir::Value val) {
if (!visited.insert(current).second)
continue;
- for (auto *user : current.getUsers()) {
- if (mlir::isa<mlir::gpu::LaunchFuncOp>(user))
+ for (mlir::OpOperand &use : current.getUses()) {
+ mlir::Operation *user = use.getOwner();
+ if (mlir::isa<mlir::gpu::LaunchFuncOp>(user) ||
+ isPassedToBindCDeviceDummy(use))
return true;
- if (auto convert = mlir::dyn_cast<fir::ConvertOp>(user))
- worklist.push_back(convert.getResult());
+ if (mlir::isa<fir::ConvertOp, mlir::arith::SelectOp>(user))
+ worklist.push_back(user->getResult(0));
+ else if (auto branch = mlir::dyn_cast<mlir::BranchOpInterface>(user))
+ if (std::optional<mlir::BlockArgument> arg =
+ branch.getSuccessorBlockArgument(use.getOperandNumber()))
+ worklist.push_back(*arg);
}
}
return false;
}
-/// Return true if any user of \p val is an OpenACC data-clause operation (an op
-/// from the `acc` dialect, e.g. `acc.present`, `acc.copyin`, `acc.create`).
-///
-/// Such a box is host-side descriptor metadata for the OpenACC data clause: the
-/// data-clause result carries the device data, so the box itself does not need
-/// to be device-accessible. It must therefore NOT be placed in managed memory
-/// even when the data it describes is device-resident: a managed descriptor
-/// created here is never freed and leaves a stale descriptor behind when its
-/// address is later reused.
-static bool isUsedByOpenACCDataClause(mlir::Value val) {
- for (auto *user : val.getUsers())
- if (mlir::isa_and_nonnull<mlir::acc::OpenACCDialect>(user->getDialect()))
- return true;
- return false;
-}
-
-static bool isDeviceAllocation(mlir::Value val, mlir::Value adaptorVal,
- llvm::StringRef cudaDescriptorAllocFunction) {
- if (val.getDefiningOp() &&
- val.getDefiningOp()->getParentOfType<mlir::gpu::GPUModuleOp>())
- return false;
- // Check if the global symbol is in the device module.
- if (auto addr = mlir::dyn_cast_or_null<fir::AddrOfOp>(val.getDefiningOp()))
- if (auto gpuMod =
- addr->getParentOfType<mlir::ModuleOp>()
- .lookupSymbol<mlir::gpu::GPUModuleOp>(cudaDeviceModuleName))
- if (gpuMod.lookupSymbol<mlir::LLVM::GlobalOp>(addr.getSymbol()) ||
- gpuMod.lookupSymbol<fir::GlobalOp>(addr.getSymbol()))
- return true;
-
- if (auto loadOp = mlir::dyn_cast_or_null<fir::LoadOp>(val.getDefiningOp()))
- return isDeviceAllocation(loadOp.getMemref(), {},
- cudaDescriptorAllocFunction);
- if (auto boxAddrOp =
- mlir::dyn_cast_or_null<fir::BoxAddrOp>(val.getDefiningOp()))
- return isDeviceAllocation(boxAddrOp.getVal(), {},
- cudaDescriptorAllocFunction);
- if (auto convertOp =
- mlir::dyn_cast_or_null<fir::ConvertOp>(val.getDefiningOp()))
- return isDeviceAllocation(convertOp.getValue(), {},
- cudaDescriptorAllocFunction);
- // fir.declare, and the fircg.ext_declare it becomes when debug info keeps it
- // alive until codegen, are pass-through on their memref operand. The adaptor
- // value is forwarded too so a declared dummy argument stays recognizable.
- if (auto declareOp =
- mlir::dyn_cast_or_null<fir::DeclareOp>(val.getDefiningOp()))
- return isDeviceAllocation(declareOp.getMemref(), adaptorVal,
- cudaDescriptorAllocFunction);
- if (auto xDeclareOp =
- mlir::dyn_cast_or_null<fir::cg::XDeclareOp>(val.getDefiningOp()))
- return isDeviceAllocation(xDeclareOp.getMemref(), adaptorVal,
- cudaDescriptorAllocFunction);
- if (!val.getDefiningOp() && adaptorVal) {
- if (auto blockArg = llvm::dyn_cast<mlir::BlockArgument>(adaptorVal)) {
- if (blockArg.getOwner() && blockArg.getOwner()->getParentOp() &&
- blockArg.getOwner()->isEntryBlock()) {
- if (auto func = mlir::dyn_cast_or_null<mlir::FunctionOpInterface>(
- *blockArg.getOwner()->getParentOp())) {
- auto argAttrs = func.getArgAttrs(blockArg.getArgNumber());
- for (auto attr : argAttrs) {
- if (attr.getName().getValue().ends_with(cuf::getDataAttrName())) {
- auto dataAttr =
- mlir::dyn_cast<cuf::DataAttributeAttr>(attr.getValue());
- if (dataAttr.getValue() != cuf::DataAttribute::Pinned &&
- dataAttr.getValue() != cuf::DataAttribute::Unified)
- return true;
- }
- }
- }
- }
- }
- }
- if (auto callOp = mlir::dyn_cast_or_null<fir::CallOp>(val.getDefiningOp()))
- 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;
-}
-
/// Create a box from a base pointer and explicit per-dimension lower bound,
/// extent, and byte stride. Unlike fir.embox, the strides are written into the
/// descriptor verbatim (no scaling by the element size), while elem_len/type
@@ -2518,10 +2453,7 @@ struct CreateBoxOpConversion : public EmboxCommonConversion<fir::CreateBoxOp> {
}
dest = insertBaseAddress(rewriter, loc, dest, adaptor.getMemref());
- bool needsDeviceAlloc =
- isDeviceAllocation(createBox.getMemref(), adaptor.getMemref(),
- this->options.cudaDescriptorAllocFunction) ||
- isUsedByGPULaunchFunc(createBox);
+ bool needsDeviceAlloc = isUsedByGPULaunchFunc(createBox);
mlir::Value result = placeInMemoryIfNotGlobalInit(rewriter, loc, boxTy,
dest, needsDeviceAlloc);
rewriter.replaceOp(createBox, result);
@@ -2717,11 +2649,7 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
dest = insertBaseAddress(rewriter, loc, dest, base);
if (fir::isDerivedTypeWithLenParams(boxTy))
TODO(loc, "fir.embox codegen of derived with length parameters");
- bool needsDeviceAlloc =
- isUsedByGPULaunchFunc(xbox) ||
- (isDeviceAllocation(xbox.getMemref(), adaptor.getMemref(),
- this->options.cudaDescriptorAllocFunction) &&
- !isUsedByOpenACCDataClause(xbox));
+ bool needsDeviceAlloc = isUsedByGPULaunchFunc(xbox);
mlir::Value result = placeInMemoryIfNotGlobalInit(rewriter, loc, boxTy,
dest, needsDeviceAlloc);
rewriter.replaceOp(xbox, result);
@@ -2839,11 +2767,7 @@ struct XReboxOpConversion : public EmboxCommonConversion<fir::cg::XReboxOp> {
dest = insertStride(rewriter, loc, dest, dim, std::get<1>(iter.value()));
}
dest = insertBaseAddress(rewriter, loc, dest, base);
- bool needsDeviceAlloc =
- isUsedByGPULaunchFunc(rebox) ||
- (isDeviceAllocation(rebox.getBox(), adaptor.getBox(),
- this->options.cudaDescriptorAllocFunction) &&
- !isUsedByOpenACCDataClause(rebox));
+ bool needsDeviceAlloc = isUsedByGPULaunchFunc(rebox);
mlir::Value result = placeInMemoryIfNotGlobalInit(
rewriter, rebox.getLoc(), destBoxTy, dest, needsDeviceAlloc);
rewriter.replaceOp(rebox, result);
diff --git a/flang/test/Fir/CUDA/cuda-code-gen.mlir b/flang/test/Fir/CUDA/cuda-code-gen.mlir
index 53492518fc4751..802bd1d6983397 100644
--- a/flang/test/Fir/CUDA/cuda-code-gen.mlir
+++ b/flang/test/Fir/CUDA/cuda-code-gen.mlir
@@ -56,7 +56,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<f80 = dense<128> : vector<2xi64>
// CHECK-LABEL: llvm.func @_QQmain()
// CHECK: llvm.call @_FortranACUFMemAlloc
-// CHECK: llvm.call @_FortranACUFAllocDescriptor
+// CHECK-NOT: llvm.call @_FortranACUFAllocDescriptor
// -----
@@ -126,7 +126,8 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<f80 = dense<128> : vector<2xi64>
}
// CHECK-LABEL: llvm.func @_QQmain()
-// CHECK-COUNT-3: llvm.call @_FortranACUFAllocDescriptor
+// CHECK-COUNT-2: llvm.call @_FortranACUFAllocDescriptor
+// CHECK-NOT: llvm.call @_FortranACUFAllocDescriptor
// -----
@@ -170,7 +171,8 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<!llvm.ptr<270> = dense<32> : vec
}
// CHECK-LABEL: llvm.func @_QQmain()
-// CHECK-COUNT-2: llvm.call @_FortranACUFAllocDescriptor
+// CHECK-COUNT-1: llvm.call @_FortranACUFAllocDescriptor
+// CHECK-NOT: llvm.call @_FortranACUFAllocDescriptor
// -----
@@ -187,9 +189,7 @@ 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
+// CHECK-NOT: _FortranACUFAllocDescriptor
// -----
@@ -200,9 +200,7 @@ 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
+// CHECK-NOT: _FortranACUFAllocDescriptor
// -----
@@ -433,6 +431,10 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<f80, dense<128> :
// CHECK: llvm.call @_FortranACUFAllocDescriptor(
// CHECK: gpu.launch_func @cuda_device_mod::@_QMtestmePmykernel
+// CUSTOM-LABEL: llvm.func @_QQmain()
+// CUSTOM: llvm.call @custom_alloc_desc(
+// CUSTOM: gpu.launch_func @cuda_device_mod::@_QMtestmePmykernel
+
// -----
// Test that an embox whose converted result is passed to gpu.launch_func gets a
@@ -465,13 +467,92 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i64, dense<64> :
// -----
-// A rebox of a device-resident global (isDeviceAllocation == true) whose result
-// is consumed by an OpenACC data-clause op (here acc.present) must NOT be placed
-// in managed memory. The box is host-side descriptor metadata for the data
-// clause; a managed descriptor created here is never freed and leaves a stale
-// descriptor behind when its address is reused. The descriptor storage should
-// therefore be a plain stack alloca and no _FortranACUFAllocDescriptor call
-// should be emitted for it.
+// A rebox that reaches gpu.launch_func through a branch to a block argument
+// (as for an OPTIONAL dummy) or through arith.select still gets a managed
+// descriptor.
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>, gpu.container_module} {
+ func.func @_QPopt_branch(%arg0: !fir.box<!fir.array<?xf32>>, %present: i1) {
+ %c0_i32 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : index
+ cf.cond_br %present, ^bb1, ^bb2
+ ^bb1:
+ %0 = fircg.ext_rebox %arg0 : (!fir.box<!fir.array<?xf32>>) -> !fir.box<!fir.array<?xf32>>
+ cf.br ^bb3(%0 : !fir.box<!fir.array<?xf32>>)
+ ^bb2:
+ %1 = fir.absent !fir.box<!fir.array<?xf32>>
+ cf.br ^bb3(%1 : !fir.box<!fir.array<?xf32>>)
+ ^bb3(%2: !fir.box<!fir.array<?xf32>>):
+ gpu.launch_func @cuda_device_mod::@kernel blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1) dynamic_shared_memory_size %c0_i32 args(%2 : !fir.box<!fir.array<?xf32>>) {cuf.proc_attr = #cuf.cuda_proc<global>}
+ return
+ }
+ func.func @_QPopt_select(%arg0: !fir.box<!fir.array<?xf32>>, %present: i1) {
+ %c0_i32 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : index
+ %0 = fircg.ext_rebox %arg0 : (!fir.box<!fir.array<?xf32>>) -> !fir.box<!fir.array<?xf32>>
+ %1 = fir.absent !fir.box<!fir.array<?xf32>>
+ %2 = arith.select %present, %0, %1 : !fir.box<!fir.array<?xf32>>
+ gpu.launch_func @cuda_device_mod::@kernel blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1) dynamic_shared_memory_size %c0_i32 args(%2 : !fir.box<!fir.array<?xf32>>) {cuf.proc_attr = #cuf.cuda_proc<global>}
+ return
+ }
+ gpu.module @cuda_device_mod {
+ gpu.func @kernel(%arg0: !fir.box<!fir.array<?xf32>>) kernel {
+ gpu.return
+ }
+ }
+}
+
+// CHECK-LABEL: llvm.func @_QPopt_branch(
+// CHECK: llvm.call @_FortranACUFAllocDescriptor(
+// CHECK: gpu.launch_func @cuda_device_mod::@kernel
+// CHECK-LABEL: llvm.func @_QPopt_select(
+// CHECK: llvm.call @_FortranACUFAllocDescriptor(
+// CHECK: gpu.launch_func @cuda_device_mod::@kernel
+
+// -----
+
+// A descriptor passed to a device dummy of a BIND(C) procedure, directly or
+// through a procedure pointer, gets a managed descriptor, since foreign code
+// may launch a kernel on it. A host dummy of a BIND(C) procedure keeps a stack
+// descriptor.
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>} {
+ func.func @_QPbindc_device(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %1 = fircg.ext_rebox %0 : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.box<!fir.array<?xf32>>
+ fir.call @claunch(%1) proc_attrs<bind_c> : (!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}) -> ()
+ return
+ }
+ func.func @_QPbindc_indirect(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, %f: (!fir.box<!fir.array<?xf32>>) -> ()) {
+ %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %1 = fircg.ext_rebox %0 : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.box<!fir.array<?xf32>>
+ fir.call %f(%1) proc_attrs<bind_c> : (!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}) -> ()
+ return
+ }
+ func.func @_QPbindc_host(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) {
+ %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %1 = fircg.ext_rebox %0 : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.box<!fir.array<?xf32>>
+ fir.call @chost(%1) proc_attrs<bind_c> : (!fir.box<!fir.array<?xf32>>) -> ()
+ return
+ }
+ func.func private @claunch(!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}) attributes {fir.proc_attrs = #fir.proc_attrs<bind_c>}
+ func.func private @chost(!fir.box<!fir.array<?xf32>>) attributes {fir.proc_attrs = #fir.proc_attrs<bind_c>}
+}
+
+// CHECK-LABEL: llvm.func @_QPbindc_device(
+// CHECK: llvm.call @_FortranACUFAllocDescriptor(
+// CHECK: llvm.call @claunch(
+// CHECK-LABEL: llvm.func @_QPbindc_indirect(
+// CHECK: llvm.call @_FortranACUFAllocDescriptor(
+// CHECK: llvm.call %{{.*}}(
+// CHECK-LABEL: llvm.func @_QPbindc_host(
+// CHECK-NOT: llvm.call @_FortranACUFAllocDescriptor
+// CHECK: llvm.call @chost(
+
+// -----
+
+// A rebox of a device-resident global consumed by an OpenACC data-clause op
+// (here acc.present) is host-side metadata and stays on the stack.
module attributes {gpu.container_module} {
fir.global @_QMm1Eda {data_attr = #cuf.cuda<device>} : !fir.box<!fir.heap<!fir.array<?x?xf32>>> {
@@ -504,10 +585,9 @@ module attributes {gpu.container_module} {
// -----
-// fircg.ext_declare, which survives until codegen when debug info is requested,
-// must not hide the device-resident memory it declares. The descriptor for an
-// embox of a managed allocation goes to managed memory whether or not the
-// allocation is reached through a declare.
+// A descriptor for device-resident data that is only passed to a host
+// procedure stays on the stack: placing it in managed memory would leak it at
+// every call, since nothing frees it.
module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>} {
func.func @_QPtesti4() {
@@ -534,7 +614,6 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i64, dense<64> :
}
// CHECK-LABEL: llvm.func @_QPtesti4()
-// CHECK: llvm.call @_FortranACUFAllocDescriptor(
-
-// CUSTOM-LABEL: llvm.func @_QPtesti4()
-// CUSTOM: llvm.call @custom_alloc_desc(
+// CHECK: llvm.alloca {{.*}} x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)>
+// CHECK-NOT: llvm.call @_FortranACUFAllocDescriptor
+// CHECK: llvm.call @_QPcallee(
diff --git a/flang/test/Lower/CUDA/cuda-bindc-call-data-attr.cuf b/flang/test/Lower/CUDA/cuda-bindc-call-data-attr.cuf
new file mode 100644
index 00000000000000..8c33a3f09cb552
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-bindc-call-data-attr.cuf
@@ -0,0 +1,45 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+
+! A call to a BIND(C) procedure carries the CUDA data attribute of its dummy
+! arguments, whether the procedure is called directly or through a procedure
+! pointer. A call to a non-BIND(C) procedure does not.
+
+module m
+ interface
+ subroutine claunch(x, inc) bind(c)
+ real, device :: x(:)
+ real, value :: inc
+ end subroutine
+ subroutine chost(x) bind(c)
+ real :: x(:)
+ end subroutine
+ end interface
+ abstract interface
+ subroutine iface(x, inc) bind(c)
+ real, device :: x(:)
+ real, value :: inc
+ end subroutine
+ end interface
+contains
+ subroutine fsub(x)
+ real, device :: x(:)
+ end subroutine
+end module
+
+subroutine test(b, h)
+ use m
+ real, device :: b(:)
+ real :: h(:)
+ procedure(iface), pointer :: pp
+ call claunch(b, 2.0)
+ pp => claunch
+ call pp(b, 2.0)
+ call chost(h)
+ call fsub(b)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtest(
+! CHECK: fir.call @claunch(%{{.*}}, %{{.*}}) proc_attrs<bind_c> {{.*}}: (!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}, f32) -> ()
+! CHECK: fir.call %{{.*}}(%{{.*}}, %{{.*}}) proc_attrs<bind_c> {{.*}}: (!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}, f32) -> ()
+! CHECK: fir.call @chost(%{{.*}}) proc_attrs<bind_c> {{.*}}: (!fir.box<!fir.array<?xf32>>) -> ()
+! CHECK: fir.call @_QMmPfsub(%{{.*}}) {{.*}}: (!fir.box<!fir.array<?xf32>>) -> ()
>From a628e8c5e5348202b9d15ddb82320fa8f459e8e8 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 21 Oct 2026 20:51:23 -0700
Subject: [PATCH 2/2] [flang] Keep call attributes when TargetRewrite rewrites
a fir.call
When TargetRewrite rebuilds a fir.call for the target ABI (for example
because an argument is a complex passed by value), the new call dropped
its procedure_attrs and the arg_attrs of every argument, keeping only the
newly computed ABI attributes.
Keep the procedure attributes, and carry the attributes of each argument
passed unchanged over to its position in the new call, merged with the
ABI attributes. Arguments that are split or passed differently only get
their ABI attributes, as before.
This keeps cuf.data_attr on a BIND(C) device dummy, which codegen uses to
place the descriptor in device-accessible memory.
---
flang/lib/Optimizer/CodeGen/TargetRewrite.cpp | 28 +++++++++++++++----
flang/test/Fir/target-rewrite-call-attrs.fir | 22 +++++++++++++++
.../Fir/target-rewrite-indirect-calls.fir | 4 +--
3 files changed, 47 insertions(+), 7 deletions(-)
create mode 100644 flang/test/Fir/target-rewrite-call-attrs.fir
diff --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
index f8eb8087a9b900..b988fd744eecfc 100644
--- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
@@ -429,6 +429,9 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
llvm::SmallVector<mlir::Value> trailingOpers;
llvm::SmallVector<mlir::Value> operands;
unsigned passArgShift = 0;
+ // (original argument index, new operand index) of the arguments passed
+ // unchanged, whose attributes carry over to the new call.
+ llvm::SmallVector<std::pair<unsigned, unsigned>> keptArgs;
if constexpr (std::is_same_v<std::decay_t<A>, mlir::gpu::LaunchFuncOp>)
operands = callOp.getKernelOperands();
else
@@ -524,6 +527,7 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
if (callOp.getPassArgPos() && *callOp.getPassArgPos() == index)
passArgShift = newOpers.size() - *callOp.getPassArgPos();
}
+ keptArgs.emplace_back(index, newOpers.size());
newInTyAndAttrs.push_back(
fir::CodeGenSpecifics::getTypeAndAttr(ty));
newOpers.push_back(oper);
@@ -571,15 +575,29 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
fir::CallOp::Properties{});
}
newCall.setFastmathAttr(callOp.getFastmathAttr());
+ newCall.setProcedureAttrsAttr(callOp.getProcedureAttrsAttr());
+ // Attributes of the arguments passed unchanged, by new argument index.
+ llvm::SmallVector<mlir::NamedAttrList> keptArgAttrs(
+ newInTyAndAttrs.size() - dropFront);
+ bool hasKeptArgAttrs = false;
+ if (mlir::ArrayAttr oldArgAttrs = callOp.getArgAttrsAttr())
+ for (auto [oldIdx, newOperIdx] : keptArgs) {
+ if (oldIdx >= oldArgAttrs.size())
+ continue;
+ auto dict = mlir::cast<mlir::DictionaryAttr>(oldArgAttrs[oldIdx]);
+ keptArgAttrs[newOperIdx - dropFront].append(dict.getValue());
+ hasKeptArgAttrs |= !dict.empty();
+ }
// Always set ABI argument attributes on call operations, even when
// direct, as required by
// https://llvm.org/docs/LangRef.html#parameter-attributes.
- if (hasAbiArgAttributes(newInTyAndAttrs)) {
+ if (hasAbiArgAttributes(newInTyAndAttrs) || hasKeptArgAttrs) {
llvm::SmallVector<mlir::Attribute> argAttrsArray;
- for (const auto &arg :
- llvm::ArrayRef<fir::CodeGenSpecifics::TypeAndAttr>(newInTyAndAttrs)
- .drop_front(dropFront)) {
- mlir::NamedAttrList argAttrs;
+ for (auto [argIdx, arg] :
+ llvm::enumerate(llvm::ArrayRef<fir::CodeGenSpecifics::TypeAndAttr>(
+ newInTyAndAttrs)
+ .drop_front(dropFront))) {
+ mlir::NamedAttrList &argAttrs = keptArgAttrs[argIdx];
const auto &attr = std::get<fir::CodeGenSpecifics::Attributes>(arg);
if (attr.isByVal()) {
mlir::Type elemType =
diff --git a/flang/test/Fir/target-rewrite-call-attrs.fir b/flang/test/Fir/target-rewrite-call-attrs.fir
new file mode 100644
index 00000000000000..d72dffc925cd3e
--- /dev/null
+++ b/flang/test/Fir/target-rewrite-call-attrs.fir
@@ -0,0 +1,22 @@
+// Test that rewriting a call keeps its procedure attributes and the attributes
+// of the arguments passed unchanged.
+// RUN: fir-opt --target-rewrite="target=x86_64-unknown-linux-gnu" %s | FileCheck %s
+
+func.func private @claunchc(!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}, complex<f32>) attributes {fir.proc_attrs = #fir.proc_attrs<bind_c>}
+
+func.func @direct(%arg0: !fir.box<!fir.array<?xf32>>, %arg1: complex<f32>) {
+ fir.call @claunchc(%arg0, %arg1) proc_attrs<bind_c> : (!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}, complex<f32>) -> ()
+ return
+}
+
+func.func @indirect(%arg: () -> (), %arg0: !fir.box<!fir.array<?xf32>>, %arg1: complex<f32>) {
+ %f = fir.convert %arg : (() -> ()) -> ((!fir.box<!fir.array<?xf32>>, complex<f32>) -> ())
+ fir.call %f(%arg0, %arg1) proc_attrs<bind_c> : (!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}, complex<f32>) -> ()
+ return
+}
+
+// CHECK-LABEL: func.func @direct(
+// CHECK: fir.call @claunchc(%{{.*}}, %{{.*}}) proc_attrs<bind_c> : (!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}, !fir.vector<2:f32>) -> ()
+
+// CHECK-LABEL: func.func @indirect(
+// CHECK: fir.call %{{.*}}(%{{.*}}, %{{.*}}) proc_attrs<bind_c> : (!fir.box<!fir.array<?xf32>> {cuf.data_attr = #cuf.cuda<device>}, !fir.vector<2:f32>) -> ()
diff --git a/flang/test/Fir/target-rewrite-indirect-calls.fir b/flang/test/Fir/target-rewrite-indirect-calls.fir
index 0e4f0c34619ba1..bc548965fc89dc 100644
--- a/flang/test/Fir/target-rewrite-indirect-calls.fir
+++ b/flang/test/Fir/target-rewrite-indirect-calls.fir
@@ -18,7 +18,7 @@ func.func @test(%arg0: () -> (), %arg1: !fir.ref<!fir.type<t{a:!fir.array<5xf64>
// CHECK-X86: %[[VAL_5:.*]] = llvm.intr.stacksave : !llvm.ptr
// CHECK-X86: %[[VAL_6:.*]] = fir.alloca !fir.type<t{a:!fir.array<5xf64>}>
// CHECK-X86: fir.store %[[VAL_3]] to %[[VAL_6]] : !fir.ref<!fir.type<t{a:!fir.array<5xf64>}>>
-// CHECK-X86: fir.call %[[VAL_4]](%[[VAL_6]], %[[VAL_2]]) : (!fir.ref<!fir.type<t{a:!fir.array<5xf64>}>> {llvm.align = 8 : i32, llvm.byval = !fir.type<t{a:!fir.array<5xf64>}>}, f64) -> ()
+// CHECK-X86: fir.call %[[VAL_4]](%[[VAL_6]], %[[VAL_2]]) proc_attrs<bind_c> : (!fir.ref<!fir.type<t{a:!fir.array<5xf64>}>> {llvm.align = 8 : i32, llvm.byval = !fir.type<t{a:!fir.array<5xf64>}>}, f64) -> ()
// CHECK-X86: llvm.intr.stackrestore %[[VAL_5]] : !llvm.ptr
// CHECK-X86: return
// CHECK-X86: }
@@ -32,7 +32,7 @@ func.func @test(%arg0: () -> (), %arg1: !fir.ref<!fir.type<t{a:!fir.array<5xf64>
// CHECK-S390: %[[VAL_5:.*]] = llvm.intr.stacksave : !llvm.ptr
// CHECK-S390: %[[VAL_6:.*]] = fir.alloca !fir.type<t{a:!fir.array<5xf64>}>
// CHECK-S390: fir.store %[[VAL_3]] to %[[VAL_6]] : !fir.ref<!fir.type<t{a:!fir.array<5xf64>}>>
-// CHECK-S390: fir.call %[[VAL_4]](%[[VAL_6]], %[[VAL_2]]) : (!fir.ref<!fir.type<t{a:!fir.array<5xf64>}>> {llvm.align = 8 : i32}, f64) -> ()
+// CHECK-S390: fir.call %[[VAL_4]](%[[VAL_6]], %[[VAL_2]]) proc_attrs<bind_c> : (!fir.ref<!fir.type<t{a:!fir.array<5xf64>}>> {llvm.align = 8 : i32}, f64) -> ()
// CHECK-S390-NOT: llvm.byval
// CHECK-S390: llvm.intr.stackrestore %[[VAL_5]] : !llvm.ptr
// CHECK-S390: return
More information about the flang-commits
mailing list