[flang-commits] [flang] [flang][cuda] Managed backing for -gpu=unified allocatables/pointers at the ALLOCATE site (PR #210149)
Vijay Kandiah via flang-commits
flang-commits at lists.llvm.org
Thu Jul 16 14:03:13 PDT 2026
https://github.com/VijayKandiah updated https://github.com/llvm/llvm-project/pull/210149
>From 7eccf5beb805aeaa767d97002d31e49bb62f215a Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Thu, 16 Jul 2026 12:25:19 -0700
Subject: [PATCH 1/2] [flang][cuda] Managed backing for -gpu=unified
allocatables/pointers at the ALLOCATE site
---
flang/lib/Lower/Allocatable.cpp | 33 +++++++-
flang/lib/Semantics/resolve-names.cpp | 73 +++++------------
flang/test/Lower/CUDA/cuda-gpu-unified.cuf | 82 +++++++++++++++++++
.../CUDA/cuda-implicit-managed-alloc.cuf | 62 --------------
4 files changed, 134 insertions(+), 116 deletions(-)
create mode 100644 flang/test/Lower/CUDA/cuda-gpu-unified.cuf
delete mode 100644 flang/test/Lower/CUDA/cuda-implicit-managed-alloc.cuf
diff --git a/flang/lib/Lower/Allocatable.cpp b/flang/lib/Lower/Allocatable.cpp
index 76108f2386738..8722a0184dac9 100644
--- a/flang/lib/Lower/Allocatable.cpp
+++ b/flang/lib/Lower/Allocatable.cpp
@@ -36,6 +36,7 @@
#include "flang/Runtime/pointer.h"
#include "flang/Semantics/tools.h"
#include "flang/Semantics/type.h"
+#include "flang/Support/Fortran-features.h"
#include "llvm/Support/CommandLine.h"
/// By default fir memory operation fir::AllocMemOp/fir::FreeMemOp are used.
@@ -512,11 +513,28 @@ class AllocateStmtHelper {
isCudaAllocate = propagateCUDAAttrsFromParent(alloc, cudaSymForAlloc);
bool isCudaDeviceContext = cuf::isCUDADeviceContext(builder.getRegion());
+ unsigned allocatorIdx = Fortran::lower::getAllocatorIdx(*cudaSymForAlloc);
+
+ // Under -gpu=mem:unified, back plain (unattributed) allocatables/pointers
+ // with managed memory by selecting the unified allocator index at the
+ // ALLOCATE site. The symbol stays unattributed, so argument passing,
+ // interfaces, and COMMON legality are unaffected.
+ bool implicitManagedBacking = false;
+ if (allocatorIdx == kDefaultAllocator && !isCudaAllocate &&
+ !isCudaDeviceContext && (box.isAllocatable() || box.isPointer()) &&
+ converter.getFoldingContext().languageFeatures().IsEnabled(
+ Fortran::common::LanguageFeature::CudaUnified)) {
+ allocatorIdx = kUnifiedAllocatorPos;
+ implicitManagedBacking = true;
+ }
+
+ // The inlined allocation path emits a plain heap allocmem that ignores the
+ // allocator index; use the runtime path when we redirected an otherwise
+ // plain allocatable to a managed allocator so the index is honored.
bool inlineAllocation = !box.isDerived() && !errorManager.hasStatSpec() &&
!alloc.type.IsPolymorphic() &&
!alloc.hasCoarraySpec() && !useAllocateRuntime &&
- !box.isPointer();
- unsigned allocatorIdx = Fortran::lower::getAllocatorIdx(*cudaSymForAlloc);
+ !box.isPointer() && !implicitManagedBacking;
if (inlineAllocation && !alloc.hasCoarraySpec() &&
((isCudaAllocate && isCudaDeviceContext) || !isCudaAllocate)) {
@@ -943,10 +961,19 @@ genDeallocate(fir::FirOpBuilder &builder,
const Fortran::semantics::Symbol *symbol = nullptr) {
bool isCudaSymbol = symbol && Fortran::semantics::HasCUDAAttr(*symbol);
bool isCudaDeviceContext = cuf::isCUDADeviceContext(builder.getRegion());
+ // A plain allocatable/pointer under -gpu=mem:unified was given the unified
+ // allocator index at ALLOCATE, so its deallocation must go through the
+ // runtime (which honors that index) rather than an inlined freemem that would
+ // call libc free() on managed memory.
+ bool implicitManagedBacking =
+ !isCudaSymbol && !isCudaDeviceContext &&
+ (box.isAllocatable() || box.isPointer()) &&
+ converter.getFoldingContext().languageFeatures().IsEnabled(
+ Fortran::common::LanguageFeature::CudaUnified);
bool inlineDeallocation =
!box.isDerived() && !box.isPolymorphic() && !box.hasAssumedRank() &&
!box.isUnlimitedPolymorphic() && !errorManager.hasStatSpec() &&
- !useAllocateRuntime && !box.isPointer();
+ !useAllocateRuntime && !box.isPointer() && !implicitManagedBacking;
bool isCoarraySymbol = symbol && Fortran::evaluate::IsCoarray(*symbol);
// Deallocate intrinsic types inline.
diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp
index cdf81e7871f36..996c05d0e193b 100644
--- a/flang/lib/Semantics/resolve-names.cpp
+++ b/flang/lib/Semantics/resolve-names.cpp
@@ -1124,7 +1124,6 @@ class DeclarationVisitor : public ArraySpecVisitor,
protected:
bool BeginDecl();
void EndDecl();
- void SetImplicitCUDADataAttr(Symbol &);
Symbol &DeclareObjectEntity(const parser::Name &, Attrs = Attrs{});
// Make sure that there's an entity in an enclosing scope called Name
Symbol &FindOrDeclareEnclosingEntity(const parser::Name &);
@@ -7126,11 +7125,6 @@ bool DeclarationVisitor::Pre(const parser::DerivedTypeDef &x) {
const auto &componentDefs{
std::get<std::list<parser::Statement<parser::ComponentDefStmt>>>(x.t)};
Walk(componentDefs);
- // Components live in the derived type's own scope, which is not visited when
- // the enclosing specification part is finished, so attribute them here.
- for (auto &pair : scope) {
- SetImplicitCUDADataAttr(*pair.second);
- }
if (derivedTypeInfo_.sequence) {
details.set_sequence(true);
if (componentDefs.empty()) {
@@ -10590,50 +10584,6 @@ void ResolveNamesVisitor::CreateGeneric(const parser::GenericSpec &x) {
info.Resolve(&MakeSymbol(symbolName, Attrs{}, std::move(genericDetails)));
}
-// Applies the implicit CUDA data attribute (managed/unified/pinned) to an
-// allocatable or pointer object entity when the matching -gpu=mem: mode is
-// enabled. Used for locals and module variables as well as derived-type
-// components.
-void DeclarationVisitor::SetImplicitCUDADataAttr(Symbol &symbol) {
- auto *object{symbol.detailsIf<ObjectEntityDetails>()};
- if (!object || object->cudaDataAttr() ||
- !(IsAllocatable(symbol) || IsPointer(symbol))) {
- return;
- }
- const bool cudaEnabled{
- context().languageFeatures().IsEnabled(common::LanguageFeature::CUDA)};
- const bool cudaManaged{context().languageFeatures().IsEnabled(
- common::LanguageFeature::CudaManaged)};
- const bool cudaUnified{context().languageFeatures().IsEnabled(
- common::LanguageFeature::CudaUnified)};
- // Implicitly treat allocatable/pointer arrays as managed when feature
- // is enabled. This is done after all explicit CUDA attributes have
- // been processed. Only applies when CUDA Fortran is enabled; otherwise
- // -gpu=mem:managed on a non-CUDA-Fortran translation unit (e.g. pure
- // OpenACC) would incorrectly route every allocatable through the CUDA
- // Fortran managed descriptor pipeline. Under -gpu=mem:unified prefer
- // the Unified attribute where it is legal (host subprogram, main
- // program, or component) so generic resolution still selects the
- // unified specific; fall back to Managed elsewhere (module scope,
- // device subprograms), which uses the same allocator.
- // COMMON objects may not carry managed/unified attributes.
- if (cudaEnabled && (cudaManaged || cudaUnified) && !object->commonBlock()) {
- const Scope &owner{symbol.owner()};
- const bool unifiedAllowed{!IsCUDADeviceContext(&owner) &&
- (owner.IsDerivedType() || owner.kind() == Scope::Kind::MainProgram ||
- owner.kind() == Scope::Kind::Subprogram)};
- object->set_cudaDataAttr(cudaUnified && unifiedAllowed
- ? common::CUDADataAttr::Unified
- : common::CUDADataAttr::Managed);
- // Implicitly treat allocatable arrays as pinned when feature is
- // enabled.
- } else if (IsAllocatable(symbol) &&
- context().languageFeatures().IsEnabled(
- common::LanguageFeature::CudaPinned)) {
- object->set_cudaDataAttr(common::CUDADataAttr::Pinned);
- }
-}
-
void ResolveNamesVisitor::FinishSpecificationPart(
const std::list<parser::DeclarationConstruct> &decls) {
misparsedStmtFuncFound_ = false;
@@ -10677,7 +10627,28 @@ void ResolveNamesVisitor::FinishSpecificationPart(
}
}
- SetImplicitCUDADataAttr(symbol);
+ if (auto *object{symbol.detailsIf<ObjectEntityDetails>()}) {
+ if ((IsAllocatable(symbol) || IsPointer(symbol)) &&
+ !object->cudaDataAttr()) {
+ // Implicitly treat allocatable/pointer arrays as managed when feature
+ // is enabled. This is done after all explicit CUDA attributes have
+ // been processed. Only applies when CUDA Fortran is enabled; otherwise
+ // -gpu=mem:managed on a non-CUDA-Fortran translation unit (e.g. pure
+ // OpenACC) would incorrectly route every allocatable through the CUDA
+ // Fortran managed descriptor pipeline.
+ if (context().languageFeatures().IsEnabled(
+ common::LanguageFeature::CudaManaged) &&
+ context().languageFeatures().IsEnabled(
+ common::LanguageFeature::CUDA))
+ object->set_cudaDataAttr(common::CUDADataAttr::Managed);
+ // Implicitly treat allocatable arrays as pinned when feature is
+ // enabled.
+ else if (IsAllocatable(symbol) &&
+ context().languageFeatures().IsEnabled(
+ common::LanguageFeature::CudaPinned))
+ object->set_cudaDataAttr(common::CUDADataAttr::Pinned);
+ }
+ }
}
currScope().InstantiateDerivedTypes();
for (const auto &decl : decls) {
diff --git a/flang/test/Lower/CUDA/cuda-gpu-unified.cuf b/flang/test/Lower/CUDA/cuda-gpu-unified.cuf
new file mode 100644
index 0000000000000..166e9230f0106
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-gpu-unified.cuf
@@ -0,0 +1,82 @@
+! RUN: bbc -emit-hlfir -fcuda -gpu=unified %s -o - | FileCheck %s
+
+! Under -gpu=unified, plain (unattributed) allocatables and pointers are backed
+! by CUDA managed memory at the ALLOCATE site (unified allocator index, idx 4).
+
+module m
+ real, allocatable :: mod_arr(:)
+contains
+ subroutine alloc_local()
+ real, allocatable :: loc_arr(:)
+ allocate(loc_arr(100))
+ deallocate(loc_arr)
+ end subroutine
+ subroutine alloc_module()
+ allocate(mod_arr(100))
+ deallocate(mod_arr)
+ end subroutine
+ subroutine alloc_pointer()
+ real, pointer :: loc_ptr(:)
+ allocate(loc_ptr(100))
+ deallocate(loc_ptr)
+ end subroutine
+ subroutine alloc_explicit_device()
+ real, allocatable, device :: dev_arr(:)
+ allocate(dev_arr(100))
+ deallocate(dev_arr)
+ end subroutine
+ ! A pointer in a COMMON block must still receive managed backing at ALLOCATE.
+ subroutine common_sub()
+ real, pointer :: com_ptr(:)
+ common /blk/ com_ptr
+ allocate(com_ptr(100))
+ deallocate(com_ptr)
+ end subroutine
+end module
+
+program p
+ real, allocatable :: prog_arr(:)
+ allocate(prog_arr(100))
+ deallocate(prog_arr)
+end program
+
+! The unified allocator index is chosen at the ALLOCATE site, and the matching
+! deallocation goes through the runtime (which honors the allocator index and
+! calls the managed deallocator) rather than an inlined plain free.
+! CHECK-LABEL: func.func @_QMmPalloc_local()
+! CHECK-NOT: cuf.alloc
+! CHECK: hlfir.declare {{.*}} {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMmFalloc_localEloc_arr"}
+! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
+! CHECK: fir.call @_FortranAAllocatableAllocate
+! CHECK-NOT: fir.freemem
+! CHECK: fir.call @_FortranAAllocatableDeallocate
+! CHECK-NOT: cuf.allocate
+
+! Module allocatable: redirected to the unified allocator at ALLOCATE.
+! CHECK-LABEL: func.func @_QMmPalloc_module()
+! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
+
+! Plain pointer: redirected to the unified allocator at ALLOCATE.
+! CHECK-LABEL: func.func @_QMmPalloc_pointer()
+! CHECK-NOT: cuf.alloc
+! CHECK: hlfir.declare {{.*}} {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QMmFalloc_pointerEloc_ptr"}
+! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
+! CHECK-NOT: cuf.allocate
+
+! Explicit device allocatable keeps the device allocator (idx 2); the implicit
+! unified redirect does not apply because the allocator index is already set.
+! CHECK-LABEL: func.func @_QMmPalloc_explicit_device()
+! CHECK: fir.embox {{.*}} {allocator_idx = 2 : i32}
+
+! COMMON-block pointer: its ALLOCATE is still redirected to the unified allocator.
+! CHECK-LABEL: func.func @_QMmPcommon_sub()
+! CHECK: hlfir.declare {{.*}} {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QMmFcommon_subEcom_ptr"}
+! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
+
+! Main-program local allocatable.
+! CHECK-LABEL: func.func @_QQmain()
+! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
+
+! The module global itself carries no CUDA data attribute.
+! CHECK: fir.global @_QMmEmod_arr : !fir.box<!fir.heap<!fir.array<?xf32>>>
+! CHECK-NOT: fir.global @_QMmEmod_arr {data_attr
diff --git a/flang/test/Lower/CUDA/cuda-implicit-managed-alloc.cuf b/flang/test/Lower/CUDA/cuda-implicit-managed-alloc.cuf
deleted file mode 100644
index 378eb8698d3c6..0000000000000
--- a/flang/test/Lower/CUDA/cuda-implicit-managed-alloc.cuf
+++ /dev/null
@@ -1,62 +0,0 @@
-! RUN: bbc -emit-hlfir -fcuda -gpu=unified %s -o - | FileCheck %s --check-prefix=UNI
-! RUN: bbc -emit-hlfir -fcuda -gpu=managed %s -o - | FileCheck %s --check-prefix=MAN
-! RUN: bbc -fcuda -gpu=unified --dump-symbols %s 2>&1 | FileCheck %s --check-prefix=UNI-SYM
-! RUN: bbc -fcuda -gpu=managed --dump-symbols %s 2>&1 | FileCheck %s --check-prefix=MAN-SYM
-
-! Under -gpu=unified and -gpu=managed, allocatables and pointers with no
-! explicit CUDA data attribute are implicitly attributed. -gpu=unified prefers
-! the Unified attribute where ATTRIBUTES(UNIFIED) is legal (host subprogram,
-! main program, or derived-type component); module-scope objects fall back to
-! Managed, which carries no such scope restriction.
-
-module m
- real, allocatable :: mod_arr(:)
- real, pointer :: mod_ptr(:)
- type :: t
- real, allocatable :: comp_arr(:)
- real, pointer :: comp_ptr(:)
- end type
-contains
- subroutine mod_sub()
- real, allocatable :: loc_arr(:)
- real, pointer :: loc_ptr(:)
- end subroutine
- subroutine common_sub()
- real, pointer :: com_ptr(:)
- common /blk/ com_ptr
- end subroutine
-end module
-
-program p
- real, allocatable :: prog_arr(:)
-end program
-
-! Host-subprogram locals: Unified under -gpu=unified, Managed under
-! -gpu=managed.
-! UNI: hlfir.declare {{.*}} {data_attr = #cuf.cuda<unified>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMmFmod_subEloc_arr"}
-! UNI: hlfir.declare {{.*}} {data_attr = #cuf.cuda<unified>, fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QMmFmod_subEloc_ptr"}
-! MAN: hlfir.declare {{.*}} {data_attr = #cuf.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMmFmod_subEloc_arr"}
-! MAN: hlfir.declare {{.*}} {data_attr = #cuf.cuda<managed>, fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QMmFmod_subEloc_ptr"}
-
-! COMMON-block pointer: unattributed, CUDA attributes may not appear in COMMON.
-! UNI: hlfir.declare {{.*}} {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QMmFcommon_subEcom_ptr"}
-! MAN: hlfir.declare {{.*}} {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QMmFcommon_subEcom_ptr"}
-
-! Main-program local: Unified under -gpu=unified, Managed under
-! -gpu=managed.
-! UNI: hlfir.declare {{.*}} {data_attr = #cuf.cuda<unified>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFEprog_arr"}
-! MAN: hlfir.declare {{.*}} {data_attr = #cuf.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFEprog_arr"}
-
-! Module-scope globals: Managed in both modes (ATTRIBUTES(UNIFIED) is not legal
-! at module scope, so the Unified case falls back to Managed).
-! UNI: fir.global @_QMmEmod_arr {data_attr = #cuf.cuda<managed>}
-! UNI: fir.global @_QMmEmod_ptr {data_attr = #cuf.cuda<managed>}
-! MAN: fir.global @_QMmEmod_arr {data_attr = #cuf.cuda<managed>}
-! MAN: fir.global @_QMmEmod_ptr {data_attr = #cuf.cuda<managed>}
-
-! Derived-type components: Unified under -gpu=unified, Managed under
-! -gpu=managed. Checked via the symbol table.
-! UNI-SYM: comp_arr, ALLOCATABLE {{.*}} cudaDataAttr: Unified
-! UNI-SYM: comp_ptr, POINTER {{.*}} cudaDataAttr: Unified
-! MAN-SYM: comp_arr, ALLOCATABLE {{.*}} cudaDataAttr: Managed
-! MAN-SYM: comp_ptr, POINTER {{.*}} cudaDataAttr: Managed
>From 7153184502dfae978745c15bbfdd9859886b913c Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Thu, 16 Jul 2026 14:02:58 -0700
Subject: [PATCH 2/2] [flang][cuda] Bake -gpu=unified managed allocator index
into the descriptor at creation
---
flang/include/flang/Lower/CUDA.h | 16 ++++++++++
flang/lib/Lower/ConvertVariable.cpp | 35 +++++++++++-----------
flang/test/Lower/CUDA/cuda-gpu-unified.cuf | 25 ++++++++++++++--
3 files changed, 57 insertions(+), 19 deletions(-)
diff --git a/flang/include/flang/Lower/CUDA.h b/flang/include/flang/Lower/CUDA.h
index eb9b311177fbd..0c71e6f4afafe 100644
--- a/flang/include/flang/Lower/CUDA.h
+++ b/flang/include/flang/Lower/CUDA.h
@@ -49,6 +49,22 @@ static inline unsigned getAllocatorIdx(const Fortran::semantics::Symbol &sym) {
return kDefaultAllocator;
}
+// Under -gpu=unified, redirect a plain (unattributed)
+// allocatable/pointer from the default allocator to the unified allocator, so
+// its managed backing is carried by the descriptor and honored by every
+// allocation path (ALLOCATE, allocate-on-assignment, SOURCE=, ...). Objects
+// with an explicit CUDA data attribute keep their own allocator.
+static inline unsigned
+getAllocatorIdxForUnified(const Fortran::semantics::Symbol &sym,
+ bool unifiedEnabled) {
+ unsigned idx = getAllocatorIdx(sym);
+ if (unifiedEnabled && idx == kDefaultAllocator &&
+ (Fortran::semantics::IsAllocatable(sym) ||
+ Fortran::semantics::IsPointer(sym)))
+ return kUnifiedAllocatorPos;
+ return idx;
+}
+
mlir::Type gatherDeviceComponentCoordinatesAndType(
fir::FirOpBuilder &builder, mlir::Location loc,
const Fortran::semantics::Symbol &sym, fir::RecordType recTy,
diff --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp
index f971feb0f01df..8031eb24b338b 100644
--- a/flang/lib/Lower/ConvertVariable.cpp
+++ b/flang/lib/Lower/ConvertVariable.cpp
@@ -495,20 +495,6 @@ createGlobalInitialization(fir::FirOpBuilder &builder, fir::GlobalOp global,
builder.restoreInsertionPoint(insertPt);
}
-static unsigned getAllocatorIdxFromDataAttr(cuf::DataAttributeAttr dataAttr) {
- if (dataAttr) {
- if (dataAttr.getValue() == cuf::DataAttribute::Pinned)
- return kPinnedAllocatorPos;
- if (dataAttr.getValue() == cuf::DataAttribute::Device)
- return kDeviceAllocatorPos;
- if (dataAttr.getValue() == cuf::DataAttribute::Managed)
- return kManagedAllocatorPos;
- if (dataAttr.getValue() == cuf::DataAttribute::Unified)
- return kUnifiedAllocatorPos;
- }
- return kDefaultAllocator;
-}
-
/// Create the global op and its init if it has one
fir::GlobalOp Fortran::lower::defineGlobal(
Fortran::lower::AbstractConverter &converter,
@@ -564,12 +550,19 @@ fir::GlobalOp Fortran::lower::defineGlobal(
fir::HasValueOp::create(b, loc, box);
});
} else {
- // Create unallocated/disassociated descriptor if no explicit init
+ // Create unallocated/disassociated descriptor if no explicit init. Under
+ // -gpu=unified, back a plain module allocatable/pointer with managed
+ // memory by baking the unified allocator index into the descriptor.
+ const bool cudaUnifiedBacking =
+ converter.getFoldingContext().languageFeatures().IsEnabled(
+ Fortran::common::LanguageFeature::CudaUnified);
+ unsigned allocatorIdx =
+ Fortran::lower::getAllocatorIdxForUnified(sym, cudaUnifiedBacking);
createGlobalInitialization(builder, global, [&](fir::FirOpBuilder &b) {
mlir::Value box = fir::factory::createUnallocatedBox(
b, loc, symTy,
/*nonDeferredParams=*/{},
- /*typeSourceBox=*/{}, getAllocatorIdxFromDataAttr(dataAttr));
+ /*typeSourceBox=*/{}, allocatorIdx);
fir::HasValueOp::create(b, loc, box);
});
}
@@ -2327,9 +2320,17 @@ void Fortran::lower::mapSymbolAttributes(
TODO(loc,
"derived type allocatable or pointer with length parameters");
}
+ // Under -gpu=unified, bake the unified allocator index into the descriptor
+ // at creation so runtime-driven allocations (allocate-on-assignment,
+ // SOURCE=, ...) get managed backing, not just explicit ALLOCATE.
+ const bool cudaUnifiedBacking =
+ converter.getFoldingContext().languageFeatures().IsEnabled(
+ Fortran::common::LanguageFeature::CudaUnified) &&
+ !cuf::isCUDADeviceContext(builder.getRegion());
fir::MutableBoxValue box = Fortran::lower::createMutableBox(
converter, loc, var, boxAlloc, nonDeferredLenParams,
- Fortran::lower::getAllocatorIdx(var.getSymbol()));
+ Fortran::lower::getAllocatorIdxForUnified(var.getSymbol(),
+ cudaUnifiedBacking));
genAllocatableOrPointerDeclare(converter, symMap, var.getSymbol(), box,
replace);
return;
diff --git a/flang/test/Lower/CUDA/cuda-gpu-unified.cuf b/flang/test/Lower/CUDA/cuda-gpu-unified.cuf
index 166e9230f0106..fe72dadf98320 100644
--- a/flang/test/Lower/CUDA/cuda-gpu-unified.cuf
+++ b/flang/test/Lower/CUDA/cuda-gpu-unified.cuf
@@ -1,7 +1,10 @@
! RUN: bbc -emit-hlfir -fcuda -gpu=unified %s -o - | FileCheck %s
! Under -gpu=unified, plain (unattributed) allocatables and pointers are backed
-! by CUDA managed memory at the ALLOCATE site (unified allocator index, idx 4).
+! by CUDA managed memory: the unified allocator index (idx 4) is baked into the
+! descriptor at creation (declaration / module global) and at the ALLOCATE site,
+! so runtime-driven allocations (allocate-on-assignment, SOURCE=, ...) also use
+! managed memory.
module m
real, allocatable :: mod_arr(:)
@@ -32,6 +35,14 @@ contains
allocate(com_ptr(100))
deallocate(com_ptr)
end subroutine
+ ! Allocate-on-assignment: assign_arr is never explicitly ALLOCATEd, so the
+ ! unified allocator index must be baked into its descriptor at declaration for
+ ! the runtime realloc to use managed memory.
+ subroutine alloc_on_assign(src)
+ real :: src(:)
+ real, allocatable :: assign_arr(:)
+ assign_arr = src
+ end subroutine
end module
program p
@@ -73,10 +84,20 @@ end program
! CHECK: hlfir.declare {{.*}} {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QMmFcommon_subEcom_ptr"}
! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
+! Allocate-on-assignment: the descriptor is created with the unified allocator
+! index at declaration, so the runtime realloc for the assignment gets managed
+! memory even though there is no explicit ALLOCATE.
+! CHECK-LABEL: func.func @_QMmPalloc_on_assign(
+! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
+! CHECK: hlfir.declare {{.*}} uniq_name = "_QMmFalloc_on_assignEassign_arr"
+! CHECK: hlfir.assign {{.*}} realloc
+
! Main-program local allocatable.
! CHECK-LABEL: func.func @_QQmain()
! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
-! The module global itself carries no CUDA data attribute.
+! The module global carries no CUDA data attribute, but its descriptor is created
+! with the unified allocator index so runtime allocations use managed memory.
! CHECK: fir.global @_QMmEmod_arr : !fir.box<!fir.heap<!fir.array<?xf32>>>
! CHECK-NOT: fir.global @_QMmEmod_arr {data_attr
+! CHECK: fir.embox {{.*}} {allocator_idx = 4 : i32}
More information about the flang-commits
mailing list