[flang-commits] [flang] 388b632 - [flang][cuda] Defined allocator for unified data (#102189)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-08-06T14:30:31-07:00
New Revision: 388b63243c6846d3bb2e9f4d8cec228756ad1d1b

URL: https://github.com/llvm/llvm-project/commit/388b63243c6846d3bb2e9f4d8cec228756ad1d1b
DIFF: https://github.com/llvm/llvm-project/commit/388b63243c6846d3bb2e9f4d8cec228756ad1d1b.diff

LOG: [flang][cuda] Defined allocator for unified data (#102189)

CUDA unified variable where set to use the same allocator than managed
variable. This patch adds a specific allocator for the unified
variables. Currently it will call the managed allocator underneath but
we want to have the flexibility to change that in the future.

Added: 
    

Modified: 
    flang/include/flang/Runtime/CUDA/allocator.h
    flang/include/flang/Runtime/allocator-registry.h
    flang/lib/Lower/ConvertVariable.cpp
    flang/runtime/CUDA/allocator.cpp

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h
index 849785cf991ff..8f5204769d7aa 100644
--- a/flang/include/flang/Runtime/CUDA/allocator.h
+++ b/flang/include/flang/Runtime/CUDA/allocator.h
@@ -36,5 +36,8 @@ void CUFFreeDevice(void *);
 void *CUFAllocManaged(std::size_t);
 void CUFFreeManaged(void *);
 
+void *CUFAllocUnified(std::size_t);
+void CUFFreeUnified(void *);
+
 } // namespace Fortran::runtime::cuda
 #endif // FORTRAN_RUNTIME_CUDA_ALLOCATOR_H_

diff  --git a/flang/include/flang/Runtime/allocator-registry.h b/flang/include/flang/Runtime/allocator-registry.h
index 209b4d2e44e9b..acfada506fafc 100644
--- a/flang/include/flang/Runtime/allocator-registry.h
+++ b/flang/include/flang/Runtime/allocator-registry.h
@@ -19,8 +19,9 @@ static constexpr unsigned kDefaultAllocator = 0;
 static constexpr unsigned kPinnedAllocatorPos = 1;
 static constexpr unsigned kDeviceAllocatorPos = 2;
 static constexpr unsigned kManagedAllocatorPos = 3;
+static constexpr unsigned kUnifiedAllocatorPos = 4;
 
-#define MAX_ALLOCATOR 5
+#define MAX_ALLOCATOR 7 // 3 bits are reserved in the descriptor.
 
 namespace Fortran::runtime {
 

diff  --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp
index 45389091b8164..ffbbea238647c 100644
--- a/flang/lib/Lower/ConvertVariable.cpp
+++ b/flang/lib/Lower/ConvertVariable.cpp
@@ -1860,9 +1860,10 @@ static unsigned getAllocatorIdx(const Fortran::semantics::Symbol &sym) {
       return kPinnedAllocatorPos;
     if (*cudaAttr == Fortran::common::CUDADataAttr::Device)
       return kDeviceAllocatorPos;
-    if (*cudaAttr == Fortran::common::CUDADataAttr::Managed ||
-        *cudaAttr == Fortran::common::CUDADataAttr::Unified)
+    if (*cudaAttr == Fortran::common::CUDADataAttr::Managed)
       return kManagedAllocatorPos;
+    if (*cudaAttr == Fortran::common::CUDADataAttr::Unified)
+      return kUnifiedAllocatorPos;
   }
   return kDefaultAllocator;
 }

diff  --git a/flang/runtime/CUDA/allocator.cpp b/flang/runtime/CUDA/allocator.cpp
index 08fae9efb3e9d..cd00d40361d28 100644
--- a/flang/runtime/CUDA/allocator.cpp
+++ b/flang/runtime/CUDA/allocator.cpp
@@ -26,6 +26,8 @@ void CUFRegisterAllocator() {
       kDeviceAllocatorPos, {&CUFAllocDevice, CUFFreeDevice});
   allocatorRegistry.Register(
       kManagedAllocatorPos, {&CUFAllocManaged, CUFFreeManaged});
+  allocatorRegistry.Register(
+      kUnifiedAllocatorPos, {&CUFAllocUnified, CUFFreeUnified});
 }
 
 void *CUFAllocPinned(std::size_t sizeInBytes) {
@@ -57,4 +59,14 @@ void CUFFreeManaged(void *p) {
   CUDA_REPORT_IF_ERROR(cuMemFree(reinterpret_cast<CUdeviceptr>(p)));
 }
 
+void *CUFAllocUnified(std::size_t sizeInBytes) {
+  // Call alloc managed for the time being.
+  return CUFAllocManaged(sizeInBytes);
+}
+
+void CUFFreeUnified(void *p) {
+  // Call free managed for the time being.
+  CUFFreeManaged(p);
+}
+
 } // namespace Fortran::runtime::cuda