MPI 4.1: add support for memory-alloc-kinds

src/mca/schizo/ompi/schizo-ompi-cli.rstxt

@@ -4,6 +4,8 @@
    Copyright (c) 2022      Cisco Systems, Inc.  All rights reserved.
    Copyright (c) 2022      IBM Corporation.  All rights reserved.
    Copyright (c) 2023      Jeffrey M. Squyres.  All rights reserved.
+   Copyright (c) 2024      Triad National Security, LLC. All rights reserved.
+
    $COPYRIGHT$
 
    Additional copyrights may follow
@@ -159,6 +161,8 @@ MPI Options
 
 * ``--initial-errhandler``: Specify the initial error handler that is
   attached to predefined communicators during the first MPI call.
+* ``--memory-alloc-kinds``: Value is a comma separated list of
+   memory allocation kinds.
 * ``--display-comm``: Display table of communication methods between
   MPI_COMM_WORLD ranks during MPI_Init
 * ``--display-comm-finalize``: Display table of communication methods

src/mca/schizo/ompi/schizo_ompi.c

@@ -19,7 +19,7 @@
  *                         and Technology (RIST). All rights reserved.
  * Copyright (c) 2018-2022 IBM Corporation.  All rights reserved.
  * Copyright (c) 2021-2024 Nanook Consulting  All rights reserved.
- * Copyright (c) 2022      Triad National Security, LLC. All rights
+ * Copyright (c) 2022-2024 Triad National Security, LLC. All rights
  *                         reserved.
  * $COPYRIGHT$
  *
@@ -194,8 +194,10 @@ static struct option ompioptions[] = {
     PMIX_OPTION_DEFINE(PRTE_CLI_CONTINUOUS, PMIX_ARG_NONE),
     PMIX_OPTION_DEFINE("with-ft", PMIX_ARG_REQD),
 
-    /* mpiexec mandated form launch key parameters */
+    /* mpiexec mandated form launch key parameters - MPI 4.0 */
     PMIX_OPTION_DEFINE("initial-errhandler", PMIX_ARG_REQD),
+    /* mpiexec mandated form launch key parameters  - MPI 4.1*/
+    PMIX_OPTION_DEFINE("memory-alloc-kinds", PMIX_ARG_REQD),
 
     /* Display Commumication Protocol : MPI_Init */
     PMIX_OPTION_DEFINE("display-comm", PMIX_ARG_NONE),
@@ -1575,6 +1577,15 @@ static int parse_env(char **srcenv, char ***dstenv,
         }
     }
 
+    if (NULL != (opt = pmix_cmd_line_get_param(results, "memory-alloc-kinds"))) {
+        rc = check_cache(&cache, &cachevals, "mpi_memory_alloc_kinds", opt->values[0]);
+        if (PRTE_SUCCESS != rc) {
+            PMIX_ARGV_FREE_COMPAT(cache);
+            PMIX_ARGV_FREE_COMPAT(cachevals);
+            return rc;
+        }
+    }
+
     if (pmix_cmd_line_is_taken(results, "display-comm") &&
         pmix_cmd_line_is_taken(results, "display-comm-finalize")) {
         PMIX_SETENV_COMPAT("OMPI_MCA_ompi_display_comm", "mpi_init,mpi_finalize", true, dstenv);

src/mca/schizo/prte/schizo_prte.c

@@ -196,6 +196,7 @@ static struct option prterunoptions[] = {
     PMIX_OPTION_SHORT_DEFINE(PRTE_CLI_PRELOAD_BIN, PMIX_ARG_NONE, 's'),
     PMIX_OPTION_DEFINE(PRTE_CLI_DO_NOT_AGG_HELP, PMIX_ARG_NONE),
     PMIX_OPTION_DEFINE(PRTE_CLI_FWD_ENVIRON, PMIX_ARG_OPTIONAL),
+    PMIX_OPTION_DEFINE(PRTE_CLI_MEM_ALLOC_KIND, PMIX_ARG_REQD),
 
     // output options
     PMIX_OPTION_DEFINE(PRTE_CLI_OUTPUT, PMIX_ARG_REQD),
@@ -310,6 +311,7 @@ static struct option prunoptions[] = {
     PMIX_OPTION_SHORT_DEFINE(PRTE_CLI_PRELOAD_BIN, PMIX_ARG_NONE, 's'),
     PMIX_OPTION_DEFINE(PRTE_CLI_DO_NOT_AGG_HELP, PMIX_ARG_NONE),
     PMIX_OPTION_DEFINE(PRTE_CLI_FWD_ENVIRON, PMIX_ARG_OPTIONAL),
+    PMIX_OPTION_DEFINE(PRTE_CLI_MEM_ALLOC_KIND, PMIX_ARG_REQD),
 
     // output options
     PMIX_OPTION_DEFINE(PRTE_CLI_OUTPUT, PMIX_ARG_REQD),

src/prted/pmix/pmix_server_queries.c

@@ -90,7 +90,8 @@ static void _query(int sd, short args, void *cbdata)
     pmix_proc_info_t *procinfo;
     pmix_data_array_t dry;
     prte_proc_t *proct;
-    pmix_proc_t *proc;
+    pmix_proc_t *proc, pproc;
+    pmix_info_t info;
     size_t sz;
     PRTE_HIDE_UNUSED_PARAMS(sd, args);
 
@@ -823,6 +824,27 @@ static void _query(int sd, short args, void *cbdata)
                 }
 #endif
 
+#ifdef PMIX_MEM_ALLOC_KIND
+            } else if (0 == strcmp(q->keys[n], PMIX_MEM_ALLOC_KIND)) {
+                pmix_value_t *value;
+                jdata = prte_get_job_data_object(jobid);
+                if (NULL == jdata) {
+                    ret = PMIX_ERR_NOT_FOUND;
+                    goto done;
+                }
+                PMIX_LOAD_PROCID(&pproc, jobid, PMIX_RANK_WILDCARD);
+                PMIX_INFO_LOAD(&info, PMIX_IMMEDIATE, NULL, PMIX_BOOL);
+                ret = PMIx_Get(&pproc, PMIX_MEM_ALLOC_KIND, &info, 1, (void**)&value);
+                if (PMIX_SUCCESS != ret) {
+                    goto done;
+                }
+                PMIX_INFO_LIST_ADD(rc, results, PMIX_MEM_ALLOC_KIND, value->data.string, PMIX_STRING);
+                PMIX_VALUE_RELEASE(value);
+                if (PMIX_SUCCESS != rc) {
+                    PMIX_ERROR_LOG(rc);
+                    goto done;
+                }
+#endif
             } else {
                 fprintf(stderr, "Query for unrecognized attribute: %s\n", q->keys[n]);
             }

src/prted/prun_common.c

@@ -661,6 +661,13 @@ int prun_common(pmix_cli_result_t *results,
         PMIX_INFO_LIST_ADD(ret, jinfo, PMIX_LOG_AGG, &flag, PMIX_BOOL);
     }
 
+#ifdef PMIX_MEM_ALLOC_KIND
+    opt = pmix_cmd_line_get_param(results, PRTE_CLI_MEM_ALLOC_KIND);
+    if (NULL != opt) {
+        PMIX_INFO_LIST_ADD(ret, jinfo, PMIX_MEM_ALLOC_KIND, opt->values[0], PMIX_STRING);
+    }
+#endif
+
     /* give the schizo components a chance to add to the job info */
     schizo->job_info(results, jinfo);
 

src/util/prte_cmd_line.h

@@ -15,7 +15,7 @@
  * Copyright (c) 2016-2017 Los Alamos National Security, LLC. All rights
  *                         reserved.
  * Copyright (c) 2017-2022 IBM Corporation.  All rights reserved.
- * Copyright (c) 2021-2023 Nanook Consulting.  All rights reserved.
+ * Copyright (c) 2021-2024 Nanook Consulting  All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -108,6 +108,7 @@ BEGIN_C_DECLS
 #define PRTE_CLI_SET_CWD_SESSION        "set-cwd-to-session-dir"    // none
 #define PRTE_CLI_ENABLE_RECOVERY        "enable-recovery"           // none
 #define PRTE_CLI_DISABLE_RECOVERY       "disable-recovery"          // none
+#define PRTE_CLI_MEM_ALLOC_KIND			"memory-alloc-kinds"        // required
 
 // Placement options
 #define PRTE_CLI_MAPBY                  "map-by"                    // required