TL/CUDA: fix step check

openucx · Jul 4, 2024 · 00f3922 · 00f3922
1 parent 9333f41
commit 00f3922
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/components/tl/cuda/bcast/bcast_linear.c b/src/components/tl/cuda/bcast/bcast_linear.c
@@ -195,7 +195,7 @@ void ucc_tl_cuda_bcast_linear_progress(ucc_coll_task_t *coll_task)
             }
             task->bcast_linear.stage = STAGE_COPY;
             // ucc_info("all others ready for next step");
-            if (task->bcast_linear.stage < task->bcast_linear.num_steps) {
+            if (task->bcast_linear.step < task->bcast_linear.num_steps) {
                 task->bcast_linear.stage = STAGE_COPY;
             } else {
                 task->bcast_linear.stage = STAGE_DONE;
@@ -241,9 +241,9 @@ void ucc_tl_cuda_bcast_linear_progress(ucc_coll_task_t *coll_task)
                     set_rank_step(task, trank, task->bcast_linear.step, 0);
                     // task->bcast_linear.stage =
                     //     STAGE_DONE; // TODO: just for debug
-                    if (task->bcast_linear.stage <
+                    if (task->bcast_linear.step <
                         task->bcast_linear.num_steps) {
-                        task->bcast_linear.stage = STAGE_COPY;
+                        task->bcast_linear.stage = STAGE_WAIT_ROOT;
                     } else {
                         task->bcast_linear.stage = STAGE_DONE;
                     }