SQUASHME: osc/ucx needs to exchange shared memory information

Signed-off-by: Joseph Schuchart <[email protected]>

Author: devreal

ompi/mca/osc/ucx/osc_ucx.h

@@ -116,6 +116,19 @@ typedef struct ompi_osc_ucx_mem_ranges {
     uint64_t tail;
 } ompi_osc_ucx_mem_ranges_t;
 
+/**
+ * Structure to hold information about shared memory regions.
+ * We store the rank, it's address, and the size of the window region.
+ * We don't store the disp_unit here, as that is stored elsewhere already.
+ */
+struct ompi_osc_ucx_shmem_info_s {
+    int peer; /* rank of the peer this information belongs to */
+    char *addr; /* address of the shared memory region */
+    size_t size; /* size of the shared memory region */
+};
+
+typedef struct ompi_osc_ucx_shmem_info_s ompi_osc_ucx_shmem_info_t;
+
 typedef struct ompi_osc_ucx_module {
     ompi_osc_base_module_t super;
     struct ompi_communicator_t *comm;
@@ -128,6 +141,7 @@ typedef struct ompi_osc_ucx_module {
                           * disp unit size; if disp_unit == -1, then we
                           * need to look at disp_units */
     ptrdiff_t *disp_units;
+    ompi_osc_ucx_shmem_info_t *shmem_info; /* shared memory info */
 
     ompi_osc_ucx_state_t state; /* remote accessible flags */
     ompi_osc_local_dynamic_win_info_t local_dynamic_win_info[OMPI_OSC_UCX_ATTACH_MAX];

ompi/mca/osc/ucx/osc_ucx_component.c

@@ -679,6 +679,13 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt
     module->acc_single_intrinsic = check_config_value_bool ("acc_single_intrinsic", info);
     module->skip_sync_check = false;
 
+    /**
+     * TODO: we need to collect the shared memory information from all processes
+     *       on the same node. This includes the size and base address, which needs
+     *       to be passed to ucp_rkey_ptr().
+     */
+    module->shmem_info = NULL;
+
     /* share everyone's displacement units. Only do an allgather if
        strictly necessary, since it requires O(p) state. */
     values[0] = disp_unit;
@@ -847,6 +854,18 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt
 
         module->size = module->sizes[ompi_comm_rank(module->comm)];
         *base = (void *)module->shmem_addrs[ompi_comm_rank(module->comm)];
+    } else {
+        /* non-shared memory: exchange sizes and addresses so they can be queried for shared memory */
+        for (i = 0; i < comm_size; i++) {
+            ompi_proc_t *peer = ompi_comm_peer_lookup(module->comm, i);
+            peer->
+            if (ompi_comm_peer_lookup(module->comm, i) == NULL) {
+                OSC_UCX_ERROR("Failed to lookup peer %d in communicator %s", i, ompi_comm_print_cid(module->comm));
+                ret = OMPI_ERR_COMM_FAILURE;
+                goto error;
+            }
+        }
+
     }
 
     void **mem_base = base;