[SYCL][UR] Implement usm memory pool for hip adapter.

sycl/plugins/hip/CMakeLists.txt

@@ -120,6 +120,7 @@ add_sycl_plugin(hip
     "../unified_runtime/ur/adapters/hip/sampler.hpp"
     "../unified_runtime/ur/adapters/hip/ur_interface_loader.cpp"
     "../unified_runtime/ur/adapters/hip/usm.cpp"
+    "../unified_runtime/ur/adapters/hip/usm.hpp"
     "../unified_runtime/ur/adapters/hip/usm_p2p.cpp"
     "${sycl_inc_dir}/sycl/detail/pi.h"
     "${sycl_inc_dir}/sycl/detail/pi.hpp"
@@ -130,6 +131,8 @@ add_sycl_plugin(hip
     ${CMAKE_CURRENT_SOURCE_DIR}/../unified_runtime
   LIBRARIES
     UnifiedRuntime-Headers
+    UnifiedRuntimeCommon
+    UnifiedMallocFramework
   HEADER
     ${CMAKE_CURRENT_SOURCE_DIR}/include/features.hpp
 )

sycl/plugins/unified_runtime/CMakeLists.txt

@@ -223,11 +223,13 @@ if ("hip" IN_LIST SYCL_ENABLE_PLUGINS)
       "ur/adapters/hip/sampler.hpp"
       "ur/adapters/hip/ur_interface_loader.cpp"
       "ur/adapters/hip/usm.cpp"
+      "ur/adapters/hip/usm.hpp"
       "ur/adapters/hip/usm_p2p.cpp"
     INCLUDE_DIRS
       ${sycl_inc_dir}
     LIBRARIES
       UnifiedRuntime-Headers
+      UnifiedRuntimeCommon
       Threads::Threads
   )
 
@@ -236,6 +238,13 @@ if ("hip" IN_LIST SYCL_ENABLE_PLUGINS)
     SOVERSION "0"
   )
 
+  if(UMF_ENABLE_POOL_TRACKING)
+    target_compile_definitions("ur_adapter_hip" PRIVATE
+      UMF_ENABLE_POOL_TRACKING)
+  else()
+    message(WARNING "HIP adapter USM pools are disabled, set UMF_ENABLE_POOL_TRACKING to enable them")
+  endif()
+
   if("${SYCL_BUILD_PI_HIP_PLATFORM}" STREQUAL "AMD")
     target_link_libraries(ur_adapter_hip PUBLIC rocmdrv)
     # Set HIP define to select AMD platform

sycl/plugins/unified_runtime/ur/adapters/hip/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.
 ///

sycl/plugins/unified_runtime/ur/adapters/hip/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 {

sycl/plugins/unified_runtime/ur/adapters/hip/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)

sycl/plugins/unified_runtime/ur/adapters/hip/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;
 }