TL/UCP: Update to use macros

openucx · Aug 20, 2024 · 6e83522 · 6e83522
1 parent 17c8cfc
commit 6e83522
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/src/components/tl/ucp/allgather/allgather_linear_xgvmi.c b/src/components/tl/ucp/allgather/allgather_linear_xgvmi.c
@@ -13,16 +13,15 @@ ucc_tl_ucp_dpu_allgather_linear_xgvmi_rdma_progress(ucc_coll_task_t *coll_task)
 {
     ucc_tl_ucp_task_t    *task           = ucc_derived_of(coll_task,
                                                           ucc_tl_ucp_task_t);
-    ucc_rank_t            size           = (ucc_rank_t)task->subset.map.ep_num;
     ucc_datatype_t        dtype          = TASK_ARGS(task).src.info.datatype;
     size_t                dt_size        = ucc_dt_size(dtype);
     ucc_count_t           count          = coll_task->bargs.args.src.info.count;
-    ucc_rank_t            host_team_size = size;
     ucc_base_team_t      *base_team      = coll_task->team;
     ucc_tl_ucp_team_t    *tl_team        = ucc_derived_of(base_team,
                                                           ucc_tl_ucp_team_t);
     ucc_coll_task_t      *allgather_task = task->dpu_xgvmi.allgather_task;
     ucc_tl_ucp_context_t *tl_ctx         = UCC_TL_UCP_TEAM_CTX(tl_team);
+    ucc_rank_t            host_team_size = UCC_TL_TEAM_SIZE(tl_team);
     ucp_request_param_t   req_param      = {0};
     int                   i              = 0;
     ucc_rank_t            rank           = UCC_TL_TEAM_RANK(tl_team);

diff --git a/src/components/tl/ucp/alltoall/alltoall_linear_xgvmi.c b/src/components/tl/ucp/alltoall/alltoall_linear_xgvmi.c
@@ -13,14 +13,13 @@ ucc_tl_ucp_dpu_alltoall_linear_xgvmi_rdma_progress(ucc_coll_task_t *coll_task)
 {
     ucc_tl_ucp_task_t    *task           = ucc_derived_of(coll_task,
                                                           ucc_tl_ucp_task_t);
-    ucc_rank_t            size           = (ucc_rank_t)task->subset.map.ep_num;
     ucc_datatype_t        dtype          = TASK_ARGS(task).src.info.datatype;
     size_t                dt_size        = ucc_dt_size(dtype);
     ucc_count_t           count          = coll_task->bargs.args.src.info.count;
-    ucc_rank_t            host_team_size = size;
     ucc_base_team_t      *base_team      = coll_task->team;
     ucc_tl_ucp_team_t    *tl_team        = ucc_derived_of(base_team,
                                                           ucc_tl_ucp_team_t);
+    ucc_rank_t            host_team_size = UCC_TL_TEAM_SIZE(tl_team);
     ucc_coll_task_t      *allgather_task = task->dpu_xgvmi.allgather_task;
     ucc_tl_ucp_context_t *tl_ctx         = UCC_TL_UCP_TEAM_CTX(tl_team);
     ucp_request_param_t   req_param      = {0};