[libomptarget] Add support for target memory allocators to cuda RTL

[jhuber6]

Summary

The allocator interface added in D97883 allows the RTL to allocate shared and
host-pinned memory from the cuda plugin. This patch adds support for these to
the runtime.

[jdoerfert]

Two nits, otherwise this looks sensible to me.

openmp/libomptarget/plugins/common/MemoryManager/MemoryManager.h

34

can we have an enum not a int32_t, also variable name Kind

openmp/libomptarget/plugins/cuda/src/rtl.cpp

363

if everything but alloc_host goes to cumemfree, only put alloc_host in the set. If it is not in the set, it's cumemfree teritory.

[jhuber6]

Making suggested changes.

[grokos]

Looks good (with one naming nit)

openmp/libomptarget/plugins/cuda/src/rtl.cpp

301

Can you rename AllocKinds to something like HostPinnedAllocs? AllocKinds does not make sense anymore after the latest change.

[xiangzhangllvm]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

327

Sorry, where is the def of "cuMemAllocHost"? I didn't find it.

[jhuber6]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

327

https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1gdd8311286d2c2691605362c689bc64e0

[JonChesterfield]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

364

cuMemFreeHost needs to be added to dynamic_cuda/cuda.h. Possibly other functions too, don't have this building successfully yet

openmp/libomptarget/plugins/common/MemoryManager/MemoryManager.h

 
 /// Base class of per-device allocator.
 class DeviceAllocatorTy {
 public:
   virtual ~DeviceAllocatorTy() = default;
 
   /// Allocate a memory of size \p Size . \p HstPtr is used to assist the
   /// allocation.
-  virtual void *allocate(size_t Size, void *HstPtr) = 0;
+  virtual void *allocate(size_t Size, void *HstPtr, TargetAllocTy Kind) = 0;
 
   /// Delete the pointer \p TgtPtr on the device
   virtual int free(void *TgtPtr) = 0;
 };
 
 /// Class of memory manager. The memory manager is per-device by using
 /// per-device allocator. Therefore, each plugin using memory manager should
 /// have an allocator for each device.
#endif
 
   /// The threshold to manage memory using memory manager. If the request size
   /// is larger than \p SizeThreshold, the allocation will not be managed by the
   /// memory manager.
   size_t SizeThreshold = 1U << 13;
 
   /// Request memory from target device
   void *allocateOnDevice(size_t Size, void *HstPtr) const {
-    return DeviceAllocator.allocate(Size, HstPtr);
+    return DeviceAllocator.allocate(Size, HstPtr, TARGET_ALLOC_DEVICE);
   }
 
   /// Deallocate data on device
   int deleteOnDevice(void *Ptr) const { return DeviceAllocator.free(Ptr); }
 
   /// This function is called when it tries to allocate memory on device but the
   /// device returns out of memory. It will first free all memory in the
   /// FreeList and try to allocate again.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "Debug.h"
 #include "omptargetplugin.h"
 
 #define TARGET_NAME CUDA
 #define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"
 
class DeviceRTLTy {
   std::vector<DeviceDataTy> DeviceData;
   std::vector<CUmodule> Modules;
 
   /// A class responsible for interacting with device native runtime library to
   /// allocate and free memory.
   class CUDADeviceAllocatorTy : public DeviceAllocatorTy {
     const int DeviceId;
     const std::vector<DeviceDataTy> &DeviceData;
+    std::unordered_map<void *, TargetAllocTy> HostPinnedAllocs;
 
   public:
     CUDADeviceAllocatorTy(int DeviceId, std::vector<DeviceDataTy> &DeviceData)
         : DeviceId(DeviceId), DeviceData(DeviceData) {}
 
-    void *allocate(size_t Size, void *) override {
+    void *allocate(size_t Size, void *, TargetAllocTy Kind) override {
       if (Size == 0)
         return nullptr;
 
       CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
       if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
         return nullptr;
 
+      void *MemAlloc = nullptr;
+      switch (Kind) {
+      case TARGET_ALLOC_DEFAULT:
+      case TARGET_ALLOC_DEVICE:
         CUdeviceptr DevicePtr;
         Err = cuMemAlloc(&DevicePtr, Size);
+        MemAlloc = (void *)DevicePtr;
         if (!checkResult(Err, "Error returned from cuMemAlloc\n"))
           return nullptr;
+        break;
+      case TARGET_ALLOC_HOST:
+        void *HostPtr;
+        Err = cuMemAllocHost(&HostPtr, Size);
+        MemAlloc = HostPtr;
+        if (!checkResult(Err, "Error returned from cuMemAllocHost\n"))
+          return nullptr;
+        HostPinnedAllocs[MemAlloc] = Kind;
+        break;
+      case TARGET_ALLOC_SHARED:
+        CUdeviceptr SharedPtr;
+        Err = cuMemAllocManaged(&SharedPtr, Size, CU_MEM_ATTACH_GLOBAL);
+        MemAlloc = (void *)SharedPtr;
+        if (!checkResult(Err, "Error returned from cuMemAllocManaged\n"))
+          return nullptr;
+        break;
+      }
 
-      return (void *)DevicePtr;
+      return MemAlloc;
     }
 
     int free(void *TgtPtr) override {
       CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
       if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
         return OFFLOAD_FAIL;
 
+      // Host pinned memory must be freed differently.
+      TargetAllocTy Kind =
+          (HostPinnedAllocs.find(TgtPtr) == HostPinnedAllocs.end())
+              ? TARGET_ALLOC_DEFAULT
+              : TARGET_ALLOC_HOST;
+      switch (Kind) {
+      case TARGET_ALLOC_DEFAULT:
+      case TARGET_ALLOC_DEVICE:
+      case TARGET_ALLOC_SHARED:
         Err = cuMemFree((CUdeviceptr)TgtPtr);
         if (!checkResult(Err, "Error returned from cuMemFree\n"))
           return OFFLOAD_FAIL;
+        break;
+      case TARGET_ALLOC_HOST:
+        Err = cuMemFreeHost(TgtPtr);
+        if (!checkResult(Err, "Error returned from cuMemFreeHost\n"))
+          return OFFLOAD_FAIL;
+        break;
+      }
 
       return OFFLOAD_SUCCESS;
     }
   };
 
   /// A vector of device allocators
   std::vector<CUDADeviceAllocatorTy> DeviceAllocators;
 
#endif
         DP("Continue, considering this is a device RTL which does not accept "
            "environment setting.\n");
       }
     }
 
     return getOffloadEntriesTable(DeviceId);
   }
 
-  void *dataAlloc(const int DeviceId, const int64_t Size) {
+  void *dataAlloc(const int DeviceId, const int64_t Size,
+                  const TargetAllocTy Kind) {
+    switch (Kind) {
+    case TARGET_ALLOC_DEFAULT:
+    case TARGET_ALLOC_DEVICE:
       if (UseMemoryManager)
         return MemoryManagers[DeviceId]->allocate(Size, nullptr);
+      else
+        return DeviceAllocators[DeviceId].allocate(Size, nullptr, Kind);
+    case TARGET_ALLOC_HOST:
+    case TARGET_ALLOC_SHARED:
+      return DeviceAllocators[DeviceId].allocate(Size, nullptr, Kind);
+    }
+
+    REPORT("Invalid target data allocation kind or requested allocator not "
+           "implemented yet\n");
 
-    return DeviceAllocators[DeviceId].allocate(Size, nullptr);
+    return nullptr;
   }
 
   int dataSubmit(const int DeviceId, const void *TgtPtr, const void *HstPtr,
                  const int64_t Size, __tgt_async_info *AsyncInfo) const {
     assert(AsyncInfo && "AsyncInfo is nullptr");
 
     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
__tgt_target_table *__tgt_rtl_load_binary(int32_t device_id,
 
   return DeviceRTL.loadBinary(device_id, image);
 }
 
 void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *,
                            int32_t kind) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
-  if (kind != TARGET_ALLOC_DEFAULT) {
-    REPORT("Invalid target data allocation kind or requested allocator not "
-           "implemented yet\n");
-    return NULL;
-  }
-
-  return DeviceRTL.dataAlloc(device_id, size);
+  return DeviceRTL.dataAlloc(device_id, size, (TargetAllocTy)kind);
 }
 
 int32_t __tgt_rtl_data_submit(int32_t device_id, void *tgt_ptr, void *hst_ptr,
                               int64_t size) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
   __tgt_async_info AsyncInfo;
   const int32_t rc = __tgt_rtl_data_submit_async(device_id, tgt_ptr, hst_ptr,

openmp/libomptarget/test/api/omp_device_managed_memory.c

+// RUN: %libomptarget-compile-run-and-check-nvptx64-nvidia-cuda
+// REQUIRES: nvptx64-nvidia-cuda
+
+#include <omp.h>
+#include <stdio.h>
+
+void *llvm_omp_target_alloc_shared(size_t, int);
+
+int main() {
+  const int N = 64;
+  const int device = omp_get_default_device();
+
+  int *shared_ptr = llvm_omp_target_alloc_shared(N * sizeof(int), device);
+
+#pragma omp target teams distribute parallel for device(device) \
+            is_device_ptr(shared_ptr)
+  for (int i = 0; i < N; ++i) {
+    shared_ptr[i] = 1;
+  }
+
+  int sum = 0;
+  for (int i = 0; i < N; ++i)
+    sum += shared_ptr[i];
+
+  omp_target_free(shared_ptr, device);
+  // CHECK: PASS
+  if (sum == N)
+    printf ("PASS\n");
+}

openmp/libomptarget/test/api/omp_host_pinned_memory.c

+// RUN: %libomptarget-compile-run-and-check-nvptx64-nvidia-cuda
+// REQUIRES: nvptx64-nvidia-cuda
+
+#include <omp.h>
+#include <stdio.h>
+
+// Allocate pinned memory on the host
+void *llvm_omp_target_alloc_host(size_t, int);
+
+int main() {
+  const int N = 64;
+  const int device = omp_get_default_device();
+  const int host = omp_get_initial_device();
+
+  int *hst_ptr = llvm_omp_target_alloc_host(N * sizeof(int), device);
+
+  for (int i = 0; i < N; ++i)
+    hst_ptr[i] = 2;
+
+#pragma omp target teams distribute parallel for device(device) \
+           map(tofrom:hst_ptr[0 : N])
+  for (int i = 0; i < N; ++i)
+    hst_ptr[i] -= 1;
+
+  int sum = 0;
+  for (int i = 0; i < N; ++i)
+    sum += hst_ptr[i];
+
+  omp_target_free(hst_ptr, device);
+  // CHECK: PASS
+  if (sum == N)
+    printf ("PASS\n");
+}