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 daec0408fb..b0ac96921f 100644 --- a/source/adapters/level_zero/usm.cpp +++ b/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(*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(*ResultPtr) % Alignment == 0, @@ -252,8 +266,14 @@ 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 = 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, @@ -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 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; @@ -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(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( 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(); } @@ -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 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 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 9e275d805e..e69de29bb2 100644 --- a/test/conformance/usm/usm_adapter_level_zero.match +++ b/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