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

Signed-off-by: Jaime Arteaga <jaime.a.arteaga.molina@intel.com>

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

@@ -185,8 +185,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,
@@ -224,8 +231,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,
@@ -252,8 +266,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,
@@ -597,6 +617,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;
@@ -746,6 +800,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);
 
@@ -829,6 +884,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();
  }
@@ -846,6 +905,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;
@@ -859,13 +920,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,36 +0,0 @@
-urUSMDeviceAllocTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMDeviceAllocTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMDeviceAllocTest.SuccessWithDescriptors/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMDeviceAllocTest.SuccessWithDescriptors/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMDeviceAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMDeviceAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMFreeTest.SuccessDeviceAlloc/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMFreeTest.SuccessHostAlloc/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMFreeTest.SuccessSharedAlloc/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMAllocInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_USM_ALLOC_INFO_TYPE
-urUSMAllocInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_USM_ALLOC_INFO_BASE_PTR
-urUSMAllocInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_USM_ALLOC_INFO_SIZE
-urUSMAllocInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_USM_ALLOC_INFO_DEVICE
-urUSMAllocInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_USM_ALLOC_INFO_POOL
-urUSMGetMemAllocInfoTest.InvalidNullHandleContext/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMGetMemAllocInfoTest.InvalidNullPointerMem/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMGetMemAllocInfoTest.InvalidEnumeration/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMGetMemAllocInfoTest.InvalidValuePropSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urUSMHostAllocTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMHostAllocTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMHostAllocTest.SuccessWithDescriptors/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMHostAllocTest.SuccessWithDescriptors/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-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.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMSharedAllocTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMSharedAllocTest.SuccessWithDescriptors/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMSharedAllocTest.SuccessWithDescriptors/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMSharedAllocTest.SuccessWithMultipleAdvices/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMSharedAllocTest.SuccessWithMultipleAdvices/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled
-urUSMSharedAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled
-urUSMSharedAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled