CORE: fix n deps overflow in pipelined schedule (openucx#1051)

* CORE: fix n deps overflow in pipelined schedule * REVIEW: fix review comments
andylin-hao · Dec 26, 2024 · d754a7e · d754a7e
1 parent aab3d69
commit d754a7e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/src/schedule/ucc_schedule.h b/src/schedule/ucc_schedule.h
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -107,9 +107,9 @@ typedef struct ucc_coll_task {
         /* used for lf mt progress queue */
         ucc_lf_queue_elem_t            lf_elem;
     };
-    uint8_t                            n_deps;
-    uint8_t                            n_deps_satisfied;
-    uint8_t                            n_deps_base;
+    uint32_t                           n_deps;
+    uint32_t                           n_deps_satisfied;
+    uint32_t                           n_deps_base;
     /* timestamp of the start time: either post or triggered_post */
     double                             start_time;
     uint32_t                           seq_num;

diff --git a/src/schedule/ucc_schedule_pipelined.c b/src/schedule/ucc_schedule_pipelined.c
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -254,9 +254,10 @@ ucc_status_t ucc_dependency_handler(ucc_coll_task_t *parent,
                                     ucc_coll_task_t *task)
 {
     ucc_status_t status;
-    uint8_t      n_deps_satisfied;
+    uint32_t     n_deps_satisfied;
 
-    n_deps_satisfied = ucc_atomic_fadd8(&task->n_deps_satisfied, 1);
+    n_deps_satisfied = ucc_atomic_fadd32(&task->n_deps_satisfied, 1);
+    ucc_assert(task->n_deps_satisfied > n_deps_satisfied);
 
     ucc_trace_req("task %p, n_deps %d, satisfied %d", task, task->n_deps,
                   n_deps_satisfied);