[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 17:58:11 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/3] [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/3] [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;
>From e5ae572a93385280bb2df498b8a2d017855e14e4 Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Mon, 29 Jun 2026 17:57:16 -0700
Subject: [PATCH 3/3] [flang][cuda] Walk DataRef chain to propagate CUDA attrs
to component allocs
---
flang/lib/Lower/Allocatable.cpp | 53 ++++++++++++++-----
.../Lower/CUDA/cuda-allocatable-component.cuf | 2 +-
flang/test/Lower/CUDA/cuda-data-transfer.cuf | 2 +-
3 files changed, 42 insertions(+), 15 deletions(-)
diff --git a/flang/lib/Lower/Allocatable.cpp b/flang/lib/Lower/Allocatable.cpp
index 3b07fac8fe15b..2b54dc29b1e14 100644
--- a/flang/lib/Lower/Allocatable.cpp
+++ b/flang/lib/Lower/Allocatable.cpp
@@ -459,22 +459,49 @@ 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.
+ /// Search a DataRef for a symbol with CUDA attributes. Returns true and
+ /// sets \p sym if found.
+ static bool findCUDAAttrInDataRef(const Fortran::parser::DataRef &ref,
+ const Fortran::semantics::Symbol *&sym) {
+ return Fortran::common::visit(
+ Fortran::common::visitors{
+ [&](const Fortran::parser::Name &name) {
+ if (name.symbol &&
+ Fortran::semantics::HasCUDAAttr(*name.symbol)) {
+ sym = name.symbol;
+ return true;
+ }
+ return false;
+ },
+ [&](const Fortran::common::Indirection<
+ Fortran::parser::StructureComponent> &sc) {
+ if (sc.value().Component().symbol &&
+ Fortran::semantics::HasCUDAAttr(
+ *sc.value().Component().symbol)) {
+ sym = sc.value().Component().symbol;
+ return true;
+ }
+ return findCUDAAttrInDataRef(sc.value().Base(), sym);
+ },
+ [&](const Fortran::common::Indirection<
+ Fortran::parser::ArrayElement> &ae) {
+ return findCUDAAttrInDataRef(ae.value().Base(), sym);
+ },
+ [&](const Fortran::common::Indirection<
+ Fortran::parser::CoindexedNamedObject> &) { return false; },
+ },
+ ref.u);
+ }
+
+ /// When the allocate object is a structure component (e.g.,
+ /// managed_var(1)%component or a%b%c), walk the DataRef chain and check
+ /// whether any parent has CUDA attributes. If so, update \p sym to that
+ /// symbol and return true so allocations use its 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)) {
- sym = baseName.symbol;
- return true;
- }
- }
+ &alloc.getAllocObj().u))
+ return findCUDAAttrInDataRef(sc->Base(), sym);
return false;
}
diff --git a/flang/test/Lower/CUDA/cuda-allocatable-component.cuf b/flang/test/Lower/CUDA/cuda-allocatable-component.cuf
index 80aa2b7255105..a5f56f157efad 100644
--- a/flang/test/Lower/CUDA/cuda-allocatable-component.cuf
+++ b/flang/test/Lower/CUDA/cuda-allocatable-component.cuf
@@ -41,7 +41,7 @@ end subroutine
! 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.
+! Managed component with source allocation.
subroutine managed_component_source()
use m
type(chunk_type), managed, allocatable :: chunks(:)
diff --git a/flang/test/Lower/CUDA/cuda-data-transfer.cuf b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
index 45ce416c084ab..f1cd5bae81419 100644
--- a/flang/test/Lower/CUDA/cuda-data-transfer.cuf
+++ b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
@@ -576,7 +576,7 @@ subroutine sub29()
end subroutine
! CHECK-LABEL: func.func @_QPsub29()
-! CHECK: %[[TMP:.*]] = fir.allocmem !fir.array<?xf16>, %24#1 {bindc_name = ".tmp", uniq_name = ""}
+! CHECK: %[[TMP:.*]] = fir.allocmem !fir.array<?xf16>, %{{.*}}#1 {bindc_name = ".tmp", uniq_name = ""}
! CHECK: %[[TMP_BUFFER:.*]]:2 = hlfir.declare %[[TMP]](%{{.*}}) {uniq_name = ".tmp"} : (!fir.heap<!fir.array<?xf16>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xf16>>, !fir.heap<!fir.array<?xf16>>)
! CHECK: cuf.data_transfer %{{.*}} to %[[TMP_BUFFER]]#0 {transfer_kind = #cuf.cuda_transfer<device_host>} : !fir.box<!fir.heap<!fir.array<?xf16>>>, !fir.box<!fir.array<?xf16>>
! CHECK: hlfir.elemental
More information about the flang-commits
mailing list