Use UMF Proxy pool manager with UMF CUDA memory provider in UR

ldorau · ldorau · commit d202b5c1338c · 2025-02-04T15:05:01.000+01:00
UMF Proxy pool manager is just a wrapper for the UMF memory provider
(CUDA memory provider in this case) plus it adds also tracking of memory
allocations.

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
@@ -100,6 +100,21 @@ CreateHostMemoryProvider(ur_device_handle_t_ *DeviceHandle,
   return UR_RESULT_SUCCESS;
 }
 
+static ur_result_t
+CreateHostMemoryPool(umf_memory_provider_handle_t MemoryProviderHost,
+                     umf_memory_pool_handle_t *MemoryPoolHost) {
+
+  umf_memory_pool_handle_t _MemoryPoolHost = nullptr;
+
+  umf_result_t UmfResult =
+      umf::createMemoryProxyPool(MemoryProviderHost, &_MemoryPoolHost);
+  UMF_RETURN_UR_ERROR(UmfResult);
+
+  *MemoryPoolHost = _MemoryPoolHost;
+
+  return UR_RESULT_SUCCESS;
+}
+
 struct ur_context_handle_t_ {
 
   struct deleter_data {
@@ -112,8 +127,10 @@ 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 CUDA memory provider and pool for the host memory
+  // (UMF_MEMORY_TYPE_HOST)
   umf_memory_provider_handle_t MemoryProviderHost = nullptr;
+  umf_memory_pool_handle_t MemoryPoolHost = nullptr;
 
   ur_context_handle_t_(const ur_device_handle_t *Devs, uint32_t NumDevices)
       : Devices{Devs, Devs + NumDevices}, RefCount{1} {
@@ -125,9 +142,13 @@ struct ur_context_handle_t_ {
     // (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_CHECK_ERROR(CreateHostMemoryPool(MemoryProviderHost, &MemoryPoolHost));
   };
 
   ~ur_context_handle_t_() {
+    if (MemoryPoolHost) {
+      umfPoolDestroy(MemoryPoolHost);
+    }
     if (MemoryProviderHost) {
       umfMemoryProviderDestroy(MemoryProviderHost);
     }
diff --git a/source/adapters/cuda/device.hpp b/source/adapters/cuda/device.hpp
@@ -11,6 +11,7 @@
 
 #include <ur/ur.hpp>
 
+#include <umf/memory_pool.h>
 #include <umf/memory_provider.h>
 
 #include "common.hpp"
@@ -84,9 +85,17 @@ struct ur_device_handle_t_ {
 
     MemoryProviderDevice = nullptr;
     MemoryProviderShared = nullptr;
+    MemoryPoolDevice = nullptr;
+    MemoryPoolShared = nullptr;
   }
 
   ~ur_device_handle_t_() {
+    if (MemoryPoolDevice) {
+      umfPoolDestroy(MemoryPoolDevice);
+    }
+    if (MemoryPoolShared) {
+      umfPoolDestroy(MemoryPoolShared);
+    }
     if (MemoryProviderDevice) {
       umfMemoryProviderDestroy(MemoryProviderDevice);
     }
@@ -131,11 +140,15 @@ 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 CUDA memory provider and pool for the device memory
+  // (UMF_MEMORY_TYPE_DEVICE)
   umf_memory_provider_handle_t MemoryProviderDevice;
+  umf_memory_pool_handle_t MemoryPoolDevice;
 
-  // UMF CUDA memory provider for the shared memory (UMF_MEMORY_TYPE_SHARED)
+  // UMF CUDA memory provider and pool for the shared memory
+  // (UMF_MEMORY_TYPE_SHARED)
   umf_memory_provider_handle_t MemoryProviderShared;
+  umf_memory_pool_handle_t MemoryPoolShared;
 };
 
 int getAttribute(ur_device_handle_t Device, CUdevice_attribute Attribute);
diff --git a/source/adapters/cuda/memory.cpp b/source/adapters/cuda/memory.cpp
@@ -50,8 +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) {
-      UMF_CHECK_ERROR(umfMemoryProviderAlloc(hContext->MemoryProviderHost, size,
-                                             0, &HostPtr));
+      HostPtr = umfPoolMalloc(hContext->MemoryPoolHost, size);
+      UMF_CHECK_PTR(HostPtr);
       AllocMode = BufferMem::AllocMode::AllocHostPtr;
     } else if (flags & UR_MEM_FLAG_ALLOC_COPY_HOST_POINTER) {
       AllocMode = BufferMem::AllocMode::CopyIn;
@@ -442,8 +442,8 @@ ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t Mem,
                                        CU_MEMHOSTALLOC_DEVICEMAP));
       UR_CHECK_ERROR(cuMemHostGetDevicePointer(&DevPtr, Buffer.HostPtr, 0));
     } else {
-      UMF_CHECK_ERROR(umfMemoryProviderAlloc(hDevice->MemoryProviderDevice,
-                                             Buffer.Size, 0, (void **)&DevPtr));
+      *(void **)&DevPtr = umfPoolMalloc(hDevice->MemoryPoolDevice, Buffer.Size);
+      UMF_CHECK_PTR(*(void **)&DevPtr);
     }
   } else {
     CUarray ImageArray{};
diff --git a/source/adapters/cuda/platform.cpp b/source/adapters/cuda/platform.cpp
@@ -20,7 +20,7 @@
 #include <sstream>
 
 static ur_result_t
-CreateDeviceMemoryProviders(ur_platform_handle_t_ *Platform) {
+CreateDeviceMemoryProvidersPools(ur_platform_handle_t_ *Platform) {
   umf_cuda_memory_provider_params_handle_t CUMemoryProviderParams = nullptr;
 
   umf_result_t UmfResult =
@@ -48,6 +48,18 @@ CreateDeviceMemoryProviders(ur_platform_handle_t_ *Platform) {
         CUMemoryProviderParamsUnique.get(), device, context,
         UMF_MEMORY_TYPE_SHARED, &device_handle->MemoryProviderShared);
     UMF_RETURN_UR_ERROR(UmfResult);
+
+    // create UMF CUDA memory pool for the device memory
+    // (UMF_MEMORY_TYPE_DEVICE)
+    UmfResult = umf::createMemoryProxyPool(device_handle->MemoryProviderDevice,
+                                           &device_handle->MemoryPoolDevice);
+    UMF_RETURN_UR_ERROR(UmfResult);
+
+    // create UMF CUDA memory pool for the shared memory
+    // (UMF_MEMORY_TYPE_SHARED)
+    UmfResult = umf::createMemoryProxyPool(device_handle->MemoryProviderShared,
+                                           &device_handle->MemoryPoolShared);
+    UMF_RETURN_UR_ERROR(UmfResult);
   }
 
   return UR_RESULT_SUCCESS;
@@ -134,7 +146,7 @@ urPlatformGet(ur_adapter_handle_t *, uint32_t, uint32_t NumEntries,
                                           static_cast<uint32_t>(i)});
             }
 
-            UR_CHECK_ERROR(CreateDeviceMemoryProviders(&Platform));
+            UR_CHECK_ERROR(CreateDeviceMemoryProvidersPools(&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
@@ -125,18 +125,17 @@ ur_result_t USMFreeImpl(ur_context_handle_t hContext, void *Pointer) {
     UR_ASSERT(DeviceOrdinal < NumDevices, UR_RESULT_ERROR_INVALID_DEVICE);
 
     ur_device_handle_t Device = Platform->Devices[DeviceOrdinal].get();
-    umf_memory_provider_handle_t MemoryProvider;
+    umf_memory_pool_handle_t MemoryPool;
 
     if (IsManaged) {
-      MemoryProvider = Device->MemoryProviderShared;
+      MemoryPool = Device->MemoryPoolShared;
     } else if (Type == CU_MEMORYTYPE_DEVICE) {
-      MemoryProvider = Device->MemoryProviderDevice;
+      MemoryPool = Device->MemoryPoolDevice;
     } else {
-      MemoryProvider = hContext->MemoryProviderHost;
+      MemoryPool = hContext->MemoryPoolHost;
     }
 
-    UMF_CHECK_ERROR(umfMemoryProviderFree(MemoryProvider, Pointer,
-                                          0 /* size is unknown */));
+    UMF_CHECK_ERROR(umfPoolFree(MemoryPool, Pointer));
   } catch (ur_result_t Err) {
     Result = Err;
   }
@@ -158,8 +157,8 @@ ur_result_t USMDeviceAllocImpl(void **ResultPtr, ur_context_handle_t,
                                uint32_t Alignment) {
   try {
     ScopedContext Active(Device);
-    UMF_CHECK_ERROR(umfMemoryProviderAlloc(Device->MemoryProviderDevice, Size,
-                                           Alignment, ResultPtr));
+    *ResultPtr = umfPoolMalloc(Device->MemoryPoolDevice, Size);
+    UMF_CHECK_PTR(*ResultPtr);
   } catch (ur_result_t Err) {
     return Err;
   }
@@ -180,8 +179,8 @@ ur_result_t USMSharedAllocImpl(void **ResultPtr, ur_context_handle_t,
                                uint32_t Alignment) {
   try {
     ScopedContext Active(Device);
-    UMF_CHECK_ERROR(umfMemoryProviderAlloc(Device->MemoryProviderShared, Size,
-                                           Alignment, ResultPtr));
+    *ResultPtr = umfPoolMalloc(Device->MemoryPoolShared, Size);
+    UMF_CHECK_PTR(*ResultPtr);
   } catch (ur_result_t Err) {
     return Err;
   }
@@ -199,8 +198,8 @@ ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t hContext,
                              ur_usm_host_mem_flags_t, size_t Size,
                              uint32_t Alignment) {
   try {
-    UMF_CHECK_ERROR(umfMemoryProviderAlloc(hContext->MemoryProviderHost, Size,
-                                           Alignment, ResultPtr));
+    *ResultPtr = umfPoolMalloc(hContext->MemoryPoolHost, Size);
+    UMF_CHECK_PTR(*ResultPtr);
   } catch (ur_result_t Err) {
     return Err;
   }
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/pools/pool_proxy.h>
 #include <umf/providers/provider_cuda.h>
 #include <ur_api.h>
 
@@ -30,6 +31,13 @@
 
 #define UMF_CHECK_ERROR(UmfResult) UR_CHECK_ERROR(umf::umf2urResult(UmfResult));
 
+#define UMF_CHECK_PTR(ptr)                                                     \
+  do {                                                                         \
+    if ((ptr) == nullptr) {                                                    \
+      UR_CHECK_ERROR(UR_RESULT_ERROR_OUT_OF_HOST_MEMORY);                      \
+    }                                                                          \
+  } while (0)
+
 #define UMF_RETURN_UMF_ERROR(UmfResult)                                        \
   do {                                                                         \
     umf_result_t UmfResult_ = (UmfResult);                                     \
@@ -328,6 +336,20 @@ inline umf_result_t createMemoryProvider(
   return UMF_RESULT_SUCCESS;
 }
 
+inline umf_result_t createMemoryProxyPool(umf_memory_provider_handle_t provider,
+                                          umf_memory_pool_handle_t *pool) {
+
+  umf_memory_pool_handle_t _pool = nullptr;
+
+  umf_result_t UmfResult =
+      umfPoolCreate(umfProxyPoolOps(), provider, nullptr, 0, &_pool);
+  UMF_RETURN_UMF_ERROR(UmfResult);
+
+  *pool = _pool;
+
+  return UMF_RESULT_SUCCESS;
+}
+
 } // namespace umf
 
 #endif /* UMF_HELPERS_H */
