diff --git a/source/adapters/level_zero/context.cpp b/source/adapters/level_zero/context.cpp index c9496dda2c..72c573a2ea 100644 --- a/source/adapters/level_zero/context.cpp +++ b/source/adapters/level_zero/context.cpp @@ -178,111 +178,105 @@ UR_APIEXPORT ur_result_t UR_APICALL urContextSetExtendedDeleter( return UR_RESULT_ERROR_UNSUPPORTED_FEATURE; } +// Template helper function for creating USM pools for given pool descriptor. +template +std::pair +createUMFPoolForDesc(usm::pool_descriptor &Desc, Args &&...args) { + umf_result_t UmfRet = UMF_RESULT_SUCCESS; + umf::provider_unique_handle_t MemProvider = nullptr; + + switch (Desc.type) { + case UR_USM_TYPE_HOST: { + std::tie(UmfRet, MemProvider) = + umf::memoryProviderMakeUnique(Desc.hContext, + Desc.hDevice); + break; + } + case UR_USM_TYPE_DEVICE: { + std::tie(UmfRet, MemProvider) = + umf::memoryProviderMakeUnique(Desc.hContext, + Desc.hDevice); + break; + } + case UR_USM_TYPE_SHARED: { + if (Desc.deviceReadOnly) { + std::tie(UmfRet, MemProvider) = + umf::memoryProviderMakeUnique( + Desc.hContext, Desc.hDevice); + } else { + std::tie(UmfRet, MemProvider) = + umf::memoryProviderMakeUnique(Desc.hContext, + Desc.hDevice); + } + break; + } + default: + assert(0 && "Invalid UR descriptor type!"); + } + + if (UmfRet) + return std::pair{ + umf::umf2urResult(UmfRet), nullptr}; + + umf::pool_unique_handle_t Pool = nullptr; + std::tie(UmfRet, Pool) = + umf::poolMakeUnique({std::move(MemProvider)}, args...); + + return std::pair{ + umf::umf2urResult(UmfRet), std::move(Pool)}; +}; + ur_result_t ur_context_handle_t_::initialize() { - // Helper lambda to create various USM allocators for a device. - // Note that the CCS devices and their respective subdevices share a - // common ze_device_handle and therefore, also share USM allocators. - auto createUSMAllocators = [this](ur_device_handle_t Device) { - auto MemProvider = umf::memoryProviderMakeUnique( - reinterpret_cast(this), Device) - .second; - DeviceMemPools.emplace( - std::piecewise_construct, std::make_tuple(Device->ZeDevice), - std::make_tuple(umf::poolMakeUnique( - {std::move(MemProvider)}, - DisjointPoolConfigInstance - .Configs[usm::DisjointPoolMemType::Device]) - .second)); - - MemProvider = umf::memoryProviderMakeUnique( - reinterpret_cast(this), Device) - .second; - SharedMemPools.emplace( - std::piecewise_construct, std::make_tuple(Device->ZeDevice), - std::make_tuple(umf::poolMakeUnique( - {std::move(MemProvider)}, - DisjointPoolConfigInstance - .Configs[usm::DisjointPoolMemType::Shared]) - .second)); - - MemProvider = umf::memoryProviderMakeUnique( - reinterpret_cast(this), Device) - .second; - SharedReadOnlyMemPools.emplace( - std::piecewise_construct, std::make_tuple(Device->ZeDevice), - std::make_tuple( - umf::poolMakeUnique( - {std::move(MemProvider)}, - DisjointPoolConfigInstance - .Configs[usm::DisjointPoolMemType::SharedReadOnly]) - .second)); - - MemProvider = umf::memoryProviderMakeUnique( - reinterpret_cast(this), Device) - .second; - DeviceMemProxyPools.emplace( - std::piecewise_construct, std::make_tuple(Device->ZeDevice), - std::make_tuple( - umf::poolMakeUnique({std::move(MemProvider)}) - .second)); - - MemProvider = umf::memoryProviderMakeUnique( - reinterpret_cast(this), Device) - .second; - SharedMemProxyPools.emplace( - std::piecewise_construct, std::make_tuple(Device->ZeDevice), - std::make_tuple( - umf::poolMakeUnique({std::move(MemProvider)}) - .second)); - - MemProvider = umf::memoryProviderMakeUnique( - reinterpret_cast(this), Device) - .second; - SharedReadOnlyMemProxyPools.emplace( - std::piecewise_construct, std::make_tuple(Device->ZeDevice), - std::make_tuple( - umf::poolMakeUnique({std::move(MemProvider)}) - .second)); + auto Context = reinterpret_cast(this); + + ur_result_t Ret; + std::tie(Ret, PoolManager) = + usm::pool_manager::create(); + if (Ret) + return Ret; + + std::vector Descs; + // Create pool descriptor for every device and subdevice + std::tie(Ret, Descs) = usm::pool_descriptor::create(nullptr, Context); + if (Ret) + return Ret; + + auto descTypeToDisjointPoolType = + [](usm::pool_descriptor &Desc) -> usm::DisjointPoolMemType { + switch (Desc.type) { + case UR_USM_TYPE_HOST: + return usm::DisjointPoolMemType::Host; + case UR_USM_TYPE_DEVICE: + return usm::DisjointPoolMemType::Device; + case UR_USM_TYPE_SHARED: + return (Desc.deviceReadOnly) ? usm::DisjointPoolMemType::SharedReadOnly + : usm::DisjointPoolMemType::Shared; + default: + assert(0 && "Invalid pool descriptor type!"); + // Added to suppress 'not all control paths return a value' warning. + return usm::DisjointPoolMemType::All; + } }; - // Recursive helper to call createUSMAllocators for all sub-devices - std::function createUSMAllocatorsRecursive; - createUSMAllocatorsRecursive = - [createUSMAllocators, - &createUSMAllocatorsRecursive](ur_device_handle_t Device) -> void { - createUSMAllocators(Device); - for (auto &SubDevice : Device->SubDevices) - createUSMAllocatorsRecursive(SubDevice); - }; + // Create USM pool for each pool descriptor and add it to pool manager. + for (auto &Desc : Descs) { + umf::pool_unique_handle_t Pool = nullptr; + auto PoolType = descTypeToDisjointPoolType(Desc); - // Create USM pool for each pair (device, context). - // - for (auto &Device : Devices) { - createUSMAllocatorsRecursive(Device); - } - // Create USM pool for host. Device and Shared USM allocations - // are device-specific. Host allocations are not device-dependent therefore - // we don't need a map with device as key. - auto MemProvider = umf::memoryProviderMakeUnique( - reinterpret_cast(this), nullptr) - .second; - HostMemPool = - umf::poolMakeUnique( - {std::move(MemProvider)}, - DisjointPoolConfigInstance.Configs[usm::DisjointPoolMemType::Host]) - .second; - - MemProvider = umf::memoryProviderMakeUnique( - reinterpret_cast(this), nullptr) - .second; - HostMemProxyPool = - umf::poolMakeUnique({std::move(MemProvider)}).second; - - // We may allocate memory to this root device so create allocators. - if (SingleRootDevice && - DeviceMemPools.find(SingleRootDevice->ZeDevice) == DeviceMemPools.end()) { - createUSMAllocators(SingleRootDevice); + std::tie(Ret, Pool) = createUMFPoolForDesc( + Desc, DisjointPoolConfigInstance.Configs[PoolType]); + if (Ret) + return Ret; + + PoolManager.addPool(Desc, Pool); + + umf::pool_unique_handle_t ProxyPool = nullptr; + std::tie(Ret, ProxyPool) = createUMFPoolForDesc(Desc); + if (Ret) + return Ret; + + ProxyPoolManager.addPool(Desc, ProxyPool); } // Create the immediate command list to be used for initializations. diff --git a/source/adapters/level_zero/context.hpp b/source/adapters/level_zero/context.hpp index 94935ee59e..cba37ea832 100644 --- a/source/adapters/level_zero/context.hpp +++ b/source/adapters/level_zero/context.hpp @@ -26,6 +26,7 @@ #include "queue.hpp" #include +#include struct ur_context_handle_t_ : _ur_object { ur_context_handle_t_(ze_context_handle_t ZeContext, uint32_t NumDevices, @@ -96,15 +97,8 @@ struct ur_context_handle_t_ : _ur_object { // Store USM pool for USM shared and device allocations. There is 1 memory // pool per each pair of (context, device) per each memory type. - std::unordered_map - DeviceMemPools; - std::unordered_map - SharedMemPools; - std::unordered_map - SharedReadOnlyMemPools; - - // Store the host memory pool. It does not depend on any device. - umf::pool_unique_handle_t HostMemPool; + usm::pool_manager PoolManager; + usm::pool_manager ProxyPoolManager; // Allocation-tracking proxy pools for direct allocations. No pooling used. std::unordered_map @@ -249,3 +243,8 @@ struct ur_context_handle_t_ : _ur_object { // mutex guarding the container with contexts because the context can be removed // from the list of tracked contexts. ur_result_t ContextReleaseHelper(ur_context_handle_t Context); + +// Template helper function for creating USM pools for given pool descriptor. +template +std::pair +createUMFPoolForDesc(usm::pool_descriptor &Desc, Args &&...args); diff --git a/source/adapters/level_zero/usm.cpp b/source/adapters/level_zero/usm.cpp index d06a0353e4..47ad5c321f 100644 --- a/source/adapters/level_zero/usm.cpp +++ b/source/adapters/level_zero/usm.cpp @@ -295,16 +295,21 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMHostAlloc( // find the allocator depending on context as we do for Shared and Device // allocations. umf_memory_pool_handle_t hPoolInternal = nullptr; + usm::pool_descriptor Desc = {nullptr, Context, nullptr, UR_USM_TYPE_HOST, + false}; if (!UseUSMAllocator || // L0 spec says that allocation fails if Alignment != 2^n, in order to // keep the same behavior for the allocator, just call L0 API directly and // return the error code. ((Align & (Align - 1)) != 0)) { - hPoolInternal = Context->HostMemProxyPool.get(); + auto hPoolInternalOpt = Context->ProxyPoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } else if (Pool) { - hPoolInternal = Pool->HostMemPool.get(); + auto hPoolInternalOpt = Pool->PoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } else { - hPoolInternal = Context->HostMemPool.get(); + auto hPoolInternalOpt = Context->PoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } *RetMem = umfPoolAlignedMalloc(hPoolInternal, Size, Align); @@ -366,24 +371,21 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMDeviceAlloc( } umf_memory_pool_handle_t hPoolInternal = nullptr; + usm::pool_descriptor Desc = {nullptr, Context, Device, UR_USM_TYPE_DEVICE, + false}; if (!UseUSMAllocator || // L0 spec says that allocation fails if Alignment != 2^n, in order to // keep the same behavior for the allocator, just call L0 API directly and // return the error code. ((Alignment & (Alignment - 1)) != 0)) { - auto It = Context->DeviceMemProxyPools.find(Device->ZeDevice); - if (It == Context->DeviceMemProxyPools.end()) - return UR_RESULT_ERROR_INVALID_VALUE; - - hPoolInternal = It->second.get(); + auto hPoolInternalOpt = Context->ProxyPoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } else if (Pool) { - hPoolInternal = Pool->DeviceMemPools[Device].get(); + auto hPoolInternalOpt = Pool->PoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } else { - auto It = Context->DeviceMemPools.find(Device->ZeDevice); - if (It == Context->DeviceMemPools.end()) - return UR_RESULT_ERROR_INVALID_VALUE; - - hPoolInternal = It->second.get(); + auto hPoolInternalOpt = Context->PoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } *RetMem = umfPoolAlignedMalloc(hPoolInternal, Size, Alignment); @@ -466,30 +468,21 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMSharedAlloc( } umf_memory_pool_handle_t hPoolInternal = nullptr; + usm::pool_descriptor Desc = {nullptr, Context, Device, UR_USM_TYPE_SHARED, + DeviceReadOnly}; if (!UseUSMAllocator || // L0 spec says that allocation fails if Alignment != 2^n, in order to // keep the same behavior for the allocator, just call L0 API directly and // return the error code. ((Alignment & (Alignment - 1)) != 0)) { - auto &Allocator = (DeviceReadOnly ? Context->SharedReadOnlyMemProxyPools - : Context->SharedMemProxyPools); - auto It = Allocator.find(Device->ZeDevice); - if (It == Allocator.end()) - return UR_RESULT_ERROR_INVALID_VALUE; - - hPoolInternal = It->second.get(); + auto hPoolInternalOpt = Context->ProxyPoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } else if (Pool) { - hPoolInternal = (DeviceReadOnly) - ? Pool->SharedReadOnlyMemPools[Device].get() - : Pool->SharedMemPools[Device].get(); + auto hPoolInternalOpt = Pool->PoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } else { - auto &Allocator = (DeviceReadOnly ? Context->SharedReadOnlyMemPools - : Context->SharedMemPools); - auto It = Allocator.find(Device->ZeDevice); - if (It == Allocator.end()) - return UR_RESULT_ERROR_INVALID_VALUE; - - hPoolInternal = It->second.get(); + auto hPoolInternalOpt = Context->PoolManager.getPool(Desc); + hPoolInternal = hPoolInternalOpt.value(); } *RetMem = umfPoolAlignedMalloc(hPoolInternal, Size, Alignment); @@ -753,50 +746,52 @@ ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t Context, pNext = const_cast(BaseDesc->pNext); } - auto MemProvider = - umf::memoryProviderMakeUnique(Context, nullptr) - .second; - - HostMemPool = - umf::poolMakeUnique( - {std::move(MemProvider)}, - this->DisjointPoolConfigs.Configs[usm::DisjointPoolMemType::Host]) - .second; - - for (auto device : Context->Devices) { - MemProvider = - umf::memoryProviderMakeUnique(Context, device) - .second; - DeviceMemPools.emplace( - std::piecewise_construct, std::make_tuple(device), - std::make_tuple(umf::poolMakeUnique( - {std::move(MemProvider)}, - this->DisjointPoolConfigs - .Configs[usm::DisjointPoolMemType::Device]) - .second)); - - MemProvider = - umf::memoryProviderMakeUnique(Context, device) - .second; - SharedMemPools.emplace( - std::piecewise_construct, std::make_tuple(device), - std::make_tuple(umf::poolMakeUnique( - {std::move(MemProvider)}, - this->DisjointPoolConfigs - .Configs[usm::DisjointPoolMemType::Shared]) - .second)); - - MemProvider = umf::memoryProviderMakeUnique( - Context, device) - .second; - SharedReadOnlyMemPools.emplace( - std::piecewise_construct, std::make_tuple(device), - std::make_tuple( - umf::poolMakeUnique( - {std::move(MemProvider)}, - this->DisjointPoolConfigs - .Configs[usm::DisjointPoolMemType::SharedReadOnly]) - .second)); + ur_result_t Ret; + std::tie(Ret, PoolManager) = + usm::pool_manager::create(); + if (Ret) { + urPrint("urUSMPoolCreate: unexpected internal error\n"); + throw UsmAllocationException(Ret); + } + + std::vector Descs; + // Create pool descriptor for every device and subdevice + std::tie(Ret, Descs) = usm::pool_descriptor::create(nullptr, Context); + if (Ret) { + urPrint("urUSMPoolCreate: unexpected internal error\n"); + throw UsmAllocationException(Ret); + } + + auto descTypeToDisjointPoolType = + [](usm::pool_descriptor &Desc) -> usm::DisjointPoolMemType { + switch (Desc.type) { + case UR_USM_TYPE_HOST: + return usm::DisjointPoolMemType::Host; + case UR_USM_TYPE_DEVICE: + return usm::DisjointPoolMemType::Device; + case UR_USM_TYPE_SHARED: + return (Desc.deviceReadOnly) ? usm::DisjointPoolMemType::SharedReadOnly + : usm::DisjointPoolMemType::Shared; + default: + assert(0 && "Invalid pool descriptor type!"); + // Added to suppress 'not all control paths return a value' warning. + return usm::DisjointPoolMemType::All; + } + }; + + // Create USM pool for each pool descriptor and add it to pool manager. + for (auto &Desc : Descs) { + umf::pool_unique_handle_t Pool = nullptr; + auto PoolType = descTypeToDisjointPoolType(Desc); + + std::tie(Ret, Pool) = createUMFPoolForDesc( + Desc, DisjointPoolConfigInstance.Configs[PoolType]); + if (Ret) { + urPrint("urUSMPoolCreate: unexpected internal error\n"); + throw UsmAllocationException(Ret); + } + + PoolManager.addPool(Desc, Pool); } } diff --git a/source/adapters/level_zero/usm.hpp b/source/adapters/level_zero/usm.hpp index 01e215c578..d9d00890d0 100644 --- a/source/adapters/level_zero/usm.hpp +++ b/source/adapters/level_zero/usm.hpp @@ -12,6 +12,7 @@ #include "common.hpp" #include +#include usm::DisjointPoolAllConfigs InitializeDisjointPoolConfig(); @@ -21,13 +22,7 @@ struct ur_usm_pool_handle_t_ : _ur_object { usm::DisjointPoolAllConfigs DisjointPoolConfigs = InitializeDisjointPoolConfig(); - std::unordered_map - DeviceMemPools; - std::unordered_map - SharedMemPools; - std::unordered_map - SharedReadOnlyMemPools; - umf::pool_unique_handle_t HostMemPool; + usm::pool_manager PoolManager; ur_usm_pool_handle_t_(ur_context_handle_t Context, ur_usm_pool_desc_t *PoolDesc);