diff --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp index 70fa32d621e2f1..f76d44f5479d32 100644 --- a/flang/lib/Lower/ConvertVariable.cpp +++ b/flang/lib/Lower/ConvertVariable.cpp @@ -478,6 +478,20 @@ void Fortran::lower::createGlobalInitialization( builder.restoreInsertionPoint(insertPt); } +static unsigned getAllocatorIdx(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 static fir::GlobalOp defineGlobal(Fortran::lower::AbstractConverter &converter, const Fortran::lower::pft::Variable &var, @@ -540,8 +554,10 @@ static fir::GlobalOp defineGlobal(Fortran::lower::AbstractConverter &converter, // Create unallocated/disassociated descriptor if no explicit init Fortran::lower::createGlobalInitialization( builder, global, [&](fir::FirOpBuilder &b) { - mlir::Value box = - fir::factory::createUnallocatedBox(b, loc, symTy, std::nullopt); + mlir::Value box = fir::factory::createUnallocatedBox( + b, loc, symTy, + /*nonDeferredParams=*/std::nullopt, + /*typeSourceBox=*/{}, getAllocatorIdx(dataAttr)); b.create(loc, box); }); } diff --git a/flang/test/Lower/CUDA/cuda-allocatable.cuf b/flang/test/Lower/CUDA/cuda-allocatable.cuf index fb72f88fe415ca..6479425c58d8be 100644 --- a/flang/test/Lower/CUDA/cuda-allocatable.cuf +++ b/flang/test/Lower/CUDA/cuda-allocatable.cuf @@ -2,6 +2,21 @@ ! Test lowering of CUDA allocatable allocate/deallocate statements. +module globals + real, device, allocatable :: a_device(:) + real, managed, allocatable :: a_managed(:) + real, pinned, allocatable :: a_pinned(:) +end module + +! CHECK-LABEL: fir.global @_QMglobalsEa_device {data_attr = #cuf.cuda} : !fir.box>> +! CHECK: %{{.*}} = fir.embox %{{.*}}(%{{.*}}) {allocator_idx = 2 : i32} : (!fir.heap>, !fir.shape<1>) -> !fir.box>> + +! CHECK-LABEL: fir.global @_QMglobalsEa_managed {data_attr = #cuf.cuda} : !fir.box>> +! CHECK: %{{.*}} = fir.embox %{{.*}}(%{{.*}}) {allocator_idx = 3 : i32} : (!fir.heap>, !fir.shape<1>) -> !fir.box>> + +! CHECK-LABEL: fir.global @_QMglobalsEa_pinned {data_attr = #cuf.cuda} : !fir.box>> +! CHECK: %{{.*}} = fir.embox %{{.*}}(%{{.*}}) {allocator_idx = 1 : i32} : (!fir.heap>, !fir.shape<1>) -> !fir.box>> + subroutine sub1() real, allocatable, device :: a(:) allocate(a(10))