[SYCL][UR] Implement usm memory pool for hip adapter. (#11065)

Based on #10758

Co-authored-by: Ronan Keryell <ronan@keryell.fr>

Author: martygrant

context.cpp

@@ -7,6 +7,28 @@
 //===----------------------------------------------------------------------===//
 
 #include "context.hpp"
+#include "usm.hpp"
+
+void ur_context_handle_t_::addPool(ur_usm_pool_handle_t Pool) {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  PoolHandles.insert(Pool);
+}
+
+void ur_context_handle_t_::removePool(ur_usm_pool_handle_t Pool) {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  PoolHandles.erase(Pool);
+}
+
+ur_usm_pool_handle_t
+ur_context_handle_t_::getOwningURPool(umf_memory_pool_t *UMFPool) {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  for (auto &Pool : PoolHandles) {
+    if (Pool->hasUMFPool(UMFPool)) {
+      return Pool;
+    }
+  }
+  return nullptr;
+}
 
 /// Create a UR HIP context.
 ///

context.hpp

@@ -7,12 +7,15 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include <set>
 #include <unordered_map>
 
 #include "common.hpp"
 #include "device.hpp"
 #include "platform.hpp"
 
+#include <umf/memory_pool.h>
+
 typedef void (*ur_context_extended_deleter_t)(void *UserData);
 
 /// UR context mapping to a HIP context object.
@@ -95,6 +98,12 @@ struct ur_context_handle_t_ {
 
   uint32_t getReferenceCount() const noexcept { return RefCount; }
 
+  void addPool(ur_usm_pool_handle_t Pool);
+
+  void removePool(ur_usm_pool_handle_t Pool);
+
+  ur_usm_pool_handle_t getOwningURPool(umf_memory_pool_t *UMFPool);
+
   /// We need to keep track of USM mappings in AMD HIP, as certain extra
   /// synchronization *is* actually required for correctness.
   /// During kernel enqueue we must dispatch a prefetch for each kernel argument
@@ -150,6 +159,7 @@ struct ur_context_handle_t_ {
   std::mutex Mutex;
   std::vector<deleter_data> ExtendedDeleters;
   std::unordered_map<const void *, size_t> USMMappings;
+  std::set<ur_usm_pool_handle_t> PoolHandles;
 };
 
 namespace {

device.hpp

@@ -24,6 +24,8 @@ struct ur_device_handle_t_ {
   ur_platform_handle_t Platform;
   hipCtx_t HIPContext;
 
+  size_t MaxAllocSize{0};
+
 public:
   ur_device_handle_t_(native_type HipDevice, hipCtx_t Context,
                       ur_platform_handle_t Platform)

ur_interface_loader.cpp

@@ -239,10 +239,10 @@ urGetUSMProcAddrTable(ur_api_version_t version, ur_usm_dditable_t *pDdiTable) {
   pDdiTable->pfnFree = urUSMFree;
   pDdiTable->pfnGetMemAllocInfo = urUSMGetMemAllocInfo;
   pDdiTable->pfnHostAlloc = urUSMHostAlloc;
-  pDdiTable->pfnPoolCreate = nullptr;
-  pDdiTable->pfnPoolRetain = nullptr;
-  pDdiTable->pfnPoolRelease = nullptr;
-  pDdiTable->pfnPoolGetInfo = nullptr;
+  pDdiTable->pfnPoolCreate = urUSMPoolCreate;
+  pDdiTable->pfnPoolRetain = urUSMPoolRetain;
+  pDdiTable->pfnPoolRelease = urUSMPoolRelease;
+  pDdiTable->pfnPoolGetInfo = urUSMPoolGetInfo;
   pDdiTable->pfnSharedAlloc = urUSMSharedAlloc;
   return UR_RESULT_SUCCESS;
 }