Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[SYCL][UR][CUDA] Access UMF pool handles through usm::pool_manager #957

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+134 −42
Open

[SYCL][UR][CUDA] Access UMF pool handles through usm::pool_manager #957

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 124 additions & 37 deletions source/adapters/cuda/usm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,10 +36,19 @@ urUSMHostAlloc(ur_context_handle_t hContext, const ur_usm_desc_t *pUSMDesc,
return USMHostAllocImpl(ppMem, hContext, nullptr, size, alignment);
}

auto UMFPool = hPool->HostMemPool.get();
*ppMem = umfPoolAlignedMalloc(UMFPool, size, alignment);
usm::pool_descriptor Desc = {hPool, hContext, nullptr, UR_USM_TYPE_HOST,
false};
auto hPoolInternalOpt = hPool->PoolManager.getPool(Desc);
if (!hPoolInternalOpt.has_value()) {
// Internal error, every usm pool should have Host, Device,
// Shared UMF pools.
return UR_RESULT_ERROR_UNKNOWN;
}

auto hPoolInternal = hPoolInternalOpt.value();
*ppMem = umfPoolAlignedMalloc(hPoolInternal, size, alignment);
if (*ppMem == nullptr) {
auto umfErr = umfPoolGetLastAllocationError(UMFPool);
auto umfErr = umfPoolGetLastAllocationError(hPoolInternal);
return umf::umf2urResult(umfErr);
}
return UR_RESULT_SUCCESS;
Expand All @@ -61,10 +70,19 @@ urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
alignment);
}

auto UMFPool = hPool->DeviceMemPool.get();
*ppMem = umfPoolAlignedMalloc(UMFPool, size, alignment);
usm::pool_descriptor Desc = {hPool, hContext, hDevice, UR_USM_TYPE_DEVICE,
false};
auto hPoolInternalOpt = hPool->PoolManager.getPool(Desc);
if (!hPoolInternalOpt.has_value()) {
// Internal error, every usm pool should have Host, Device,
// Shared UMF pools.
return UR_RESULT_ERROR_UNKNOWN;
}

auto hPoolInternal = hPoolInternalOpt.value();
*ppMem = umfPoolAlignedMalloc(hPoolInternal, size, alignment);
if (*ppMem == nullptr) {
auto umfErr = umfPoolGetLastAllocationError(UMFPool);
auto umfErr = umfPoolGetLastAllocationError(hPoolInternal);
return umf::umf2urResult(umfErr);
}
return UR_RESULT_SUCCESS;
Expand All @@ -86,10 +104,19 @@ urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
alignment);
}

auto UMFPool = hPool->SharedMemPool.get();
*ppMem = umfPoolAlignedMalloc(UMFPool, size, alignment);
usm::pool_descriptor Desc = {hPool, hContext, hDevice, UR_USM_TYPE_SHARED,
false};
auto hPoolInternalOpt = hPool->PoolManager.getPool(Desc);
if (!hPoolInternalOpt.has_value()) {
// Internal error, every usm pool should have Host, Device,
// Shared UMF pools.
return UR_RESULT_ERROR_UNKNOWN;
}

auto hPoolInternal = hPoolInternalOpt.value();
*ppMem = umfPoolAlignedMalloc(hPoolInternal, size, alignment);
if (*ppMem == nullptr) {
auto umfErr = umfPoolGetLastAllocationError(UMFPool);
auto umfErr = umfPoolGetLastAllocationError(hPoolInternal);
return umf::umf2urResult(umfErr);
}
return UR_RESULT_SUCCESS;
Expand Down Expand Up @@ -373,6 +400,48 @@ ur_result_t USMHostMemoryProvider::allocateImpl(void **ResultPtr, size_t Size,
return USMHostAllocImpl(ResultPtr, Context, nullptr, Size, Alignment);
}

// Template helper function for creating USM pools for given pool descriptor.
template <typename P, typename... Args>
std::pair<ur_result_t, umf::pool_unique_handle_t>
createUMFPoolForDesc(usm::pool_descriptor &Desc, Args &&...args) {
kswiecicki marked this conversation as resolved.
Show resolved Hide resolved
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<USMHostMemoryProvider>(Desc.hContext,
Desc.hDevice);
break;
}
case UR_USM_TYPE_DEVICE: {
std::tie(UmfRet, MemProvider) =
umf::memoryProviderMakeUnique<USMDeviceMemoryProvider>(Desc.hContext,
Desc.hDevice);
break;
}
case UR_USM_TYPE_SHARED: {
std::tie(UmfRet, MemProvider) =
umf::memoryProviderMakeUnique<USMSharedMemoryProvider>(Desc.hContext,
Desc.hDevice);
break;
}
default:
UmfRet = UMF_RESULT_ERROR_INVALID_ARGUMENT;
}

if (UmfRet)
return std::pair<ur_result_t, umf::pool_unique_handle_t>{
umf::umf2urResult(UmfRet), nullptr};

umf::pool_unique_handle_t Pool = nullptr;
std::tie(UmfRet, Pool) =
umf::poolMakeUnique<P, 1>({std::move(MemProvider)}, args...);

return std::pair<ur_result_t, umf::pool_unique_handle_t>{
umf::umf2urResult(UmfRet), std::move(Pool)};
};

ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t Context,
ur_usm_pool_desc_t *PoolDesc)
: Context(Context) {
Expand All @@ -396,40 +465,58 @@ ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t Context,
pNext = BaseDesc->pNext;
}

auto MemProvider =
umf::memoryProviderMakeUnique<USMHostMemoryProvider>(Context, nullptr)
.second;

HostMemPool =
umf::poolMakeUnique<usm::DisjointPool, 1>(
{std::move(MemProvider)},
this->DisjointPoolConfigs.Configs[usm::DisjointPoolMemType::Host])
.second;
ur_result_t Ret;
std::tie(Ret, PoolManager) =
usm::pool_manager<usm::pool_descriptor>::create();
if (Ret) {
throw UsmAllocationException(Ret);
}

auto Device = Context->DeviceID;
MemProvider =
umf::memoryProviderMakeUnique<USMDeviceMemoryProvider>(Context, Device)
.second;
DeviceMemPool =
umf::poolMakeUnique<usm::DisjointPool, 1>(
{std::move(MemProvider)},
this->DisjointPoolConfigs.Configs[usm::DisjointPoolMemType::Device])
.second;

MemProvider =
umf::memoryProviderMakeUnique<USMSharedMemoryProvider>(Context, Device)
.second;
SharedMemPool =
umf::poolMakeUnique<usm::DisjointPool, 1>(
{std::move(MemProvider)},
this->DisjointPoolConfigs.Configs[usm::DisjointPoolMemType::Shared])
.second;
auto UrUSMPool = reinterpret_cast<ur_usm_pool_handle_t>(this);

// TODO: Replace this with appropriate usm::pool_descriptor 'create' static
aarongreig marked this conversation as resolved.
Show resolved Hide resolved
// function.
usm::pool_descriptor Descs[] = {
{UrUSMPool, Context, nullptr, UR_USM_TYPE_HOST, false},
{UrUSMPool, Context, Device, UR_USM_TYPE_DEVICE, false},
{UrUSMPool, Context, Device, UR_USM_TYPE_SHARED, false}};

// Helper lambda function matching USM type to DisjointPoolMemType
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:
// Should not be reachable.
ur::unreachable();
}
};

for (auto &Desc : Descs) {
umf::pool_unique_handle_t Pool = nullptr;
auto PoolType = descTypeToDisjointPoolType(Desc);

std::tie(Ret, Pool) = createUMFPoolForDesc<usm::DisjointPool>(
Desc, this->DisjointPoolConfigs.Configs[PoolType]);
if (Ret) {
throw UsmAllocationException(Ret);
}

PoolManager.addPool(Desc, Pool);
}

Context->addPool(this);
}

bool ur_usm_pool_handle_t_::hasUMFPool(umf_memory_pool_t *umf_pool) {
return DeviceMemPool.get() == umf_pool || SharedMemPool.get() == umf_pool ||
HostMemPool.get() == umf_pool;
return PoolManager.hasPool(umf_pool);
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolCreate(
Expand Down
5 changes: 2 additions & 3 deletions source/adapters/cuda/usm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

#include <umf_helpers.hpp>
#include <umf_pools/disjoint_pool_config_parser.hpp>
#include <ur_pool_manager.hpp>

usm::DisjointPoolAllConfigs InitializeDisjointPoolConfig();

Expand All @@ -23,9 +24,7 @@ struct ur_usm_pool_handle_t_ {
usm::DisjointPoolAllConfigs DisjointPoolConfigs =
usm::DisjointPoolAllConfigs();

umf::pool_unique_handle_t DeviceMemPool;
umf::pool_unique_handle_t SharedMemPool;
umf::pool_unique_handle_t HostMemPool;
usm::pool_manager<usm::pool_descriptor> PoolManager;

ur_usm_pool_handle_t_(ur_context_handle_t Context,
ur_usm_pool_desc_t *PoolDesc);
Expand Down
10 changes: 8 additions & 2 deletions source/common/ur_pool_manager.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -249,13 +249,19 @@ template <typename D> struct pool_manager {
std::optional<umf_memory_pool_handle_t> getPool(const D &desc) noexcept {
auto it = descToPoolMap.find(desc);
if (it == descToPoolMap.end()) {
logger::error("Pool descriptor doesn't match any existing pool: {}",
desc);
logger::error(
"Pool descriptor: {}, doesn't match any existing pool", desc);
return std::nullopt;
}

return it->second.get();
}

bool hasPool(umf_memory_pool_handle_t hPool) noexcept {
return std::any_of(
descToPoolMap.begin(), descToPoolMap.end(),
[&hPool](const auto &pair) { return hPool == pair.second.get(); });
}
};

} // namespace usm
Expand Down
Loading