Skip to content

Commit a5aa0b3

Browse files
bratpiorkabratpiorka
committed
add support for CUDA allocation flags
1 parent 1fa3f8a commit a5aa0b3

File tree

8 files changed

+178
-12
lines changed

8 files changed

+178
-12
lines changed

‎include/umf/providers/provider_cuda.h‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,22 @@ umf_result_t umfCUDAMemoryProviderParamsSetMemoryType(
5353
umf_cuda_memory_provider_params_handle_t hParams,
5454
umf_usm_memory_type_t memoryType);
5555

56+
/// @brief Set the host allocation flags in the parameters struct.
57+
/// @param hParams handle to the parameters of the CUDA Memory Provider.
58+
/// @param flags combination of CU_MEMHOSTALLOC_PORTABLE,
59+
/// CU_MEMHOSTALLOC_DEVICEMAP and CU_MEMHOSTALLOC_WRITECOMBINED flags
60+
/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
61+
umf_result_t umfCUDAMemoryProviderParamsSetHostAllocFlags(
62+
umf_cuda_memory_provider_params_handle_t hParams, unsigned int flags);
63+
64+
/// @brief Set the managed allocation flags in the parameters struct.
65+
/// @param hParams handle to the parameters of the CUDA Memory Provider.
66+
/// @param flags must be one of CU_MEM_ATTACH_GLOBAL or CU_MEM_ATTACH_HOST
67+
/// flags
68+
/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
69+
umf_result_t umfCUDAMemoryProviderParamsSetManagedAllocFlags(
70+
umf_cuda_memory_provider_params_handle_t hParams, unsigned int flags);
71+
5672
umf_memory_provider_ops_t *umfCUDAMemoryProviderOps(void);
5773

5874
#ifdef __cplusplus

‎src/libumf.def‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,8 @@ EXPORTS
118118
umfScalablePoolParamsSetGranularity
119119
umfScalablePoolParamsSetKeepAllMemory
120120
; Added in UMF_0.11
121+
umfCUDAMemoryProviderParamsSetHostAllocFlags
122+
umfCUDAMemoryProviderParamsSetManagedAllocFlags
121123
umfFixedMemoryProviderOps
122124
umfFixedMemoryProviderParamsCreate
123125
umfFixedMemoryProviderParamsDestroy

‎src/libumf.map‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,8 @@ UMF_0.10 {
116116
};
117117

118118
UMF_0.11 {
119+
umfCUDAMemoryProviderParamsSetHostAllocFlags;
120+
umfCUDAMemoryProviderParamsSetManagedAllocFlags;
119121
umfFixedMemoryProviderOps;
120122
umfFixedMemoryProviderParamsCreate;
121123
umfFixedMemoryProviderParamsDestroy;

‎src/provider/provider_cuda.c‎

Lines changed: 72 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,22 @@ umf_result_t umfCUDAMemoryProviderParamsSetMemoryType(
5555
return UMF_RESULT_ERROR_NOT_SUPPORTED;
5656
}
5757

58+
umf_result_t umfCUDAMemoryProviderParamsSetHostAllocFlags(
59+
umf_cuda_memory_provider_params_handle_t hParams, unsigned int flags) {
60+
(void)hParams;
61+
(void)flags;
62+
LOG_ERR("CUDA provider is disabled (UMF_BUILD_CUDA_PROVIDER is OFF)!");
63+
return UMF_RESULT_ERROR_NOT_SUPPORTED;
64+
}
65+
66+
umf_result_t umfCUDAMemoryProviderParamsSetManagedAllocFlags(
67+
umf_cuda_memory_provider_params_handle_t hParams, unsigned int flags) {
68+
(void)hParams;
69+
(void)flags;
70+
LOG_ERR("CUDA provider is disabled (UMF_BUILD_CUDA_PROVIDER is OFF)!");
71+
return UMF_RESULT_ERROR_NOT_SUPPORTED;
72+
}
73+
5874
umf_memory_provider_ops_t *umfCUDAMemoryProviderOps(void) {
5975
// not supported
6076
LOG_ERR("CUDA provider is disabled (UMF_BUILD_CUDA_PROVIDER is OFF)!");
@@ -89,13 +105,17 @@ typedef struct cu_memory_provider_t {
89105
CUdevice device;
90106
umf_usm_memory_type_t memory_type;
91107
size_t min_alignment;
108+
unsigned int host_alloc_flags;
109+
unsigned int managed_alloc_flags;
92110
} cu_memory_provider_t;
93111

94112
// CUDA Memory Provider settings struct
95113
typedef struct umf_cuda_memory_provider_params_t {
96-
void *cuda_context_handle; ///< Handle to the CUDA context
97-
int cuda_device_handle; ///< Handle to the CUDA device
98-
umf_usm_memory_type_t memory_type; ///< Allocation memory type
114+
void *cuda_context_handle; // Handle to the CUDA context
115+
int cuda_device_handle; // Handle to the CUDA device
116+
umf_usm_memory_type_t memory_type; // Allocation memory type
117+
unsigned int host_alloc_flags; // Allocation flags for cuMemHostAlloc
118+
unsigned int managed_alloc_flags; // Allocation flags for cuMemAllocManaged
99119
} umf_cuda_memory_provider_params_t;
100120

101121
typedef struct cu_ops_t {
@@ -104,6 +124,7 @@ typedef struct cu_ops_t {
104124
CUmemAllocationGranularity_flags option);
105125
CUresult (*cuMemAlloc)(CUdeviceptr *dptr, size_t bytesize);
106126
CUresult (*cuMemAllocHost)(void **pp, size_t bytesize);
127+
CUresult (*cuMemHostAlloc)(void **pp, size_t bytesize, unsigned int flags);
107128
CUresult (*cuMemAllocManaged)(CUdeviceptr *dptr, size_t bytesize,
108129
unsigned int flags);
109130
CUresult (*cuMemFree)(CUdeviceptr dptr);
@@ -175,6 +196,8 @@ static void init_cu_global_state(void) {
175196
utils_get_symbol_addr(0, "cuMemAlloc_v2", lib_name);
176197
*(void **)&g_cu_ops.cuMemAllocHost =
177198
utils_get_symbol_addr(0, "cuMemAllocHost_v2", lib_name);
199+
*(void **)&g_cu_ops.cuMemHostAlloc =
200+
utils_get_symbol_addr(0, "cuMemHostAlloc", lib_name);
178201
*(void **)&g_cu_ops.cuMemAllocManaged =
179202
utils_get_symbol_addr(0, "cuMemAllocManaged", lib_name);
180203
*(void **)&g_cu_ops.cuMemFree =
@@ -197,7 +220,7 @@ static void init_cu_global_state(void) {
197220
utils_get_symbol_addr(0, "cuIpcCloseMemHandle", lib_name);
198221

199222
if (!g_cu_ops.cuMemGetAllocationGranularity || !g_cu_ops.cuMemAlloc ||
200-
!g_cu_ops.cuMemAllocHost || !g_cu_ops.cuMemAllocManaged ||
223+
!g_cu_ops.cuMemAllocHost || !g_cu_ops.cuMemHostAlloc || !g_cu_ops.cuMemAllocManaged ||
201224
!g_cu_ops.cuMemFree || !g_cu_ops.cuMemFreeHost ||
202225
!g_cu_ops.cuGetErrorName || !g_cu_ops.cuGetErrorString ||
203226
!g_cu_ops.cuCtxGetCurrent || !g_cu_ops.cuCtxSetCurrent ||
@@ -226,6 +249,8 @@ umf_result_t umfCUDAMemoryProviderParamsCreate(
226249
params_data->cuda_context_handle = NULL;
227250
params_data->cuda_device_handle = -1;
228251
params_data->memory_type = UMF_MEMORY_TYPE_UNKNOWN;
252+
params_data->host_alloc_flags = 0;
253+
params_data->managed_alloc_flags = CU_MEM_ATTACH_GLOBAL;
229254

230255
*hParams = params_data;
231256

@@ -276,6 +301,44 @@ umf_result_t umfCUDAMemoryProviderParamsSetMemoryType(
276301
return UMF_RESULT_SUCCESS;
277302
}
278303

304+
umf_result_t umfCUDAMemoryProviderParamsSetHostAllocFlags(
305+
umf_cuda_memory_provider_params_handle_t hParams, unsigned int flags) {
306+
if (!hParams) {
307+
LOG_ERR("CUDA Memory Provider params handle is NULL");
308+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
309+
}
310+
311+
// mask out valid flags and check if there are bits left
312+
if (flags &~
313+
(CU_MEMHOSTALLOC_PORTABLE | CU_MEMHOSTALLOC_DEVICEMAP |
314+
CU_MEMHOSTALLOC_WRITECOMBINED)
315+
) {
316+
LOG_ERR("Invalid host allocation flags");
317+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
318+
}
319+
320+
hParams->host_alloc_flags = flags;
321+
322+
return UMF_RESULT_SUCCESS;
323+
}
324+
325+
umf_result_t umfCUDAMemoryProviderParamsSetManagedAllocFlags(
326+
umf_cuda_memory_provider_params_handle_t hParams, unsigned int flags) {
327+
if (!hParams) {
328+
LOG_ERR("CUDA Memory Provider params handle is NULL");
329+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
330+
}
331+
332+
if (flags != CU_MEM_ATTACH_GLOBAL && flags != CU_MEM_ATTACH_HOST) {
333+
LOG_ERR("Invalid managed allocation flags");
334+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
335+
}
336+
337+
hParams->managed_alloc_flags = flags;
338+
339+
return UMF_RESULT_SUCCESS;
340+
}
341+
279342
static umf_result_t cu_memory_provider_initialize(void *params,
280343
void **provider) {
281344
if (params == NULL) {
@@ -325,6 +388,8 @@ static umf_result_t cu_memory_provider_initialize(void *params,
325388
cu_provider->device = cu_params->cuda_device_handle;
326389
cu_provider->memory_type = cu_params->memory_type;
327390
cu_provider->min_alignment = min_alignment;
391+
cu_provider->host_alloc_flags = cu_params->host_alloc_flags;
392+
cu_provider->managed_alloc_flags = cu_params->managed_alloc_flags;
328393

329394
*provider = cu_provider;
330395

@@ -382,7 +447,8 @@ static umf_result_t cu_memory_provider_alloc(void *provider, size_t size,
382447
CUresult cu_result = CUDA_SUCCESS;
383448
switch (cu_provider->memory_type) {
384449
case UMF_MEMORY_TYPE_HOST: {
385-
cu_result = g_cu_ops.cuMemAllocHost(resultPtr, size);
450+
cu_result = g_cu_ops.cuMemHostAlloc(resultPtr, size,
451+
cu_provider->host_alloc_flags);
386452
break;
387453
}
388454
case UMF_MEMORY_TYPE_DEVICE: {
@@ -391,7 +457,7 @@ static umf_result_t cu_memory_provider_alloc(void *provider, size_t size,
391457
}
392458
case UMF_MEMORY_TYPE_SHARED: {
393459
cu_result = g_cu_ops.cuMemAllocManaged((CUdeviceptr *)resultPtr, size,
394-
CU_MEM_ATTACH_GLOBAL);
460+
cu_provider->managed_alloc_flags);
395461
break;
396462
}
397463
default:

‎test/providers/cuda_helpers.cpp‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@ struct libcu_ops {
2222
CUresult (*cuDeviceGet)(CUdevice *device, int ordinal);
2323
CUresult (*cuMemAlloc)(CUdeviceptr *dptr, size_t size);
2424
CUresult (*cuMemFree)(CUdeviceptr dptr);
25-
CUresult (*cuMemAllocHost)(void **pp, size_t size);
25+
CUresult (*cuMemAllocHost)(void **pp, size_t size);
26+
CUresult (*cuMemHostAlloc)(void **pp, size_t size, unsigned int flags);
2627
CUresult (*cuMemAllocManaged)(CUdeviceptr *dptr, size_t bytesize,
2728
unsigned int flags);
2829
CUresult (*cuMemFreeHost)(void *p);
@@ -34,6 +35,7 @@ struct libcu_ops {
3435
CUresult (*cuPointerGetAttributes)(unsigned int numAttributes,
3536
CUpointer_attribute *attributes,
3637
void **data, CUdeviceptr ptr);
38+
CUresult (*cuMemHostGetFlags)(unsigned int *pFlags, void *p);
3739
CUresult (*cuStreamSynchronize)(CUstream hStream);
3840
CUresult (*cuCtxSynchronize)(void);
3941
} libcu_ops;
@@ -72,6 +74,9 @@ struct DlHandleCloser {
7274
libcu_ops.cuMemAllocHost = [](auto... args) {
7375
return noop_stub(args...);
7476
};
77+
libcu_ops.cuMemHostAlloc = [](auto... args) {
78+
return noop_stub(args...);
79+
};
7580
libcu_ops.cuMemAllocManaged = [](auto... args) {
7681
return noop_stub(args...);
7782
};
@@ -90,6 +95,9 @@ struct DlHandleCloser {
9095
libcu_ops.cuPointerGetAttributes = [](auto... args) {
9196
return noop_stub(args...);
9297
};
98+
libcu_ops.cuMemHostGetFlags = [](auto... args) {
99+
return noop_stub(args...);
100+
};
93101
libcu_ops.cuStreamSynchronize = [](auto... args) {
94102
return noop_stub(args...);
95103
};
@@ -170,6 +178,12 @@ int InitCUDAOps() {
170178
fprintf(stderr, "cuMemAllocHost_v2 symbol not found in %s\n", lib_name);
171179
return -1;
172180
}
181+
*(void **)&libcu_ops.cuMemHostAlloc =
182+
utils_get_symbol_addr(cuDlHandle.get(), "cuMemHostAlloc", lib_name);
183+
if (libcu_ops.cuMemHostAlloc == nullptr) {
184+
fprintf(stderr, "cuMemHostAlloc symbol not found in %s\n", lib_name);
185+
return -1;
186+
}
173187
*(void **)&libcu_ops.cuMemAllocManaged =
174188
utils_get_symbol_addr(cuDlHandle.get(), "cuMemAllocManaged", lib_name);
175189
if (libcu_ops.cuMemAllocManaged == nullptr) {
@@ -208,6 +222,13 @@ int InitCUDAOps() {
208222
lib_name);
209223
return -1;
210224
}
225+
*(void **)&libcu_ops.cuMemHostGetFlags = utils_get_symbol_addr(
226+
cuDlHandle.get(), "cuMemHostGetFlags", lib_name);
227+
if (libcu_ops.cuMemHostGetFlags == nullptr) {
228+
fprintf(stderr, "cuMemHostGetFlags symbol not found in %s\n",
229+
lib_name);
230+
return -1;
231+
}
211232
*(void **)&libcu_ops.cuStreamSynchronize = utils_get_symbol_addr(
212233
cuDlHandle.get(), "cuStreamSynchronize", lib_name);
213234
if (libcu_ops.cuStreamSynchronize == nullptr) {
@@ -237,13 +258,15 @@ int InitCUDAOps() {
237258
libcu_ops.cuDeviceGet = cuDeviceGet;
238259
libcu_ops.cuMemAlloc = cuMemAlloc;
239260
libcu_ops.cuMemAllocHost = cuMemAllocHost;
261+
libcu_ops.cuMemHostAlloc = cuMemHostAlloc;
240262
libcu_ops.cuMemAllocManaged = cuMemAllocManaged;
241263
libcu_ops.cuMemFree = cuMemFree;
242264
libcu_ops.cuMemFreeHost = cuMemFreeHost;
243265
libcu_ops.cuMemsetD32 = cuMemsetD32;
244266
libcu_ops.cuMemcpy = cuMemcpy;
245267
libcu_ops.cuPointerGetAttribute = cuPointerGetAttribute;
246268
libcu_ops.cuPointerGetAttributes = cuPointerGetAttributes;
269+
libcu_ops.cuMemHostGetFlags = cuMemHostGetFlags;
247270
libcu_ops.cuStreamSynchronize = cuStreamSynchronize;
248271
libcu_ops.cuCtxSynchronize = cuCtxSynchronize;
249272

@@ -373,6 +396,17 @@ umf_usm_memory_type_t get_mem_type(CUcontext context, void *ptr) {
373396
return UMF_MEMORY_TYPE_UNKNOWN;
374397
}
375398

399+
unsigned int get_mem_host_alloc_flags(void *ptr) {
400+
unsigned int flags;
401+
CUresult res = libcu_ops.cuMemHostGetFlags(&flags, ptr);
402+
if (res != CUDA_SUCCESS) {
403+
fprintf(stderr, "cuPointerGetAttribute() failed!\n");
404+
return 0;
405+
}
406+
407+
return flags;
408+
}
409+
376410
CUcontext get_mem_context(void *ptr) {
377411
CUcontext context;
378412
CUresult res = libcu_ops.cuPointerGetAttribute(

‎test/providers/cuda_helpers.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ int cuda_copy(CUcontext context, CUdevice device, void *dst_ptr,
4242

4343
umf_usm_memory_type_t get_mem_type(CUcontext context, void *ptr);
4444

45+
unsigned int get_mem_host_alloc_flags(void *ptr);
46+
4547
CUcontext get_mem_context(void *ptr);
4648

4749
CUcontext get_current_context();

‎test/providers/provider_cuda.cpp‎

Lines changed: 41 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ typedef umf_result_t (*pfnProviderParamsDestroy)(void *);
119119

120120
using CUDAProviderTestParams =
121121
std::tuple<pfnProviderParamsCreate, pfnProviderParamsDestroy, CUcontext,
122-
umf_usm_memory_type_t, MemoryAccessor *>;
122+
CUdevice, umf_usm_memory_type_t, MemoryAccessor *>;
123123

124124
struct umfCUDAProviderTest
125125
: umf_test::test,
@@ -128,7 +128,7 @@ struct umfCUDAProviderTest
128128
void SetUp() override {
129129
test::SetUp();
130130

131-
auto [params_create, params_destroy, cu_context, memory_type,
131+
auto [params_create, params_destroy, cu_context, cu_device, memory_type,
132132
accessor] = this->GetParam();
133133

134134
params = nullptr;
@@ -139,6 +139,7 @@ struct umfCUDAProviderTest
139139

140140
memAccessor = accessor;
141141
expected_context = cu_context;
142+
expected_device = cu_device;
142143
expected_memory_type = memory_type;
143144
}
144145

@@ -155,6 +156,7 @@ struct umfCUDAProviderTest
155156

156157
MemoryAccessor *memAccessor = nullptr;
157158
CUcontext expected_context;
159+
CUdevice expected_device;
158160
umf_usm_memory_type_t expected_memory_type;
159161
};
160162

@@ -320,6 +322,40 @@ TEST_P(umfCUDAProviderTest, cudaProviderNullParams) {
320322
EXPECT_EQ(res, UMF_RESULT_ERROR_INVALID_ARGUMENT);
321323
}
322324

325+
TEST_P(umfCUDAProviderTest, cudaHostAllocFlags) {
326+
umf_cuda_memory_provider_params_handle_t params = nullptr;
327+
umf_result_t res = umfCUDAMemoryProviderParamsCreate(&params);
328+
EXPECT_EQ(res, UMF_RESULT_ERROR_INVALID_ARGUMENT);
329+
330+
umfCUDAMemoryProviderParamsSetContext(params, expected_context);
331+
umfCUDAMemoryProviderParamsSetDevice(params, expected_device );
332+
umfCUDAMemoryProviderParamsSetMemoryType(params, expected_memory_type);
333+
umfCUDAMemoryProviderParamsSetHostAllocFlags(params,CU_MEMHOSTALLOC_WRITECOMBINED);
334+
335+
umf_memory_provider_handle_t provider = nullptr;
336+
umf_result_t umf_result =
337+
umfMemoryProviderCreate(umfCUDAMemoryProviderOps(), params, &provider);
338+
ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
339+
ASSERT_NE(provider, nullptr);
340+
341+
if (expected_memory_type == UMF_MEMORY_TYPE_HOST) {
342+
void *ptr = nullptr;
343+
umf_result = umfMemoryProviderAlloc(provider, 128, 0, &ptr);
344+
ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
345+
ASSERT_NE(ptr, nullptr);
346+
347+
// check if the memory allocation flags are set correctly
348+
unsigned int flags = get_mem_host_alloc_flags(ptr);
349+
ASSERT_EQ(flags, CU_MEMHOSTALLOC_WRITECOMBINED);
350+
351+
umf_result = umfMemoryProviderFree(provider, ptr, 128);
352+
ASSERT_EQ(umf_result, UMF_RESULT_SUCCESS);
353+
}
354+
355+
umfMemoryProviderDestroy(provider);
356+
umfCUDAMemoryProviderParamsDestroy(params);
357+
}
358+
323359
TEST_P(umfCUDAProviderTest, multiContext) {
324360
CUdevice device;
325361
int ret = get_cuda_device(&device);
@@ -421,13 +457,13 @@ INSTANTIATE_TEST_SUITE_P(
421457
umfCUDAProviderTestSuite, umfCUDAProviderTest,
422458
::testing::Values(
423459
CUDAProviderTestParams{createCuParamsDeviceMemory, destroyCuParams,
424-
cudaTestHelper.get_test_context(),
460+
cudaTestHelper.get_test_context(),cudaTestHelper.get_test_device(),
425461
UMF_MEMORY_TYPE_DEVICE, &cuAccessor},
426462
CUDAProviderTestParams{createCuParamsSharedMemory, destroyCuParams,
427-
cudaTestHelper.get_test_context(),
463+
cudaTestHelper.get_test_context(),cudaTestHelper.get_test_device(),
428464
UMF_MEMORY_TYPE_SHARED, &hostAccessor},
429465
CUDAProviderTestParams{createCuParamsHostMemory, destroyCuParams,
430-
cudaTestHelper.get_test_context(),
466+
cudaTestHelper.get_test_context(),cudaTestHelper.get_test_device(),
431467
UMF_MEMORY_TYPE_HOST, &hostAccessor}));
432468

433469
// TODO: add IPC API

0 commit comments

Comments
 (0)