oshmem: Fix mkey release for remote segments

The bug was introduced in PR #9058 and caused rmkey_free to fail when
cleaning service data for remote PEs.

Enable memheap_find_segnum to find a segment by va for both local and remote
PEs.

Co-authored-by: Artem Y. Polyakov <artemp@nvidia.com>
Signed-off-by: Subhadeep Bhattacharya <subhadeepb@nvidia.com>
Signed-off-by: Artem Polyakov <artemp@nvidia.com>

Author: sb17v

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

@@ -571,7 +571,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) {
@@ -673,7 +679,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 */
@@ -747,13 +758,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);
     }