[flang-commits] [flang] 4c1dbbe - [flang][cuda] Make CUFRegisterAllocator callable from C/Fortran (#102543)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-08-08T17:09:53-07:00
New Revision: 4c1dbbe7aaeb1cb3f991e1de9c7d0dd312e565f5

URL: https://github.com/llvm/llvm-project/commit/4c1dbbe7aaeb1cb3f991e1de9c7d0dd312e565f5
DIFF: https://github.com/llvm/llvm-project/commit/4c1dbbe7aaeb1cb3f991e1de9c7d0dd312e565f5.diff

LOG: [flang][cuda] Make CUFRegisterAllocator callable from C/Fortran (#102543)

Added: 
    

Modified: 
    flang/include/flang/Runtime/CUDA/allocator.h
    flang/runtime/CUDA/allocator.cpp
    flang/unittests/Runtime/CUDA/AllocatorCUF.cpp

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h
index 8f5204769d7aa..f0bfc1548e645 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 cd00d40361d28..bd657b800c61e 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 4f53e654034cb..9f5ec289ee8f7 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)};