Skip to content

[UR][CUDA] Cleanup UMF pools creation code #18361

New issue

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

#

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

Already on GitHub? # to your account

Merged
merged 2 commits into from
May 12, 2025
Merged

[UR][CUDA] Cleanup UMF pools creation code #18361

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c658d22
[UR][CUDA] Cleanup UMF pools creation code
npmiller May 8, 2025
0ea21b1
Merge branch 'sycl' into umf-cleanup
npmiller May 12, 2025
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
31 changes: 31 additions & 0 deletions unified-runtime/source/adapters/cuda/common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
#include "common.hpp"
#include "logger/ur_logger.hpp"

#include "umf_helpers.hpp"

#include <cuda.h>
#include <nvml.h>

Expand Down Expand Up @@ -169,4 +171,33 @@ ur_result_t getProviderNativeError(const char *providerName, int32_t error) {

return UR_RESULT_ERROR_UNKNOWN;
}

ur_result_t CreateProviderPool(int cuDevice, void *cuContext,
umf_usm_memory_type_t type,
umf_memory_provider_handle_t *provider,
umf_memory_pool_handle_t *pool) {
umf_cuda_memory_provider_params_handle_t CUMemoryProviderParams = nullptr;
UMF_RETURN_UR_ERROR(
umfCUDAMemoryProviderParamsCreate(&CUMemoryProviderParams));
OnScopeExit Cleanup(
[=]() { umfCUDAMemoryProviderParamsDestroy(CUMemoryProviderParams); });

// Setup memory provider parameters
UMF_RETURN_UR_ERROR(
umfCUDAMemoryProviderParamsSetContext(CUMemoryProviderParams, cuContext));
UMF_RETURN_UR_ERROR(
umfCUDAMemoryProviderParamsSetDevice(CUMemoryProviderParams, cuDevice));
UMF_RETURN_UR_ERROR(
umfCUDAMemoryProviderParamsSetMemoryType(CUMemoryProviderParams, type));

// Create memory provider
UMF_RETURN_UR_ERROR(umfMemoryProviderCreate(
umfCUDAMemoryProviderOps(), CUMemoryProviderParams, provider));

// Create memory pool
UMF_RETURN_UR_ERROR(
umfPoolCreate(umfProxyPoolOps(), *provider, nullptr, 0, pool));

return UR_RESULT_SUCCESS;
}
} // namespace umf
34 changes: 6 additions & 28 deletions unified-runtime/source/adapters/cuda/common.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,10 @@
#include <ur/ur.hpp>

#include <umf/base.h>
#include <umf/memory_pool.h>
#include <umf/memory_provider.h>
#include <umf/providers/provider_cuda.h>

#define UMF_RETURN_UMF_ERROR(UmfResult) \
do { \
umf_result_t UmfResult_ = (UmfResult); \
if (UmfResult_ != UMF_RESULT_SUCCESS) { \
return UmfResult_; \
} \
} while (0)

ur_result_t mapErrorUR(CUresult Result);

/// Converts CUDA error into UR error codes, and outputs error information
Expand Down Expand Up @@ -57,24 +51,8 @@ extern thread_local char ErrorMessage[MaxMessageSize];
void setPluginSpecificMessage(CUresult cu_res);

namespace umf {

inline umf_result_t setCUMemoryProviderParams(
umf_cuda_memory_provider_params_handle_t CUMemoryProviderParams,
int cuDevice, void *cuContext, umf_usm_memory_type_t memType) {

umf_result_t UmfResult =
umfCUDAMemoryProviderParamsSetContext(CUMemoryProviderParams, cuContext);
UMF_RETURN_UMF_ERROR(UmfResult);

UmfResult =
umfCUDAMemoryProviderParamsSetDevice(CUMemoryProviderParams, cuDevice);
UMF_RETURN_UMF_ERROR(UmfResult);

UmfResult =
umfCUDAMemoryProviderParamsSetMemoryType(CUMemoryProviderParams, memType);
UMF_RETURN_UMF_ERROR(UmfResult);

return UMF_RESULT_SUCCESS;
}

ur_result_t CreateProviderPool(int cuDevice, void *cuContext,
umf_usm_memory_type_t type,
umf_memory_provider_handle_t *provider,
umf_memory_pool_handle_t *pool);
} // namespace umf
37 changes: 3 additions & 34 deletions unified-runtime/source/adapters/cuda/context.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,38 +78,6 @@ typedef void (*ur_context_extended_deleter_t)(void *user_data);
///
///

static ur_result_t
CreateHostMemoryProviderPool(ur_device_handle_t_ *DeviceHandle,
umf_memory_provider_handle_t *MemoryProviderHost,
umf_memory_pool_handle_t *MemoryPoolHost) {

*MemoryProviderHost = nullptr;
CUcontext context = DeviceHandle->getNativeContext();

umf_cuda_memory_provider_params_handle_t CUMemoryProviderParams = nullptr;
umf_result_t UmfResult =
umfCUDAMemoryProviderParamsCreate(&CUMemoryProviderParams);
UMF_RETURN_UR_ERROR(UmfResult);
OnScopeExit Cleanup(
[=]() { umfCUDAMemoryProviderParamsDestroy(CUMemoryProviderParams); });

UmfResult = umf::setCUMemoryProviderParams(
CUMemoryProviderParams, 0 /* cuDevice */, context, UMF_MEMORY_TYPE_HOST);
UMF_RETURN_UR_ERROR(UmfResult);

// create UMF CUDA memory provider and pool for the host memory
// (UMF_MEMORY_TYPE_HOST)
UmfResult = umfMemoryProviderCreate(
umfCUDAMemoryProviderOps(), CUMemoryProviderParams, MemoryProviderHost);
UMF_RETURN_UR_ERROR(UmfResult);

UmfResult = umfPoolCreate(umfProxyPoolOps(), *MemoryProviderHost, nullptr, 0,
MemoryPoolHost);
UMF_RETURN_UR_ERROR(UmfResult);

return UR_RESULT_SUCCESS;
}

struct ur_context_handle_t_ {

struct deleter_data {
Expand All @@ -132,8 +100,9 @@ struct ur_context_handle_t_ {
// Create UMF CUDA memory provider for the host memory
// (UMF_MEMORY_TYPE_HOST) from any device (Devices[0] is used here, because
// it is guaranteed to exist).
UR_CHECK_ERROR(CreateHostMemoryProviderPool(Devices[0], &MemoryProviderHost,
&MemoryPoolHost));
UR_CHECK_ERROR(umf::CreateProviderPool(
0, Devices[0]->getNativeContext(), UMF_MEMORY_TYPE_HOST,
&MemoryProviderHost, &MemoryPoolHost));
UR_CHECK_ERROR(urAdapterRetain(ur::cuda::adapter));
};

Expand Down
68 changes: 9 additions & 59 deletions unified-runtime/source/adapters/cuda/platform.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,62 +19,6 @@
#include <cuda.h>
#include <sstream>

static ur_result_t
CreateDeviceMemoryProvidersPools(ur_platform_handle_t_ *Platform) {
umf_cuda_memory_provider_params_handle_t CUMemoryProviderParams = nullptr;

umf_result_t UmfResult =
umfCUDAMemoryProviderParamsCreate(&CUMemoryProviderParams);
UMF_RETURN_UR_ERROR(UmfResult);

OnScopeExit Cleanup(
[=]() { umfCUDAMemoryProviderParamsDestroy(CUMemoryProviderParams); });

for (auto &Device : Platform->Devices) {
ur_device_handle_t_ *device_handle = Device.get();
CUdevice device = device_handle->get();
CUcontext context = device_handle->getNativeContext();

// create UMF CUDA memory provider for the device memory
// (UMF_MEMORY_TYPE_DEVICE)
UmfResult = umf::setCUMemoryProviderParams(CUMemoryProviderParams, device,
context, UMF_MEMORY_TYPE_DEVICE);
UMF_RETURN_UR_ERROR(UmfResult);

UmfResult = umfMemoryProviderCreate(umfCUDAMemoryProviderOps(),
CUMemoryProviderParams,
&device_handle->MemoryProviderDevice);
UMF_RETURN_UR_ERROR(UmfResult);

// create UMF CUDA memory provider for the shared memory
// (UMF_MEMORY_TYPE_SHARED)
UmfResult = umf::setCUMemoryProviderParams(CUMemoryProviderParams, device,
context, UMF_MEMORY_TYPE_SHARED);
UMF_RETURN_UR_ERROR(UmfResult);

UmfResult = umfMemoryProviderCreate(umfCUDAMemoryProviderOps(),
CUMemoryProviderParams,
&device_handle->MemoryProviderShared);
UMF_RETURN_UR_ERROR(UmfResult);

// create UMF CUDA memory pool for the device memory
// (UMF_MEMORY_TYPE_DEVICE)
UmfResult =
umfPoolCreate(umfProxyPoolOps(), device_handle->MemoryProviderDevice,
nullptr, 0, &device_handle->MemoryPoolDevice);
UMF_RETURN_UR_ERROR(UmfResult);

// create UMF CUDA memory pool for the shared memory
// (UMF_MEMORY_TYPE_SHARED)
UmfResult =
umfPoolCreate(umfProxyPoolOps(), device_handle->MemoryProviderShared,
nullptr, 0, &device_handle->MemoryPoolShared);
UMF_RETURN_UR_ERROR(UmfResult);
}

return UR_RESULT_SUCCESS;
}

UR_APIEXPORT ur_result_t UR_APICALL
urPlatformGetInfo(ur_platform_handle_t, ur_platform_info_t PlatformInfoType,
size_t Size, void *pPlatformInfo, size_t *pSizeRet) {
Expand Down Expand Up @@ -148,10 +92,16 @@ urPlatformGet(ur_adapter_handle_t, uint32_t, ur_platform_handle_t *phPlatforms,
new ur_device_handle_t_{Device, Context, EvBase,
ur::cuda::adapter->Platform.get(),
static_cast<uint32_t>(i)});
}

UR_CHECK_ERROR(CreateDeviceMemoryProvidersPools(
ur::cuda::adapter->Platform.get()));
// Create UMF memory providers and pools
auto &dev = ur::cuda::adapter->Platform->Devices.back();
UR_CHECK_ERROR(umf::CreateProviderPool(
Device, Context, UMF_MEMORY_TYPE_DEVICE,
&dev->MemoryProviderDevice, &dev->MemoryPoolDevice));
UR_CHECK_ERROR(umf::CreateProviderPool(
Device, Context, UMF_MEMORY_TYPE_SHARED,
&dev->MemoryProviderShared, &dev->MemoryPoolShared));
}
} catch (const std::bad_alloc &) {
// Signal out-of-memory situation
for (int i = 0; i < NumDevices; ++i) {
Expand Down