Merge pull request #9058 from sb17v/topic/v5.x/oshmem/mem_footprint_fix_v2

oshmem: Reduce the service memory footprint

Author: awlauria

oshmem/mca/atomic/ucx/atomic_ucx_cswap.c

@@ -48,7 +48,8 @@ int mca_atomic_ucx_cswap(shmem_ctx_t ctx,
     assert(NULL != prev);
 
     *prev      = value;
-    ucx_mkey   = mca_spml_ucx_get_mkey(ctx, pe, target, (void *)&rva, mca_spml_self);
+    ucx_mkey   = mca_spml_ucx_ctx_mkey_by_va(ctx, pe, target, (void *)&rva, mca_spml_self);
+    assert(NULL != ucx_mkey);
 #if HAVE_DECL_UCP_ATOMIC_OP_NBX
     status_ptr = ucp_atomic_op_nbx(ucx_ctx->ucp_peers[pe].ucp_conn,
                                    UCP_ATOMIC_OP_CSWAP, &cond, 1, rva,

oshmem/mca/atomic/ucx/atomic_ucx_module.c

@@ -65,8 +65,8 @@ int mca_atomic_ucx_op(shmem_ctx_t ctx,
 
     assert((8 == size) || (4 == size));
 
-    ucx_mkey = mca_spml_ucx_get_mkey(ctx, pe, target, (void *)&rva, mca_spml_self);
-
+    ucx_mkey = mca_spml_ucx_ctx_mkey_by_va(ctx, pe, target, (void *)&rva, mca_spml_self);
+    assert(NULL != ucx_mkey);
 #if HAVE_DECL_UCP_ATOMIC_OP_NBX
     status_ptr = ucp_atomic_op_nbx(ucx_ctx->ucp_peers[pe].ucp_conn,
                                    op, &value, 1, rva, ucx_mkey->rkey,
@@ -115,7 +115,8 @@ int mca_atomic_ucx_fop(shmem_ctx_t ctx,
 
     assert((8 == size) || (4 == size));
 
-    ucx_mkey = mca_spml_ucx_get_mkey(ctx, pe, target, (void *)&rva, mca_spml_self);
+    ucx_mkey = mca_spml_ucx_ctx_mkey_by_va(ctx, pe, target, (void *)&rva, mca_spml_self);
+    assert(NULL != ucx_mkey);
 #if HAVE_DECL_UCP_ATOMIC_OP_NBX
     status_ptr = ucp_atomic_op_nbx(ucx_ctx->ucp_peers[pe].ucp_conn, op, &value, 1,
                                    rva, ucx_mkey->rkey, &param);

oshmem/mca/memheap/base/base.h

@@ -200,22 +200,6 @@ static inline void *map_segment_va2rva(mkey_segment_t *seg, void *va)
     return memheap_va2rva(va, seg->super.va_base, seg->rva_base);
 }
 
-static inline map_base_segment_t *map_segment_find_va(map_base_segment_t *segs,
-                                                      size_t elem_size, void *va)
-{
-    map_base_segment_t *rseg;
-    int i;
-
-    for (i = 0; i < MCA_MEMHEAP_MAX_SEGMENTS; i++) {
-        rseg = (map_base_segment_t *)((char *)segs + elem_size * i);
-        if (OPAL_LIKELY(map_segment_is_va_in(rseg, va))) {
-            return rseg;
-        }
-    }
-
-    return NULL;
-}
-
 void mkey_segment_init(mkey_segment_t *seg, sshmem_mkey_t *mkey, uint32_t segno);
 
 static inline map_segment_t *memheap_find_va(void* va)

oshmem/mca/memheap/base/memheap_base_mkey.c

@@ -775,7 +775,6 @@ void mkey_segment_init(mkey_segment_t *seg, sshmem_mkey_t *mkey, uint32_t segno)
 
     s = memheap_find_seg(segno);
     assert(NULL != s);
-
     seg->super.va_base = s->super.va_base;
     seg->super.va_end  = s->super.va_end;
     seg->rva_base      = mkey->va_base;

oshmem/mca/memheap/base/memheap_base_register.c

@@ -86,7 +86,7 @@ static int _dereg_segment(map_segment_t *s)
                 continue;
             if (s->mkeys_cache[j]) {
                 if (s->mkeys_cache[j]->len) {
-                    MCA_SPML_CALL(rmkey_free(s->mkeys_cache[j]));
+                    MCA_SPML_CALL(rmkey_free(s->mkeys_cache[j], j));
                     free(s->mkeys_cache[j]->u.data);
                     s->mkeys_cache[j]->len = 0;
                 }

oshmem/mca/spml/base/base.h

@@ -78,7 +78,7 @@ OSHMEM_DECLSPEC int mca_spml_base_oob_get_mkeys(shmem_ctx_t ctx,
                                                 sshmem_mkey_t *mkeys);
 
 OSHMEM_DECLSPEC void mca_spml_base_rmkey_unpack(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t seg, int pe, int tr_id);
-OSHMEM_DECLSPEC void mca_spml_base_rmkey_free(sshmem_mkey_t *mkey);
+OSHMEM_DECLSPEC void mca_spml_base_rmkey_free(sshmem_mkey_t *mkey, int pe);
 OSHMEM_DECLSPEC void *mca_spml_base_rmkey_ptr(const void *dst_addr, sshmem_mkey_t *mkey, int pe);
 
 OSHMEM_DECLSPEC int mca_spml_base_put_nb(void *dst_addr,

oshmem/mca/spml/base/spml_base.c

@@ -256,7 +256,7 @@ void mca_spml_base_rmkey_unpack(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t s
 {
 }
 
-void mca_spml_base_rmkey_free(sshmem_mkey_t *mkey)
+void mca_spml_base_rmkey_free(sshmem_mkey_t *mkey, int pe)
 {
 }
 

oshmem/mca/spml/spml.h

@@ -150,7 +150,7 @@ typedef void * (*mca_spml_base_module_mkey_ptr_fn_t)(const void *dst_addr, sshme
  *
  * @param mkey remote mkey
  */
-typedef void (*mca_spml_base_module_mkey_free_fn_t)(sshmem_mkey_t *);
+typedef void (*mca_spml_base_module_mkey_free_fn_t)(sshmem_mkey_t *, int pe);
 
 /**
  * Register (Pinn) a buffer of 'size' bits starting in address addr