From 62b9f0d3a72595a442d8fd3fe8fa502d704ab095 Mon Sep 17 00:00:00 2001
From: gilbertlee-amd <44450918+gilbertlee-amd@users.noreply.github.com>
Date: Thu, 4 Apr 2024 15:11:47 -0600
Subject: [PATCH] [topo_expl] Adding -n option to override number of nodes
 (#1134)

[ROCm/rccl commit: 93982533d7506dc6a406f2b7b7e94152d5ecc277]
---
 projects/rccl/tools/topo_expl/topo_expl.cpp | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/projects/rccl/tools/topo_expl/topo_expl.cpp b/projects/rccl/tools/topo_expl/topo_expl.cpp
index f6963f7830..2d8d18312b 100644
--- a/projects/rccl/tools/topo_expl/topo_expl.cpp
+++ b/projects/rccl/tools/topo_expl/topo_expl.cpp
@@ -171,7 +171,7 @@ int main(int argc,char* argv[])
   int maxCTAsEnv;
 
   if (!cmdOptionExists(argv, argv + argc, "-m")) {
-    printf("Usage: ./topo_expl -m model_id\n");
+    printf("Usage: ./topo_expl -m model_id [-n numNodes=1]\n");
     printf("List of model_id:\n");
     for (int i = 0; i < num_models; i++)
       printf("  %d: %s\n", i, model_descs[i].description);
@@ -194,7 +194,13 @@ int main(int argc,char* argv[])
   initCollNet();
 
   NodeModelDesc *desc = &model_descs[model_id];
-  for (int i=0; i<desc->num_nodes; i++) {
+  int numNodes = desc->num_nodes;
+  if (cmdOptionExists(argv, argv + argc, "-n")) {
+    char *numNodesStr = getCmdOption(argv, argv + argc, "-n");
+    if (numNodesStr)
+      numNodes = atol(numNodesStr);
+  }
+  for (int i=0; i < numNodes; i++) {
       node = new NodeModel(desc->filename);
       network.AddNode(node);
   }