Merge pull request #9336 from sb17v/topic/master/oshmem/fix_assertion_error_on_oshmem_local_vpids_cleanup_v2

oshmem: fix assertion error on oshmem local vpids cleanup

Author: artpol84

oshmem/mca/memheap/base/base.h

@@ -176,13 +176,35 @@ static inline int memheap_is_va_in_segment(void *va, int segno)
     return map_segment_is_va_in(&memheap_find_seg(segno)->super, va);
 }
 
-static inline int memheap_find_segnum(void *va)
+static inline int memheap_find_segnum(void *va, int pe)
 {
     int i;
-
-    for (i = 0; i < mca_memheap_base_map.n_segments; i++) {
-        if (memheap_is_va_in_segment(va, i)) {
-            return i;
+    int my_pe = oshmem_my_proc_id();
+
+    if (pe == my_pe) {
+        /* Find segment number for local segment using va_base
+         * TODO: Merge local and remote segment information in mkeys_cache
+         */
+        for (i = 0; i < mca_memheap_base_map.n_segments; i++) {
+            if (memheap_is_va_in_segment(va, i)) {
+                return i;
+            }
+        }
+    } else {
+        /* Find segment number for remote segments using va_base */
+        for (i = 0; i < mca_memheap_base_map.n_segments; i++) {
+            map_segment_t *seg = memheap_find_seg(i);
+            if (seg) {
+                sshmem_mkey_t **mkeys_cache = seg->mkeys_cache;
+                if (mkeys_cache) {
+                    if (mkeys_cache[pe]) {
+                        if ((va >= mkeys_cache[pe]->va_base) &&
+                            (va < mkeys_cache[pe]->va_base + mkeys_cache[pe]->len)) {
+                            return i;
+                        }
+                    }
+                }
+            }
         }
     }
     return MEMHEAP_SEG_INVALID;

oshmem/mca/spml/ucx/spml_ucx.c

@@ -574,7 +574,13 @@ void mca_spml_ucx_rmkey_free(sshmem_mkey_t *mkey, int pe)
     if (!mkey->spml_context) {
         return;
     }
-    segno = memheap_find_segnum(mkey->va_base);
+    segno = memheap_find_segnum(mkey->va_base, pe);
+    if (MEMHEAP_SEG_INVALID == segno) {
+        SPML_UCX_ERROR("mca_spml_ucx_rmkey_free failed because of invalid "
+            "segment number: %d\n", segno);
+        return;
+    }
+
     ucx_mkey = (spml_ucx_mkey_t *)(mkey->spml_context);
     rc = mca_spml_ucx_ctx_mkey_del(&mca_spml_ucx_ctx_default, pe, segno, ucx_mkey);
     if (OSHMEM_SUCCESS != rc) {
@@ -676,7 +682,12 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
         return NULL;
     }
 
-    segno   = memheap_find_segnum(addr);
+    segno   = memheap_find_segnum(addr, my_pe);
+    if (MEMHEAP_SEG_INVALID == segno) {
+        SPML_UCX_ERROR("mca_spml_ucx_register failed because of invalid "
+            "segment number: %d\n", segno);
+        return NULL;
+    }
     mem_seg = memheap_find_seg(segno);
 
     /* if possible use mem handle already created by ucx allocator */
@@ -750,13 +761,19 @@ int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
         return OSHMEM_SUCCESS;
 
     mem_seg  = memheap_find_va(mkeys[SPML_UCX_TRANSP_IDX].va_base);
-    ucx_mkey = (spml_ucx_mkey_t*)mkeys[SPML_UCX_TRANSP_IDX].spml_context;
-    segno = memheap_find_segnum(mkeys[SPML_UCX_TRANSP_IDX].va_base);
-
     if (OPAL_UNLIKELY(NULL == mem_seg)) {
         return OSHMEM_ERROR;
     }
 
+    segno = memheap_find_segnum(mkeys[SPML_UCX_TRANSP_IDX].va_base, my_pe);
+    if (MEMHEAP_SEG_INVALID == segno) {
+        SPML_UCX_ERROR("mca_spml_ucx_deregister failed because of invalid "
+            "segment number: %d\n", segno);
+        return OSHMEM_ERROR;
+    }
+
+    ucx_mkey = (spml_ucx_mkey_t*)mkeys[SPML_UCX_TRANSP_IDX].spml_context;
+
     if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
         ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
     }

oshmem/proc/proc.c

@@ -179,6 +179,7 @@ int oshmem_proc_group_finalize(void)
         }
     }
 
+    OBJ_DESTRUCT(&_oshmem_local_vpids);
     OBJ_DESTRUCT(&oshmem_group_array);
 
     oshmem_group_cache_destroy();
@@ -265,8 +266,6 @@ oshmem_proc_group_destroy_internal(oshmem_group_t* group, int scoll_unselect)
         mca_scoll_base_group_unselect(group);
     }
 
-    /* Destroy proc array */
-    OBJ_DESTRUCT(&_oshmem_local_vpids);
     if (group->proc_vpids) {
         free(group->proc_vpids);
     }