Skip to content

Commit e6d5043

Browse files
lplewalplewa
committed
ctl devresidency
1 parent cde4cb8 commit e6d5043

File tree

1 file changed

+195
-3
lines changed

1 file changed

+195
-3
lines changed

‎src/provider/provider_level_zero.c‎

Lines changed: 195 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ void fini_ze_global_state(void) {
3636
#include "utils_concurrency.h"
3737
#include "utils_log.h"
3838
#include "utils_sanitizers.h"
39+
#include "critnib.h"
40+
#include "provider_ctl_stats_type.h"
3941
#include "ze_api.h"
4042

4143
// Level Zero Memory Provider settings struct
@@ -73,8 +75,18 @@ typedef struct ze_memory_provider_t {
7375
size_t min_page_size;
7476

7577
uint32_t device_ordinal;
78+
79+
critnib *allocations;
80+
ctl_stats_t stats;
81+
utils_rwlock_t lock; // lock for residency updates and allocations
7682
} ze_memory_provider_t;
7783

84+
#define CTL_PROVIDER_TYPE ze_memory_provider_t
85+
#include "provider_ctl_stats_impl.h"
86+
87+
struct ctl level_zero_ctl_root;
88+
static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
89+
7890
typedef struct ze_ops_t {
7991
ze_result_t (*zeMemAllocHost)(ze_context_handle_t,
8092
const ze_host_mem_alloc_desc_t *, size_t,
@@ -97,6 +109,8 @@ typedef struct ze_ops_t {
97109
ze_result_t (*zeContextMakeMemoryResident)(ze_context_handle_t,
98110
ze_device_handle_t, void *,
99111
size_t);
112+
ze_result_t (*zeContextEvictMemory)(ze_context_handle_t, ze_device_handle_t,
113+
void *, size_t);
100114
ze_result_t (*zeDeviceGetProperties)(ze_device_handle_t,
101115
ze_device_properties_t *);
102116
ze_result_t (*zeMemFreeExt)(ze_context_handle_t,
@@ -190,6 +204,8 @@ static void init_ze_global_state(void) {
190204
utils_get_symbol_addr(lib_handle, "zeMemCloseIpcHandle", lib_name);
191205
*(void **)&g_ze_ops.zeContextMakeMemoryResident = utils_get_symbol_addr(
192206
lib_handle, "zeContextMakeMemoryResident", lib_name);
207+
*(void **)&g_ze_ops.zeContextEvictMemory =
208+
utils_get_symbol_addr(lib_handle, "zeContextEvictMemory", lib_name);
193209
*(void **)&g_ze_ops.zeDeviceGetProperties =
194210
utils_get_symbol_addr(lib_handle, "zeDeviceGetProperties", lib_name);
195211
*(void **)&g_ze_ops.zeMemFreeExt =
@@ -202,7 +218,8 @@ static void init_ze_global_state(void) {
202218
!g_ze_ops.zeMemGetIpcHandle || !g_ze_ops.zeMemOpenIpcHandle ||
203219
!g_ze_ops.zeMemCloseIpcHandle ||
204220
!g_ze_ops.zeContextMakeMemoryResident ||
205-
!g_ze_ops.zeDeviceGetProperties || !g_ze_ops.zeMemGetAllocProperties) {
221+
!g_ze_ops.zeContextEvictMemory || !g_ze_ops.zeDeviceGetProperties ||
222+
!g_ze_ops.zeMemGetAllocProperties) {
206223
// g_ze_ops.zeMemPutIpcHandle can be NULL because it was introduced
207224
// starting from Level Zero 1.6
208225
LOG_FATAL("Required Level Zero symbols not found.");
@@ -213,6 +230,134 @@ static void init_ze_global_state(void) {
213230
ze_lib_handle = lib_handle;
214231
}
215232

233+
struct device_action_args {
234+
ze_memory_provider_t *provider;
235+
ze_device_handle_t device;
236+
umf_result_t result;
237+
};
238+
239+
static int make_resident_iter(uintptr_t key, void *value, void *priv) {
240+
struct device_action_args *args = priv;
241+
size_t size = (size_t)value;
242+
ze_result_t zr = g_ze_ops.zeContextMakeMemoryResident(
243+
args->provider->context, args->device, (void *)key, size);
244+
if (zr != ZE_RESULT_SUCCESS) {
245+
args->result = ze2umf_result(zr);
246+
return 1;
247+
}
248+
return 0;
249+
}
250+
251+
static int evict_memory_iter(uintptr_t key, void *value, void *priv) {
252+
struct device_action_args *args = priv;
253+
size_t size = (size_t)value;
254+
ze_result_t zr = g_ze_ops.zeContextEvictMemory(
255+
args->provider->context, args->device, (void *)key, size);
256+
if (zr != ZE_RESULT_SUCCESS) {
257+
args->result = ze2umf_result(zr);
258+
return 1;
259+
}
260+
return 0;
261+
}
262+
263+
static umf_result_t CTL_RUNNABLE_HANDLER(add_resident_device)(
264+
void *ctx, umf_ctl_query_source_t source, void *arg, size_t size,
265+
umf_ctl_index_utlist_t *indexes, const char *extra_name,
266+
umf_ctl_query_type_t query_type) {
267+
(void)indexes;
268+
(void)source;
269+
(void)extra_name;
270+
(void)query_type;
271+
272+
if (!arg || size != sizeof(ze_device_handle_t)) {
273+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
274+
}
275+
276+
ze_memory_provider_t *provider = (ze_memory_provider_t *)ctx;
277+
ze_device_handle_t device = *(ze_device_handle_t *)arg;
278+
279+
utils_write_lock(&provider->lock);
280+
281+
ze_device_handle_t *new_list = umf_ba_global_alloc(
282+
sizeof(ze_device_handle_t) * (provider->resident_device_count + 1));
283+
if (!new_list) {
284+
utils_write_unlock(&provider->lock);
285+
return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
286+
}
287+
288+
for (uint32_t i = 0; i < provider->resident_device_count; i++) {
289+
new_list[i] = provider->resident_device_handles[i];
290+
}
291+
new_list[provider->resident_device_count] = device;
292+
293+
umf_ba_global_free(provider->resident_device_handles);
294+
provider->resident_device_handles = new_list;
295+
provider->resident_device_count++;
296+
297+
struct device_action_args args_iter = {provider, device,
298+
UMF_RESULT_SUCCESS};
299+
critnib_iter(provider->allocations, 0, UINTPTR_MAX, make_resident_iter,
300+
&args_iter);
301+
302+
utils_write_unlock(&provider->lock);
303+
304+
return args_iter.result;
305+
}
306+
307+
static umf_result_t CTL_RUNNABLE_HANDLER(remove_resident_device)(
308+
void *ctx, umf_ctl_query_source_t source, void *arg, size_t size,
309+
umf_ctl_index_utlist_t *indexes, const char *extra_name,
310+
umf_ctl_query_type_t query_type) {
311+
(void)indexes;
312+
(void)source;
313+
(void)extra_name;
314+
(void)query_type;
315+
316+
if (!arg || size != sizeof(ze_device_handle_t)) {
317+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
318+
}
319+
320+
ze_memory_provider_t *provider = (ze_memory_provider_t *)ctx;
321+
ze_device_handle_t device = *(ze_device_handle_t *)arg;
322+
323+
utils_write_lock(&provider->lock);
324+
325+
bool found = false;
326+
for (uint32_t i = 0; i < provider->resident_device_count; i++) {
327+
if (provider->resident_device_handles[i] == device) {
328+
found = true;
329+
for (uint32_t j = i + 1; j < provider->resident_device_count; j++) {
330+
provider->resident_device_handles[j - 1] =
331+
provider->resident_device_handles[j];
332+
}
333+
provider->resident_device_count--;
334+
break;
335+
}
336+
}
337+
if (!found) {
338+
utils_write_unlock(&provider->lock);
339+
return UMF_RESULT_SUCCESS;
340+
}
341+
342+
struct device_action_args args_iter = {provider, device,
343+
UMF_RESULT_SUCCESS};
344+
critnib_iter(provider->allocations, 0, UINTPTR_MAX, evict_memory_iter,
345+
&args_iter);
346+
347+
utils_write_unlock(&provider->lock);
348+
349+
return args_iter.result;
350+
}
351+
352+
static const umf_ctl_node_t CTL_NODE(residency)[] = {
353+
CTL_LEAF_RUNNABLE(add_resident_device),
354+
CTL_LEAF_RUNNABLE(remove_resident_device), CTL_NODE_END};
355+
356+
static void initialize_level_zero_ctl(void) {
357+
CTL_REGISTER_MODULE(&level_zero_ctl_root, stats);
358+
CTL_REGISTER_MODULE(&level_zero_ctl_root, residency);
359+
}
360+
216361
umf_result_t umfLevelZeroMemoryProviderParamsCreate(
217362
umf_level_zero_memory_provider_params_handle_t *hParams) {
218363
libumfInit();
@@ -368,6 +513,8 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
368513
void **resultPtr) {
369514
ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
370515

516+
utils_read_lock(&ze_provider->lock);
517+
371518
ze_result_t ze_result = ZE_RESULT_SUCCESS;
372519
switch (ze2umf_memory_type(ze_provider->memory_type)) {
373520
case UMF_MEMORY_TYPE_HOST: {
@@ -417,6 +564,7 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
417564
}
418565

419566
if (ze_result != ZE_RESULT_SUCCESS) {
567+
utils_read_unlock(&ze_provider->lock);
420568
return ze2umf_result(ze_result);
421569
}
422570

@@ -425,10 +573,16 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
425573
ze_provider->context, ze_provider->resident_device_handles[i],
426574
*resultPtr, size);
427575
if (ze_result != ZE_RESULT_SUCCESS) {
576+
utils_read_unlock(&ze_provider->lock);
428577
return ze2umf_result(ze_result);
429578
}
430579
}
431580

581+
critnib_insert(ze_provider->allocations, (uintptr_t)(*resultPtr),
582+
(void *)(uintptr_t)size, 0);
583+
provider_ctl_stats_alloc(ze_provider, size);
584+
utils_read_unlock(&ze_provider->lock);
585+
432586
return ze2umf_result(ze_result);
433587
}
434588

@@ -441,18 +595,32 @@ static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
441595
}
442596

443597
ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
598+
utils_read_lock(&ze_provider->lock);
444599

445600
if (ze_provider->freePolicyFlags == 0) {
446-
return ze2umf_result(g_ze_ops.zeMemFree(ze_provider->context, ptr));
601+
umf_result_t ret =
602+
ze2umf_result(g_ze_ops.zeMemFree(ze_provider->context, ptr));
603+
if (ret == UMF_RESULT_SUCCESS) {
604+
critnib_remove_release(ze_provider->allocations, (uintptr_t)ptr);
605+
provider_ctl_stats_free(ze_provider, bytes);
606+
}
607+
utils_read_unlock(&ze_provider->lock);
608+
return ret;
447609
}
448610

449611
ze_memory_free_ext_desc_t desc = {
450612
.stype = ZE_STRUCTURE_TYPE_MEMORY_FREE_EXT_DESC,
451613
.pNext = NULL,
452614
.freePolicy = ze_provider->freePolicyFlags};
453615

454-
return ze2umf_result(
616+
umf_result_t ret = ze2umf_result(
455617
g_ze_ops.zeMemFreeExt(ze_provider->context, &desc, ptr));
618+
if (ret == UMF_RESULT_SUCCESS) {
619+
critnib_remove_release(ze_provider->allocations, (uintptr_t)ptr);
620+
provider_ctl_stats_free(ze_provider, bytes);
621+
}
622+
utils_read_unlock(&ze_provider->lock);
623+
return ret;
456624
}
457625

458626
static umf_result_t query_min_page_size(ze_memory_provider_t *ze_provider,
@@ -482,6 +650,8 @@ static umf_result_t query_min_page_size(ze_memory_provider_t *ze_provider,
482650
static umf_result_t ze_memory_provider_finalize(void *provider) {
483651
ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
484652
umf_ba_global_free(ze_provider->resident_device_handles);
653+
utils_rwlock_destroy_not_free(&ze_provider->lock);
654+
critnib_delete(ze_provider->allocations);
485655

486656
umf_ba_global_free(provider);
487657
return UMF_RESULT_SUCCESS;
@@ -534,6 +704,17 @@ static umf_result_t ze_memory_provider_initialize(const void *params,
534704
umfFreePolicyToZePolicy(ze_params->freePolicy);
535705
ze_provider->min_page_size = 0;
536706
ze_provider->device_ordinal = ze_params->device_ordinal;
707+
ze_provider->allocations = critnib_new(NULL, NULL);
708+
if (!ze_provider->allocations) {
709+
umf_ba_global_free(ze_provider);
710+
return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
711+
}
712+
if (utils_rwlock_init(&ze_provider->lock) == NULL) {
713+
umf_ba_global_free(ze_provider);
714+
return UMF_RESULT_ERROR_UNKNOWN;
715+
}
716+
ze_provider->stats.allocated_memory = 0;
717+
ze_provider->stats.peak_memory = 0;
537718

538719
memset(&ze_provider->device_properties, 0,
539720
sizeof(ze_provider->device_properties));
@@ -577,6 +758,8 @@ static umf_result_t ze_memory_provider_initialize(const void *params,
577758
return result;
578759
}
579760

761+
utils_init_once(&ctl_initialized, initialize_level_zero_ctl);
762+
580763
*provider = ze_provider;
581764

582765
return UMF_RESULT_SUCCESS;
@@ -787,6 +970,14 @@ ze_memory_provider_close_ipc_handle(void *provider, void *ptr, size_t size) {
787970
return UMF_RESULT_SUCCESS;
788971
}
789972

973+
static umf_result_t ze_ctl(void *provider, int operationType, const char *name,
974+
void *arg, size_t size,
975+
umf_ctl_query_type_t query_type) {
976+
utils_init_once(&ctl_initialized, initialize_level_zero_ctl);
977+
return ctl_query(&level_zero_ctl_root, provider, operationType, name,
978+
query_type, arg, size);
979+
}
980+
790981
static umf_memory_provider_ops_t UMF_LEVEL_ZERO_MEMORY_PROVIDER_OPS = {
791982
.version = UMF_PROVIDER_OPS_VERSION_CURRENT,
792983
.initialize = ze_memory_provider_initialize,
@@ -806,6 +997,7 @@ static umf_memory_provider_ops_t UMF_LEVEL_ZERO_MEMORY_PROVIDER_OPS = {
806997
.ext_put_ipc_handle = ze_memory_provider_put_ipc_handle,
807998
.ext_open_ipc_handle = ze_memory_provider_open_ipc_handle,
808999
.ext_close_ipc_handle = ze_memory_provider_close_ipc_handle,
1000+
.ext_ctl = ze_ctl,
8091001
};
8101002

8111003
const umf_memory_provider_ops_t *umfLevelZeroMemoryProviderOps(void) {

0 commit comments

Comments
 (0)