TL/CUDA: moved get/set rank step

openucx · Aug 2, 2024 · f6d7536 · f6d7536
1 parent e43c6b5
commit f6d7536
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 65 deletions.
diff --git a/src/components/tl/cuda/allgatherv/allgatherv_linear.c b/src/components/tl/cuda/allgatherv/allgatherv_linear.c
@@ -55,22 +55,6 @@ enum
  * other ranks to finish */
 };
 
-static inline int get_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- return sync->seq_num[step_id];
-}
-
-static inline void set_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step, int step_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- sync->seq_num[step_id] = step;
-}
-
 ucc_status_t ucc_tl_cuda_allgatherv_linear_finalize(ucc_coll_task_t *coll_task)
 {
  ucc_tl_cuda_task_t *task = ucc_derived_of(coll_task, ucc_tl_cuda_task_t);

diff --git a/src/components/tl/cuda/bcast/bcast_linear.c b/src/components/tl/cuda/bcast/bcast_linear.c
@@ -21,23 +21,6 @@ enum
  STAGE_CLIENT_COPY_WAIT,
 };
 
-// TODO: move out to common with allgather
-static inline int get_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- return sync->seq_num[step_id];
-}
-
-static inline void set_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step, int step_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- sync->seq_num[step_id] = step;
-}
-
 ucc_status_t ucc_tl_cuda_bcast_linear_setup_start(ucc_tl_cuda_task_t *task)
 {
  ucc_tl_cuda_team_t *team = TASK_TEAM(task);

diff --git a/src/components/tl/cuda/reduce_scatterv/reduce_scatterv_linear.c b/src/components/tl/cuda/reduce_scatterv/reduce_scatterv_linear.c
@@ -59,22 +59,6 @@ enum
  * other ranks to finish */
 };
 
-static inline int get_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- return sync->seq_num[step_id];
-}
-
-static inline void set_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step, int step_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- sync->seq_num[step_id] = step;
-}
-
 ucc_status_t
 ucc_tl_cuda_reduce_scatterv_linear_finalize(ucc_coll_task_t *coll_task)
 {

diff --git a/src/components/tl/cuda/tl_cuda_coll.h b/src/components/tl/cuda/tl_cuda_coll.h
@@ -156,4 +156,21 @@ ucc_status_t ucc_tl_cuda_alg_id_to_init(int alg_id, const char *alg_id_str,
  ucc_memory_type_t mem_type,
  ucc_base_coll_init_fn_t *init);
 
+// common utils function for collectives:
+static inline int get_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
+ int step_id)
+{
+ ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
+
+ return sync->seq_num[step_id];
+}
+
+static inline void set_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
+ int step, int step_id)
+{
+ ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
+
+ sync->seq_num[step_id] = step;
+}
+
 #endif
diff --git a/src/components/tl/cuda/tl_cuda_ring.h b/src/components/tl/cuda/tl_cuda_ring.h
@@ -83,20 +83,4 @@ static inline ucc_rank_t get_recv_block(ucc_tl_cuda_team_t *team,
  return ring->ring[(ring->iring[trank] + tsize - step - 1) % tsize];
 }
 
-static inline int get_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int ring_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- return sync->seq_num[ring_id];
-}
-
-static inline void set_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step, int ring_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- sync->seq_num[ring_id] = step;
-}
-
 #endif