From c5cae38bb87c85ca1386e3c6f52cd1b03959ef37 Mon Sep 17 00:00:00 2001
From: Rahul Vaidya <ravaidya@amd.com>
Date: Fri, 14 Jun 2024 11:46:08 -0500
Subject: [PATCH] Fix --root all issue. (#83)

Signed-off-by: rahulvaidya20 <ravaidya@amd.com>
---
 src/common.cu | 4 ++--
 src/common.h  | 7 +++++++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/common.cu b/src/common.cu
index f7ab2f4372..02c03bace5 100644
--- a/src/common.cu
+++ b/src/common.cu
@@ -1019,7 +1019,7 @@ int main(int argc, char* argv[]) {
         nccltype = ncclstringtotype(optarg);
         break;
       case 'r':
-        ncclroot = strtol(optarg, NULL, 0);
+        ncclroot = ncclstringtoroot(optarg);
         break;
       case 'z':
         blocking_coll = strtol(optarg, NULL, 0);
@@ -1099,7 +1099,7 @@ int main(int argc, char* argv[]) {
             "[-o,--op <sum/prod/min/max/all>] \n\t"
 #endif
             "[-d,--datatype <nccltype/all>] \n\t"
-            "[-r,--root <root>] \n\t"
+            "[-r,--root <root/all>] \n\t"
             "[-z,--blocking <0/1>] \n\t"
             "[-Y,--memory_type <coarse/fine/host/managed>] \n\t"
             "[-s,--stress_cycles <number of cycles>] \n\t"
diff --git a/src/common.h b/src/common.h
index e14648dc97..e5c8b0fbe0 100644
--- a/src/common.h
+++ b/src/common.h
@@ -294,6 +294,13 @@ static int ncclstringtoop (char *str) {
     return ncclSum;
 }
 
+static int ncclstringtoroot (char *str) {
+    if (strcmp(str, "all") == 0) {
+      return -1;
+    }
+    return strtol(str, NULL, 0);
+}
+
 static int ncclstringtomtype (char *str) {
     for (int o=0; o<nccl_NUM_MTYPES; o++) {
       if (strcmp(str, test_memorytypes[o]) == 0) {