osc/ucx: Fix return values and exchange of disp_unit/sizes

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

Author: devreal

ompi/mca/osc/ucx/osc_ucx_component.c

@@ -479,11 +479,11 @@ static int ompi_osc_ucx_shared_query_peer(ompi_osc_ucx_module_t *module, int pee
     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;
+        return OMPI_ERR_NOT_SUPPORTED;
     }
     void *addr_p;
     if (UCS_OK != ucp_rkey_ptr(rkey, module->addrs[peer], &addr_p)) {
-        return OMPI_ERR_NOT_AVAILABLE;
+        return OMPI_ERR_NOT_SUPPORTED;
     }
     *size = ompi_osc_ucx_get_size(module, peer);
     *((void**) baseptr) = addr_p;
@@ -719,18 +719,21 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt
     if (same_disp_unit) { /* everyone has the same disp_unit, we do not need O(p) space */
         module->disp_unit = disp_unit;
         module->disp_units = NULL;
+    } else {
         values[val_count++] = disp_unit;
     }
 
     if (same_size) {
         module->same_size = true;
         module->sizes = NULL;
+    } else {
         values[val_count++] = size;
     }
 
     if (!same_disp_unit || !same_size) {
+        long* peer_values = malloc(comm_size * val_count * sizeof(long));
         ret = module->comm->c_coll->coll_allgather(values, val_count * sizeof(long), MPI_BYTE,
-                                                    (void *)my_info, sizeof(long) * val_count, MPI_BYTE,
+                                                    peer_values, sizeof(long) * val_count, MPI_BYTE,
                                                     module->comm,
                                                     module->comm->c_coll->coll_allgather_module);
         if (OMPI_SUCCESS != ret) {
@@ -745,7 +748,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt
                 goto error;
             }
             for (i = 0; i < comm_size; i++) {
-                module->disp_units[i] = (ptrdiff_t)values[i*val_count];
+                module->disp_units[i] = (ptrdiff_t)peer_values[i*val_count];
             }
         }
 
@@ -758,9 +761,10 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt
             }
 
             for (i = 0; i < comm_size; i++) {
-                module->sizes[i] = (size_t)values[i*val_count + val_count-1];
+                module->sizes[i] = (size_t)peer_values[(i+1)*val_count - 1];
             }
         }
+        free(peer_values);
     }
 
     ret = opal_common_ucx_wpctx_create(mca_osc_ucx_component.wpool, comm_size,