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

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/102238

Allocation/deallocation from device code can use standard malloc/free since the cuda functions cannot be called. Force the default allocator/deallocator in the device runtime build so we don't need to register extra allocators. 

>From 89ba16b84a2142da229bc9e887e62c8604d5d743 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 6 Aug 2024 15:04:53 -0700
Subject: [PATCH] [flang][cuda] Force default allocator in device code

---
 flang/runtime/descriptor.cpp | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/flang/runtime/descriptor.cpp b/flang/runtime/descriptor.cpp
index 34f7a02ea8c7b..6d43bacaed697 100644
--- a/flang/runtime/descriptor.cpp
+++ b/flang/runtime/descriptor.cpp
@@ -162,11 +162,17 @@ RT_API_ATTRS int Descriptor::Allocate() {
     elementBytes = raw_.elem_len = 0;
   }
   std::size_t byteSize{Elements() * elementBytes};
+
+  // Force default allocator in device code.
+#ifdef RT_DEVICE_COMPILATION
+  AllocFct alloc{allocatorRegistry.GetAllocator(kDefaultAllocator)};
+#else
+  AllocFct alloc{allocatorRegistry.GetAllocator(GetAllocIdx())};
+#endif
+
   // 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 +215,12 @@ RT_API_ATTRS int Descriptor::Deallocate() {
   if (!descriptor.base_addr) {
     return CFI_ERROR_BASE_ADDR_NULL;
   } else {
+    // Force default deallocator in device code.
+#ifdef RT_DEVICE_COMPILATION
+    FreeFct free{allocatorRegistry.GetDeallocator(kDefaultAllocator)};
+#else
     FreeFct free{allocatorRegistry.GetDeallocator(GetAllocIdx())};
+#endif
     free(descriptor.base_addr);
     descriptor.base_addr = nullptr;
     return CFI_SUCCESS;