[flang][cuda] Make sure allocator id is set for pointer allocate

Author: clementval

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<cuf::KernelOp>())
     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<Fortran::common::CUDADataAttr> 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

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,15 @@ 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<fir::StoreOp>(loc, nullPointer, box.getAddr());
     } else {
       assert(box.isAllocatable() && "must be an allocatable");
@@ -612,11 +614,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

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<Fortran::common::CUDADataAttr> 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;