Fixing bug with ExtractSubDataset function not fully initializing subdataset (#390)

test/CorrectnessTest.hpp

@@ -268,6 +268,7 @@ namespace CorrectnessTests
             subDataset.numElements = lastElement - startElement + 1;
             subDataset.dataType    = dataType;
             subDataset.inPlace     = inPlace;
+            subDataset.function    = function;
 
             subDataset.inputs.resize(numDevices);
             subDataset.outputs.resize(numDevices);

test/test_AllGather.cpp

@@ -14,7 +14,7 @@ namespace CorrectnessTests
 
         // Prepare input / output / expected results
         Dataset dataset;
-        dataset.Initialize(numDevices, numElements, dataType, inPlace);
+        dataset.Initialize(numDevices, numElements, dataType, inPlace, ncclCollAllGather);
         FillDatasetWithPattern(dataset);
         ComputeExpectedResults(dataset);
 
@@ -46,7 +46,7 @@ namespace CorrectnessTests
 
         // Allocate dataset
         Dataset dataset;
-        dataset.Initialize(numDevices, numElements, dataType, inPlace);
+        dataset.Initialize(numDevices, numElements, dataType, inPlace, ncclCollAllGather);
 
         // Loop over several offsets (so that device pointers are not aligned)
         for (int firstElement = 1; firstElement <= 11; firstElement += 2)

test/test_AllReduce.cpp

@@ -14,7 +14,7 @@ namespace CorrectnessTests
 
     // Prepare input / output / expected results
     Dataset dataset;
-    dataset.Initialize(numDevices, numElements, dataType, inPlace);
+    dataset.Initialize(numDevices, numElements, dataType, inPlace, ncclCollAllReduce);
     FillDatasetWithPattern(dataset);
     ComputeExpectedResults(dataset, op);
 

test/test_AllReduceAbort.cpp

@@ -30,7 +30,7 @@ namespace CorrectnessTests
 
         // Prepare input / output / expected results
         Dataset dataset;
-        dataset.Initialize(numDevices, numElements, dataType, inPlace);
+        dataset.Initialize(numDevices, numElements, dataType, inPlace, ncclCollAllReduce);
         FillDatasetWithPattern(dataset);
 
         int gpu = 0; // GPU number to trigger abort

test/test_AllReduceGroup.cpp

@@ -15,9 +15,9 @@ namespace CorrectnessTests
 
     // Prepare input / output / expected results
     Dataset dataset1, dataset2, dataset3;
-    dataset1.Initialize(numDevices, numElements, dataType, inPlace);
-    dataset2.Initialize(numDevices, numElements, dataType, inPlace);
-    dataset3.Initialize(numDevices, numElements, dataType, inPlace);
+    dataset1.Initialize(numDevices, numElements, dataType, inPlace, ncclCollAllReduce);
+    dataset2.Initialize(numDevices, numElements, dataType, inPlace, ncclCollAllReduce);
+    dataset3.Initialize(numDevices, numElements, dataType, inPlace, ncclCollAllReduce);
     FillDatasetWithPattern(dataset1);
     FillDatasetWithPattern(dataset2);
     FillDatasetWithPattern(dataset3);

test/test_Broadcast.cpp

@@ -14,7 +14,7 @@ namespace CorrectnessTests
 
         // Allocate data
         Dataset dataset;
-        dataset.Initialize(numDevices, numElements, dataType, inPlace);
+        dataset.Initialize(numDevices, numElements, dataType, inPlace, ncclCollBroadcast);
 
         // Test each possible root
         for (int root = 0; root < numDevices; root++)

test/test_BroadcastAbort.cpp

@@ -30,7 +30,7 @@ namespace CorrectnessTests
 
         // Prepare input / output / expected results
         Dataset dataset;
-        dataset.Initialize(numDevices, numElements, dataType, inPlace);
+        dataset.Initialize(numDevices, numElements, dataType, inPlace, ncclCollBroadcast);
         FillDatasetWithPattern(dataset);
 
         int root = 0;

test/test_CombinedCalls.cpp

@@ -20,9 +20,16 @@ namespace CorrectnessTests
 
         // Create multiple datasets for combined operation
         std::vector<Dataset> datasets(5);
+        std::vector<ncclFunc_t> ncclFuncs(5);
+        ncclFuncs.push_back(ncclCollAllGather);
+        ncclFuncs.push_back(ncclCollAllReduce);
+        ncclFuncs.push_back(ncclCollBroadcast);
+        ncclFuncs.push_back(ncclCollReduce);
+        ncclFuncs.push_back(ncclCollReduceScatter);
+
         for (int i = 0; i < datasets.size(); i++)
         {
-            datasets[i].Initialize(numDevices, numElements, dataType, inPlace);
+            datasets[i].Initialize(numDevices, numElements, dataType, inPlace, ncclFuncs[i]);
             FillDatasetWithPattern(datasets[i]);
         }
 

test/test_GroupCalls.cpp

@@ -19,9 +19,16 @@ namespace CorrectnessTests
 
         // Create multiple datasets for group operation
         std::vector<Dataset> datasets(5);
+        std::vector<ncclFunc_t> ncclFuncs(5);
+        ncclFuncs.push_back(ncclCollAllGather);
+        ncclFuncs.push_back(ncclCollAllReduce);
+        ncclFuncs.push_back(ncclCollBroadcast);
+        ncclFuncs.push_back(ncclCollReduce);
+        ncclFuncs.push_back(ncclCollReduceScatter);
+
         for (int i = 0; i < datasets.size(); i++)
         {
-            datasets[i].Initialize(numDevices, numElements, dataType, inPlace);
+            datasets[i].Initialize(numDevices, numElements, dataType, inPlace, ncclFuncs[i]);
             FillDatasetWithPattern(datasets[i]);
         }
 

test/test_Reduce.cpp

@@ -14,7 +14,7 @@ namespace CorrectnessTests
 
         // Allocate data
         Dataset dataset;
-        dataset.Initialize(numDevices, numElements, dataType, inPlace);
+        dataset.Initialize(numDevices, numElements, dataType, inPlace, ncclCollReduce);
 
         // Test each possible root
         for (int root = 0; root < numDevices; root++)

test/test_ReduceScatter.cpp

@@ -15,7 +15,7 @@ namespace CorrectnessTests
 
         // Prepare input / output / expected results
         Dataset dataset;
-        dataset.Initialize(numDevices, numElements, dataType, inPlace);
+        dataset.Initialize(numDevices, numElements, dataType, inPlace, ncclCollReduceScatter);
         FillDatasetWithPattern(dataset);
         ComputeExpectedResults(dataset, op);