TL/MLX5: a2a part 2 -- full coll

openucx · Jun 29, 2023 · cc68112 · cc68112
1 parent 22dc12d
commit cc68112
Show file tree

Hide file tree

Showing 9 changed files with 1,176 additions and 72 deletions.
diff --git a/src/components/tl/mlx5/Makefile.am b/src/components/tl/mlx5/Makefile.am
@@ -4,11 +4,13 @@
 
 if TL_MLX5_ENABLED
 
-alltoall =                    \
-	alltoall/alltoall.h       \
-	alltoall/alltoall.c       \
-	alltoall/alltoall_mkeys.h \
-	alltoall/alltoall_mkeys.c
+alltoall =                      \
+	alltoall/alltoall.h         \
+	alltoall/alltoall.c         \
+	alltoall/alltoall_mkeys.h   \
+	alltoall/alltoall_mkeys.c   \
+	alltoall/alltoall_inline.h  \
+	alltoall/alltoall_coll.c
 
 mcast =                           \
 	mcast/tl_mlx5_mcast_context.c \

diff --git a/src/components/tl/mlx5/alltoall/alltoall.c b/src/components/tl/mlx5/alltoall/alltoall.c
@@ -65,7 +65,7 @@ static ucc_status_t build_rank_map(ucc_tl_mlx5_alltoall_t *a2a,
     return UCC_OK;
 }
 
-ucc_status_t ucc_tl_mlx5_alltoall_init_start(ucc_tl_mlx5_team_t *team)
+ucc_status_t ucc_tl_mlx5_team_alltoall_init_start(ucc_tl_mlx5_team_t *team)
 {
     ucc_tl_mlx5_context_t  *ctx = UCC_TL_MLX5_TEAM_CTX(team);
     ucc_tl_mlx5_alltoall_t *a2a = NULL;
@@ -270,7 +270,8 @@ static ucc_status_t ucc_tl_mlx5_alltoall_barrier_alloc(ucc_tl_mlx5_team_t *team)
     return UCC_OK;
 }
 
-ucc_status_t ucc_tl_mlx5_alltoall_init_progress(ucc_tl_mlx5_team_t *tl_team)
+ucc_status_t
+ucc_tl_mlx5_team_alltoall_init_progress(ucc_tl_mlx5_team_t *tl_team)
 {
     ucc_tl_mlx5_team_t        *team        = ucc_derived_of(tl_team,
                                                            ucc_tl_mlx5_team_t);

diff --git a/src/components/tl/mlx5/alltoall/alltoall.h b/src/components/tl/mlx5/alltoall/alltoall.h
@@ -138,8 +138,11 @@ typedef struct ucc_tl_mlx5_alltoall {
     ucc_tl_mlx5_a2a_bcast_data_t bcast_data;
 } ucc_tl_mlx5_alltoall_t;
 
-ucc_status_t ucc_tl_mlx5_alltoall_init_start(ucc_tl_mlx5_team_t *team);
-ucc_status_t ucc_tl_mlx5_alltoall_init_progress(ucc_tl_mlx5_team_t *team);
+ucc_status_t ucc_tl_mlx5_team_alltoall_init_start(ucc_tl_mlx5_team_t *team);
+ucc_status_t ucc_tl_mlx5_team_alltoall_init_progress(ucc_tl_mlx5_team_t *team);
+ucc_status_t ucc_tl_mlx5_alltoall_init(ucc_base_coll_args_t *coll_args,
+                                       ucc_base_team_t *     team,
+                                       ucc_coll_task_t **    task_h);
 void         ucc_tl_mlx5_alltoall_cleanup(ucc_tl_mlx5_team_t *team);
 
 static inline ucc_tl_mlx5_alltoall_ctrl_t*