extending the unit-tests for multi-rank support

2022-03-15 18:53:52 -04:00
parent 0336ffdf70
commit a87d61db2b
6 changed files with 90 additions and 36 deletions
@@ -94,6 +94,9 @@ namespace RcclUnitTesting
    std::vector<ncclDataType_t> const& GetAllSupportedDataTypes();

    // Helper function that splits up GPUs to the given number of processes
+    static std::vector<std::vector<int>> GetDeviceIdsList(int const numProcesses,
+                                                          int const numGpus,
+							  int const ranksPerGpu);
    static std::vector<std::vector<int>> GetDeviceIdsList(int const numProcesses,
                                                          int const numGpus);

@@ -105,7 +108,8 @@ namespace RcclUnitTesting
                                       ncclRedOp_t    const redOp,
                                       int            const root,
                                       bool           const inPlace,
-                                       bool           const managedMem);
+                                       bool           const managedMem,
+                                       int            const ranksPerProc=1);

    // Run a simple sweep
    void RunSimpleSweep(std::vector<ncclFunc_t>     const& funcTypes,