[flang][cuda] Force default allocator in device code

[clementval]

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.

[llvmbot]

@llvm/pr-subscribers-flang-runtime

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/102238.diff

1 Files Affected:

• (modified) flang/runtime/descriptor.cpp (+13-2)

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;

[clementval]

@vzakhari @klausler I'm not a big fan of GetNormalizedAllocIdx name. Do you have a better idea.

I was thinking to rename GetAllocIdx to DecodeAllocIdx and use GetAllocIdx to have the ifdef inside.

[klausler]

Use two functions, GetAllocIdx and MapAllocIdx. The second one can be static to the file.

[clementval]

@klausler Does it look good for you?

[klausler]

Add inline, please.