Adding support for alignment tests via sub-datasets

Added sample alignment test for AllGather
Datasets no longer free memory on destruction so Release() must be used


[ROCm/rccl commit: a50c852851]

projects/rccl/test/CorrectnessTest.hpp

@@ -60,6 +60,7 @@ namespace CorrectnessTests
         size_t              numElements; // Number of elements per array
         ncclDataType_t      dataType;    // Data type of each input/output pointer
         bool                inPlace;     // Whether or not output pointers are same as input pointers
+
         std::vector<void *> inputs;      // Input pointers (1 per device)
         std::vector<void *> outputs;     // Output pointers (1 per device)
                                          // May be identical to input pointers for in-place tests
@@ -100,7 +101,8 @@ namespace CorrectnessTests
             }
         }
 
-        ~Dataset()
+        // Explicit memory release to avoid double-free from subDatasets
+        void Release()
         {
             for (int i = 0; i < outputs.size(); i++)
             {
@@ -108,6 +110,35 @@ namespace CorrectnessTests
                 hipFree(inputs[i]);
                 free(expected[i]);
             }
+
+            outputs.clear();
+        }
+
+        // Creates a dataset by pointing to an existing dataset
+        // Primarily to allow for testing with different starting byte-alignments
+        void ExtractSubDataset(size_t const startElement,
+                               size_t const lastElement,
+                               Dataset& subDataset)
+        {
+            ASSERT_LE(startElement, lastElement);
+            ASSERT_LT(lastElement, numElements);
+
+            subDataset.numDevices  = numDevices;
+            subDataset.numElements = lastElement - startElement + 1;
+            subDataset.dataType    = dataType;
+            subDataset.inPlace     = inPlace;
+
+            subDataset.inputs.resize(numDevices);
+            subDataset.outputs.resize(numDevices);
+            subDataset.expected.resize(numDevices);
+
+            size_t const byteOffset = (startElement * DataTypeToBytes(dataType));
+            for (int i = 0; i < numDevices; i++)
+            {
+                subDataset.inputs[i]   = (int8_t *)inputs[i] + byteOffset;
+                subDataset.outputs[i]  = (int8_t *)outputs[i] + byteOffset;
+                subDataset.expected[i] = (int8_t *)expected[i] + byteOffset;
+            }
         }
     };
 

projects/rccl/test/test_AllGather.cpp

@@ -36,8 +36,57 @@ namespace CorrectnessTests
 
         // Check results
         ValidateResults(dataset);
+        dataset.Release();
     }
 
+    TEST_P(AllGatherCorrectnessTest, Alignment)
+    {
+        if (numDevices > numDevicesAvailable) return;
+        if (numElements % numDevices != 0) return;
+
+        // Allocate dataset
+        Dataset dataset;
+        dataset.Initialize(numDevices, numElements, dataType, inPlace);
+
+        // Loop over several offsets (so that device pointers are not aligned)
+        for (int firstElement = 1; firstElement <= 11; firstElement += 2)
+        {
+            if (firstElement < numElements)
+            {
+                // Select last element so that total number of elements is multiple of numDevices
+                int const lastElement = firstElement + ((numElements - firstElement) / numDevices) * numDevices - 1;
+                if (lastElement >= numElements) break;
+
+                Dataset subDataset;
+                dataset.ExtractSubDataset(firstElement, lastElement, subDataset);
+
+                // Compute reference results for sub-dataset
+                FillDatasetWithPattern(subDataset);
+                ComputeExpectedResults(subDataset);
+
+                size_t const byteCount = subDataset.NumBytes() / subDataset.numDevices;
+                size_t const sendCount = subDataset.numElements / subDataset.numDevices;
+
+                // Launch the reduction (1 thread per GPU)
+                #pragma omp parallel for num_threads(numDevices)
+                for (int i = 0; i < numDevices; i++)
+                {
+                    ncclAllGather((int8_t *)subDataset.inputs[i] + (i * byteCount),
+                                  subDataset.outputs[i], sendCount,
+                                  dataType, comms[i], streams[i]);
+                }
+
+                // Wait for reduction to complete
+                Synchronize();
+
+                // Check results
+                ValidateResults(subDataset);
+            }
+        }
+        dataset.Release();
+    }
+
+
     INSTANTIATE_TEST_CASE_P(AllGatherCorrectnessSweep,
                             AllGatherCorrectnessTest,
                             testing::Combine(

projects/rccl/test/test_AllReduce.cpp

@@ -32,6 +32,8 @@ namespace CorrectnessTests
 
         // Check results
         ValidateResults(dataset);
+
+        dataset.Release();
     }
 
     INSTANTIATE_TEST_CASE_P(AllReduceCorrectnessSweep,

projects/rccl/test/test_Broadcast.cpp

@@ -41,6 +41,8 @@ namespace CorrectnessTests
             // Check results
             ValidateResults(dataset);
         }
+
+        dataset.Release();
     }
 
     INSTANTIATE_TEST_CASE_P(BroadcastCorrectnessSweep,

projects/rccl/test/test_GroupCalls.cpp

@@ -92,6 +92,7 @@ namespace CorrectnessTests
         for (int i = 0; i < 5; i++)
         {
             ValidateResults(datasets[i]);
+            datasets[i].Release();
         }
     }
 

projects/rccl/test/test_Reduce.cpp

@@ -40,6 +40,8 @@ namespace CorrectnessTests
             // Check results
             ValidateResults(dataset);
         }
+
+        dataset.Release();
     }
 
     INSTANTIATE_TEST_CASE_P(ReduceCorrectnessSweep,

projects/rccl/test/test_ReduceScatter.cpp

@@ -39,6 +39,8 @@ namespace CorrectnessTests
 
         // Check results
         ValidateResults(dataset);
+
+        dataset.Release();
     }
 
     INSTANTIATE_TEST_CASE_P(ReduceScatterCorrectnessSweep,