Skip to content

Commit

Permalink
Merge pull request #1569 from lbushi25/fix_usm_allocation
Browse files Browse the repository at this point in the history
[L0] Fix usm allocation functions when alignment is not a power of 2
  • Loading branch information
kbenzie authored May 23, 2024
2 parents 6d676f2 + 2d02c21 commit fb3cbd1
Showing 1 changed file with 9 additions and 18 deletions.
27 changes: 9 additions & 18 deletions source/adapters/level_zero/usm.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMHostAlloc(
uint32_t Align = USMDesc ? USMDesc->align : 0;
// L0 supports alignment up to 64KB and silently ignores higher values.
// We flag alignment > 64KB as an invalid value.
if (Align > 65536)
// L0 spec says that alignment values that are not powers of 2 are invalid.
if (Align > 65536 || (Align & (Align - 1)) != 0)
return UR_RESULT_ERROR_INVALID_VALUE;

ur_platform_handle_t Plt = Context->getPlatform();
Expand Down Expand Up @@ -335,11 +336,7 @@ 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;
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)) {
if (!UseUSMAllocator) {
hPoolInternal = Context->HostMemProxyPool.get();
} else if (Pool) {
hPoolInternal = Pool->HostMemPool.get();
Expand Down Expand Up @@ -379,7 +376,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMDeviceAlloc(

// L0 supports alignment up to 64KB and silently ignores higher values.
// We flag alignment > 64KB as an invalid value.
if (Alignment > 65536)
// L0 spec says that alignment values that are not powers of 2 are invalid.
if (Alignment > 65536 || (Alignment & (Alignment - 1)) != 0)
return UR_RESULT_ERROR_INVALID_VALUE;

ur_platform_handle_t Plt = Device->Platform;
Expand All @@ -406,11 +404,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMDeviceAlloc(
}

umf_memory_pool_handle_t hPoolInternal = nullptr;
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)) {
if (!UseUSMAllocator) {
auto It = Context->DeviceMemProxyPools.find(Device->ZeDevice);
if (It == Context->DeviceMemProxyPools.end())
return UR_RESULT_ERROR_INVALID_VALUE;
Expand Down Expand Up @@ -483,7 +477,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMSharedAlloc(

// L0 supports alignment up to 64KB and silently ignores higher values.
// We flag alignment > 64KB as an invalid value.
if (Alignment > 65536)
// L0 spec says that alignment values that are not powers of 2 are invalid.
if (Alignment > 65536 || (Alignment & (Alignment - 1)) != 0)
return UR_RESULT_ERROR_INVALID_VALUE;

ur_platform_handle_t Plt = Device->Platform;
Expand All @@ -506,11 +501,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMSharedAlloc(
}

umf_memory_pool_handle_t hPoolInternal = nullptr;
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)) {
if (!UseUSMAllocator) {
auto &Allocator = (DeviceReadOnly ? Context->SharedReadOnlyMemProxyPools
: Context->SharedMemProxyPools);
auto It = Allocator.find(Device->ZeDevice);
Expand Down

0 comments on commit fb3cbd1

Please sign in to comment.