Skip to content

[UR][HIP] Implement USM memory pool experimental API #22673

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

Open
zjin-lcf wants to merge 2 commits into
base: sycl
Choose a base branch
from
Open

[UR][HIP] Implement USM memory pool experimental API #22673

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
b515c67
[UR][HIP] Implement async device allocation
zjin-lcf Jul 18, 2026
f7b51d0
[UR][HIP] Implement USM memory pool experimental API
zjin-lcf Jul 18, 2026
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
96 changes: 82 additions & 14 deletions unified-runtime/source/adapters/hip/async_alloc.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
//===--------- async_alloc.cpp - CUDA Adapter -----------------------------===//
//===--------- async_alloc.cpp - HIP Adapter ------------------------------===//
//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM
Expand All @@ -9,31 +9,99 @@

#include <unified-runtime/ur_api.h>

UR_APIEXPORT ur_result_t urEnqueueUSMDeviceAllocExp(
ur_queue_handle_t, ur_usm_pool_handle_t, const size_t,
const ur_exp_async_usm_alloc_properties_t *, uint32_t,
const ur_event_handle_t *, void **, ur_event_handle_t *) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
#include "common.hpp"
#include "context.hpp"
#include "enqueue.hpp"
#include "event.hpp"
#include "queue.hpp"
#include "usm.hpp"

UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMDeviceAllocExp(
ur_queue_handle_t hQueue, ur_usm_pool_handle_t hPool, const size_t size,
const ur_exp_async_usm_alloc_properties_t *, uint32_t numEventsInWaitList,
const ur_event_handle_t *phEventWaitList, void **ppMem,
ur_event_handle_t *phEvent) try {
std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};

ScopedDevice Active(hQueue->getDevice());
uint32_t StreamToken;
ur_stream_guard Guard;
hipStream_t HIPStream = hQueue->getNextComputeStream(
numEventsInWaitList, phEventWaitList, Guard, &StreamToken);

UR_CHECK_ERROR(enqueueEventsWait(hQueue, HIPStream, numEventsInWaitList,
phEventWaitList));

if (phEvent) {
RetImplEvent = std::make_unique<ur_event_handle_t_>(
UR_COMMAND_ENQUEUE_USM_DEVICE_ALLOC_EXP, hQueue, HIPStream,
StreamToken);
UR_CHECK_ERROR(RetImplEvent->start());
}

// Allocate from the specified device pool if provided. Otherwise, allocate
// from the current device's default pool.
if (hPool) {
assert(hPool->usesHipPool());
UR_CHECK_ERROR(
hipMallocFromPoolAsync(ppMem, size, hPool->getHipPool(), HIPStream));
} else {
UR_CHECK_ERROR(hipMallocAsync(ppMem, size, HIPStream));
}

if (phEvent) {
UR_CHECK_ERROR(RetImplEvent->record());
*phEvent = RetImplEvent.release();
}

return UR_RESULT_SUCCESS;
} catch (ur_result_t Err) {
return Err;
}

UR_APIEXPORT ur_result_t urEnqueueUSMSharedAllocExp(
UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMSharedAllocExp(
ur_queue_handle_t, ur_usm_pool_handle_t, const size_t,
const ur_exp_async_usm_alloc_properties_t *, uint32_t,
const ur_event_handle_t *, void **, ur_event_handle_t *) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
}

UR_APIEXPORT ur_result_t urEnqueueUSMHostAllocExp(
UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMHostAllocExp(
ur_queue_handle_t, ur_usm_pool_handle_t, const size_t,
const ur_exp_async_usm_alloc_properties_t *, uint32_t,
const ur_event_handle_t *, void **, ur_event_handle_t *) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
}

UR_APIEXPORT ur_result_t urEnqueueUSMFreeExp(ur_queue_handle_t,
ur_usm_pool_handle_t, void *,
uint32_t,
const ur_event_handle_t *,
ur_event_handle_t *) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMFreeExp(
ur_queue_handle_t hQueue, [[maybe_unused]] ur_usm_pool_handle_t hPool,
void *pMem, uint32_t numEventsInWaitList,
const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) try {
std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};

ScopedDevice Active(hQueue->getDevice());
uint32_t StreamToken;
ur_stream_guard Guard;
hipStream_t HIPStream = hQueue->getNextComputeStream(
numEventsInWaitList, phEventWaitList, Guard, &StreamToken);

UR_CHECK_ERROR(enqueueEventsWait(hQueue, HIPStream, numEventsInWaitList,
phEventWaitList));

if (phEvent) {
RetImplEvent = std::make_unique<ur_event_handle_t_>(
UR_COMMAND_ENQUEUE_USM_FREE_EXP, hQueue, HIPStream, StreamToken);
UR_CHECK_ERROR(RetImplEvent->start());
}

UR_CHECK_ERROR(hipFreeAsync(pMem, HIPStream));

if (phEvent) {
UR_CHECK_ERROR(RetImplEvent->record());
*phEvent = RetImplEvent.release();
}

return UR_RESULT_SUCCESS;
} catch (ur_result_t Err) {
return Err;
}
5 changes: 5 additions & 0 deletions unified-runtime/source/adapters/hip/device.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -999,6 +999,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
return ReturnValue(false);
case UR_DEVICE_INFO_ASYNC_BARRIER:
return ReturnValue(false);
case UR_DEVICE_INFO_ASYNC_USM_ALLOCATIONS_SUPPORT_EXP: {
int MemPoolsSupported =
getAttribute(hDevice, hipDeviceAttributeMemoryPoolsSupported);
return ReturnValue(static_cast<ur_bool_t>(MemPoolsSupported));
}
case UR_DEVICE_INFO_IL_VERSION:
return ReturnValue("");

Expand Down
4 changes: 4 additions & 0 deletions unified-runtime/source/adapters/hip/ur_interface_loader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -499,6 +499,10 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetEnqueueExpProcAddrTable(

pDdiTable->pfnTimestampRecordingExp = urEnqueueTimestampRecordingExp;
pDdiTable->pfnNativeCommandExp = urEnqueueNativeCommandExp;
pDdiTable->pfnUSMDeviceAllocExp = urEnqueueUSMDeviceAllocExp;
pDdiTable->pfnUSMSharedAllocExp = urEnqueueUSMSharedAllocExp;
pDdiTable->pfnUSMHostAllocExp = urEnqueueUSMHostAllocExp;
pDdiTable->pfnUSMFreeExp = urEnqueueUSMFreeExp;
pDdiTable->pfnCommandBufferExp = urEnqueueCommandBufferExp;
pDdiTable->pfnKernelLaunchWithArgsExp = urEnqueueKernelLaunchWithArgsExp;
pDdiTable->pfnGraphExp = urEnqueueGraphExp;
Expand Down
208 changes: 184 additions & 24 deletions unified-runtime/source/adapters/hip/usm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -413,6 +413,52 @@ bool ur_usm_pool_handle_t_::hasUMFPool(umf_memory_pool_t *umf_pool) {
HostMemPool.get() == umf_pool;
}

ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t Context,
ur_device_handle_t Device,
ur_usm_pool_desc_t *PoolDesc)
: handle_base(), Context{Context}, Device{Device} {
if (!(PoolDesc->flags & UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP))
throw UR_RESULT_ERROR_INVALID_ARGUMENT;

hipMemPoolProps MemPoolProps{};
size_t threshold = 0;

const void *pNext = PoolDesc->pNext;
while (pNext != nullptr) {
const ur_base_desc_t *BaseDesc = static_cast<const ur_base_desc_t *>(pNext);
switch (BaseDesc->stype) {
case UR_STRUCTURE_TYPE_USM_POOL_LIMITS_DESC: {
const ur_usm_pool_limits_desc_t *Limits =
reinterpret_cast<const ur_usm_pool_limits_desc_t *>(BaseDesc);
MemPoolProps.maxSize = Limits->maxPoolableSize;
maxSize = Limits->maxPoolableSize;
threshold = Limits->minDriverAllocSize;
break;
}
default: {
throw UR_RESULT_ERROR_INVALID_ARGUMENT;
}
}
pNext = BaseDesc->pNext;
}

MemPoolProps.allocType = hipMemAllocationTypePinned;
MemPoolProps.location.id = Device->getIndex();
MemPoolProps.location.type = hipMemLocationTypeDevice;
UR_CHECK_ERROR(hipMemPoolCreate(&HIPMemPool, &MemPoolProps));

// Release threshold is not a property when creating a pool.
// It must be set separately.
UR_CHECK_ERROR(urUSMPoolSetInfoExp(this,
UR_USM_POOL_INFO_RELEASE_THRESHOLD_EXP,
&threshold, 8 /*uint64_t*/));
}

ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t Context,
ur_device_handle_t Device,
hipMemPool_t HIPMemPool)
: handle_base(), Context{Context}, Device{Device}, HIPMemPool(HIPMemPool) {}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolCreate(
/// [in] handle of the context object
ur_context_handle_t Context,
Expand Down Expand Up @@ -488,34 +534,135 @@ bool checkUSMAlignment(uint32_t &alignment, const ur_usm_desc_t *pUSMDesc) {
(alignment == 0 || ((alignment & (alignment - 1)) == 0)));
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolCreateExp(ur_context_handle_t,
ur_device_handle_t,
ur_usm_pool_desc_t *,
ur_usm_pool_handle_t *) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
UR_APIEXPORT ur_result_t UR_APICALL
urUSMPoolCreateExp(ur_context_handle_t Context, ur_device_handle_t Device,
ur_usm_pool_desc_t *pPoolDesc, ur_usm_pool_handle_t *pPool) {
// This entry point only supports native mem pools.
if (!(pPoolDesc->flags & UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP))
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
try {
*pPool = reinterpret_cast<ur_usm_pool_handle_t>(
new ur_usm_pool_handle_t_(Context, Device, pPoolDesc));
} catch (ur_result_t e) {
return e;
} catch (...) {
return UR_RESULT_ERROR_UNKNOWN;
}
return UR_RESULT_SUCCESS;
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolDestroyExp(ur_context_handle_t,
ur_device_handle_t,
ur_usm_pool_handle_t) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
UR_APIEXPORT ur_result_t UR_APICALL
urUSMPoolDestroyExp(ur_context_handle_t hContext, ur_device_handle_t hDevice,
ur_usm_pool_handle_t hPool) {
UR_ASSERT(std::find(hContext->getDevices().begin(),
hContext->getDevices().end(),
hDevice) != hContext->getDevices().end(),
UR_RESULT_ERROR_INVALID_CONTEXT);
ScopedDevice Active(hDevice);

try {
UR_CHECK_ERROR(hipMemPoolDestroy(hPool->getHipPool()));
} catch (ur_result_t Err) {
return Err;
} catch (...) {
return UR_RESULT_ERROR_UNKNOWN;
}

return UR_RESULT_SUCCESS;
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolSetInfoExp(ur_usm_pool_handle_t,
ur_usm_pool_info_t,
void *, size_t) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolSetInfoExp(
ur_usm_pool_handle_t hPool, ur_usm_pool_info_t propName, void *pPropValue,
size_t) {
hipMemPoolAttr attr;

// All current values are expected to be of size uint64_t.
switch (propName) {
case UR_USM_POOL_INFO_RELEASE_THRESHOLD_EXP:
attr = hipMemPoolAttrReleaseThreshold;
break;
case UR_USM_POOL_INFO_RESERVED_HIGH_EXP:
attr = hipMemPoolAttrReservedMemHigh;
break;
case UR_USM_POOL_INFO_USED_HIGH_EXP:
attr = hipMemPoolAttrUsedMemHigh;
break;
default:
return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
}
try {
UR_CHECK_ERROR(
hipMemPoolSetAttribute(hPool->getHipPool(), attr, pPropValue));
} catch (ur_result_t Err) {
return Err;
} catch (...) {
return UR_RESULT_ERROR_UNKNOWN;
}
return UR_RESULT_SUCCESS;
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolGetDefaultDevicePoolExp(
ur_context_handle_t, ur_device_handle_t, ur_usm_pool_handle_t *) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
ur_context_handle_t hContext, ur_device_handle_t hDevice,
ur_usm_pool_handle_t *pPool) {
UR_ASSERT(std::find(hContext->getDevices().begin(),
hContext->getDevices().end(),
hDevice) != hContext->getDevices().end(),
UR_RESULT_ERROR_INVALID_CONTEXT);
ScopedDevice Active(hDevice);

try {
hipMemPool_t HipPool;
UR_CHECK_ERROR(
hipDeviceGetDefaultMemPool(&HipPool, hDevice->getIndex()));

*pPool = reinterpret_cast<ur_usm_pool_handle_t>(
new ur_usm_pool_handle_t_(hContext, hDevice, HipPool));
} catch (ur_result_t Err) {
return Err;
} catch (...) {
return UR_RESULT_ERROR_UNKNOWN;
}

return UR_RESULT_SUCCESS;
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolGetInfoExp(ur_usm_pool_handle_t,
ur_usm_pool_info_t,
void *, size_t *) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
UR_APIEXPORT ur_result_t UR_APICALL
urUSMPoolGetInfoExp(ur_usm_pool_handle_t hPool, ur_usm_pool_info_t propName,
void *pPropValue, size_t *pPropSizeRet) {
hipMemPoolAttr attr;

switch (propName) {
case UR_USM_POOL_INFO_RELEASE_THRESHOLD_EXP:
attr = hipMemPoolAttrReleaseThreshold;
break;
case UR_USM_POOL_INFO_RESERVED_CURRENT_EXP:
attr = hipMemPoolAttrReservedMemCurrent;
break;
case UR_USM_POOL_INFO_RESERVED_HIGH_EXP:
attr = hipMemPoolAttrReservedMemHigh;
break;
case UR_USM_POOL_INFO_USED_CURRENT_EXP:
attr = hipMemPoolAttrUsedMemCurrent;
break;
case UR_USM_POOL_INFO_USED_HIGH_EXP:
attr = hipMemPoolAttrUsedMemHigh;
break;
default:
return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
}

uint64_t value = 0;
UR_CHECK_ERROR(
hipMemPoolGetAttribute(hPool->getHipPool(), attr, (void *)&value));

if (pPropValue) {
*(size_t *)pPropValue = value;
}
if (pPropSizeRet) {
*(size_t *)pPropSizeRet = sizeof(size_t);
}

return UR_RESULT_SUCCESS;
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolGetDevicePoolExp(
Expand All @@ -528,11 +675,24 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolSetDevicePoolExp(
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolTrimToExp(ur_context_handle_t,
ur_device_handle_t,
ur_usm_pool_handle_t,
size_t) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
UR_APIEXPORT ur_result_t UR_APICALL
urUSMPoolTrimToExp(ur_context_handle_t hContext, ur_device_handle_t hDevice,
ur_usm_pool_handle_t hPool, size_t minBytesToKeep) {
UR_ASSERT(std::find(hContext->getDevices().begin(),
hContext->getDevices().end(),
hDevice) != hContext->getDevices().end(),
UR_RESULT_ERROR_INVALID_CONTEXT);
ScopedDevice Active(hDevice);

try {
UR_CHECK_ERROR(hipMemPoolTrimTo(hPool->getHipPool(), minBytesToKeep));
} catch (ur_result_t Err) {
return Err;
} catch (...) {
return UR_RESULT_ERROR_UNKNOWN;
}

return UR_RESULT_SUCCESS;
}

UR_APIEXPORT ur_result_t UR_APICALL urUSMContextMemcpyExp(ur_context_handle_t,
Expand Down
Loading
Loading