Merge pull request #6492 from janjust/oshmem-multiple-contexts-master

Oshmem multiple contexts

Author: jladd-mlnx

opal/mca/common/ucx/common_ucx.c

@@ -153,6 +153,11 @@ void opal_common_ucx_mca_proc_added(void)
 #endif
 }
 
+OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence_nb(int *fenced)
+{
+    return opal_pmix.fence_nb(NULL, 0, opal_common_ucx_mca_fence_complete_cb, (void *)fenced);
+}
+
 OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence(ucp_worker_h worker)
 {
     volatile int fenced = 0;
@@ -182,9 +187,8 @@ static void opal_common_ucx_wait_all_requests(void **reqs, int count, ucp_worker
     }
 }
 
-OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, size_t count,
-                                            size_t my_rank, size_t max_disconnect, ucp_worker_h worker)
-{
+OPAL_DECLSPEC int opal_common_ucx_del_procs_nofence(opal_common_ucx_del_proc_t *procs, size_t count,
+                                            size_t my_rank, size_t max_disconnect, ucp_worker_h worker) {
     size_t num_reqs;
     size_t max_reqs;
     void *dreq, **dreqs;
@@ -232,10 +236,14 @@ OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, s
     opal_common_ucx_wait_all_requests(dreqs, num_reqs, worker);
     free(dreqs);
 
-    if (OPAL_SUCCESS != (ret = opal_common_ucx_mca_pmix_fence(worker))) {
-        return ret;
-    }
-
     return OPAL_SUCCESS;
 }
 
+OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, size_t count,
+                                            size_t my_rank, size_t max_disconnect, ucp_worker_h worker)
+{
+    opal_common_ucx_del_procs_nofence(procs, count, my_rank, max_disconnect, worker);
+
+    return opal_common_ucx_mca_pmix_fence(worker);
+}
+

opal/mca/common/ucx/common_ucx.h

@@ -105,8 +105,11 @@ OPAL_DECLSPEC void opal_common_ucx_mca_deregister(void);
 OPAL_DECLSPEC void opal_common_ucx_mca_proc_added(void);
 OPAL_DECLSPEC void opal_common_ucx_empty_complete_cb(void *request, ucs_status_t status);
 OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence(ucp_worker_h worker);
+OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence_nb(int *fenced);
 OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, size_t count,
                                             size_t my_rank, size_t max_disconnect, ucp_worker_h worker);
+OPAL_DECLSPEC int opal_common_ucx_del_procs_nofence(opal_common_ucx_del_proc_t *procs, size_t count,
+                                               size_t my_rank, size_t max_disconnect, ucp_worker_h worker);
 OPAL_DECLSPEC void opal_common_ucx_mca_var_register(const mca_base_component_t *component);
 
 static inline

oshmem/mca/spml/ikrit/spml_ikrit.c

@@ -149,7 +149,7 @@ int mca_spml_ikrit_put_simple(void* dst_addr,
                               void* src_addr,
                               int dst);
 
-static void mca_spml_ikrit_cache_mkeys(shmem_ctx_t ctx, sshmem_mkey_t *, uint32_t seg, int remote_pe, int tr_id);
+static void mca_spml_ikrit_cache_mkeys(sshmem_mkey_t *, uint32_t seg, int remote_pe, int tr_id);
 
 static mxm_mem_key_t *mca_spml_ikrit_get_mkey_slow(int pe, void *va, int ptl_id, void **rva);
 
@@ -185,7 +185,7 @@ mca_spml_ikrit_t mca_spml_ikrit = {
     mca_spml_ikrit_get_mkey_slow
 };
 
-static void mca_spml_ikrit_cache_mkeys(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t seg, int dst_pe, int tr_id)
+static void mca_spml_ikrit_cache_mkeys(sshmem_mkey_t *mkey, uint32_t seg, int dst_pe, int tr_id)
 {
     mxm_peer_t *peer;
 
@@ -504,7 +504,7 @@ sshmem_mkey_t *mca_spml_ikrit_register(void* addr,
                      my_rank, i, addr, (unsigned long long)size,
                      mca_spml_base_mkey2str(&mkeys[i]));
 
-        mca_spml_ikrit_cache_mkeys(NULL, &mkeys[i], memheap_find_segnum(addr), my_rank, i);
+        mca_spml_ikrit_cache_mkeys(&mkeys[i], memheap_find_segnum(addr), my_rank, i);
     }
     *count = MXM_PTL_LAST;
 
@@ -548,7 +548,7 @@ int mca_spml_ikrit_deregister(sshmem_mkey_t *mkeys)
 
 }
 
-int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe, uint32_t seg, sshmem_mkey_t *mkeys)
+int mca_spml_ikrit_oob_get_mkeys(int pe, uint32_t seg, sshmem_mkey_t *mkeys)
 {
     int ptl;
 
@@ -567,7 +567,7 @@ int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe, uint32_t seg, sshmem_m
         mkeys[ptl].len     = 0;
         mkeys[ptl].va_base = mca_memheap_seg2base_va(seg);
         mkeys[ptl].u.key   = MAP_SEGMENT_SHM_INVALID;
-        mca_spml_ikrit_cache_mkeys(NULL, &mkeys[ptl], seg, pe, ptl);
+        mca_spml_ikrit_cache_mkeys(&mkeys[ptl], seg, pe, ptl);
         return OSHMEM_SUCCESS;
     }
 

oshmem/mca/spml/ikrit/spml_ikrit.h

@@ -182,7 +182,7 @@ extern sshmem_mkey_t *mca_spml_ikrit_register(void* addr,
                                                 uint64_t shmid,
                                                 int *count);
 extern int mca_spml_ikrit_deregister(sshmem_mkey_t *mkeys);
-extern int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe,
+extern int mca_spml_ikrit_oob_get_mkeys(int pe,
                                         uint32_t segno,
                                         sshmem_mkey_t *mkeys);