[flang-commits] [flang] [flang][cuda] Add RT_API_ATTRS to allocator functions (PR #102235)

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

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

Allocators might be called from device code. Add RT_API_ATTRS so it can be build for the device. 

>From 41a9d122948266a55582dde3338cb73ed993adf4 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 6 Aug 2024 14:38:51 -0700
Subject: [PATCH] [flang][cuda] Add RT_API_ATTRS to allocator functions

---
 flang/include/flang/Runtime/CUDA/allocator.h | 18 +++++++++---------
 flang/runtime/CUDA/allocator.cpp             | 20 +++++++++++---------
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h
index 8f5204769d7aa..309582c296969 100644
--- a/flang/include/flang/Runtime/CUDA/allocator.h
+++ b/flang/include/flang/Runtime/CUDA/allocator.h
@@ -25,19 +25,19 @@
 
 namespace Fortran::runtime::cuda {
 
-void CUFRegisterAllocator();
+RT_API_ATTRS void CUFRegisterAllocator();
 
-void *CUFAllocPinned(std::size_t);
-void CUFFreePinned(void *);
+RT_API_ATTRS void *CUFAllocPinned(std::size_t);
+RT_API_ATTRS void CUFFreePinned(void *);
 
-void *CUFAllocDevice(std::size_t);
-void CUFFreeDevice(void *);
+RT_API_ATTRS void *CUFAllocDevice(std::size_t);
+RT_API_ATTRS void CUFFreeDevice(void *);
 
-void *CUFAllocManaged(std::size_t);
-void CUFFreeManaged(void *);
+RT_API_ATTRS void *CUFAllocManaged(std::size_t);
+RT_API_ATTRS void CUFFreeManaged(void *);
 
-void *CUFAllocUnified(std::size_t);
-void CUFFreeUnified(void *);
+RT_API_ATTRS void *CUFAllocUnified(std::size_t);
+RT_API_ATTRS void CUFFreeUnified(void *);
 
 } // namespace Fortran::runtime::cuda
 #endif // FORTRAN_RUNTIME_CUDA_ALLOCATOR_H_
diff --git a/flang/runtime/CUDA/allocator.cpp b/flang/runtime/CUDA/allocator.cpp
index cd00d40361d28..ff645e27f4f6d 100644
--- a/flang/runtime/CUDA/allocator.cpp
+++ b/flang/runtime/CUDA/allocator.cpp
@@ -19,7 +19,7 @@
 
 namespace Fortran::runtime::cuda {
 
-void CUFRegisterAllocator() {
+RT_API_ATTRS void CUFRegisterAllocator() {
   allocatorRegistry.Register(
       kPinnedAllocatorPos, {&CUFAllocPinned, CUFFreePinned});
   allocatorRegistry.Register(
@@ -30,41 +30,43 @@ void CUFRegisterAllocator() {
       kUnifiedAllocatorPos, {&CUFAllocUnified, CUFFreeUnified});
 }
 
-void *CUFAllocPinned(std::size_t sizeInBytes) {
+RT_API_ATTRS void *CUFAllocPinned(std::size_t sizeInBytes) {
   void *p;
   CUDA_REPORT_IF_ERROR(cuMemAllocHost(&p, sizeInBytes));
   return p;
 }
 
-void CUFFreePinned(void *p) { CUDA_REPORT_IF_ERROR(cuMemFreeHost(p)); }
+RT_API_ATTRS void CUFFreePinned(void *p) {
+  CUDA_REPORT_IF_ERROR(cuMemFreeHost(p));
+}
 
-void *CUFAllocDevice(std::size_t sizeInBytes) {
+RT_API_ATTRS void *CUFAllocDevice(std::size_t sizeInBytes) {
   CUdeviceptr p = 0;
   CUDA_REPORT_IF_ERROR(cuMemAlloc(&p, sizeInBytes));
   return reinterpret_cast<void *>(p);
 }
 
-void CUFFreeDevice(void *p) {
+RT_API_ATTRS void CUFFreeDevice(void *p) {
   CUDA_REPORT_IF_ERROR(cuMemFree(reinterpret_cast<CUdeviceptr>(p)));
 }
 
-void *CUFAllocManaged(std::size_t sizeInBytes) {
+RT_API_ATTRS void *CUFAllocManaged(std::size_t sizeInBytes) {
   CUdeviceptr p = 0;
   CUDA_REPORT_IF_ERROR(
       cuMemAllocManaged(&p, sizeInBytes, CU_MEM_ATTACH_GLOBAL));
   return reinterpret_cast<void *>(p);
 }
 
-void CUFFreeManaged(void *p) {
+RT_API_ATTRS void CUFFreeManaged(void *p) {
   CUDA_REPORT_IF_ERROR(cuMemFree(reinterpret_cast<CUdeviceptr>(p)));
 }
 
-void *CUFAllocUnified(std::size_t sizeInBytes) {
+RT_API_ATTRS void *CUFAllocUnified(std::size_t sizeInBytes) {
   // Call alloc managed for the time being.
   return CUFAllocManaged(sizeInBytes);
 }
 
-void CUFFreeUnified(void *p) {
+RT_API_ATTRS void CUFFreeUnified(void *p) {
   // Call free managed for the time being.
   CUFFreeManaged(p);
 }