Allow MPI_WIN_SHARED_QUERY on regular windows

MPI 4.0 introduced allows applications to query regular windows for shared
memory. This patch enables it for osc/rdma and osc/ucx and otherwise makes
sure we fail gracefully if the component does not provide the query callback.
For osc/rdma, this is currently supported only for allocated windows but could
later be extended to windows with application-provided memory through xpmem.

Signed-off-by: Joseph Schuchart <joseph.schuchart@stonybrook.edu>

Author: devreal

ompi/mca/osc/rdma/osc_rdma_component.c

@@ -69,6 +69,8 @@
 #include "ompi/mca/bml/base/base.h"
 #include "ompi/mca/mtl/base/base.h"
 
+static int ompi_osc_rdma_shared_query(struct ompi_win_t *win, int rank, size_t *size,
+                               ptrdiff_t *disp_unit, void *baseptr);
 static int ompi_osc_rdma_component_register (void);
 static int ompi_osc_rdma_component_init (bool enable_progress_threads, bool enable_mpi_threads);
 static int ompi_osc_rdma_component_finalize (void);
@@ -112,6 +114,7 @@ ompi_osc_rdma_component_t mca_osc_rdma_component = {
 };
 
 ompi_osc_base_module_t ompi_osc_rdma_module_rdma_template = {
+    .osc_win_shared_query = ompi_osc_rdma_shared_query,
     .osc_win_attach = ompi_osc_rdma_attach,
     .osc_win_detach  = ompi_osc_rdma_detach,
     .osc_free = ompi_osc_rdma_free,
@@ -876,7 +879,7 @@ static void ompi_osc_rdma_ensure_local_add_procs (void)
             /* this will cause add_proc to get called if it has not already been called */
             (void) mca_bml_base_get_endpoint (proc);
         }
-    } 
+    }
 
     free(procs);
 }
@@ -1610,3 +1613,58 @@ ompi_osc_rdma_set_no_lock_info(opal_infosubscriber_t *obj, const char *key, cons
      */
     return module->no_locks ? "true" : "false";
 }
+
+int ompi_osc_rdma_shared_query(
+    struct ompi_win_t *win, int rank, size_t *size,
+    ptrdiff_t *disp_unit, void *baseptr)
+{
+    ompi_osc_rdma_peer_t *peer;
+    int actual_rank = rank;
+    ompi_osc_rdma_module_t *module = GET_MODULE(win);
+
+    peer = ompi_osc_rdma_module_peer (module, actual_rank);
+    if (NULL == peer) {
+        return OMPI_ERR_BAD_PARAM;
+    }
+
+    /* currently only supported for allocated windows */
+    if (MPI_WIN_FLAVOR_ALLOCATE != module->flavor) {
+        return OMPI_ERR_NOT_SUPPORTED;
+    }
+
+    if (!ompi_osc_rdma_peer_local_base(peer)) {
+        return OMPI_ERR_NOT_SUPPORTED;
+    }
+
+    if (MPI_PROC_NULL == rank) {
+        /* iterate until we find a rank that has a non-zero size */
+        for (int i = 0 ; i < ompi_comm_size(module->comm) ; ++i) {
+            peer = ompi_osc_rdma_module_peer (module, i);
+            ompi_osc_rdma_peer_extended_t *ex_peer = (ompi_osc_rdma_peer_extended_t *) peer;
+            if (!ompi_osc_rdma_peer_local_base(peer)) {
+                continue;
+            } else if (module->same_size && ex_peer->super.base) {
+                break;
+            } else if (ex_peer->size > 0) {
+                break;
+            }
+        }
+    }
+
+    if (module->same_size && module->same_disp_unit) {
+        *size = module->size;
+        *disp_unit = module->disp_unit;
+        ompi_osc_rdma_peer_basic_t *ex_peer = (ompi_osc_rdma_peer_basic_t *) peer;
+        *((void**) baseptr) = (void *) (intptr_t)ex_peer->base;
+    } else {
+        ompi_osc_rdma_peer_extended_t *ex_peer = (ompi_osc_rdma_peer_extended_t *) peer;
+        if (ex_peer->super.base != 0) {
+            /* we know the base of the peer */
+            *((void**) baseptr) = (void *) (intptr_t)ex_peer->super.base;
+            *size = ex_peer->size;
+            *disp_unit = ex_peer->disp_unit;
+            return OMPI_SUCCESS;
+        }
+    }
+    return OMPI_ERR_NOT_SUPPORTED;
+}

ompi/mca/osc/ucx/osc_ucx_component.c

@@ -465,30 +465,70 @@ static const char* ompi_osc_ucx_set_no_lock_info(opal_infosubscriber_t *obj, con
     return module->no_locks ? "true" : "false";
 }
 
+static int ompi_osc_ucx_shared_query_peer(ompi_osc_ucx_module_t *module, int rank, size_t *size,
+    ptrdiff_t *disp_unit, void *baseptr) {
+
+    ucp_ep_h *dflt_ep;
+    ucp_ep_h ep; // ignored
+    OSC_UCX_GET_DEFAULT_EP(dflt_ep, module, peer); // TODO: needed?
+    ucs_status_t status;
+    opal_common_ucx_winfo_t *winfo; // ignored
+    rc = opal_common_ucx_tlocal_fetch(module->mem, peer, &ep, &rkey, &winfo, dflt_ep);
+    if (OMPI_SUCCESS != rc) {
+        return rc;
+    }
+    uint64_t raddr;
+    void *addr_p;
+    if (UCS_OK != ucp_rkey_ptr(rkey, module->addrs[peer], &addr_p)) {
+        return OMPI_ERR_NOT_AVAILABLE;
+    }
+    *size = module->sizes[i];
+    *((void**) baseptr) = (void *)module->shmem_addrs[i];
+    *disp_unit = module->disp_units[i];
+
+    return OMPI_SUCCESS;
+}
+
 int ompi_osc_ucx_shared_query(struct ompi_win_t *win, int rank, size_t *size,
         ptrdiff_t *disp_unit, void *baseptr)
 {
     ompi_osc_ucx_module_t *module =
         (ompi_osc_ucx_module_t*) win->w_osc_module;
 
+    *size = 0;
+    *((void**) baseptr) = NULL;
+    *disp_unit = 0;
+
     if (module->flavor != MPI_WIN_FLAVOR_SHARED) {
-        return MPI_ERR_WIN;
-    }
 
-    if (MPI_PROC_NULL != rank) {
+        if (MPI_PROC_NULL == rank) {
+            for (int i = 0 ; i < ompi_comm_size(module->comm) ; ++i) {
+                if (0 != module->sizes[i]) {
+                    if (OMPI_SUCCESS == ompi_osc_ucx_shared_query_peer(module, i, size, disp_unit, baseptr)) {
+                        return OMPI_SUCCESS;
+                    }
+                }
+            }
+        } else {
+            if (0 != module->sizes[i]) {
+                if (OMPI_SUCCESS == ompi_osc_ucx_shared_query_peer(module, i, size, disp_unit, baseptr)) {
+                    return OMPI_SUCCESS;
+                }
+            }
+        }
+        return OMPI_ERR_NOT_SUPPORTED;
+
+    } else if (MPI_PROC_NULL != rank) { // shared memory window with given rank
         *size = module->sizes[rank];
         *((void**) baseptr) = (void *)module->shmem_addrs[rank];
         if (module->disp_unit == -1) {
             *disp_unit = module->disp_units[rank];
         } else {
             *disp_unit = module->disp_unit;
         }
-    } else {
+    } else { // shared memory window with MPI_PROC_NULL
         int i = 0;
 
-        *size = 0;
-        *((void**) baseptr) = NULL;
-        *disp_unit = 0;
         for (i = 0 ; i < ompi_comm_size(module->comm) ; ++i) {
             if (0 != module->sizes[i]) {
                 *size = module->sizes[i];

ompi/mpi/c/win_shared_query.c

@@ -34,9 +34,9 @@ static const char FUNC_NAME[] = "MPI_Win_shared_query";
 
 int MPI_Win_shared_query(MPI_Win win, int rank, MPI_Aint *size, int *disp_unit, void *baseptr)
 {
-    int rc;
     size_t tsize;
     ptrdiff_t du;
+    int rc = OMPI_SUCCESS;
 
     if (MPI_PARAM_CHECK) {
         OMPI_ERR_INIT_FINALIZE(FUNC_NAME);
@@ -48,12 +48,23 @@ int MPI_Win_shared_query(MPI_Win win, int rank, MPI_Aint *size, int *disp_unit,
          }
     }
 
+    rc = OMPI_ERR_NOT_SUPPORTED;
+
     if (NULL != win->w_osc_module->osc_win_shared_query) {
         rc = win->w_osc_module->osc_win_shared_query(win, rank, &tsize, &du, baseptr);
-        *size = tsize;
-        *disp_unit = du;
-    } else {
-        rc = MPI_ERR_RMA_FLAVOR;
+        if (OMPI_SUCCESS == rc) {
+            *size = tsize;
+            *disp_unit = du;
+        }
+    }
+
+    if (OMPI_ERR_NOT_SUPPORTED == rc) {
+        /* gracefully bail out */
+        *size = 0;
+        *disp_unit = 0;
+        *(void**) baseptr = NULL;
+        rc = MPI_SUCCESS; // don't raise an error if the function is not supported
     }
+
     OMPI_ERRHANDLER_RETURN(rc, win, rc, FUNC_NAME);
 }