diff --git a/flang/include/flang/Lower/Cuda.h b/flang/include/flang/Lower/Cuda.h index e446b505ee6e4..d97045383d195 100644 --- a/flang/include/flang/Lower/Cuda.h +++ b/flang/include/flang/Lower/Cuda.h @@ -25,7 +25,7 @@ namespace Fortran::lower { // for it. // If the insertion point is inside an OpenACC region op, it is considered // device context. -static bool isCudaDeviceContext(fir::FirOpBuilder &builder) { +static bool inline isCudaDeviceContext(fir::FirOpBuilder &builder) { if (builder.getRegion().getParentOfType()) return true; if (builder.getRegion() @@ -41,6 +41,23 @@ static bool isCudaDeviceContext(fir::FirOpBuilder &builder) { } return false; } + +static inline unsigned getAllocatorIdx(const Fortran::semantics::Symbol &sym) { + std::optional cudaAttr = + Fortran::semantics::GetCUDADataAttr(&sym.GetUltimate()); + if (cudaAttr) { + if (*cudaAttr == Fortran::common::CUDADataAttr::Pinned) + return kPinnedAllocatorPos; + if (*cudaAttr == Fortran::common::CUDADataAttr::Device) + return kDeviceAllocatorPos; + if (*cudaAttr == Fortran::common::CUDADataAttr::Managed) + return kManagedAllocatorPos; + if (*cudaAttr == Fortran::common::CUDADataAttr::Unified) + return kUnifiedAllocatorPos; + } + return kDefaultAllocator; +} + } // end namespace Fortran::lower #endif // FORTRAN_LOWER_CUDA_H diff --git a/flang/lib/Lower/Allocatable.cpp b/flang/lib/Lower/Allocatable.cpp index dc135543fafc7..3d21e7a3fa8d5 100644 --- a/flang/lib/Lower/Allocatable.cpp +++ b/flang/lib/Lower/Allocatable.cpp @@ -475,6 +475,7 @@ class AllocateStmtHelper { !alloc.type.IsPolymorphic() && !alloc.hasCoarraySpec() && !useAllocateRuntime && !box.isPointer(); + unsigned allocatorIdx = Fortran::lower::getAllocatorIdx(alloc.getSymbol()); if (inlineAllocation && ((isCudaSymbol && isCudaDeviceContext) || !isCudaSymbol)) { @@ -488,7 +489,7 @@ class AllocateStmtHelper { // Generate a sequence of runtime calls. errorManager.genStatCheck(builder, loc); - genAllocateObjectInit(box); + genAllocateObjectInit(box, allocatorIdx); if (alloc.hasCoarraySpec()) TODO(loc, "coarray: allocation of a coarray object"); if (alloc.type.IsPolymorphic()) @@ -549,14 +550,16 @@ class AllocateStmtHelper { TODO(loc, "derived type length parameters in allocate"); } - void genAllocateObjectInit(const fir::MutableBoxValue &box) { + void genAllocateObjectInit(const fir::MutableBoxValue &box, + unsigned allocatorIdx) { if (box.isPointer()) { // For pointers, the descriptor may still be uninitialized (see Fortran // 2018 19.5.2.2). The allocation runtime needs to be given a descriptor // with initialized rank, types and attributes. Initialize the descriptor // here to ensure these constraints are fulfilled. mlir::Value nullPointer = fir::factory::createUnallocatedBox( - builder, loc, box.getBoxTy(), box.nonDeferredLenParams()); + builder, loc, box.getBoxTy(), box.nonDeferredLenParams(), + /*typeSourceBox=*/{}, allocatorIdx); builder.create(loc, nullPointer, box.getAddr()); } else { assert(box.isAllocatable() && "must be an allocatable"); @@ -612,11 +615,12 @@ class AllocateStmtHelper { void genSourceMoldAllocation(const Allocation &alloc, const fir::MutableBoxValue &box, bool isSource) { + unsigned allocatorIdx = Fortran::lower::getAllocatorIdx(alloc.getSymbol()); fir::ExtendedValue exv = isSource ? sourceExv : moldExv; - ; + // Generate a sequence of runtime calls. errorManager.genStatCheck(builder, loc); - genAllocateObjectInit(box); + genAllocateObjectInit(box, allocatorIdx); if (alloc.hasCoarraySpec()) TODO(loc, "coarray: allocation of a coarray object"); // Set length of the allocate object if it has. Otherwise, get the length diff --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp index 3cf2988a8b48c..8a7be19f8040b 100644 --- a/flang/lib/Lower/ConvertVariable.cpp +++ b/flang/lib/Lower/ConvertVariable.cpp @@ -19,6 +19,7 @@ #include "flang/Lower/ConvertExpr.h" #include "flang/Lower/ConvertExprToHLFIR.h" #include "flang/Lower/ConvertProcedureDesignator.h" +#include "flang/Lower/Cuda.h" #include "flang/Lower/Mangler.h" #include "flang/Lower/PFTBuilder.h" #include "flang/Lower/StatementContext.h" @@ -1985,22 +1986,6 @@ static void genBoxDeclare(Fortran::lower::AbstractConverter &converter, replace); } -static unsigned getAllocatorIdx(const Fortran::semantics::Symbol &sym) { - std::optional cudaAttr = - Fortran::semantics::GetCUDADataAttr(&sym.GetUltimate()); - if (cudaAttr) { - if (*cudaAttr == Fortran::common::CUDADataAttr::Pinned) - return kPinnedAllocatorPos; - if (*cudaAttr == Fortran::common::CUDADataAttr::Device) - return kDeviceAllocatorPos; - if (*cudaAttr == Fortran::common::CUDADataAttr::Managed) - return kManagedAllocatorPos; - if (*cudaAttr == Fortran::common::CUDADataAttr::Unified) - return kUnifiedAllocatorPos; - } - return kDefaultAllocator; -} - /// Lower specification expressions and attributes of variable \p var and /// add it to the symbol map. For a global or an alias, the address must be /// pre-computed and provided in \p preAlloc. A dummy argument for the current @@ -2091,7 +2076,7 @@ void Fortran::lower::mapSymbolAttributes( converter, loc, var, boxAlloc, nonDeferredLenParams, /*alwaysUseBox=*/ converter.getLoweringOptions().getLowerToHighLevelFIR(), - getAllocatorIdx(var.getSymbol())); + Fortran::lower::getAllocatorIdx(var.getSymbol())); genAllocatableOrPointerDeclare(converter, symMap, var.getSymbol(), box, replace); return; diff --git a/flang/test/Lower/CUDA/cuda-pointer.cuf b/flang/test/Lower/CUDA/cuda-pointer.cuf new file mode 100644 index 0000000000000..2a9dbe54c2922 --- /dev/null +++ b/flang/test/Lower/CUDA/cuda-pointer.cuf @@ -0,0 +1,11 @@ +! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s + +! Test lowering of CUDA pointers. + +subroutine allocate_pointer + real, device, pointer :: pr(:) + allocate(pr(10)) +end + +! CHECK-LABEL: func.func @_QPallocate_pointer() +! CHECK-COUNT-2: fir.embox %{{.*}} {allocator_idx = 2 : i32} : (!fir.ptr>, !fir.shape<1>) -> !fir.box>>