[UR][L0] Add several fixes to L0 adapter for test-usm

Signed-off-by: Jaime Arteaga <[email protected]>

Author: Jaime Arteaga

source/adapters/level_zero/context.hpp

@@ -115,6 +115,9 @@ struct ur_context_handle_t_ : _ur_object {
       SharedReadOnlyMemProxyPools;
   umf::pool_unique_handle_t HostMemProxyPool;
 
+  // Map associating pools created with urUsmPoolCreate and internal pools
+  std::list<ur_usm_pool_handle_t> UsmPoolHandles{};
+
   // We need to store all memory allocations in the context because there could
   // be kernels with indirect access. Kernels with indirect access start to
   // reference all existing memory allocations at the time when they are

source/adapters/level_zero/usm.cpp

@@ -187,8 +187,15 @@ static ur_result_t USMDeviceAllocImpl(void **ResultPtr,
     ZeDesc.pNext = &RelaxedDesc;
   }
 
-  ZE2UR_CALL(zeMemAllocDevice, (Context->ZeContext, &ZeDesc, Size, Alignment,
-                                Device->ZeDevice, ResultPtr));
+  ze_result_t ZeResult =
+      zeMemAllocDevice(Context->ZeContext, &ZeDesc, Size, Alignment,
+                       Device->ZeDevice, ResultPtr);
+  if (ZeResult != ZE_RESULT_SUCCESS) {
+    if (ZeResult == ZE_RESULT_ERROR_UNSUPPORTED_SIZE) {
+      return UR_RESULT_ERROR_INVALID_USM_SIZE;
+    }
+    return ze2urResult(ZeResult);
+  }
 
   UR_ASSERT(Alignment == 0 ||
                 reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0,
@@ -226,8 +233,15 @@ static ur_result_t USMSharedAllocImpl(void **ResultPtr,
     ZeDevDesc.pNext = &RelaxedDesc;
   }
 
-  ZE2UR_CALL(zeMemAllocShared, (Context->ZeContext, &ZeDevDesc, &ZeHostDesc,
-                                Size, Alignment, Device->ZeDevice, ResultPtr));
+  ze_result_t ZeResult =
+      zeMemAllocShared(Context->ZeContext, &ZeDevDesc, &ZeHostDesc, Size,
+                       Alignment, Device->ZeDevice, ResultPtr);
+  if (ZeResult != ZE_RESULT_SUCCESS) {
+    if (ZeResult == ZE_RESULT_ERROR_UNSUPPORTED_SIZE) {
+      return UR_RESULT_ERROR_INVALID_USM_SIZE;
+    }
+    return ze2urResult(ZeResult);
+  }
 
   UR_ASSERT(Alignment == 0 ||
                 reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0,
@@ -254,8 +268,14 @@ static ur_result_t USMHostAllocImpl(void **ResultPtr,
   // TODO: translate PI properties to Level Zero flags
   ZeStruct<ze_host_mem_alloc_desc_t> ZeHostDesc;
   ZeHostDesc.flags = 0;
-  ZE2UR_CALL(zeMemAllocHost,
-             (Context->ZeContext, &ZeHostDesc, Size, Alignment, ResultPtr));
+  ze_result_t ZeResult = zeMemAllocHost(Context->ZeContext, &ZeHostDesc, Size,
+                                        Alignment, ResultPtr);
+  if (ZeResult != ZE_RESULT_SUCCESS) {
+    if (ZeResult == ZE_RESULT_ERROR_UNSUPPORTED_SIZE) {
+      return UR_RESULT_ERROR_INVALID_USM_SIZE;
+    }
+    return ze2urResult(ZeResult);
+  }
 
   UR_ASSERT(Alignment == 0 ||
                 reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0,
@@ -599,6 +619,40 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMGetMemAllocInfo(
     ZE2UR_CALL(zeMemGetAddressRange, (Context->ZeContext, Ptr, nullptr, &Size));
     return ReturnValue(Size);
   }
+  case UR_USM_ALLOC_INFO_POOL: {
+    auto UMFPool = umfPoolByPtr(Ptr);
+    if (!UMFPool) {
+      return UR_RESULT_ERROR_INVALID_VALUE;
+    }
+
+    std::shared_lock<ur_shared_mutex> ContextLock(Context->Mutex);
+
+    auto SearchMatchingPool =
+        [](std::unordered_map<ur_device_handle_t, umf::pool_unique_handle_t>
+               &PoolMap,
+           umf_memory_pool_handle_t UMFPool) {
+          for (auto &PoolPair : PoolMap) {
+            if (PoolPair.second.get() == UMFPool) {
+              return true;
+            }
+          }
+          return false;
+        };
+
+    for (auto &Pool : Context->UsmPoolHandles) {
+      if (SearchMatchingPool(Pool->DeviceMemPools, UMFPool)) {
+        return ReturnValue(Pool);
+      }
+      if (SearchMatchingPool(Pool->SharedMemPools, UMFPool)) {
+        return ReturnValue(Pool);
+      }
+      if (Pool->HostMemPool.get() == UMFPool) {
+        return ReturnValue(Pool);
+      }
+    }
+
+    return UR_RESULT_ERROR_INVALID_VALUE;
+  }
   default:
     urPrint("urUSMGetMemAllocInfo: unsupported ParamName\n");
     return UR_RESULT_ERROR_INVALID_VALUE;
@@ -748,6 +802,7 @@ ur_result_t L0HostMemoryProvider::allocateImpl(void **ResultPtr, size_t Size,
 ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t Context,
                                              ur_usm_pool_desc_t *PoolDesc) {
 
+  this->Context = Context;
   zeroInit = static_cast<uint32_t>(PoolDesc->flags &
                                    UR_USM_POOL_FLAG_ZERO_INITIALIZE_BLOCK);
 
@@ -831,6 +886,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolCreate(
   try {
     *Pool = reinterpret_cast<ur_usm_pool_handle_t>(
         new ur_usm_pool_handle_t_(Context, PoolDesc));
+
+    std::shared_lock<ur_shared_mutex> ContextLock(Context->Mutex);
+    Context->UsmPoolHandles.insert(Context->UsmPoolHandles.cend(), *Pool);
+
   } catch (const UsmAllocationException &Ex) {
     return Ex.getError();
   }
@@ -848,6 +907,8 @@ ur_result_t
 urUSMPoolRelease(ur_usm_pool_handle_t Pool ///< [in] pointer to USM memory pool
 ) {
   if (Pool->RefCount.decrementAndTest()) {
+    std::shared_lock<ur_shared_mutex> ContextLock(Pool->Context->Mutex);
+    Pool->Context->UsmPoolHandles.remove(Pool);
     delete Pool;
   }
   return UR_RESULT_SUCCESS;
@@ -861,13 +922,19 @@ ur_result_t urUSMPoolGetInfo(
                      ///< property
     size_t *PropSizeRet ///< [out] size in bytes returned in pool property value
 ) {
-  std::ignore = Pool;
-  std::ignore = PropName;
-  std::ignore = PropSize;
-  std::ignore = PropValue;
-  std::ignore = PropSizeRet;
-  urPrint("[UR][L0] %s function not implemented!\n", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  UrReturnHelper ReturnValue(PropSize, PropValue, PropSizeRet);
+
+  switch (PropName) {
+  case UR_USM_POOL_INFO_REFERENCE_COUNT: {
+    return ReturnValue(Pool->RefCount.load());
+  }
+  case UR_USM_POOL_INFO_CONTEXT: {
+    return ReturnValue(Pool->Context);
+  }
+  default: {
+    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
+  }
+  }
 }
 
 // If indirect access tracking is not enabled then this functions just performs

source/adapters/level_zero/usm.hpp

@@ -29,6 +29,8 @@ struct ur_usm_pool_handle_t_ : _ur_object {
       SharedReadOnlyMemPools;
   umf::pool_unique_handle_t HostMemPool;
 
+  ur_context_handle_t Context{};
+
   ur_usm_pool_handle_t_(ur_context_handle_t Context,
                         ur_usm_pool_desc_t *PoolDesc);
 };

test/conformance/usm/usm_adapter_level_zero.match

@@ -1,11 +0,0 @@
-urUSMDeviceAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMDeviceAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMAllocInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_USM_ALLOC_INFO_POOL
-urUSMHostAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMHostAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMPoolGetInfoTestWithInfoParam.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_USM_POOL_INFO_CONTEXT
-urUSMPoolGetInfoTestWithInfoParam.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_USM_POOL_INFO_REFERENCE_COUNT
-urUSMPoolGetInfoTest.InvalidSizeTooSmall/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMPoolRetainTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMSharedAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMSharedAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled