From 887d46c826a31f1eb0535e03d5f6792ab6da0b83 Mon Sep 17 00:00:00 2001
From: Ralph Castain <rhc@pmix.org>
Date: Mon, 12 Apr 2021 15:48:44 -0700
Subject: [PATCH] Retrieve cpuset when configured with pmix rte

When configured `--with-ompi-pmix-rte`, ensure the `pmix_process_info_t`
structure has a `cpuset` entry and that it gets set.

Signed-off-by: Ralph Castain <rhc@pmix.org>
---
 ompi/mca/rte/pmix/rte_pmix.h        |  2 ++
 ompi/mca/rte/pmix/rte_pmix_module.c | 12 ++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/ompi/mca/rte/pmix/rte_pmix.h b/ompi/mca/rte/pmix/rte_pmix.h
index c2df3fe8bf6..bb083ab4eea 100644
--- a/ompi/mca/rte/pmix/rte_pmix.h
+++ b/ompi/mca/rte/pmix/rte_pmix.h
@@ -7,6 +7,7 @@
  *                         and Technology (RIST). All rights reserved.
  * Copyright (c) 2015      Intel, Inc. All rights reserved.
  * Copyright (c) 2016-2017 IBM Corporation. All rights reserved.
+ * Copyright (c) 2021      Nanook Consulting.  All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -115,6 +116,7 @@ typedef struct {
     int32_t num_local_peers;
     uint32_t num_procs;
     uint32_t app_num;
+    char *cpuset;                       /**< String-representation of bitmap where we are bound */
 } pmix_process_info_t;
 OMPI_DECLSPEC extern pmix_process_info_t pmix_process_info;
 #define ompi_process_info pmix_process_info
diff --git a/ompi/mca/rte/pmix/rte_pmix_module.c b/ompi/mca/rte/pmix/rte_pmix_module.c
index cd393e3a0c1..a16e0895e14 100644
--- a/ompi/mca/rte/pmix/rte_pmix_module.c
+++ b/ompi/mca/rte/pmix/rte_pmix_module.c
@@ -6,6 +6,7 @@
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
  * Copyright (c) 2014      Cisco Systems, Inc.  All rights reserved.
+ * Copyright (c) 2021      Nanook Consulting.  All rights reserved.
  * $COPYRIGHT$
  */
 #include "ompi_config.h"
@@ -674,6 +675,17 @@ int ompi_rte_init(int *pargc, char ***pargv)
         peers = NULL;
     }
 
+    /* get our cpuset */
+    val = NULL;
+    OPAL_MODEX_RECV_VALUE_OPTIONAL(ret, OPAL_PMIX_CPUSET,
+                                   &pmix_process_info.my_name,
+                                   &val, OPAL_STRING);
+    if (OPAL_SUCCESS == ret && NULL != val) {
+        pmix_process_info.cpuset = val;
+    } else {
+        pmix_process_info.cpuset = NULL;
+    }
+
     /* set the locality */
     if (NULL != peers) {
         /* identify our location */