11/*
2- * Copyright (C) 2024 Intel Corporation
2+ * Copyright (C) 2024-2025 Intel Corporation
33 *
44 * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
55 * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -118,6 +118,8 @@ typedef struct umf_level_zero_memory_provider_params_t {
118118
119119 umf_level_zero_memory_provider_free_policy_t
120120 freePolicy ; ///< Memory free policy
121+
122+ uint32_t device_ordinal ;
121123} umf_level_zero_memory_provider_params_t ;
122124
123125typedef struct ze_memory_provider_t {
@@ -131,6 +133,10 @@ typedef struct ze_memory_provider_t {
131133 ze_device_properties_t device_properties ;
132134
133135 ze_driver_memory_free_policy_ext_flags_t freePolicyFlags ;
136+
137+ size_t min_page_size ;
138+
139+ uint32_t device_ordinal ;
134140} ze_memory_provider_t ;
135141
136142typedef struct ze_ops_t {
@@ -159,6 +165,9 @@ typedef struct ze_ops_t {
159165 ze_device_properties_t * );
160166 ze_result_t (* zeMemFreeExt )(ze_context_handle_t ,
161167 ze_memory_free_ext_desc_t * , void * );
168+ ze_result_t (* zeMemGetAllocProperties )(ze_context_handle_t , const void * ,
169+ ze_memory_allocation_properties_t * ,
170+ ze_device_handle_t * );
162171} ze_ops_t ;
163172
164173static ze_ops_t g_ze_ops ;
@@ -214,13 +223,15 @@ static void init_ze_global_state(void) {
214223 utils_get_symbol_addr (0 , "zeDeviceGetProperties" , lib_name );
215224 * (void * * )& g_ze_ops .zeMemFreeExt =
216225 utils_get_symbol_addr (0 , "zeMemFreeExt" , lib_name );
226+ * (void * * )& g_ze_ops .zeMemGetAllocProperties =
227+ utils_get_symbol_addr (0 , "zeMemGetAllocProperties" , lib_name );
217228
218229 if (!g_ze_ops .zeMemAllocHost || !g_ze_ops .zeMemAllocDevice ||
219230 !g_ze_ops .zeMemAllocShared || !g_ze_ops .zeMemFree ||
220231 !g_ze_ops .zeMemGetIpcHandle || !g_ze_ops .zeMemOpenIpcHandle ||
221232 !g_ze_ops .zeMemCloseIpcHandle ||
222233 !g_ze_ops .zeContextMakeMemoryResident ||
223- !g_ze_ops .zeDeviceGetProperties ) {
234+ !g_ze_ops .zeDeviceGetProperties || ! g_ze_ops . zeMemGetAllocProperties ) {
224235 // g_ze_ops.zeMemPutIpcHandle can be NULL because it was introduced
225236 // starting from Level Zero 1.6
226237 LOG_ERR ("Required Level Zero symbols not found." );
@@ -250,6 +261,7 @@ umf_result_t umfLevelZeroMemoryProviderParamsCreate(
250261 params -> resident_device_handles = NULL ;
251262 params -> resident_device_count = 0 ;
252263 params -> freePolicy = UMF_LEVEL_ZERO_MEMORY_PROVIDER_FREE_POLICY_DEFAULT ;
264+ params -> device_ordinal = 0 ;
253265
254266 * hParams = params ;
255267
@@ -307,6 +319,18 @@ umf_result_t umfLevelZeroMemoryProviderParamsSetMemoryType(
307319 return UMF_RESULT_SUCCESS ;
308320}
309321
322+ umf_result_t umfLevelZeroMemoryProviderSetDeviceOrdinal (
323+ umf_level_zero_memory_provider_params_handle_t hParams ,
324+ uint32_t deviceOrdinal ) {
325+ if (!hParams ) {
326+ LOG_ERR ("Level zero memory provider params handle is NULL" );
327+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
328+ }
329+ hParams -> device_ordinal = deviceOrdinal ;
330+
331+ return UMF_RESULT_SUCCESS ;
332+ }
333+
310334umf_result_t umfLevelZeroMemoryProviderParamsSetResidentDevices (
311335 umf_level_zero_memory_provider_params_handle_t hParams ,
312336 ze_device_handle_t * hDevices , uint32_t deviceCount ) {
@@ -351,100 +375,6 @@ umfFreePolicyToZePolicy(umf_level_zero_memory_provider_free_policy_t policy) {
351375 return 0 ;
352376 }
353377}
354-
355- static umf_result_t ze_memory_provider_initialize (void * params ,
356- void * * provider ) {
357- if (params == NULL ) {
358- return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
359- }
360-
361- umf_level_zero_memory_provider_params_handle_t ze_params =
362- (umf_level_zero_memory_provider_params_handle_t )params ;
363-
364- if (!ze_params -> level_zero_context_handle ) {
365- LOG_ERR ("Level Zero context handle is NULL" );
366- return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
367- }
368-
369- if ((ze_params -> memory_type == UMF_MEMORY_TYPE_HOST ) ==
370- (ze_params -> level_zero_device_handle != NULL )) {
371- LOG_ERR ("Level Zero device handle is NULL" );
372- return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
373- }
374-
375- if ((bool )ze_params -> resident_device_count &&
376- (ze_params -> resident_device_handles == NULL )) {
377- LOG_ERR ("Resident devices handles array is NULL, but device_count is "
378- "not zero" );
379- return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
380- }
381-
382- utils_init_once (& ze_is_initialized , init_ze_global_state );
383- if (Init_ze_global_state_failed ) {
384- LOG_ERR ("Loading Level Zero symbols failed" );
385- return UMF_RESULT_ERROR_UNKNOWN ;
386- }
387-
388- ze_memory_provider_t * ze_provider =
389- umf_ba_global_alloc (sizeof (ze_memory_provider_t ));
390- if (!ze_provider ) {
391- LOG_ERR ("Cannot allocate memory for Level Zero Memory Provider" );
392- return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY ;
393- }
394-
395- ze_provider -> context = ze_params -> level_zero_context_handle ;
396- ze_provider -> device = ze_params -> level_zero_device_handle ;
397- ze_provider -> memory_type = (ze_memory_type_t )ze_params -> memory_type ;
398- ze_provider -> freePolicyFlags =
399- umfFreePolicyToZePolicy (ze_params -> freePolicy );
400-
401- memset (& ze_provider -> device_properties , 0 ,
402- sizeof (ze_provider -> device_properties ));
403- ze_provider -> device_properties .stype = ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES ;
404-
405- if (ze_provider -> device ) {
406- umf_result_t ret = ze2umf_result (g_ze_ops .zeDeviceGetProperties (
407- ze_provider -> device , & ze_provider -> device_properties ));
408-
409- if (ret != UMF_RESULT_SUCCESS ) {
410- LOG_ERR ("Cannot get device properties" );
411- umf_ba_global_free (ze_provider );
412- return ret ;
413- }
414- }
415-
416- if (ze_params -> resident_device_count ) {
417- ze_provider -> resident_device_handles = umf_ba_global_alloc (
418- sizeof (ze_device_handle_t ) * ze_params -> resident_device_count );
419- if (!ze_provider -> resident_device_handles ) {
420- LOG_ERR ("Cannot allocate memory for resident devices" );
421- umf_ba_global_free (ze_provider );
422- return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY ;
423- }
424-
425- ze_provider -> resident_device_count = ze_params -> resident_device_count ;
426-
427- for (uint32_t i = 0 ; i < ze_provider -> resident_device_count ; i ++ ) {
428- ze_provider -> resident_device_handles [i ] =
429- ze_params -> resident_device_handles [i ];
430- }
431- } else {
432- ze_provider -> resident_device_handles = NULL ;
433- ze_provider -> resident_device_count = 0 ;
434- }
435-
436- * provider = ze_provider ;
437-
438- return UMF_RESULT_SUCCESS ;
439- }
440-
441- static void ze_memory_provider_finalize (void * provider ) {
442- ze_memory_provider_t * ze_provider = (ze_memory_provider_t * )provider ;
443- umf_ba_global_free (ze_provider -> resident_device_handles );
444-
445- umf_ba_global_free (provider );
446- }
447-
448378static bool use_relaxed_allocation (ze_memory_provider_t * ze_provider ,
449379 size_t size ) {
450380 assert (ze_provider );
@@ -482,8 +412,7 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
482412 ? & relaxed_device_allocation_desc
483413 : NULL ,
484414 .flags = 0 ,
485- .ordinal = 0 // TODO
486- };
415+ .ordinal = ze_provider -> device_ordinal };
487416 ze_result = g_ze_ops .zeMemAllocDevice (ze_provider -> context , & dev_desc ,
488417 size , alignment ,
489418 ze_provider -> device , resultPtr );
@@ -500,8 +429,7 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
500429 ? & relaxed_device_allocation_desc
501430 : NULL ,
502431 .flags = 0 ,
503- .ordinal = 0 // TODO
504- };
432+ .ordinal = ze_provider -> device_ordinal };
505433 ze_result = g_ze_ops .zeMemAllocShared (ze_provider -> context , & dev_desc ,
506434 & host_desc , size , alignment ,
507435 ze_provider -> device , resultPtr );
@@ -553,6 +481,132 @@ static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
553481 g_ze_ops .zeMemFreeExt (ze_provider -> context , & desc , ptr ));
554482}
555483
484+ static umf_result_t query_min_page_size (ze_memory_provider_t * ze_provider ,
485+ size_t * min_page_size ) {
486+ assert (min_page_size );
487+
488+ LOG_DEBUG ("Querying minimum page size" );
489+
490+ void * ptr ;
491+ umf_result_t result = ze_memory_provider_alloc (ze_provider , 1 , 0 , & ptr );
492+ if (result != UMF_RESULT_SUCCESS ) {
493+ return result ;
494+ }
495+
496+ ze_memory_allocation_properties_t properties = {
497+ .stype = ZE_STRUCTURE_TYPE_MEMORY_ALLOCATION_PROPERTIES };
498+ ze_result_t ze_result = g_ze_ops .zeMemGetAllocProperties (
499+ ze_provider -> context , ptr , & properties , NULL );
500+
501+ * min_page_size = properties .pageSize ;
502+
503+ ze_memory_provider_free (ze_provider , ptr , 1 );
504+
505+ return ze2umf_result (ze_result );
506+ }
507+
508+ static void ze_memory_provider_finalize (void * provider ) {
509+ ze_memory_provider_t * ze_provider = (ze_memory_provider_t * )provider ;
510+ umf_ba_global_free (ze_provider -> resident_device_handles );
511+
512+ umf_ba_global_free (provider );
513+ }
514+
515+ static umf_result_t ze_memory_provider_initialize (void * params ,
516+ void * * provider ) {
517+ if (params == NULL ) {
518+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
519+ }
520+
521+ umf_level_zero_memory_provider_params_handle_t ze_params =
522+ (umf_level_zero_memory_provider_params_handle_t )params ;
523+
524+ if (!ze_params -> level_zero_context_handle ) {
525+ LOG_ERR ("Level Zero context handle is NULL" );
526+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
527+ }
528+
529+ if ((ze_params -> memory_type == UMF_MEMORY_TYPE_HOST ) ==
530+ (ze_params -> level_zero_device_handle != NULL )) {
531+ LOG_ERR ("Level Zero device handle is NULL" );
532+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
533+ }
534+
535+ if ((bool )ze_params -> resident_device_count &&
536+ (ze_params -> resident_device_handles == NULL )) {
537+ LOG_ERR ("Resident devices handles array is NULL, but device_count is "
538+ "not zero" );
539+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
540+ }
541+
542+ utils_init_once (& ze_is_initialized , init_ze_global_state );
543+ if (Init_ze_global_state_failed ) {
544+ LOG_ERR ("Loading Level Zero symbols failed" );
545+ return UMF_RESULT_ERROR_UNKNOWN ;
546+ }
547+
548+ ze_memory_provider_t * ze_provider =
549+ umf_ba_global_alloc (sizeof (ze_memory_provider_t ));
550+ if (!ze_provider ) {
551+ LOG_ERR ("Cannot allocate memory for Level Zero Memory Provider" );
552+ return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY ;
553+ }
554+
555+ ze_provider -> context = ze_params -> level_zero_context_handle ;
556+ ze_provider -> device = ze_params -> level_zero_device_handle ;
557+ ze_provider -> memory_type = (ze_memory_type_t )ze_params -> memory_type ;
558+ ze_provider -> freePolicyFlags =
559+ umfFreePolicyToZePolicy (ze_params -> freePolicy );
560+ ze_provider -> min_page_size = 0 ;
561+ ze_provider -> device_ordinal = ze_params -> device_ordinal ;
562+
563+ memset (& ze_provider -> device_properties , 0 ,
564+ sizeof (ze_provider -> device_properties ));
565+ ze_provider -> device_properties .stype = ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES ;
566+
567+ if (ze_provider -> device ) {
568+ umf_result_t ret = ze2umf_result (g_ze_ops .zeDeviceGetProperties (
569+ ze_provider -> device , & ze_provider -> device_properties ));
570+
571+ if (ret != UMF_RESULT_SUCCESS ) {
572+ LOG_ERR ("Cannot get device properties" );
573+ umf_ba_global_free (ze_provider );
574+ return ret ;
575+ }
576+ }
577+
578+ if (ze_params -> resident_device_count ) {
579+ ze_provider -> resident_device_handles = umf_ba_global_alloc (
580+ sizeof (ze_device_handle_t ) * ze_params -> resident_device_count );
581+ if (!ze_provider -> resident_device_handles ) {
582+ LOG_ERR ("Cannot allocate memory for resident devices" );
583+ umf_ba_global_free (ze_provider );
584+ return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY ;
585+ }
586+
587+ ze_provider -> resident_device_count = ze_params -> resident_device_count ;
588+
589+ for (uint32_t i = 0 ; i < ze_provider -> resident_device_count ; i ++ ) {
590+ ze_provider -> resident_device_handles [i ] =
591+ ze_params -> resident_device_handles [i ];
592+ }
593+ } else {
594+ ze_provider -> resident_device_handles = NULL ;
595+ ze_provider -> resident_device_count = 0 ;
596+ }
597+
598+ umf_result_t result =
599+ query_min_page_size (ze_provider , & ze_provider -> min_page_size );
600+ if (result != UMF_RESULT_SUCCESS ) {
601+ ze_memory_provider_finalize (provider );
602+ return result ;
603+ }
604+
605+ * provider = ze_provider ;
606+
607+ return UMF_RESULT_SUCCESS ;
608+ }
609+
556610static void ze_memory_provider_get_last_native_error (void * provider ,
557611 const char * * ppMessage ,
558612 int32_t * pError ) {
@@ -569,12 +623,19 @@ static void ze_memory_provider_get_last_native_error(void *provider,
569623static umf_result_t ze_memory_provider_get_min_page_size (void * provider ,
570624 void * ptr ,
571625 size_t * pageSize ) {
572- (void )provider ;
573- (void )ptr ;
626+ ze_memory_provider_t * ze_provider = (ze_memory_provider_t * )provider ;
574627
575- // TODO
576- * pageSize = 1024 * 64 ;
577- return UMF_RESULT_SUCCESS ;
628+ if (!ptr ) {
629+ * pageSize = ze_provider -> min_page_size ;
630+ return UMF_RESULT_SUCCESS ;
631+ }
632+
633+ ze_memory_allocation_properties_t properties = {
634+ .stype = ZE_STRUCTURE_TYPE_MEMORY_ALLOCATION_PROPERTIES };
635+ ze_result_t ze_result = g_ze_ops .zeMemGetAllocProperties (
636+ ze_provider -> context , ptr , & properties , NULL );
637+
638+ return ze2umf_result (ze_result );
578639}
579640
580641static umf_result_t ze_memory_provider_purge_lazy (void * provider , void * ptr ,
@@ -600,12 +661,8 @@ static umf_result_t ze_memory_provider_purge_force(void *provider, void *ptr,
600661static umf_result_t
601662ze_memory_provider_get_recommended_page_size (void * provider , size_t size ,
602663 size_t * pageSize ) {
603- (void )provider ;
604664 (void )size ;
605-
606- // TODO
607- * pageSize = 1024 * 64 ;
608- return UMF_RESULT_SUCCESS ;
665+ return ze_memory_provider_get_min_page_size (provider , NULL , pageSize );
609666}
610667
611668static const char * ze_memory_provider_get_name (void * provider ) {
0 commit comments