TL/SHARP: Option to enable SHARP multi-channel

openucx · Nov 21, 2024 · 9dbecd8 · 9dbecd8
1 parent 7f3c2f6
commit 9dbecd8
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/src/components/tl/sharp/tl_sharp.c b/src/components/tl/sharp/tl_sharp.c
@@ -72,6 +72,11 @@ static ucc_config_field_t ucc_tl_sharp_context_config_table[] = {
      ucc_offsetof(ucc_tl_sharp_context_config_t, team_max_ppn),
      UCC_CONFIG_TYPE_UINT},
 
+    {"MULTI_CHANNEL", "0",
+     "Use SHARP Multi-channel feature. Options: 0-disable 1-enable",
+     ucc_offsetof(ucc_tl_sharp_context_config_t, multi_channel),
+     UCC_CONFIG_TYPE_UINT},
+
     {NULL}};
 
 UCC_CLASS_DEFINE_NEW_FUNC(ucc_tl_sharp_lib_t, ucc_base_lib_t,

diff --git a/src/components/tl/sharp/tl_sharp.h b/src/components/tl/sharp/tl_sharp.h
@@ -53,6 +53,7 @@ typedef struct ucc_tl_sharp_context_config {
     int                      context_per_team;
     int                      enable_lazy_group_alloc;
     int                      team_max_ppn;
+    int                      multi_channel;
 } ucc_tl_sharp_context_config_t;
 
 typedef struct ucc_tl_sharp_lib {

diff --git a/src/components/tl/sharp/tl_sharp_context.c b/src/components/tl/sharp/tl_sharp_context.c
@@ -305,7 +305,11 @@ ucc_status_t ucc_tl_sharp_context_init(ucc_tl_sharp_context_t *sharp_ctx,
 
     init_spec.progress_func                  = NULL;
     init_spec.world_local_rank               = local_rank;
-    init_spec.group_channel_idx              = 0;
+    if (sharp_ctx->cfg.multi_channel) {
+        init_spec.group_channel_idx          = local_rank;
+    } else {
+        init_spec.group_channel_idx          = 0;
+    }
     init_spec.oob_ctx                        = oob_ctx;
     init_spec.config                         = sharp_coll_default_config;
     init_spec.config.user_progress_num_polls = sharp_ctx->cfg.uprogress_num_polls;