[flang-commits] [flang] 5c016bf - [flang][cuda] Force default allocator in device code (#102238)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-08-09T09:18:34-07:00
New Revision: 5c016bf40ea38bc0198b53b1eff743d305e16707

URL: https://github.com/llvm/llvm-project/commit/5c016bf40ea38bc0198b53b1eff743d305e16707
DIFF: https://github.com/llvm/llvm-project/commit/5c016bf40ea38bc0198b53b1eff743d305e16707.diff

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

Added: 
    

Modified: 
    flang/runtime/descriptor.cpp

Removed: 
    


################################################################################
diff  --git a/flang/runtime/descriptor.cpp b/flang/runtime/descriptor.cpp
index 34f7a02ea8c7bd..32f43e89dc7a36 100644
--- a/flang/runtime/descriptor.cpp
+++ b/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;