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

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-runtime

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

<details>
<summary>Changes</summary>

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
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;

``````````

</details>


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