Use UMF CUDA provider in UR

ldorau · ldorau · commit f7e21bbc1ac2 · 2025-02-04T15:28:44.000+01:00
Signed-off-by: Lukasz Dorau &lt;lukasz.dorau@intel.com&gt;
diff --git a/source/adapters/cuda/context.hpp b/source/adapters/cuda/context.hpp
@@ -19,6 +19,7 @@
 
 #include "common.hpp"
 #include "device.hpp"
+#include "umf_helpers.hpp"
 
 #include <umf/memory_pool.h>
 
@@ -74,6 +75,31 @@ typedef void (*ur_context_extended_deleter_t)(void *user_data);
 ///  if necessary.
 ///
 ///
+
+static ur_result_t
+CreateHostMemoryProvider(ur_device_handle_t_ *DeviceHandle,
+                         umf_memory_provider_handle_t *MemoryProviderHost) {
+  umf_cuda_memory_provider_params_handle_t CUMemoryProviderParams = nullptr;
+
+  *MemoryProviderHost = nullptr;
+  CUcontext context = DeviceHandle->getNativeContext();
+
+  umf_result_t UmfResult =
+      umfCUDAMemoryProviderParamsCreate(&CUMemoryProviderParams);
+  UMF_RETURN_UR_ERROR(UmfResult);
+
+  umf::cuda_params_unique_handle_t CUMemoryProviderParamsUnique(
+      CUMemoryProviderParams, umfCUDAMemoryProviderParamsDestroy);
+
+  // create UMF CUDA memory provider for the host memory (UMF_MEMORY_TYPE_HOST)
+  UmfResult = umf::createMemoryProvider(
+      CUMemoryProviderParamsUnique.get(), 0 /* cuDevice */, context,
+      UMF_MEMORY_TYPE_HOST, MemoryProviderHost);
+  UMF_RETURN_UR_ERROR(UmfResult);
+
+  return UR_RESULT_SUCCESS;
+}
+
 struct ur_context_handle_t_ {
 
   struct deleter_data {
@@ -86,14 +112,25 @@ struct ur_context_handle_t_ {
   std::vector<ur_device_handle_t> Devices;
   std::atomic_uint32_t RefCount;
 
+  // UMF CUDA memory provider for the host memory (UMF_MEMORY_TYPE_HOST)
+  umf_memory_provider_handle_t MemoryProviderHost = nullptr;
+
   ur_context_handle_t_(const ur_device_handle_t *Devs, uint32_t NumDevices)
       : Devices{Devs, Devs + NumDevices}, RefCount{1} {
     for (auto &Dev : Devices) {
       urDeviceRetain(Dev);
     }
+
+    // Create UMF CUDA memory provider for the host memory
+    // (UMF_MEMORY_TYPE_HOST) from any device (Devices[0] is used here, because
+    // it is guaranteed to exist).
+    UR_CHECK_ERROR(CreateHostMemoryProvider(Devices[0], &MemoryProviderHost));
   };
 
   ~ur_context_handle_t_() {
+    if (MemoryProviderHost) {
+      umfMemoryProviderDestroy(MemoryProviderHost);
+    }
     for (auto &Dev : Devices) {
       urDeviceRelease(Dev);
     }
diff --git a/source/adapters/cuda/device.hpp b/source/adapters/cuda/device.hpp
@@ -11,6 +11,8 @@
 
 #include <ur/ur.hpp>
 
+#include <umf/memory_provider.h>
+
 #include "common.hpp"
 
 struct ur_device_handle_t_ {
@@ -79,9 +81,20 @@ struct ur_device_handle_t_ {
     // CUDA doesn't really have this concept, and could allow almost 100% of
     // global memory in one allocation, but is dependent on device usage.
     UR_CHECK_ERROR(cuDeviceTotalMem(&MaxAllocSize, cuDevice));
+
+    MemoryProviderDevice = nullptr;
+    MemoryProviderShared = nullptr;
   }
 
-  ~ur_device_handle_t_() { cuDevicePrimaryCtxRelease(CuDevice); }
+  ~ur_device_handle_t_() {
+    if (MemoryProviderDevice) {
+      umfMemoryProviderDestroy(MemoryProviderDevice);
+    }
+    if (MemoryProviderShared) {
+      umfMemoryProviderDestroy(MemoryProviderShared);
+    }
+    cuDevicePrimaryCtxRelease(CuDevice);
+  }
 
   native_type get() const noexcept { return CuDevice; };
 
@@ -117,6 +130,12 @@ struct ur_device_handle_t_ {
 
   // bookkeeping for mipmappedArray leaks in Mapping external Memory
   std::map<CUarray, CUmipmappedArray> ChildCuarrayFromMipmapMap;
+
+  // UMF CUDA memory provider for the device memory (UMF_MEMORY_TYPE_DEVICE)
+  umf_memory_provider_handle_t MemoryProviderDevice;
+
+  // UMF CUDA memory provider for the shared memory (UMF_MEMORY_TYPE_SHARED)
+  umf_memory_provider_handle_t MemoryProviderShared;
 };
 
 int getAttribute(ur_device_handle_t Device, CUdevice_attribute Attribute);
diff --git a/source/adapters/cuda/memory.cpp b/source/adapters/cuda/memory.cpp
@@ -14,6 +14,7 @@
 #include "context.hpp"
 #include "enqueue.hpp"
 #include "memory.hpp"
+#include "umf_helpers.hpp"
 
 /// Creates a UR Memory object using a CUDA memory allocation.
 /// Can trigger a manual copy depending on the mode.
@@ -49,7 +50,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemBufferCreate(
           cuMemHostRegister(HostPtr, size, CU_MEMHOSTREGISTER_DEVICEMAP));
       AllocMode = BufferMem::AllocMode::UseHostPtr;
     } else if (flags & UR_MEM_FLAG_ALLOC_HOST_POINTER) {
-      UR_CHECK_ERROR(cuMemAllocHost(&HostPtr, size));
+      UMF_CHECK_ERROR(umfMemoryProviderAlloc(hContext->MemoryProviderHost, size,
+                                             0, &HostPtr));
       AllocMode = BufferMem::AllocMode::AllocHostPtr;
     } else if (flags & UR_MEM_FLAG_ALLOC_COPY_HOST_POINTER) {
       AllocMode = BufferMem::AllocMode::CopyIn;
@@ -440,7 +442,8 @@ ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t Mem,
                                        CU_MEMHOSTALLOC_DEVICEMAP));
       UR_CHECK_ERROR(cuMemHostGetDevicePointer(&DevPtr, Buffer.HostPtr, 0));
     } else {
-      UR_CHECK_ERROR(cuMemAlloc(&DevPtr, Buffer.Size));
+      UMF_CHECK_ERROR(umfMemoryProviderAlloc(hDevice->MemoryProviderDevice,
+                                             Buffer.Size, 0, (void **)&DevPtr));
     }
   } else {
     CUarray ImageArray{};
diff --git a/source/adapters/cuda/platform.cpp b/source/adapters/cuda/platform.cpp
@@ -13,11 +13,46 @@
 #include "common.hpp"
 #include "context.hpp"
 #include "device.hpp"
+#include "umf_helpers.hpp"
 
 #include <cassert>
 #include <cuda.h>
 #include <sstream>
 
+static ur_result_t
+CreateDeviceMemoryProviders(ur_platform_handle_t_ *Platform) {
+  umf_cuda_memory_provider_params_handle_t CUMemoryProviderParams = nullptr;
+
+  umf_result_t UmfResult =
+      umfCUDAMemoryProviderParamsCreate(&CUMemoryProviderParams);
+  UMF_RETURN_UR_ERROR(UmfResult);
+
+  umf::cuda_params_unique_handle_t CUMemoryProviderParamsUnique(
+      CUMemoryProviderParams, umfCUDAMemoryProviderParamsDestroy);
+
+  for (auto &Device : Platform->Devices) {
+    ur_device_handle_t_ *device_handle = Device.get();
+    CUdevice device = device_handle->get();
+    CUcontext context = device_handle->getNativeContext();
+
+    // create UMF CUDA memory provider for the device memory
+    // (UMF_MEMORY_TYPE_DEVICE)
+    UmfResult = umf::createMemoryProvider(
+        CUMemoryProviderParamsUnique.get(), device, context,
+        UMF_MEMORY_TYPE_DEVICE, &device_handle->MemoryProviderDevice);
+    UMF_RETURN_UR_ERROR(UmfResult);
+
+    // create UMF CUDA memory provider for the shared memory
+    // (UMF_MEMORY_TYPE_SHARED)
+    UmfResult = umf::createMemoryProvider(
+        CUMemoryProviderParamsUnique.get(), device, context,
+        UMF_MEMORY_TYPE_SHARED, &device_handle->MemoryProviderShared);
+    UMF_RETURN_UR_ERROR(UmfResult);
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urPlatformGetInfo(
     ur_platform_handle_t hPlatform, ur_platform_info_t PlatformInfoType,
     size_t Size, void *pPlatformInfo, size_t *pSizeRet) {
@@ -98,6 +133,8 @@ urPlatformGet(ur_adapter_handle_t *, uint32_t, uint32_t NumEntries,
                   new ur_device_handle_t_{Device, Context, EvBase, &Platform,
                                           static_cast<uint32_t>(i)});
             }
+
+            UR_CHECK_ERROR(CreateDeviceMemoryProviders(&Platform));
           } catch (const std::bad_alloc &) {
             // Signal out-of-memory situation
             for (int i = 0; i < NumDevices; ++i) {
diff --git a/source/adapters/cuda/usm.cpp b/source/adapters/cuda/usm.cpp
@@ -102,26 +102,41 @@ urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t USMFreeImpl(ur_context_handle_t, void *Pointer) {
+ur_result_t USMFreeImpl(ur_context_handle_t hContext, void *Pointer) {
   ur_result_t Result = UR_RESULT_SUCCESS;
   try {
     unsigned int IsManaged;
     unsigned int Type;
-    void *AttributeValues[2] = {&IsManaged, &Type};
-    CUpointer_attribute Attributes[2] = {CU_POINTER_ATTRIBUTE_IS_MANAGED,
-                                         CU_POINTER_ATTRIBUTE_MEMORY_TYPE};
-    UR_CHECK_ERROR(cuPointerGetAttributes(2, Attributes, AttributeValues,
-                                          (CUdeviceptr)Pointer));
+    unsigned int DeviceOrdinal;
+    const int NumAttributes = 3;
+    void *AttributeValues[NumAttributes] = {&IsManaged, &Type, &DeviceOrdinal};
+
+    CUpointer_attribute Attributes[NumAttributes] = {
+        CU_POINTER_ATTRIBUTE_IS_MANAGED, CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
+        CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL};
+    UR_CHECK_ERROR(cuPointerGetAttributes(
+        NumAttributes, Attributes, AttributeValues, (CUdeviceptr)Pointer));
     UR_ASSERT(Type == CU_MEMORYTYPE_DEVICE || Type == CU_MEMORYTYPE_HOST,
               UR_RESULT_ERROR_INVALID_MEM_OBJECT);
-    if (IsManaged || Type == CU_MEMORYTYPE_DEVICE) {
-      // Memory allocated with cuMemAlloc and cuMemAllocManaged must be freed
-      // with cuMemFree
-      UR_CHECK_ERROR(cuMemFree((CUdeviceptr)Pointer));
+
+    std::vector<ur_device_handle_t> ContextDevices = hContext->getDevices();
+    ur_platform_handle_t Platform = ContextDevices[0]->getPlatform();
+    unsigned int NumDevices = Platform->Devices.size();
+    UR_ASSERT(DeviceOrdinal < NumDevices, UR_RESULT_ERROR_INVALID_DEVICE);
+
+    ur_device_handle_t Device = Platform->Devices[DeviceOrdinal].get();
+    umf_memory_provider_handle_t MemoryProvider;
+
+    if (IsManaged) {
+      MemoryProvider = Device->MemoryProviderShared;
+    } else if (Type == CU_MEMORYTYPE_DEVICE) {
+      MemoryProvider = Device->MemoryProviderDevice;
     } else {
-      // Memory allocated with cuMemAllocHost must be freed with cuMemFreeHost
-      UR_CHECK_ERROR(cuMemFreeHost(Pointer));
+      MemoryProvider = hContext->MemoryProviderHost;
     }
+
+    UMF_CHECK_ERROR(umfMemoryProviderFree(MemoryProvider, Pointer,
+                                          0 /* size is unknown */));
   } catch (ur_result_t Err) {
     Result = Err;
   }
@@ -143,7 +158,8 @@ ur_result_t USMDeviceAllocImpl(void **ResultPtr, ur_context_handle_t,
                                uint32_t Alignment) {
   try {
     ScopedContext Active(Device);
-    UR_CHECK_ERROR(cuMemAlloc((CUdeviceptr *)ResultPtr, Size));
+    UMF_CHECK_ERROR(umfMemoryProviderAlloc(Device->MemoryProviderDevice, Size,
+                                           Alignment, ResultPtr));
   } catch (ur_result_t Err) {
     return Err;
   }
@@ -164,8 +180,8 @@ ur_result_t USMSharedAllocImpl(void **ResultPtr, ur_context_handle_t,
                                uint32_t Alignment) {
   try {
     ScopedContext Active(Device);
-    UR_CHECK_ERROR(cuMemAllocManaged((CUdeviceptr *)ResultPtr, Size,
-                                     CU_MEM_ATTACH_GLOBAL));
+    UMF_CHECK_ERROR(umfMemoryProviderAlloc(Device->MemoryProviderShared, Size,
+                                           Alignment, ResultPtr));
   } catch (ur_result_t Err) {
     return Err;
   }
@@ -179,11 +195,12 @@ ur_result_t USMSharedAllocImpl(void **ResultPtr, ur_context_handle_t,
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t,
+ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t hContext,
                              ur_usm_host_mem_flags_t, size_t Size,
                              uint32_t Alignment) {
   try {
-    UR_CHECK_ERROR(cuMemAllocHost(ResultPtr, Size));
+    UMF_CHECK_ERROR(umfMemoryProviderAlloc(hContext->MemoryProviderHost, Size,
+                                           Alignment, ResultPtr));
   } catch (ur_result_t Err) {
     return Err;
   }
diff --git a/source/common/CMakeLists.txt b/source/common/CMakeLists.txt
@@ -64,7 +64,7 @@ else()
   set(UMF_BUILD_EXAMPLES OFF CACHE INTERNAL "Build UMF examples")
   set(UMF_BUILD_SHARED_LIBRARY ${UMF_BUILD_SHARED_LIBRARY} CACHE INTERNAL "Build UMF shared library")
   set(UMF_BUILD_LIBUMF_POOL_DISJOINT ON CACHE INTERNAL "Build Disjoint Pool")
-  set(UMF_BUILD_CUDA_PROVIDER OFF CACHE INTERNAL "Build UMF CUDA provider")
+  set(UMF_BUILD_CUDA_PROVIDER ON CACHE INTERNAL "Build UMF CUDA provider")
 
   FetchContent_MakeAvailable(unified-memory-framework)
   FetchContent_GetProperties(unified-memory-framework)
diff --git a/source/common/umf_helpers.hpp b/source/common/umf_helpers.hpp
@@ -16,6 +16,7 @@
 #include <umf/memory_pool_ops.h>
 #include <umf/memory_provider.h>
 #include <umf/memory_provider_ops.h>
+#include <umf/providers/provider_cuda.h>
 #include <ur_api.h>
 
 #include "logger/ur_logger.hpp"
@@ -27,6 +28,24 @@
 #include <tuple>
 #include <utility>
 
+#define UMF_CHECK_ERROR(UmfResult) UR_CHECK_ERROR(umf::umf2urResult(UmfResult));
+
+#define UMF_RETURN_UMF_ERROR(UmfResult)                                        \
+  do {                                                                         \
+    umf_result_t UmfResult_ = (UmfResult);                                     \
+    if (UmfResult_ != UMF_RESULT_SUCCESS) {                                    \
+      return UmfResult_;                                                       \
+    }                                                                          \
+  } while (0)
+
+#define UMF_RETURN_UR_ERROR(UmfResult)                                         \
+  do {                                                                         \
+    umf_result_t UmfResult_ = (UmfResult);                                     \
+    if (UmfResult_ != UMF_RESULT_SUCCESS) {                                    \
+      return umf::umf2urResult(UmfResult_);                                    \
+    }                                                                          \
+  } while (0)
+
 namespace umf {
 
 using pool_unique_handle_t =
@@ -35,6 +54,9 @@ using pool_unique_handle_t =
 using provider_unique_handle_t =
     std::unique_ptr<umf_memory_provider_t,
                     std::function<void(umf_memory_provider_handle_t)>>;
+using cuda_params_unique_handle_t = std::unique_ptr<
+    umf_cuda_memory_provider_params_t,
+    std::function<umf_result_t(umf_cuda_memory_provider_params_handle_t)>>;
 
 #define DEFINE_CHECK_OP(op)                                                    \
   template <typename T> class HAS_OP_##op {                                    \
@@ -279,6 +301,33 @@ inline ur_result_t umf2urResult(umf_result_t umfResult) {
   };
 }
 
+inline umf_result_t createMemoryProvider(
+    umf_cuda_memory_provider_params_handle_t CUMemoryProviderParams,
+    int cuDevice, void *cuContext, umf_usm_memory_type_t memType,
+    umf_memory_provider_handle_t *provider) {
+
+  umf_result_t UmfResult =
+      umfCUDAMemoryProviderParamsSetContext(CUMemoryProviderParams, cuContext);
+  UMF_RETURN_UMF_ERROR(UmfResult);
+
+  UmfResult =
+      umfCUDAMemoryProviderParamsSetDevice(CUMemoryProviderParams, cuDevice);
+  UMF_RETURN_UMF_ERROR(UmfResult);
+
+  UmfResult =
+      umfCUDAMemoryProviderParamsSetMemoryType(CUMemoryProviderParams, memType);
+  UMF_RETURN_UMF_ERROR(UmfResult);
+
+  umf_memory_provider_handle_t umfCUDAprovider = nullptr;
+  UmfResult = umfMemoryProviderCreate(umfCUDAMemoryProviderOps(),
+                                      CUMemoryProviderParams, &umfCUDAprovider);
+  UMF_RETURN_UMF_ERROR(UmfResult);
+
+  *provider = umfCUDAprovider;
+
+  return UMF_RESULT_SUCCESS;
+}
+
 } // namespace umf
 
 #endif /* UMF_HELPERS_H */
