int ompi_osc_ucx_put(const void *origin_addr, int origin_count, struct ompi_datatype_t *origin_dt,
int target, ptrdiff_t target_disp, int target_count,
struct ompi_datatype_t *target_dt, struct ompi_win_t *win) {
ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module;
ucp_ep_h ep = OSC_UCX_GET_EP(module->comm, target);
uint64_t remote_addr = (module->win_info_array[target]).addr + target_disp * OSC_UCX_GET_DISP(module, target);
ucp_rkey_h rkey;
bool is_origin_contig = false, is_target_contig = false;
ptrdiff_t origin_lb, origin_extent, target_lb, target_extent;
ucs_status_t status;
int ret = OMPI_SUCCESS;
ret = check_sync_state(module, target, false);
if (ret != OMPI_SUCCESS) {
return ret;
}
if (module->flavor == MPI_WIN_FLAVOR_DYNAMIC) {
status = get_dynamic_win_info(remote_addr, module, ep, target);
if (status != UCS_OK) {
return OMPI_ERROR;
}
}
rkey = (module->win_info_array[target]).rkey;
ompi_datatype_get_true_extent(origin_dt, &origin_lb, &origin_extent);
ompi_datatype_get_true_extent(target_dt, &target_lb, &target_extent);
is_origin_contig = ompi_datatype_is_contiguous_memory_layout(origin_dt, origin_count);
is_target_contig = ompi_datatype_is_contiguous_memory_layout(target_dt, target_count);
if (is_origin_contig && is_target_contig) {
/* fast path */
size_t origin_len;
ompi_datatype_type_size(origin_dt, &origin_len);
origin_len *= origin_count;
status = ucp_put_nbi(ep, (void *)((intptr_t)origin_addr + origin_lb), origin_len,
remote_addr + target_lb, rkey);
if (status != UCS_OK && status != UCS_INPROGRESS) {
opal_output_verbose(1, ompi_osc_base_framework.framework_output,
"%s:%d: ucp_put_nbi failed: %d\n",
__FILE__, __LINE__, status);
return OMPI_ERROR;
}
return incr_and_check_ops_num(module, target, ep);
} else {
return ddt_put_get(module, origin_addr, origin_count, origin_dt, is_origin_contig,
origin_lb, target, ep, remote_addr, rkey, target_count, target_dt,
is_target_contig, target_lb, false);
}
}
int mca_spml_ucx_put_nb(void* dst_addr, size_t size, void* src_addr, int dst, void **handle)
{
void *rva;
ucs_status_t status;
spml_ucx_mkey_t *ucx_mkey;
ucx_mkey = mca_spml_ucx_get_mkey(dst, dst_addr, &rva);
status = ucp_put_nbi(mca_spml_ucx.ucp_peers[dst].ucp_conn, src_addr, size,
(uint64_t)rva, ucx_mkey->rkey);
return ucx_status_to_oshmem(status);
}
int mca_spml_ucx_put_nb(shmem_ctx_t ctx, void* dst_addr, size_t size, void* src_addr, int dst, void **handle)
{
void *rva;
ucs_status_t status;
spml_ucx_mkey_t *ucx_mkey;
mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
ucx_mkey = mca_spml_ucx_get_mkey(ctx, dst, dst_addr, &rva, &mca_spml_ucx);
status = ucp_put_nbi(ucx_ctx->ucp_peers[dst].ucp_conn, src_addr, size,
(uint64_t)rva, ucx_mkey->rkey);
return ucx_status_to_oshmem_nb(status);
}
static inline int ddt_put_get(ompi_osc_ucx_module_t *module,
const void *origin_addr, int origin_count,
struct ompi_datatype_t *origin_dt,
bool is_origin_contig, ptrdiff_t origin_lb,
int target, ucp_ep_h ep, uint64_t remote_addr, ucp_rkey_h rkey,
int target_count, struct ompi_datatype_t *target_dt,
bool is_target_contig, ptrdiff_t target_lb, bool is_get) {
ucx_iovec_t *origin_ucx_iov = NULL, *target_ucx_iov = NULL;
uint32_t origin_ucx_iov_count = 0, target_ucx_iov_count = 0;
uint32_t origin_ucx_iov_idx = 0, target_ucx_iov_idx = 0;
ucs_status_t status;
int ret = OMPI_SUCCESS;
if (!is_origin_contig) {
ret = create_iov_list(origin_addr, origin_count, origin_dt,
&origin_ucx_iov, &origin_ucx_iov_count);
if (ret != OMPI_SUCCESS) {
return ret;
}
}
if (!is_target_contig) {
ret = create_iov_list(NULL, target_count, target_dt,
&target_ucx_iov, &target_ucx_iov_count);
if (ret != OMPI_SUCCESS) {
return ret;
}
}
if (!is_origin_contig && !is_target_contig) {
size_t curr_len = 0;
while (origin_ucx_iov_idx < origin_ucx_iov_count) {
curr_len = MIN(origin_ucx_iov[origin_ucx_iov_idx].len,
target_ucx_iov[target_ucx_iov_idx].len);
if (!is_get) {
status = ucp_put_nbi(ep, origin_ucx_iov[origin_ucx_iov_idx].addr, curr_len,
remote_addr + (uint64_t)(target_ucx_iov[target_ucx_iov_idx].addr), rkey);
if (status != UCS_OK && status != UCS_INPROGRESS) {
opal_output_verbose(1, ompi_osc_base_framework.framework_output,
"%s:%d: ucp_put_nbi failed: %d\n",
__FILE__, __LINE__, status);
return OMPI_ERROR;
}
} else {
status = ucp_get_nbi(ep, origin_ucx_iov[origin_ucx_iov_idx].addr, curr_len,
remote_addr + (uint64_t)(target_ucx_iov[target_ucx_iov_idx].addr), rkey);
if (status != UCS_OK && status != UCS_INPROGRESS) {
opal_output_verbose(1, ompi_osc_base_framework.framework_output,
"%s:%d: ucp_get_nbi failed: %d\n",
__FILE__, __LINE__, status);
return OMPI_ERROR;
}
}
ret = incr_and_check_ops_num(module, target, ep);
if (ret != OMPI_SUCCESS) {
return ret;
}
origin_ucx_iov[origin_ucx_iov_idx].addr = (void *)((intptr_t)origin_ucx_iov[origin_ucx_iov_idx].addr + curr_len);
target_ucx_iov[target_ucx_iov_idx].addr = (void *)((intptr_t)target_ucx_iov[target_ucx_iov_idx].addr + curr_len);
origin_ucx_iov[origin_ucx_iov_idx].len -= curr_len;
if (origin_ucx_iov[origin_ucx_iov_idx].len == 0) {
origin_ucx_iov_idx++;
}
target_ucx_iov[target_ucx_iov_idx].len -= curr_len;
if (target_ucx_iov[target_ucx_iov_idx].len == 0) {
target_ucx_iov_idx++;
}
}
assert(origin_ucx_iov_idx == origin_ucx_iov_count &&
target_ucx_iov_idx == target_ucx_iov_count);
} else if (!is_origin_contig) {
size_t prev_len = 0;
while (origin_ucx_iov_idx < origin_ucx_iov_count) {
if (!is_get) {
status = ucp_put_nbi(ep, origin_ucx_iov[origin_ucx_iov_idx].addr,
origin_ucx_iov[origin_ucx_iov_idx].len,
remote_addr + target_lb + prev_len, rkey);
if (status != UCS_OK && status != UCS_INPROGRESS) {
opal_output_verbose(1, ompi_osc_base_framework.framework_output,
"%s:%d: ucp_put_nbi failed: %d\n",
__FILE__, __LINE__, status);
return OMPI_ERROR;
}
} else {
status = ucp_get_nbi(ep, origin_ucx_iov[origin_ucx_iov_idx].addr,
origin_ucx_iov[origin_ucx_iov_idx].len,
remote_addr + target_lb + prev_len, rkey);
if (status != UCS_OK && status != UCS_INPROGRESS) {
opal_output_verbose(1, ompi_osc_base_framework.framework_output,
"%s:%d: ucp_get_nbi failed: %d\n",
__FILE__, __LINE__, status);
return OMPI_ERROR;
}
}
ret = incr_and_check_ops_num(module, target, ep);
if (ret != OMPI_SUCCESS) {
return ret;
}
prev_len += origin_ucx_iov[origin_ucx_iov_idx].len;
origin_ucx_iov_idx++;
}
} else {
size_t prev_len = 0;
while (target_ucx_iov_idx < target_ucx_iov_count) {
if (!is_get) {
status = ucp_put_nbi(ep, (void *)((intptr_t)origin_addr + origin_lb + prev_len),
target_ucx_iov[target_ucx_iov_idx].len,
remote_addr + (uint64_t)(target_ucx_iov[target_ucx_iov_idx].addr), rkey);
if (status != UCS_OK && status != UCS_INPROGRESS) {
opal_output_verbose(1, ompi_osc_base_framework.framework_output,
"%s:%d: ucp_put_nbi failed: %d\n",
__FILE__, __LINE__, status);
return OMPI_ERROR;
}
} else {
status = ucp_get_nbi(ep, (void *)((intptr_t)origin_addr + origin_lb + prev_len),
target_ucx_iov[target_ucx_iov_idx].len,
remote_addr + (uint64_t)(target_ucx_iov[target_ucx_iov_idx].addr), rkey);
if (status != UCS_OK && status != UCS_INPROGRESS) {
opal_output_verbose(1, ompi_osc_base_framework.framework_output,
"%s:%d: ucp_get_nbi failed: %d\n",
__FILE__, __LINE__, status);
return OMPI_ERROR;
}
}
ret = incr_and_check_ops_num(module, target, ep);
if (ret != OMPI_SUCCESS) {
return ret;
}
prev_len += target_ucx_iov[target_ucx_iov_idx].len;
target_ucx_iov_idx++;
}
}
if (origin_ucx_iov != NULL) {
free(origin_ucx_iov);
}
if (target_ucx_iov != NULL) {
free(target_ucx_iov);
}
return ret;
}
