accelerator: add interface to retrieve memkind

add an API to the accelerator component to retrieve the
memory_alloc_kind information that is supported by the component.
The values stored/returned are based on the side document that is
about to be ratified, see

https://github.com/mpi-forum/mem-alloc/blob/main/mem_alloc.tex

Signed-off-by: Edgar Gabriel <[email protected]>

Author: edgargabriel

ompi/info/info_memkind.c

@@ -166,11 +166,9 @@ static int ompi_info_memkind_get_available(int *num_memkinds, ompi_memkind_t **m
     }
 
     int tmp_num = 2;
-#if 0
     if (0 != strcmp(opal_accelerator_base_selected_component.base_version.mca_component_name, "null")) {
         tmp_num++;
     }
-#endif
 
     ompi_info_memkind_available = (ompi_memkind_t *) malloc (tmp_num * sizeof(ompi_memkind_t));
     if (NULL == ompi_info_memkind_available) {
@@ -190,11 +188,12 @@ static int ompi_info_memkind_get_available(int *num_memkinds, ompi_memkind_t **m
     ompi_info_memkind_available[1].im_restrictors[1] = strdup ("win_allocate");
     ompi_info_memkind_available[1].im_restrictors[2] = strdup ("win_allocate_shared");
 
-#if 0
     if (tmp_num > 2) {
-        opal_accelerator.get_memkind_info (&ompi_info_memkind_available[2]);
+        ompi_info_memkind_available[2].im_num_restrictors = OMPI_MAX_NUM_MEMKIND_RESTRICTORS;
+        opal_accelerator.get_memkind (&(ompi_info_memkind_available[2].im_name),
+                                      &(ompi_info_memkind_available[2].im_num_restrictors),
+                                      (char**)ompi_info_memkind_available[2].im_restrictors);
     }
-#endif
     ompi_info_memkind_num_available = tmp_num;
 
  exit:

opal/mca/accelerator/accelerator.h

@@ -4,7 +4,7 @@
  *                         reserved.
  * Copyright (c)           Amazon.com, Inc. or its affiliates.
  *                         All Rights reserved.
- * Copyright (c) 2023      Advanced Micro Devices, Inc. All Rights reserved.
+ * Copyright (c) 2023-2024 Advanced Micro Devices, Inc. All Rights reserved.
  * Copyright (c) 2024      The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
@@ -654,6 +654,22 @@ typedef int (*opal_accelerator_base_module_get_num_devices_fn_t)(int *num_device
  */
 typedef int (*opal_accelerator_base_module_get_mem_bw_fn_t)(int device, float *bw);
 
+/**
+ * Get the memkind information of the accelerator component.
+ *
+ * @param[OUT]   name             Name of memory alloc kinds supported by component.
+ *                                This field will have to be released by the calling function.
+ * @param[INOUT] num_restrictors  As input, this parameter represents the lenght of the
+ *                                restrictors array allocated by the caller.
+ *                                At return, this variable will indicate the number of
+ *                                restrictors set by the function
+ * @param[OUT]   restrictors      Array of restrictors supported by the component.
+ *                                The array of char* pointers has been allocated by the caller.
+ *                                The elements of the array will have to be released by the caller.
+ *
+ */
+typedef void (*opal_accelerator_base_module_get_memkind_fn_t)(char **name, int *num_restrictors,
+                                                              char **restrictors);
 
 /*
  * the standard public API data structure
@@ -700,6 +716,7 @@ typedef struct {
 
     opal_accelerator_base_module_get_num_devices_fn_t num_devices;
     opal_accelerator_base_module_get_mem_bw_fn_t get_mem_bw;
+    opal_accelerator_base_module_get_memkind_fn_t get_memkind;
 } opal_accelerator_base_module_t;
 
 /**

opal/mca/accelerator/cuda/accelerator_cuda.c

@@ -26,6 +26,7 @@
 #include "opal/mca/rcache/rcache.h"
 #include "opal/util/show_help.h"
 #include "opal/util/proc.h"
+
 /* Accelerator API's */
 static int accelerator_cuda_check_addr(const void *addr, int *dev_id, uint64_t *flags);
 static int accelerator_cuda_create_stream(int dev_id, opal_accelerator_stream_t **stream);
@@ -80,10 +81,14 @@ static int accelerator_cuda_get_buffer_id(int dev_id, const void *addr, opal_acc
 static int accelerator_cuda_sync_stream(opal_accelerator_stream_t *stream);
 static int accelerator_cuda_get_num_devices(int *num_devices);
 static int accelerator_cuda_get_mem_bw(int device, float *bw);
+static void accelerator_cuda_get_memkind(char **name, int *num_restrictors, char **restrictors);
 
 #define GET_STREAM(_stream) \
     ((_stream) == MCA_ACCELERATOR_STREAM_DEFAULT ? 0 : *((CUstream *) (_stream)->stream))
 
+// This value is based on the memory kind MPI side document
+#define MCA_ACCELERATOR_CUDA_NUM_RESTRICTORS 3
+
 opal_accelerator_base_module_t opal_accelerator_cuda_module =
 {
     accelerator_cuda_check_addr,
@@ -125,7 +130,8 @@ opal_accelerator_base_module_t opal_accelerator_cuda_module =
     accelerator_cuda_get_buffer_id,
 
     accelerator_cuda_get_num_devices,
-    accelerator_cuda_get_mem_bw
+    accelerator_cuda_get_mem_bw,
+    accelerator_cuda_get_memkind
 };
 
 static inline int opal_accelerator_cuda_delayed_init_check(void)
@@ -1218,3 +1224,24 @@ static int accelerator_cuda_get_mem_bw(int device, float *bw)
     *bw = opal_accelerator_cuda_mem_bw[device];
     return OPAL_SUCCESS;
 }
+
+static void accelerator_cuda_get_memkind (char **name, int *num_restrictors, char **restrictors)
+{
+    int n_restrictors = *num_restrictors > MCA_ACCELERATOR_CUDA_NUM_RESTRICTORS ?
+	MCA_ACCELERATOR_CUDA_NUM_RESTRICTORS : *num_restrictors;
+
+    *name = strdup("cuda");
+
+    if (n_restrictors > 0) {
+	restrictors[0] = strdup("host");
+    }
+    if (n_restrictors > 1) {
+	restrictors[1] = strdup("device");
+    }
+    if (n_restrictors > 2) {
+	restrictors[2] = strdup("managed");
+    }
+    *num_restrictors = n_restrictors;
+
+  return;
+}

opal/mca/accelerator/null/accelerator_null_component.c

@@ -94,6 +94,7 @@ static int accelerator_null_sync_stream(opal_accelerator_stream_t *stream);
 static int accelerator_null_get_num_devices(int *num_devices);
 
 static int accelerator_null_get_mem_bw(int device, float *bw);
+static void accelerator_null_get_memkind(char **name, int *num_restrictors, char **restrictors);
 
 /*
  * Instantiate the public struct with all of our public information
@@ -174,7 +175,8 @@ opal_accelerator_base_module_t opal_accelerator_null_module =
     accelerator_null_get_buffer_id,
 
     accelerator_null_get_num_devices,
-    accelerator_null_get_mem_bw
+    accelerator_null_get_mem_bw,
+    accelerator_null_get_memkind
 };
 
 static int accelerator_null_open(void)
@@ -393,3 +395,11 @@ static int accelerator_null_get_mem_bw(int device, float *bw)
     *bw = 1.0; // return something that is not 0
     return OPAL_SUCCESS;
 }
+
+static void accelerator_null_get_memkind (char **name, int *num_restrictors, char **restrictors)
+{
+  *name = NULL;
+  *num_restrictors = 0;
+
+  return;
+}

opal/mca/accelerator/rocm/accelerator_rocm_module.c

@@ -74,9 +74,13 @@ static int mca_accelerator_rocm_sync_stream(opal_accelerator_stream_t *stream);
 static int mca_accelerator_rocm_get_num_devices(int *num_devices);
 
 static int mca_accelerator_rocm_get_mem_bw(int device, float *bw);
+static void mca_accelerator_rocm_get_memkind(char **name, int *num_restrictors, char **restrictors);
 
 #define GET_STREAM(_stream) (_stream == MCA_ACCELERATOR_STREAM_DEFAULT ? 0 : *((hipStream_t *)_stream->stream))
 
+// This value is based on the memory kind MPI side document
+#define MCA_ACCELERATOR_ROCM_NUM_RESTRICTORS 3
+
 opal_accelerator_base_module_t opal_accelerator_rocm_module =
 {
     mca_accelerator_rocm_check_addr,
@@ -118,7 +122,8 @@ opal_accelerator_base_module_t opal_accelerator_rocm_module =
     mca_accelerator_rocm_get_buffer_id,
 
     mca_accelerator_rocm_get_num_devices,
-    mca_accelerator_rocm_get_mem_bw
+    mca_accelerator_rocm_get_mem_bw,
+    mca_accelerator_rocm_get_memkind
 };
 
 
@@ -946,3 +951,24 @@ static int mca_accelerator_rocm_get_mem_bw(int device, float *bw)
     *bw = opal_accelerator_rocm_mem_bw[device];
     return OPAL_SUCCESS;
 }
+
+static void mca_accelerator_rocm_get_memkind (char **name, int *num_restrictors, char **restrictors)
+{
+    int n_restrictors = *num_restrictors > MCA_ACCELERATOR_ROCM_NUM_RESTRICTORS ?
+	MCA_ACCELERATOR_ROCM_NUM_RESTRICTORS : *num_restrictors;
+
+    *name = strdup("rocm");
+
+    if (n_restrictors > 0) {
+	restrictors[0] = strdup("host");
+    }
+    if (n_restrictors > 1) {
+	restrictors[1] = strdup("device");
+    }
+    if (n_restrictors > 2) {
+	restrictors[2] = strdup("managed");
+    }
+    *num_restrictors = n_restrictors;
+
+    return;
+}

opal/mca/accelerator/ze/accelerator_ze_module.c

@@ -77,6 +77,10 @@ static int mca_accelerator_ze_sync_stream(opal_accelerator_stream_t *stream);
 static int mca_accelerator_ze_get_num_devices(int *num_devices);
 
 static int mca_accelerator_ze_get_mem_bw(int device, float *bw);
+static void mca_accelerator_ze_get_memkind(char **name, int *num_restrictors, char **restrictors);
+
+// This value is based on the memory kind MPI side document
+#define MCA_ACCELERATOR_ZE_NUM_RESTRICTORS 3
 
 opal_accelerator_base_module_t opal_accelerator_ze_module =
 {
@@ -118,7 +122,8 @@ opal_accelerator_base_module_t opal_accelerator_ze_module =
 
     .get_buffer_id = mca_accelerator_ze_get_buffer_id,
     .num_devices = mca_accelerator_ze_get_num_devices,
-    .get_mem_bw = mca_accelerator_ze_get_mem_bw
+    .get_mem_bw = mca_accelerator_ze_get_mem_bw,
+    .get_memkind = mca_accelerator_ze_get_memkind
 };
 
 static int accelerator_ze_dev_handle_to_dev_id(ze_device_handle_t hDevice)
@@ -872,4 +877,25 @@ static int mca_accelerator_ze_get_mem_bw(int device, float *bw)
      * TODO
      */
     return OPAL_ERR_NOT_IMPLEMENTED;
-}
+}
+
+static void mca_accelerator_ze_get_memkind (char **name, int *num_restrictors, char **restrictors)
+{
+    int n_restrictors = *num_restrictors > MCA_ACCELERATOR_ZE_NUM_RESTRICTORS ?
+	MCA_ACCELERATOR_ZE_NUM_RESTRICTORS : *num_restrictors;
+
+    *name = strdup("level_zero");
+
+    if (n_restrictors > 0) {
+	restrictors[0] = strdup("host");
+    }
+    if (n_restrictors > 1) {
+	restrictors[1] = strdup("device");
+    }
+    if (n_restrictors > 2) {
+	restrictors[2] = strdup("shared");
+    }
+    *num_restrictors = n_restrictors;
+
+    return;
+}