[flang-commits] [flang] [flang][cuda] Propagate CUDA attrs from parent variable to component allocations (PR #206614)
Vijay Kandiah via flang-commits
flang-commits at lists.llvm.org
Mon Jun 29 16:50:19 PDT 2026
https://github.com/VijayKandiah updated https://github.com/llvm/llvm-project/pull/206614
>From 3e0da1fcbf93318688fb0180bde2928fca6e9cc5 Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Mon, 29 Jun 2026 16:31:57 -0700
Subject: [PATCH 1/2] [flang][cuda] Propagate CUDA attrs from parent variable
to component allocations
---
flang/lib/Lower/Allocatable.cpp | 57 ++++++++++++++----
.../Lower/CUDA/cuda-allocatable-component.cuf | 59 +++++++++++++++++++
2 files changed, 104 insertions(+), 12 deletions(-)
create mode 100644 flang/test/Lower/CUDA/cuda-allocatable-component.cuf
diff --git a/flang/lib/Lower/Allocatable.cpp b/flang/lib/Lower/Allocatable.cpp
index 3d2eea6a23830..3ed4dcb2e9d1e 100644
--- a/flang/lib/Lower/Allocatable.cpp
+++ b/flang/lib/Lower/Allocatable.cpp
@@ -459,17 +459,41 @@ class AllocateStmtHelper {
fir::StoreOp::create(builder, loc, falseConv, pinned);
}
+ /// When the allocate object is a structure component
+ /// (e.g., managed_var(1)%component), check whether the base variable has
+ /// CUDA attributes. If so, update \p sym to the base symbol and return true
+ /// so allocations use the parent's memory space.
+ bool propagateCUDAAttrsFromParent(const Allocation &alloc,
+ const Fortran::semantics::Symbol *&sym) {
+ if (const auto *sc = std::get_if<Fortran::parser::StructureComponent>(
+ &alloc.getAllocObj().u)) {
+ const Fortran::parser::Name &baseName =
+ Fortran::parser::GetFirstName(*sc);
+ if (baseName.symbol &&
+ (Fortran::semantics::HasCUDAAttr(*baseName.symbol) ||
+ Fortran::semantics::HasCUDAComponent(*baseName.symbol))) {
+ sym = baseName.symbol;
+ return true;
+ }
+ }
+ return false;
+ }
+
void genSimpleAllocation(const Allocation &alloc,
const fir::MutableBoxValue &box) {
bool isCudaAllocate =
Fortran::semantics::HasCUDAAttr(alloc.getSymbol()) ||
Fortran::semantics::HasCUDAComponent(alloc.getSymbol());
+ const Fortran::semantics::Symbol *cudaSymForAlloc = &alloc.getSymbol();
+ if (!isCudaAllocate)
+ isCudaAllocate = propagateCUDAAttrsFromParent(alloc, cudaSymForAlloc);
+
bool isCudaDeviceContext = cuf::isCUDADeviceContext(builder.getRegion());
bool inlineAllocation = !box.isDerived() && !errorManager.hasStatSpec() &&
!alloc.type.IsPolymorphic() &&
!alloc.hasCoarraySpec() && !useAllocateRuntime &&
!box.isPointer();
- unsigned allocatorIdx = Fortran::lower::getAllocatorIdx(alloc.getSymbol());
+ unsigned allocatorIdx = Fortran::lower::getAllocatorIdx(*cudaSymForAlloc);
if (inlineAllocation && !alloc.hasCoarraySpec() &&
((isCudaAllocate && isCudaDeviceContext) || !isCudaAllocate)) {
@@ -508,8 +532,7 @@ class AllocateStmtHelper {
stat = genRuntimeAllocate(builder, loc, box, errorManager);
setPinnedToFalse();
} else {
- stat =
- genCudaAllocate(builder, loc, box, errorManager, alloc.getSymbol());
+ stat = genCudaAllocate(builder, loc, box, errorManager, *cudaSymForAlloc);
}
fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
postAllocationAction(alloc, box);
@@ -570,8 +593,16 @@ class AllocateStmtHelper {
fir::StoreOp::create(builder, loc, nullPointer, box.getAddr());
} else {
assert(box.isAllocatable() && "must be an allocatable");
- // For allocatables, sync the MutableBoxValue and descriptor before the
- // calls in case it is tracked locally by a set of variables.
+ // For allocatables, re-establish the descriptor with the correct
+ // allocator index so that AllocatableAllocate uses the right memory
+ // space.
+ if (allocatorIdx != kDefaultAllocator) {
+ fir::factory::disassociateMutableBox(builder, loc, box,
+ /*polymorphicSetType=*/false,
+ allocatorIdx);
+ }
+ // Sync the MutableBoxValue and descriptor before the calls in case
+ // it is tracked locally by a set of variables.
fir::factory::getMutableIRBox(builder, loc, box);
}
}
@@ -623,11 +654,15 @@ class AllocateStmtHelper {
void genSourceMoldAllocation(const Allocation &alloc,
const fir::MutableBoxValue &box, bool isSource) {
- unsigned allocatorIdx = Fortran::lower::getAllocatorIdx(alloc.getSymbol());
+ bool isCudaAllocate = Fortran::semantics::HasCUDAAttr(alloc.getSymbol());
+ const Fortran::semantics::Symbol *cudaSymForAlloc = &alloc.getSymbol();
+ if (!isCudaAllocate)
+ isCudaAllocate = propagateCUDAAttrsFromParent(alloc, cudaSymForAlloc);
+ unsigned allocatorIdx = Fortran::lower::getAllocatorIdx(*cudaSymForAlloc);
fir::ExtendedValue exv = isSource ? sourceExv : moldExv;
bool sourceIsDevice = false;
- if (const Fortran::semantics::Symbol *sym{GetLastSymbol(sourceExpr)})
+ if (const Fortran::semantics::Symbol * sym{GetLastSymbol(sourceExpr)})
if (Fortran::semantics::IsCUDADevice(*sym))
sourceIsDevice = true;
@@ -652,10 +687,8 @@ class AllocateStmtHelper {
converter, loc, alloc.getSymbol(), box.getAddr(),
alloc.getCoarraySpec(), errorManager.errMsgAddr,
errorManager.hasStatSpec());
- } else if (Fortran::semantics::HasCUDAAttr(alloc.getSymbol()) ||
- sourceIsDevice) {
- stat =
- genCudaAllocate(builder, loc, box, errorManager, alloc.getSymbol());
+ } else if (isCudaAllocate || sourceIsDevice) {
+ stat = genCudaAllocate(builder, loc, box, errorManager, *cudaSymForAlloc);
} else {
if (isSource)
stat = genRuntimeAllocateSource(builder, loc, box, exv, errorManager);
@@ -792,7 +825,7 @@ class AllocateStmtHelper {
mlir::Type retTy = fir::runtime::getModel<int>()(builder.getContext());
bool isSourceDevice = false;
- if (const Fortran::semantics::Symbol *sym{GetLastSymbol(sourceExpr)})
+ if (const Fortran::semantics::Symbol * sym{GetLastSymbol(sourceExpr)})
if (Fortran::semantics::IsCUDADevice(*sym))
isSourceDevice = true;
diff --git a/flang/test/Lower/CUDA/cuda-allocatable-component.cuf b/flang/test/Lower/CUDA/cuda-allocatable-component.cuf
new file mode 100644
index 0000000000000..80aa2b7255105
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-allocatable-component.cuf
@@ -0,0 +1,59 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+
+! Test that allocating a component of a CUDA-attributed derived type variable
+! routes through the CUF allocator, not the default (host) allocator.
+! This ensures the component data lives in the same memory space as its parent.
+
+module m
+ type :: chunk_type
+ real(8), dimension(:,:), allocatable :: arr
+ end type
+end module
+
+! Explicit managed attribute on variable with allocatable component.
+subroutine managed_component()
+ use m
+ type(chunk_type), managed, allocatable :: chunks(:)
+ allocate(chunks(1))
+ allocate(chunks(1)%arr(10,10))
+ deallocate(chunks)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPmanaged_component()
+! CHECK: cuf.allocate %{{.*}} : {{.*}} {data_attr = #cuf.cuda<managed>} -> i32
+! The component descriptor must be re-established with the managed allocator.
+! CHECK: hlfir.designate %{{.*}}{"arr"} {{.*}} : {{.*}} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf64>>>>
+! CHECK: fir.embox %{{.*}} {allocator_idx = 3 : i32}
+! CHECK: cuf.allocate %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf64>>>> {data_attr = #cuf.cuda<managed>} -> i32
+
+! Device attribute on variable with allocatable component.
+subroutine device_component()
+ use m
+ type(chunk_type), device, allocatable :: chunks(:)
+ allocate(chunks(1))
+ allocate(chunks(1)%arr(10,10))
+ deallocate(chunks)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPdevice_component()
+! CHECK: cuf.allocate %{{.*}} : {{.*}} {data_attr = #cuf.cuda<device>} -> i32
+! CHECK: hlfir.designate %{{.*}}{"arr"} {{.*}} : {{.*}} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf64>>>>
+! CHECK: fir.embox %{{.*}} {allocator_idx = 2 : i32}
+! CHECK: cuf.allocate %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf64>>>> {data_attr = #cuf.cuda<device>} -> i32
+
+! SOURCE= allocation of a managed component.
+subroutine managed_component_source()
+ use m
+ type(chunk_type), managed, allocatable :: chunks(:)
+ real(8), allocatable :: src(:,:)
+ allocate(chunks(1))
+ allocate(src(10,10))
+ allocate(chunks(1)%arr, source=src)
+ deallocate(chunks)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPmanaged_component_source()
+! CHECK: cuf.allocate %{{.*}} : {{.*}} {data_attr = #cuf.cuda<managed>} -> i32
+! CHECK: hlfir.designate %{{.*}}{"arr"} {{.*}} : {{.*}} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf64>>>>
+! CHECK: fir.embox %{{.*}} {allocator_idx = 3 : i32}
+! CHECK: cuf.allocate %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf64>>>> source({{.*}}) {data_attr = #cuf.cuda<managed>} -> i32
>From 7014a7f6f4146f1537a908d378ad02177e247d4d Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Mon, 29 Jun 2026 16:50:07 -0700
Subject: [PATCH 2/2] [flang][cuda] Fix formatting and address feedback
---
flang/lib/Lower/Allocatable.cpp | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/flang/lib/Lower/Allocatable.cpp b/flang/lib/Lower/Allocatable.cpp
index 3ed4dcb2e9d1e..3b07fac8fe15b 100644
--- a/flang/lib/Lower/Allocatable.cpp
+++ b/flang/lib/Lower/Allocatable.cpp
@@ -470,8 +470,7 @@ class AllocateStmtHelper {
const Fortran::parser::Name &baseName =
Fortran::parser::GetFirstName(*sc);
if (baseName.symbol &&
- (Fortran::semantics::HasCUDAAttr(*baseName.symbol) ||
- Fortran::semantics::HasCUDAComponent(*baseName.symbol))) {
+ Fortran::semantics::HasCUDAAttr(*baseName.symbol)) {
sym = baseName.symbol;
return true;
}
@@ -662,7 +661,7 @@ class AllocateStmtHelper {
fir::ExtendedValue exv = isSource ? sourceExv : moldExv;
bool sourceIsDevice = false;
- if (const Fortran::semantics::Symbol * sym{GetLastSymbol(sourceExpr)})
+ if (const Fortran::semantics::Symbol *sym{GetLastSymbol(sourceExpr)})
if (Fortran::semantics::IsCUDADevice(*sym))
sourceIsDevice = true;
@@ -825,7 +824,7 @@ class AllocateStmtHelper {
mlir::Type retTy = fir::runtime::getModel<int>()(builder.getContext());
bool isSourceDevice = false;
- if (const Fortran::semantics::Symbol * sym{GetLastSymbol(sourceExpr)})
+ if (const Fortran::semantics::Symbol *sym{GetLastSymbol(sourceExpr)})
if (Fortran::semantics::IsCUDADevice(*sym))
isSourceDevice = true;
More information about the flang-commits
mailing list