Merge pull request #13055 from edgargabriel/topic/memkind-support

ompi/info: introduce support for the mpi_memory_alloc_kinds info object (II)

Author: edgargabriel

ompi/communicator/comm.c

@@ -26,7 +26,7 @@
  * Copyright (c) 2021      Nanook Consulting.  All rights reserved.
  * Copyright (c) 2018-2024 Triad National Security, LLC. All rights
  *                         reserved.
- * Copyright (c) 2023      Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2023-2025 Advanced Micro Devices, Inc. All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -56,6 +56,7 @@
 #include "ompi/communicator/communicator.h"
 #include "ompi/mca/pml/pml.h"
 #include "ompi/request/request.h"
+#include "ompi/info/info_memkind.h"
 
 #include "ompi/runtime/params.h"
 
@@ -447,6 +448,7 @@ int ompi_comm_create_w_info (ompi_communicator_t *comm, ompi_group_t *group, opa
     if (info) {
         opal_info_dup(info, &(newcomp->super.s_info));
     }
+    ompi_info_memkind_copy_or_set (&comm->instance->super, &newcomp->super, info);
 
     /* Set name for debugging purposes */
     snprintf(newcomp->c_name, MPI_MAX_OBJECT_NAME, "MPI COMMUNICATOR %s CREATE FROM %s",
@@ -699,10 +701,11 @@ int ompi_comm_split_with_info( ompi_communicator_t* comm, int color, int key,
 	     ompi_comm_print_cid (newcomp), ompi_comm_print_cid (comm));
 
     /* Copy info if there is one */
+    newcomp->super.s_info = OBJ_NEW(opal_info_t);
     if (info) {
-        newcomp->super.s_info = OBJ_NEW(opal_info_t);
         opal_info_dup(info, &(newcomp->super.s_info));
     }
+    ompi_info_memkind_copy_or_set (&comm->instance->super, &newcomp->super, info);
 
     /* Activate the communicator and init coll-component */
     rc = ompi_comm_activate (&newcomp, comm, NULL, NULL, NULL, false, mode);
@@ -994,6 +997,7 @@ static int ompi_comm_split_type_core(ompi_communicator_t *comm,
     if (info) {
         opal_infosubscribe_change_info(&newcomp->super, info);
     }
+    ompi_info_memkind_copy_or_set (&comm->instance->super, &newcomp->super, info);
 
     /* Activate the communicator and init coll-component */
     rc = ompi_comm_activate (&newcomp, comm, NULL, NULL, NULL, false, mode);
@@ -1347,6 +1351,7 @@ int ompi_comm_dup_with_info ( ompi_communicator_t * comm, opal_info_t *info, omp
     if (info) {
         opal_infosubscribe_change_info(&newcomp->super, info);
     }
+    ompi_info_memkind_copy_or_set (&comm->instance->super, &newcomp->super, info);
 
     /* activate communicator and init coll-module */
     rc = ompi_comm_activate (&newcomp, comm, NULL, NULL, NULL, false, mode);
@@ -1437,6 +1442,7 @@ static int ompi_comm_idup_internal (ompi_communicator_t *comm, ompi_group_t *gro
         if (info) {
             opal_info_dup(info, &(newcomp->super.s_info));
         }
+        ompi_info_memkind_copy_or_set (&comm->super, &newcomp->super, info);
     }
 
     ompi_comm_request_schedule_append (request, ompi_comm_idup_getcid, subreq, subreq[0] ? 1 : 0);
@@ -1588,6 +1594,7 @@ int ompi_comm_create_from_group (ompi_group_t *group, const char *tag, opal_info
     if (NULL == newcomp->super.s_info) {
         return OMPI_ERR_OUT_OF_RESOURCE;
     }
+    ompi_info_memkind_copy_or_set (&group->grp_instance->super, &newcomp->super, info);
 
     /* activate communicator and init coll-module. use the group allreduce implementation as
      * no collective module has yet been selected. the tag does not matter as any tag will
@@ -1727,6 +1734,10 @@ int ompi_intercomm_create (ompi_communicator_t *local_comm, int local_leader, om
         return rc;
     }
 
+    // Copy info if there is one.
+    newcomp->super.s_info = OBJ_NEW(opal_info_t);
+    ompi_info_memkind_copy_or_set (&local_comm->instance->super, &newcomp->super, &ompi_mpi_info_null.info.super);
+
     *newintercomm = newcomp;
 
     return OMPI_SUCCESS;
@@ -1791,6 +1802,7 @@ int ompi_intercomm_create_from_groups (ompi_group_t *local_group, int local_lead
                 ompi_comm_free (&local_comm);
                 return OMPI_ERR_OUT_OF_RESOURCE;
             }
+            leader_group->grp_instance = local_group->grp_instance;
 
             /* create a unique tag for allocating the leader communicator. we can eliminate this step
              * if we take a CID from the newly allocated block belonging to local_comm. this is
@@ -1888,6 +1900,7 @@ int ompi_intercomm_create_from_groups (ompi_group_t *local_group, int local_lead
     if (info) {
         opal_info_dup(info, &(newcomp->super.s_info));
     }
+    ompi_info_memkind_copy_or_set (&local_group->grp_instance->super, &newcomp->super, info);
 
     /* activate communicator and init coll-module */
     rc = ompi_comm_activate (&newcomp, local_comm, leader_comm, &local_leader, &leader_comm_remote_leader,

ompi/communicator/comm_init.c

@@ -25,7 +25,7 @@
  * Copyright (c) 2016-2017 IBM Corporation. All rights reserved.
  * Copyright (c) 2018-2024 Triad National Security, LLC. 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) 2023      NVIDIA Corporation.  All rights reserved.
  * $COPYRIGHT$
  *
@@ -53,6 +53,7 @@
 #include "ompi/dpm/dpm.h"
 #include "ompi/memchecker.h"
 #include "ompi/instance/instance.h"
+#include "ompi/info/info_memkind.h"
 
 /*
 ** Table for Fortran <-> C communicator handle conversion
@@ -266,6 +267,7 @@ int ompi_comm_init_mpi3 (void)
             free(str);
         }
     }
+
     /* Setup MPI_COMM_SELF */
     OBJ_CONSTRUCT(&ompi_mpi_comm_self, ompi_communicator_t);
     assert(ompi_mpi_comm_self.comm.c_f_to_c_index == 1);
@@ -300,6 +302,17 @@ int ompi_comm_init_mpi3 (void)
        MPI_COMM_SELF, the keyhash will automatically be created. */
     ompi_mpi_comm_self.comm.c_keyhash = NULL;
 
+    char *memkind_requested = getenv ("OMPI_MCA_mpi_memory_alloc_kinds");
+    if (NULL != memkind_requested) {
+        char *memkind_provided;
+
+        ompi_info_memkind_process (memkind_requested, &memkind_provided);
+        opal_infosubscribe_subscribe (&ompi_mpi_comm_world.comm.super, "mpi_memory_alloc_kinds", memkind_provided, ompi_info_memkind_cb);
+        opal_infosubscribe_subscribe (&ompi_mpi_comm_self.comm.super, "mpi_memory_alloc_kinds", memkind_provided, ompi_info_memkind_cb);
+        opal_infosubscribe_subscribe (&ompi_mpi_comm_world.comm.instance->super, "mpi_memory_alloc_kinds", memkind_provided, ompi_info_memkind_cb);
+        free (memkind_provided);
+    }
+
     /*
      * finally here we set the predefined attribute keyvals
      */

ompi/file/file.c

@@ -18,6 +18,7 @@
  * Copyright (c) 2016-2017 IBM Corporation. All rights reserved.
  * Copyright (c) 2024      Triad National Security, LLC. All rights
  *                         reserved.
+ * Copyright (c) 2024      Advanced Micro Devices, Inc. All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -34,7 +35,7 @@
 #include "ompi/runtime/params.h"
 #include "ompi/mca/io/base/base.h"
 #include "ompi/info/info.h"
-
+#include "ompi/info/info_memkind.h"
 
 opal_mutex_t ompi_mpi_file_bootstrap_mutex = OPAL_MUTEX_STATIC_INIT;
 
@@ -123,6 +124,7 @@ int ompi_file_open(struct ompi_communicator_t *comm, const char *filename,
     if (info) {
         opal_info_dup(info, &(file->super.s_info));
     }
+    ompi_info_memkind_copy_or_set (&comm->instance->super, &file->super, info);
 
     file->f_amode = amode;
     file->f_filename = strdup(filename);

ompi/info/Makefile.am

@@ -21,7 +21,9 @@
 # This makefile.am does not stand on its own - it is included from ompi/Makefile.am
 
 headers += \
-        info/info.h
+        info/info.h                 \
+        info/info_memkind.h
 
 lib@OMPI_LIBMPI_NAME@_la_SOURCES += \
-        info/info.c
+        info/info.c                 \
+        info/info_memkind.c

ompi/info/info.c

@@ -51,6 +51,7 @@
 #include "opal/util/info.h"
 
 #include "ompi/info/info.h"
+#include "ompi/info/info_memkind.h"
 #include "ompi/runtime/mpiruntime.h"
 #include "ompi/runtime/params.h"
 #include "ompi/runtime/ompi_rte.h"
@@ -351,6 +352,9 @@ int ompi_mpiinfo_finalize(void)
         }
     }
 
+    /* Release the array of available memkind objects */
+    ompi_info_memkind_free_available();
+
     /* All done -- destroy the table */
 
     OBJ_DESTRUCT(&ompi_info_f_to_c_table);