@@ -426,3 +426,153 @@ umfMemspaceMemtargetRemove(umf_memspace_handle_t hMemspace,
426426 hMemspace -> size -- ;
427427 return UMF_RESULT_SUCCESS ;
428428}
429+
430+ // Helper function - returns zero on success, negative in case of error in filter function
431+ // and positive error code, in case of other errors.
432+ static int umfMemspaceFilterHelper (umf_memspace_handle_t memspace ,
433+ umf_memspace_filter_func_t filter ,
434+ void * args ) {
435+
436+ if (!memspace || !filter ) {
437+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
438+ }
439+
440+ size_t idx = 0 ;
441+ int ret ;
442+ umf_memtarget_handle_t * nodesToRemove =
443+ umf_ba_global_alloc (sizeof (* nodesToRemove ) * memspace -> size );
444+ if (!nodesToRemove ) {
445+ return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY ;
446+ }
447+
448+ for (size_t i = 0 ; i < memspace -> size ; i ++ ) {
449+ ret = filter (memspace , memspace -> nodes [i ], args );
450+ if (ret < 0 ) {
451+ LOG_ERR ("filter function failed" );
452+ goto free_mem ;
453+ } else if (ret == 0 ) {
454+ nodesToRemove [idx ++ ] = memspace -> nodes [i ];
455+ }
456+ }
457+
458+ size_t i = 0 ;
459+ for (; i < idx ; i ++ ) {
460+ ret = umfMemspaceMemtargetRemove (memspace , nodesToRemove [i ]);
461+ if (ret != UMF_RESULT_SUCCESS ) {
462+ goto re_add ;
463+ }
464+ }
465+
466+ umf_ba_global_free (nodesToRemove );
467+ return UMF_RESULT_SUCCESS ;
468+
469+ re_add :
470+ // If target removal failed, add back previously removed targets.
471+ for (size_t j = 0 ; j < i ; j ++ ) {
472+ umf_result_t ret2 = umfMemspaceMemtargetAdd (memspace , nodesToRemove [j ]);
473+ if (ret2 != UMF_RESULT_SUCCESS ) {
474+ ret =
475+ UMF_RESULT_ERROR_UNKNOWN ; // indicate that memspace is corrupted
476+ break ;
477+ }
478+ }
479+ free_mem :
480+ umf_ba_global_free (nodesToRemove );
481+ return ret ;
482+ }
483+
484+ umf_result_t umfMemspaceUserFilter (umf_memspace_handle_t memspace ,
485+ umf_memspace_filter_func_t filter ,
486+ void * args ) {
487+
488+ if (!memspace || !filter ) {
489+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
490+ }
491+
492+ int ret = umfMemspaceFilterHelper (memspace , filter , args );
493+ if (ret < 0 ) {
494+ return UMF_RESULT_ERROR_USER_SPECIFIC ;
495+ }
496+
497+ return ret ;
498+ }
499+
500+ typedef struct filter_by_id_args {
501+ unsigned * ids ; // array of numa nodes ids
502+ size_t size ; // size of the array
503+ } filter_by_id_args_t ;
504+
505+ /*
506+ * The following predefined filter callbacks returns umf_result_t codes as negative value
507+ * because only negative values are treated as errors. umfMemspaceFilterHelper() will pass
508+ * this error code through and umfMemspaceFilterBy*() functions will translate this code to positive
509+ * umf_result_t code.
510+ */
511+
512+ static int filterById (umf_const_memspace_handle_t memspace ,
513+ umf_const_memtarget_handle_t target , void * args ) {
514+ if (!memspace || !target || !args ) {
515+ return - UMF_RESULT_ERROR_INVALID_ARGUMENT ;
516+ }
517+
518+ filter_by_id_args_t * filterArgs = args ;
519+ for (size_t i = 0 ; i < filterArgs -> size ; i ++ ) {
520+ unsigned id ;
521+ umf_result_t ret = umfMemtargetGetId (target , & id );
522+ if (ret != UMF_RESULT_SUCCESS ) {
523+ return - ret ;
524+ }
525+
526+ if (id == filterArgs -> ids [i ]) {
527+ return 1 ;
528+ }
529+ }
530+ return 0 ;
531+ }
532+
533+ static int filterByCapacity (umf_const_memspace_handle_t memspace ,
534+ umf_const_memtarget_handle_t target , void * args ) {
535+ if (!memspace || !target || !args ) {
536+ return - UMF_RESULT_ERROR_INVALID_ARGUMENT ;
537+ }
538+
539+ size_t capacity ;
540+ umf_result_t ret = umfMemtargetGetCapacity (target , & capacity );
541+ if (ret != UMF_RESULT_SUCCESS ) {
542+ return - ret ;
543+ }
544+
545+ size_t * targetCapacity = args ;
546+ return (capacity >= * targetCapacity ) ? 1 : 0 ;
547+ }
548+
549+ umf_result_t umfMemspaceFilterById (umf_memspace_handle_t memspace ,
550+ unsigned * ids , size_t size ) {
551+ if (!memspace || !ids || size == 0 ) {
552+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
553+ }
554+
555+ filter_by_id_args_t args = {ids , size };
556+ int ret = umfMemspaceFilterHelper (memspace , & filterById , & args );
557+
558+ // if umfMemspaceFilter() returned negative umf_result_t change it to positive
559+ return ret < 0 ? - ret : ret ;
560+ }
561+
562+ umf_result_t umfMemspaceFilterByCapacity (umf_memspace_handle_t memspace ,
563+ int64_t capacity ) {
564+ if (!memspace ) {
565+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
566+ }
567+ // TODO: At this moment this function filters out memory targets that capacity is
568+ // less than specified size. We can extend this function to support reverse filter,
569+ // by using negative values of capacity parameter.
570+ // For now we just return invalid argument.
571+ if (capacity < 0 ) {
572+ return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
573+ }
574+ int ret = umfMemspaceFilterHelper (memspace , & filterByCapacity , & capacity );
575+
576+ // if umfMemspaceFilter() returned negative umf_result_t change it to positive
577+ return ret < 0 ? - ret : ret ;
578+ }
0 commit comments