diff --git a/ompi/runtime/ompi_rte.c b/ompi/runtime/ompi_rte.c index a2feca8285..b408b4ed77 100644 --- a/ompi/runtime/ompi_rte.c +++ b/ompi/runtime/ompi_rte.c @@ -12,6 +12,8 @@ * reserved. * Copyright (c) 2019 Research Organization for Information Science * and Technology (RIST). All rights reserved. + * Copyright (c) 2020 Amazon.com, Inc. or its affiliates. All Rights + * reserved. * $COPYRIGHT$ */ #include "ompi_config.h" @@ -70,6 +72,7 @@ pmix_process_info_t pmix_process_info = { .proc_session_dir = NULL, .my_local_rank = 0, .my_node_rank = 0, + .my_numa_rank = UINT16_MAX, /* Assume invalid NUMA rank, set to UINT16_MAX */ .num_local_peers = 0, .num_procs = 0, .app_num = 0, @@ -777,6 +780,28 @@ int ompi_rte_init(int *pargc, char ***pargv) pmix_proc_is_bound = false; } + /* get our numa rank from PMIx */ + if (pmix_proc_is_bound) { + OPAL_MODEX_RECV_VALUE_OPTIONAL(rc, PMIX_NUMA_RANK, + &pmix_process_info.my_name, &u16ptr, PMIX_UINT16); + if (PMIX_SUCCESS != rc) { + if (ompi_singleton) { + /* just assume the numa_rank is invalid, set to UINT16_MAX */ + u16 = UINT16_MAX; + } else { + ret = opal_pmix_convert_status(rc); + error = "numa rank"; + goto error; + } + } + pmix_process_info.my_numa_rank = u16; + } else { + /* If processes are not bound, the numa_rank is not available + * Assign UINT16_MAX to the numa_rank to indicate an invalid value + */ + pmix_process_info.my_numa_rank = UINT16_MAX; + } + /* get our local peers */ if (0 < pmix_process_info.num_local_peers) { /* if my local rank if too high, then that's an error */ @@ -866,6 +891,7 @@ int ompi_rte_init(int *pargc, char ***pargv) opal_process_info.proc_session_dir = pmix_process_info.proc_session_dir; opal_process_info.num_local_peers = (int32_t)pmix_process_info.num_local_peers; opal_process_info.my_local_rank = (int32_t)pmix_process_info.my_local_rank; + opal_process_info.my_numa_rank = pmix_process_info.my_numa_rank; opal_process_info.cpuset = pmix_process_info.cpuset; return OPAL_SUCCESS; diff --git a/ompi/runtime/ompi_rte.h b/ompi/runtime/ompi_rte.h index 616df8af34..981152ed9a 100644 --- a/ompi/runtime/ompi_rte.h +++ b/ompi/runtime/ompi_rte.h @@ -8,6 +8,8 @@ * and Technology (RIST). All rights reserved. * Copyright (c) 2020 Triad National Security, LLC. All rights * reserved. + * Copyright (c) 2020 Amazon.com, Inc. or its affiliates. All Rights + * reserved. * * $COPYRIGHT$ * @@ -253,6 +255,8 @@ typedef struct { char *proc_session_dir; uint16_t my_local_rank; uint16_t my_node_rank; + /* process rank on local NUMA node. Set to UINT16_MAX if NUMA rank is unavailable */ + uint16_t my_numa_rank; int32_t num_local_peers; uint32_t num_procs; uint32_t app_num; diff --git a/opal/util/proc.c b/opal/util/proc.c index 16b1968bd8..8bb9443400 100644 --- a/opal/util/proc.c +++ b/opal/util/proc.c @@ -9,6 +9,8 @@ * and Technology (RIST). All rights reserved. * Copyright (c) 2015 Los Alamos National Security, LLC. All rights * reserved. + * Copyright (c) 2020 Amazon.com, Inc. or its affiliates. All Rights + * reserved. * $COPYRIGHT$ * * Additional copyrights may follow @@ -35,6 +37,7 @@ opal_process_info_t opal_process_info = { .proc_session_dir = NULL, .num_local_peers = 0, /* there is nobody else but me */ .my_local_rank = 0, /* I'm the only process around here */ + .my_numa_rank = UINT16_MAX, /* Assume numa_rank is unavailable, set to UINT16_MAX */ .cpuset = NULL, }; diff --git a/opal/util/proc.h b/opal/util/proc.h index 3bd86b6292..1cff4e0b6d 100644 --- a/opal/util/proc.h +++ b/opal/util/proc.h @@ -7,6 +7,8 @@ * Copyright (c) 2014-2016 Research Organization for Information Science * and Technology (RIST). All rights reserved. * Copyright (c) 2017 Cisco Systems, Inc. All rights reserved + * Copyright (c) 2020 Amazon.com, Inc. or its affiliates. All Rights + * reserved. * $COPYRIGHT$ * * Additional copyrights may follow @@ -108,6 +110,7 @@ typedef struct opal_process_info_t { char *proc_session_dir; /**< Session directory for the process */ int32_t num_local_peers; /**< number of procs from my job that share my node with me */ int32_t my_local_rank; /**< local rank on this node within my job */ + int16_t my_numa_rank; /**< rank on this processes NUMA node. A value of UINT16_MAX indicates unavailable numa_rank */ char *cpuset; /**< String-representation of bitmap where we are bound */ } opal_process_info_t; OPAL_DECLSPEC extern opal_process_info_t opal_process_info;