[flang][cuda] Force default allocator in device code (#102238)

Author: clementval

flang/runtime/descriptor.cpp

@@ -154,6 +154,15 @@ RT_API_ATTRS std::size_t Descriptor::Elements() const {
   return elements;
 }
 
+RT_API_ATTRS static inline int MapAllocIdx(const Descriptor &desc) {
+#ifdef RT_DEVICE_COMPILATION
+  // Force default allocator in device code.
+  return kDefaultAllocator;
+#else
+  return desc.GetAllocIdx();
+#endif
+}
+
 RT_API_ATTRS int Descriptor::Allocate() {
   std::size_t elementBytes{ElementBytes()};
   if (static_cast<std::int64_t>(elementBytes) < 0) {
@@ -162,11 +171,10 @@ RT_API_ATTRS int Descriptor::Allocate() {
     elementBytes = raw_.elem_len = 0;
   }
   std::size_t byteSize{Elements() * elementBytes};
+  AllocFct alloc{allocatorRegistry.GetAllocator(MapAllocIdx(*this))};
   // Zero size allocation is possible in Fortran and the resulting
   // descriptor must be allocated/associated. Since std::malloc(0)
   // result is implementation defined, always allocate at least one byte.
-
-  AllocFct alloc{allocatorRegistry.GetAllocator(GetAllocIdx())};
   void *p{alloc(byteSize ? byteSize : 1)};
   if (!p) {
     return CFI_ERROR_MEM_ALLOCATION;
@@ -209,7 +217,7 @@ RT_API_ATTRS int Descriptor::Deallocate() {
   if (!descriptor.base_addr) {
     return CFI_ERROR_BASE_ADDR_NULL;
   } else {
-    FreeFct free{allocatorRegistry.GetDeallocator(GetAllocIdx())};
+    FreeFct free{allocatorRegistry.GetDeallocator(MapAllocIdx(*this))};
     free(descriptor.base_addr);
     descriptor.base_addr = nullptr;
     return CFI_SUCCESS;