enable UT to test with channels greater than 64

2024-06-12 12:04:58 -05:00
parent 955347bab4
commit 392a73fdef
4 changed files with 93 additions and 0 deletions
@@ -85,4 +85,28 @@ namespace RcclUnitTesting
                           inPlaceList, managedMemList, useHipGraphList);
    testBed.Finalize();
  }
+
+    TEST(AllToAll, Channels)
+  {
+    TestBed testBed;
+    if(testBed.ev.isGfx94) {
+      // Configuration
+      std::vector<ncclFunc_t>     const funcTypes       = {ncclCollAllToAll};
+      std::vector<ncclDataType_t> const dataTypes       = {ncclBfloat16, ncclHalf};
+      std::vector<ncclRedOp_t>    const redOps          = {ncclSum};
+      std::vector<int>            const roots           = {0};
+      std::vector<int>            const numElements     = {64 * 1024 * 1024, 1024};
+      std::vector<bool>           const inPlaceList     = {false};
+      std::vector<bool>           const managedMemList  = {false};
+      std::vector<bool>           const useHipGraphList = {false, true};
+      std::vector<char *>         const channelList     = {"56", "84", "112"}; 
+      for (auto channel : channelList) {
+        setenv("NCCL_MIN_NCHANNELS", channel, 1);
+        testBed.RunSimpleSweep(funcTypes, dataTypes, redOps, roots, numElements,
+                              inPlaceList, managedMemList, useHipGraphList);
+        testBed.Finalize();
+        unsetenv("NCCL_MIN_NCHANNELS");
+      }
+    }
+  }
 }