Skip to content

Commit a57191a

Browse files
RossBruntonRossBrunton
committed
Add size validation for USM device/shared allocs
The spec states that an allocation size higher than `DEVICE_INFO_MAX_MEM_ALLOC_SIZE` should report an error. This has been added to the Cuda, Hip and Native CPU backends now. The tests have also been updated, which sadly introduces a new regression for level zero.
1 parent ad88f0a commit a57191a

File tree

9 files changed

+92
-32
lines changed

9 files changed

+92
-32
lines changed

‎source/adapters/cuda/usm.cpp

+8
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,10 @@ urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
6363
(alignment == 0 || ((alignment & (alignment - 1)) == 0)),
6464
UR_RESULT_ERROR_INVALID_VALUE);
6565

66+
if (size > hDevice->getMaxAllocSize()) {
67+
return UR_RESULT_ERROR_INVALID_USM_SIZE;
68+
}
69+
6670
if (!hPool) {
6771
return USMDeviceAllocImpl(ppMem, hContext, hDevice, /* flags */ 0, size,
6872
alignment);
@@ -88,6 +92,10 @@ urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
8892
(alignment == 0 || ((alignment & (alignment - 1)) == 0)),
8993
UR_RESULT_ERROR_INVALID_VALUE);
9094

95+
if (size > hDevice->getMaxAllocSize()) {
96+
return UR_RESULT_ERROR_INVALID_USM_SIZE;
97+
}
98+
9199
if (!hPool) {
92100
return USMSharedAllocImpl(ppMem, hContext, hDevice, /*host flags*/ 0,
93101
/*device flags*/ 0, size, alignment);

‎source/adapters/hip/device.cpp

+2-16
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
#include "context.hpp"
1414
#include "event.hpp"
1515
#include "logger/ur_logger.hpp"
16+
#include "usm.hpp"
1617

1718
#include <sstream>
1819

@@ -206,22 +207,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
206207
return ReturnValue(Bits);
207208
}
208209
case UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE: {
209-
// Max size of memory object allocation in bytes.
210-
// The minimum value is max(min(1024 × 1024 ×
211-
// 1024, 1/4th of CL_DEVICE_GLOBAL_MEM_SIZE),
212-
// 32 × 1024 × 1024) for devices that are not of type
213-
// CL_DEVICE_TYPE_CUSTOM.
214-
215-
size_t Global = 0;
216-
detail::ur::assertion(hipDeviceTotalMem(&Global, hDevice->get()) ==
217-
hipSuccess);
218-
219-
auto QuarterGlobal = static_cast<uint32_t>(Global / 4u);
220-
221-
auto MaxAlloc = std::max(std::min(1024u * 1024u * 1024u, QuarterGlobal),
222-
32u * 1024u * 1024u);
223-
224-
return ReturnValue(uint64_t{MaxAlloc});
210+
return ReturnValue(uint64_t{maxUSMAllocationSize(hDevice)});
225211
}
226212
case UR_DEVICE_INFO_IMAGE_SUPPORTED: {
227213
bool Enabled = false;

‎source/adapters/hip/usm.cpp

+26
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,10 @@ urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
4949
UR_ASSERT(checkUSMAlignment(alignment, pUSMDesc),
5050
UR_RESULT_ERROR_INVALID_VALUE);
5151

52+
if (size > maxUSMAllocationSize(hDevice)) {
53+
return UR_RESULT_ERROR_INVALID_USM_SIZE;
54+
}
55+
5256
if (!hPool) {
5357
return USMDeviceAllocImpl(ppMem, hContext, hDevice, /* flags */ 0, size,
5458
alignment);
@@ -66,6 +70,10 @@ urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
6670
UR_ASSERT(checkUSMAlignment(alignment, pUSMDesc),
6771
UR_RESULT_ERROR_INVALID_VALUE);
6872

73+
if (size > maxUSMAllocationSize(hDevice)) {
74+
return UR_RESULT_ERROR_INVALID_USM_SIZE;
75+
}
76+
6977
if (!hPool) {
7078
return USMSharedAllocImpl(ppMem, hContext, hDevice, /*host flags*/ 0,
7179
/*device flags*/ 0, size, alignment);
@@ -454,6 +462,24 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolGetInfo(
454462
}
455463
}
456464

465+
uint64_t maxUSMAllocationSize(const ur_device_handle_t &Device) {
466+
// Max size of memory object allocation in bytes.
467+
// The minimum value is max(min(1024 × 1024 ×
468+
// 1024, 1/4th of CL_DEVICE_GLOBAL_MEM_SIZE),
469+
// 32 × 1024 × 1024) for devices that are not of type
470+
// CL_DEVICE_TYPE_CUSTOM.
471+
size_t Global = 0;
472+
detail::ur::assertion(hipDeviceTotalMem(&Global, Device->get()) ==
473+
hipSuccess);
474+
475+
auto QuarterGlobal = static_cast<uint32_t>(Global / 4u);
476+
477+
auto MaxAlloc = std::max(std::min(1024u * 1024u * 1024u, QuarterGlobal),
478+
32u * 1024u * 1024u);
479+
480+
return MaxAlloc;
481+
}
482+
457483
bool checkUSMAlignment(uint32_t &alignment, const ur_usm_desc_t *pUSMDesc) {
458484
alignment = pUSMDesc ? pUSMDesc->align : 0u;
459485
return (!pUSMDesc ||

‎source/adapters/hip/usm.hpp

+2
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,8 @@ ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t Context,
137137
ur_usm_host_mem_flags_t Flags, size_t Size,
138138
uint32_t Alignment);
139139

140+
uint64_t maxUSMAllocationSize(const ur_device_handle_t &Device);
141+
140142
bool checkUSMAlignment(uint32_t &alignment, const ur_usm_desc_t *pUSMDesc);
141143

142144
bool checkUSMImplAlignment(uint32_t Alignment, void **ResultPtr);

‎source/adapters/native_cpu/device.cpp

+3-8
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
#include "common.hpp"
1414
#include "platform.hpp"
15+
#include "usm.hpp"
1516

1617
#if defined(_MSC_VER) || defined(__MINGW32__) || defined(__MINGW64__)
1718
#ifndef NOMINMAX
@@ -321,14 +322,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
321322
case UR_DEVICE_INFO_PARTITION_AFFINITY_DOMAIN:
322323
return ReturnValue(ur_device_affinity_domain_flags_t{0});
323324
case UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE: {
324-
size_t Global = hDevice->mem_size;
325-
326-
auto QuarterGlobal = static_cast<uint32_t>(Global / 4u);
327-
328-
auto MaxAlloc = std::max(std::min(1024u * 1024u * 1024u, QuarterGlobal),
329-
32u * 1024u * 1024u);
330-
331-
return ReturnValue(uint64_t{MaxAlloc});
325+
return ReturnValue(
326+
uint64_t{native_cpu::detail::maxUSMAllocationSize(hDevice)});
332327
}
333328
case UR_DEVICE_INFO_EXECUTION_CAPABILITIES:
334329
// TODO : CHECK

‎source/adapters/native_cpu/usm.cpp

+16-3
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313

1414
#include "common.hpp"
1515
#include "context.hpp"
16+
#include "usm.hpp"
1617
#include <cstdlib>
1718

1819
namespace umf {
@@ -27,7 +28,8 @@ static ur_result_t alloc_helper(ur_context_handle_t hContext,
2728
auto alignment = (pUSMDesc && pUSMDesc->align) ? pUSMDesc->align : 1u;
2829
UR_ASSERT(isPowerOf2(alignment), UR_RESULT_ERROR_UNSUPPORTED_ALIGNMENT);
2930
UR_ASSERT(ppMem, UR_RESULT_ERROR_INVALID_NULL_POINTER);
30-
// TODO: Check Max size when UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE is implemented
31+
// TODO: Check Max size for host allocations when
32+
// UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE is implemented
3133
UR_ASSERT(size > 0, UR_RESULT_ERROR_INVALID_USM_SIZE);
3234

3335
auto *ptr = hContext->add_alloc(alignment, type, size, nullptr);
@@ -49,8 +51,9 @@ UR_APIEXPORT ur_result_t UR_APICALL
4951
urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
5052
const ur_usm_desc_t *pUSMDesc, ur_usm_pool_handle_t pool,
5153
size_t size, void **ppMem) {
52-
std::ignore = hDevice;
5354
std::ignore = pool;
55+
UR_ASSERT(size < native_cpu::detail::maxUSMAllocationSize(hDevice),
56+
UR_RESULT_ERROR_INVALID_USM_SIZE);
5457

5558
return alloc_helper(hContext, pUSMDesc, size, ppMem, UR_USM_TYPE_DEVICE);
5659
}
@@ -59,8 +62,9 @@ UR_APIEXPORT ur_result_t UR_APICALL
5962
urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
6063
const ur_usm_desc_t *pUSMDesc, ur_usm_pool_handle_t pool,
6164
size_t size, void **ppMem) {
62-
std::ignore = hDevice;
6365
std::ignore = pool;
66+
UR_ASSERT(size < native_cpu::detail::maxUSMAllocationSize(hDevice),
67+
UR_RESULT_ERROR_INVALID_USM_SIZE);
6468

6569
return alloc_helper(hContext, pUSMDesc, size, ppMem, UR_USM_TYPE_SHARED);
6670
}
@@ -155,3 +159,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMReleaseExp(ur_context_handle_t Context,
155159
std::ignore = HostPtr;
156160
DIE_NO_IMPLEMENTATION;
157161
}
162+
163+
uint64_t maxUSMAllocationSize(const ur_device_handle_t &Device) {
164+
size_t Global = Device->mem_size;
165+
166+
auto QuarterGlobal = static_cast<uint32_t>(Global / 4u);
167+
168+
return std::max(std::min(1024u * 1024u * 1024u, QuarterGlobal),
169+
32u * 1024u * 1024u);
170+
}

‎source/adapters/native_cpu/usm.hpp

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
//===--------- usm.hpp - Native CPU Adapter --------------------------===//
2+
//
3+
// Copyright (C) 2025 Intel Corporation
4+
//
5+
// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
6+
// Exceptions. See LICENSE.TXT
7+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
8+
//
9+
//===-----------------------------------------------------------------===//
10+
11+
#include "common.hpp"
12+
13+
namespace native_cpu {
14+
namespace detail {
15+
uint64_t maxUSMAllocationSize(const ur_device_handle_t &Device);
16+
} // namespace detail
17+
} // namespace native_cpu

‎test/conformance/usm/urUSMDeviceAlloc.cpp

+9-3
Original file line numberDiff line numberDiff line change
@@ -118,13 +118,19 @@ TEST_P(urUSMDeviceAllocTest, InvalidNullPtrResult) {
118118
}
119119

120120
TEST_P(urUSMDeviceAllocTest, InvalidUSMSize) {
121-
UUR_KNOWN_FAILURE_ON(uur::CUDA{}, uur::HIP{}, uur::LevelZero{},
122-
uur::NativeCPU{});
121+
UUR_KNOWN_FAILURE_ON(uur::LevelZero{}, uur::LevelZeroV2{});
122+
123+
size_t max_size;
124+
ASSERT_SUCCESS(urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE,
125+
sizeof(max_size), &max_size, nullptr));
126+
if (max_size == std::numeric_limits<size_t>::max()) {
127+
GTEST_SKIP() << "Device has no max allocation size";
128+
}
123129

124130
void *ptr = nullptr;
125131
ASSERT_EQ_RESULT(
126132
UR_RESULT_ERROR_INVALID_USM_SIZE,
127-
urUSMDeviceAlloc(context, device, nullptr, pool, -1, &ptr));
133+
urUSMDeviceAlloc(context, device, nullptr, pool, max_size + 1, &ptr));
128134
}
129135

130136
TEST_P(urUSMDeviceAllocTest, InvalidValueAlignPowerOfTwo) {

‎test/conformance/usm/urUSMSharedAlloc.cpp

+9-2
Original file line numberDiff line numberDiff line change
@@ -141,11 +141,18 @@ TEST_P(urUSMSharedAllocTest, InvalidNullPtrMem) {
141141
}
142142

143143
TEST_P(urUSMSharedAllocTest, InvalidUSMSize) {
144-
UUR_KNOWN_FAILURE_ON(uur::CUDA{}, uur::HIP{}, uur::NativeCPU{});
144+
UUR_KNOWN_FAILURE_ON(uur::LevelZero{}, uur::LevelZeroV2{});
145+
146+
size_t max_size;
147+
ASSERT_SUCCESS(urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE,
148+
sizeof(max_size), &max_size, nullptr));
149+
if (max_size == std::numeric_limits<size_t>::max()) {
150+
GTEST_SKIP() << "Device has no max allocation size";
151+
}
145152

146153
ASSERT_EQ_RESULT(
147154
UR_RESULT_ERROR_INVALID_USM_SIZE,
148-
urUSMSharedAlloc(context, device, nullptr, pool, -1, &ptr));
155+
urUSMSharedAlloc(context, device, nullptr, pool, max_size + 1, &ptr));
149156
}
150157

151158
TEST_P(urUSMSharedAllocTest, InvalidValueAlignPowerOfTwo) {

0 commit comments

Comments
 (0)