Skip to content

Commit da6cb6b

Browse files
authored
Merge pull request #11442 from wzamazon/coll_han_fix_t3_bcast_count
coll/han: fix count of allreduce t3 task's bcast operations
2 parents bd8f772 + 006c4e9 commit da6cb6b

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

ompi/mca/coll/han/coll_han_allreduce.c

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -354,7 +354,7 @@ int mca_coll_han_allreduce_t3_task(void *task_args)
354354
if (t->cur_seg == t->num_segments - 2 && t->last_seg_count != t->seg_count) {
355355
tmp_count = t->last_seg_count;
356356
}
357-
t->up_comm->c_coll->coll_ibcast((char *) t->rbuf + extent * t->seg_count, t->seg_count,
357+
t->up_comm->c_coll->coll_ibcast((char *) t->rbuf + extent * t->seg_count, tmp_count,
358358
t->dtype, t->root_up_rank, t->up_comm, &(reqs[0]),
359359
t->up_comm->c_coll->coll_ibcast_module);
360360
req_count++;
@@ -391,7 +391,13 @@ int mca_coll_han_allreduce_t3_task(void *task_args)
391391
t->low_comm->c_coll->coll_reduce_module);
392392
}
393393
/* lb of cur_seg */
394-
t->low_comm->c_coll->coll_bcast((char *) t->rbuf, t->seg_count, t->dtype, t->root_low_rank,
394+
if (t->cur_seg == t->num_segments - 1 && t->last_seg_count != t->seg_count) {
395+
tmp_count = t->last_seg_count;
396+
} else {
397+
tmp_count = t->seg_count;
398+
}
399+
400+
t->low_comm->c_coll->coll_bcast((char *) t->rbuf, tmp_count, t->dtype, t->root_low_rank,
395401
t->low_comm, t->low_comm->c_coll->coll_bcast_module);
396402
if (!t->noop && req_count > 0) {
397403
ompi_request_wait_all(req_count, reqs, MPI_STATUSES_IGNORE);

0 commit comments

Comments
 (0)