[flang][cuda] Make CUFRegisterAllocator callable from C/Fortran

[clementval]

The registration of the CUDA Fortran allocators is meant to be done from C code or Fortran code. Use extern "C" so we can do so. Also use the macro to make the name coherent with the rest of the runtime calls.

[llvmbot]

@llvm/pr-subscribers-flang-runtime

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

Changes

The registration of the CUDA Fortran allocators is meant to be done from C code or Fortran code. Use extern "C" so we can do so. Also use the macro to make the name coherent with the rest of the runtime calls.

---

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

3 Files Affected:

• (modified) flang/include/flang/Runtime/CUDA/allocator.h (+5-1)
• (modified) flang/runtime/CUDA/allocator.cpp (+3-1)
• (modified) flang/unittests/Runtime/CUDA/AllocatorCUF.cpp (+3-3)

diff --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h
index 8f5204769d7aa7..f0bfc1548e6458 100644
--- a/flang/include/flang/Runtime/CUDA/allocator.h
+++ b/flang/include/flang/Runtime/CUDA/allocator.h
@@ -10,6 +10,7 @@
 #define FORTRAN_RUNTIME_CUDA_ALLOCATOR_H_
 
 #include "flang/Runtime/descriptor.h"
+#include "flang/Runtime/entry-names.h"
 
 #define CUDA_REPORT_IF_ERROR(expr) \
   [](CUresult result) { \
@@ -25,7 +26,10 @@
 
 namespace Fortran::runtime::cuda {
 
-void CUFRegisterAllocator();
+extern "C" {
+
+void RTDECL(CUFRegisterAllocator)();
+}
 
 void *CUFAllocPinned(std::size_t);
 void CUFFreePinned(void *);
diff --git a/flang/runtime/CUDA/allocator.cpp b/flang/runtime/CUDA/allocator.cpp
index cd00d40361d28b..bd657b800c61e8 100644
--- a/flang/runtime/CUDA/allocator.cpp
+++ b/flang/runtime/CUDA/allocator.cpp
@@ -18,8 +18,9 @@
 #include "cuda.h"
 
 namespace Fortran::runtime::cuda {
+extern "C" {
 
-void CUFRegisterAllocator() {
+void RTDEF(CUFRegisterAllocator)() {
   allocatorRegistry.Register(
       kPinnedAllocatorPos, {&CUFAllocPinned, CUFFreePinned});
   allocatorRegistry.Register(
@@ -29,6 +30,7 @@ void CUFRegisterAllocator() {
   allocatorRegistry.Register(
       kUnifiedAllocatorPos, {&CUFAllocUnified, CUFFreeUnified});
 }
+}
 
 void *CUFAllocPinned(std::size_t sizeInBytes) {
   void *p;
diff --git a/flang/unittests/Runtime/CUDA/AllocatorCUF.cpp b/flang/unittests/Runtime/CUDA/AllocatorCUF.cpp
index 4f53e654034cb7..9f5ec289ee8f74 100644
--- a/flang/unittests/Runtime/CUDA/AllocatorCUF.cpp
+++ b/flang/unittests/Runtime/CUDA/AllocatorCUF.cpp
@@ -55,7 +55,7 @@ class ScopedContext {
 
 TEST(AllocatableCUFTest, SimpleDeviceAllocate) {
   using Fortran::common::TypeCategory;
-  Fortran::runtime::cuda::CUFRegisterAllocator();
+  RTNAME(CUFRegisterAllocator)();
   ScopedContext ctx;
   // REAL(4), DEVICE, ALLOCATABLE :: a(:)
   auto a{createAllocatable(TypeCategory::Real, 4)};
@@ -73,7 +73,7 @@ TEST(AllocatableCUFTest, SimpleDeviceAllocate) {
 
 TEST(AllocatableCUFTest, SimplePinnedAllocate) {
   using Fortran::common::TypeCategory;
-  Fortran::runtime::cuda::CUFRegisterAllocator();
+  RTNAME(CUFRegisterAllocator)();
   ScopedContext ctx;
   // INTEGER(4), PINNED, ALLOCATABLE :: a(:)
   auto a{createAllocatable(TypeCategory::Integer, 4)};
@@ -92,7 +92,7 @@ TEST(AllocatableCUFTest, SimplePinnedAllocate) {
 
 TEST(AllocatableCUFTest, DescriptorAllocationTest) {
   using Fortran::common::TypeCategory;
-  Fortran::runtime::cuda::CUFRegisterAllocator();
+  RTNAME(CUFRegisterAllocator)();
   ScopedContext ctx;
   // REAL(4), DEVICE, ALLOCATABLE :: a(:)
   auto a{createAllocatable(TypeCategory::Real, 4)};