diff --git a/source/adapters/level_zero/context.hpp b/source/adapters/level_zero/context.hpp index 94935ee59e..96935d470e 100644 --- a/source/adapters/level_zero/context.hpp +++ b/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 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 diff --git a/source/adapters/level_zero/usm.cpp b/source/adapters/level_zero/usm.cpp index c6d98855e7..d2dfc9b37d 100644 --- a/source/adapters/level_zero/usm.cpp +++ b/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 = ZE_CALL_NOCHECK( + 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(*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 = ZE_CALL_NOCHECK( + 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(*ResultPtr) % Alignment == 0, @@ -254,8 +268,15 @@ static ur_result_t USMHostAllocImpl(void **ResultPtr, // TODO: translate PI properties to Level Zero flags ZeStruct ZeHostDesc; ZeHostDesc.flags = 0; - ZE2UR_CALL(zeMemAllocHost, - (Context->ZeContext, &ZeHostDesc, Size, Alignment, ResultPtr)); + ze_result_t ZeResult = + ZE_CALL_NOCHECK(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(*ResultPtr) % Alignment == 0, @@ -599,6 +620,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 ContextLock(Context->Mutex); + + auto SearchMatchingPool = + [](std::unordered_map + &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 +803,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(PoolDesc->flags & UR_USM_POOL_FLAG_ZERO_INITIALIZE_BLOCK); @@ -831,6 +887,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolCreate( try { *Pool = reinterpret_cast( new ur_usm_pool_handle_t_(Context, PoolDesc)); + + std::shared_lock ContextLock(Context->Mutex); + Context->UsmPoolHandles.insert(Context->UsmPoolHandles.cend(), *Pool); + } catch (const UsmAllocationException &Ex) { return Ex.getError(); } @@ -848,6 +908,8 @@ ur_result_t urUSMPoolRelease(ur_usm_pool_handle_t Pool ///< [in] pointer to USM memory pool ) { if (Pool->RefCount.decrementAndTest()) { + std::shared_lock ContextLock(Pool->Context->Mutex); + Pool->Context->UsmPoolHandles.remove(Pool); delete Pool; } return UR_RESULT_SUCCESS; @@ -861,13 +923,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 diff --git a/source/adapters/level_zero/usm.hpp b/source/adapters/level_zero/usm.hpp index 01e215c578..958fca9354 100644 --- a/source/adapters/level_zero/usm.hpp +++ b/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); }; diff --git a/test/conformance/usm/usm_adapter_level_zero.match b/test/conformance/usm/usm_adapter_level_zero.match index bf45b83ec2..c036fa785c 100644 --- a/test/conformance/usm/usm_adapter_level_zero.match +++ b/test/conformance/usm/usm_adapter_level_zero.match @@ -1,11 +1,2 @@ -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 +{{OPT}}urUSMDeviceAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolEnabled +{{OPT}}urUSMDeviceAllocTest.InvalidUSMSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UsePoolDisabled