From 9dbecd84c478ca9765b04f95760f54603609301e Mon Sep 17 00:00:00 2001
From: Devendar Bureddy <devendar@nvidia.com>
Date: Fri, 22 Nov 2024 01:28:22 +0200
Subject: [PATCH] TL/SHARP: Option to enable SHARP multi-channel

---
 src/components/tl/sharp/tl_sharp.c         | 5 +++++
 src/components/tl/sharp/tl_sharp.h         | 1 +
 src/components/tl/sharp/tl_sharp_context.c | 6 +++++-
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/components/tl/sharp/tl_sharp.c b/src/components/tl/sharp/tl_sharp.c
index 464ef50478..511dac63b0 100644
--- 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
index adfbc86036..4570c01eeb 100644
--- 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
index 42d10f8d87..150c684090 100644
--- 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;