diff --git a/projects/rccl/src/collectives/device/reduce_kernel.h b/projects/rccl/src/collectives/device/reduce_kernel.h index f6d8f95e03..00dd5e4d18 100644 --- a/projects/rccl/src/collectives/device/reduce_kernel.h +++ b/projects/rccl/src/collectives/device/reduce_kernel.h @@ -520,8 +520,8 @@ struct FuncAvg: FuncSum { // Obviously, the only invalid combination is both true. An improvement would be // make this parameterized as a build time setting and passed here through // preprocessor definitions. - static constexpr bool IsPreOpIdentity = false; - static constexpr bool IsPostOpIdentity = true; + static constexpr bool IsPreOpIdentity = true; + static constexpr bool IsPostOpIdentity = false; #if __CUDA_ARCH__ >= 800 __device__ FuncAvg(int n) { diff --git a/projects/rccl/test/CorrectnessTest.hpp b/projects/rccl/test/CorrectnessTest.hpp index b294e607ee..368f06ae38 100644 --- a/projects/rccl/test/CorrectnessTest.hpp +++ b/projects/rccl/test/CorrectnessTest.hpp @@ -22,6 +22,7 @@ #include #include +#include #include "rccl.h" #include "../include/rccl_bfloat16.h" @@ -558,7 +559,8 @@ dropback: {ncclSum, "sum"}, {ncclProd, "prod"}, {ncclMax, "max"}, - {ncclMin, "min"} + {ncclMin, "min"}, + {ncclAvg, "avg"} }; std::map dataTypeStrings { @@ -735,6 +737,36 @@ dropback: } } + static void Average(Dataset const& dataset, int8_t* resultI1) + { + uint8_t* resultU1 = (uint8_t *)resultI1; + int32_t* resultI4 = (int32_t *)resultI1; + uint32_t* resultU4 = (uint32_t *)resultI1; + int64_t* resultI8 = (int64_t *)resultI1; + uint64_t* resultU8 = (uint64_t *)resultI1; + float* resultF4 = (float *)resultI1; + double* resultF8 = (double *)resultI1; + rccl_bfloat16* resultB2 = (rccl_bfloat16 *)resultI1; + for (int j = 0; j < dataset.numElements; j++) + { + switch (dataset.dataType) + { + case ncclInt8: resultI1[j] = resultI1[j]/dataset.numDevices; break; + case ncclUint8: resultU1[j] = resultU1[j]/dataset.numDevices; break; + case ncclInt32: resultI4[j] = resultI4[j]/dataset.numDevices; break; + case ncclUint32: resultU4[j] = resultU4[j]/dataset.numDevices; break; + case ncclInt64: resultI8[j] = resultI8[j]/dataset.numDevices; break; + case ncclUint64: resultU8[j] = resultU8[j]/dataset.numDevices; break; + case ncclFloat32: resultF4[j] = resultF4[j]/dataset.numDevices; break; + case ncclFloat64: resultF8[j] = resultF8[j]/dataset.numDevices; break; + case ncclBfloat16: resultB2[j] = rccl_bfloat16((float)(resultB2[j])/dataset.numDevices); break; + default: + fprintf(stderr, "[ERROR] Unsupported datatype\n"); + exit(0); + } + } + } + void ValidateResults(Dataset const& dataset, int root = 0) const { int8_t* outputI1 = (int8_t *)malloc(dataset.NumBytes(ncclOutputBuffer)); @@ -778,9 +810,9 @@ dropback: case ncclUint32: isMatch &= (outputU4[j] == expectedU4[j]); break; case ncclInt64: isMatch &= (outputI8[j] == expectedI8[j]); break; case ncclUint64: isMatch &= (outputU8[j] == expectedU8[j]); break; - case ncclFloat32: isMatch &= (outputF4[j] == expectedF4[j]); break; - case ncclFloat64: isMatch &= (outputF8[j] == expectedF8[j]); break; - case ncclBfloat16: isMatch &= (outputB2[j] == expectedB2[j]); break; + case ncclFloat32: isMatch &= (fabs(outputF4[j] - expectedF4[j]) < 1e-5); break; + case ncclFloat64: isMatch &= (fabs(outputF8[j] - expectedF8[j]) < 1e-12); break; + case ncclBfloat16: isMatch &= (fabs((float)outputB2[j] - (float)expectedB2[j]) < 1e-2); break; default: fprintf(stderr, "[ERROR] Unsupported datatype\n"); exit(0); diff --git a/projects/rccl/test/test_AllReduce.cpp b/projects/rccl/test/test_AllReduce.cpp index 74e4017c19..36bae9721a 100644 --- a/projects/rccl/test/test_AllReduce.cpp +++ b/projects/rccl/test/test_AllReduce.cpp @@ -40,7 +40,7 @@ namespace CorrectnessTests AllReduceCorrectnessTest, testing::Combine( // Reduction operator - testing::Values(ncclSum, ncclProd, ncclMax, ncclMin), + testing::Values(ncclSum, ncclProd, ncclMax, ncclMin, ncclAvg), // Data types testing::Values(ncclInt8, ncclUint8, diff --git a/projects/rccl/test/test_AllReduce.hpp b/projects/rccl/test/test_AllReduce.hpp index b5d783b4af..220877207b 100644 --- a/projects/rccl/test/test_AllReduce.hpp +++ b/projects/rccl/test/test_AllReduce.hpp @@ -33,6 +33,7 @@ namespace CorrectnessTests // Initialize the result with the first device's array memcpy(resultI1, dataset.expected[0], dataset.NumBytes()); + ncclRedOp_t red_op = ((op == ncclAvg) ? ncclSum : op); // Perform reduction on the other device arrays for (int i = 1; i < dataset.numDevices; i++) @@ -51,15 +52,15 @@ namespace CorrectnessTests { switch (dataset.dataType) { - case ncclInt8: resultI1[j] = ReduceOp(op, resultI1[j], arrayI1[j]); break; - case ncclUint8: resultU1[j] = ReduceOp(op, resultU1[j], arrayU1[j]); break; - case ncclInt32: resultI4[j] = ReduceOp(op, resultI4[j], arrayI4[j]); break; - case ncclUint32: resultU4[j] = ReduceOp(op, resultU4[j], arrayU4[j]); break; - case ncclInt64: resultI8[j] = ReduceOp(op, resultI8[j], arrayI8[j]); break; - case ncclUint64: resultU8[j] = ReduceOp(op, resultU8[j], arrayU8[j]); break; - case ncclFloat32: resultF4[j] = ReduceOp(op, resultF4[j], arrayF4[j]); break; - case ncclFloat64: resultF8[j] = ReduceOp(op, resultF8[j], arrayF8[j]); break; - case ncclBfloat16: resultB2[j] = ReduceOp(op, resultB2[j], arrayB2[j]); break; + case ncclInt8: resultI1[j] = ReduceOp(red_op, resultI1[j], arrayI1[j]); break; + case ncclUint8: resultU1[j] = ReduceOp(red_op, resultU1[j], arrayU1[j]); break; + case ncclInt32: resultI4[j] = ReduceOp(red_op, resultI4[j], arrayI4[j]); break; + case ncclUint32: resultU4[j] = ReduceOp(red_op, resultU4[j], arrayU4[j]); break; + case ncclInt64: resultI8[j] = ReduceOp(red_op, resultI8[j], arrayI8[j]); break; + case ncclUint64: resultU8[j] = ReduceOp(red_op, resultU8[j], arrayU8[j]); break; + case ncclFloat32: resultF4[j] = ReduceOp(red_op, resultF4[j], arrayF4[j]); break; + case ncclFloat64: resultF8[j] = ReduceOp(red_op, resultF8[j], arrayF8[j]); break; + case ncclBfloat16: resultB2[j] = ReduceOp(red_op, resultB2[j], arrayB2[j]); break; default: fprintf(stderr, "[ERROR] Unsupported datatype\n"); exit(0); @@ -67,6 +68,9 @@ namespace CorrectnessTests } } + if (op == ncclAvg) + Average(dataset, resultI1); + // Copy results into expected arrays for (int i = 0; i < dataset.numDevices; i++) memcpy(dataset.expected[i], resultI1, dataset.NumBytes()); diff --git a/projects/rccl/test/test_Reduce.cpp b/projects/rccl/test/test_Reduce.cpp index c433bffe87..8927ba59c2 100644 --- a/projects/rccl/test/test_Reduce.cpp +++ b/projects/rccl/test/test_Reduce.cpp @@ -48,7 +48,7 @@ namespace CorrectnessTests ReduceCorrectnessTest, testing::Combine( // Reduction operator - testing::Values(ncclSum, ncclProd, ncclMax, ncclMin), + testing::Values(ncclSum, ncclProd, ncclMax, ncclMin, ncclAvg), // Data types testing::Values(ncclInt8, ncclUint8, diff --git a/projects/rccl/test/test_Reduce.hpp b/projects/rccl/test/test_Reduce.hpp index 4ac6dd9486..b4afdab4aa 100644 --- a/projects/rccl/test/test_Reduce.hpp +++ b/projects/rccl/test/test_Reduce.hpp @@ -33,6 +33,7 @@ namespace CorrectnessTests // Initialize the result with the first device's array memcpy(resultI1, dataset.expected[0], dataset.NumBytes()); + ncclRedOp_t red_op = ((op == ncclAvg) ? ncclSum : op); // Perform reduction on the other device arrays for (int i = 1; i < dataset.numDevices; i++) @@ -51,15 +52,15 @@ namespace CorrectnessTests { switch (dataset.dataType) { - case ncclInt8: resultI1[j] = ReduceOp(op, resultI1[j], arrayI1[j]); break; - case ncclUint8: resultU1[j] = ReduceOp(op, resultU1[j], arrayU1[j]); break; - case ncclInt32: resultI4[j] = ReduceOp(op, resultI4[j], arrayI4[j]); break; - case ncclUint32: resultU4[j] = ReduceOp(op, resultU4[j], arrayU4[j]); break; - case ncclInt64: resultI8[j] = ReduceOp(op, resultI8[j], arrayI8[j]); break; - case ncclUint64: resultU8[j] = ReduceOp(op, resultU8[j], arrayU8[j]); break; - case ncclFloat32: resultF4[j] = ReduceOp(op, resultF4[j], arrayF4[j]); break; - case ncclFloat64: resultF8[j] = ReduceOp(op, resultF8[j], arrayF8[j]); break; - case ncclBfloat16: resultB2[j] = ReduceOp(op, resultB2[j], arrayB2[j]); break; + case ncclInt8: resultI1[j] = ReduceOp(red_op, resultI1[j], arrayI1[j]); break; + case ncclUint8: resultU1[j] = ReduceOp(red_op, resultU1[j], arrayU1[j]); break; + case ncclInt32: resultI4[j] = ReduceOp(red_op, resultI4[j], arrayI4[j]); break; + case ncclUint32: resultU4[j] = ReduceOp(red_op, resultU4[j], arrayU4[j]); break; + case ncclInt64: resultI8[j] = ReduceOp(red_op, resultI8[j], arrayI8[j]); break; + case ncclUint64: resultU8[j] = ReduceOp(red_op, resultU8[j], arrayU8[j]); break; + case ncclFloat32: resultF4[j] = ReduceOp(red_op, resultF4[j], arrayF4[j]); break; + case ncclFloat64: resultF8[j] = ReduceOp(red_op, resultF8[j], arrayF8[j]); break; + case ncclBfloat16: resultB2[j] = ReduceOp(red_op, resultB2[j], arrayB2[j]); break; default: fprintf(stderr, "[ERROR] Unsupported datatype\n"); exit(0); @@ -67,6 +68,9 @@ namespace CorrectnessTests } } + if (op == ncclAvg) + Average(dataset, resultI1); + // Copy results into expected arrays for (int i = 0; i < dataset.numDevices; i++) { diff --git a/projects/rccl/test/test_ReduceScatter.cpp b/projects/rccl/test/test_ReduceScatter.cpp index 9979e29954..bed3a5f5f3 100644 --- a/projects/rccl/test/test_ReduceScatter.cpp +++ b/projects/rccl/test/test_ReduceScatter.cpp @@ -46,7 +46,7 @@ namespace CorrectnessTests ReduceScatterCorrectnessTest, testing::Combine( // Reduction operator - testing::Values(ncclSum, ncclProd, ncclMax, ncclMin), + testing::Values(ncclSum, ncclProd, ncclMax, ncclMin, ncclAvg), // Data types testing::Values(ncclInt8, ncclUint8, diff --git a/projects/rccl/test/test_ReduceScatter.hpp b/projects/rccl/test/test_ReduceScatter.hpp index 9aff579b13..3ba54b2983 100644 --- a/projects/rccl/test/test_ReduceScatter.hpp +++ b/projects/rccl/test/test_ReduceScatter.hpp @@ -33,6 +33,7 @@ namespace CorrectnessTests // Initialize the result with the first device's array memcpy(resultI1, dataset.expected[0], dataset.NumBytes()); + ncclRedOp_t red_op = ((op == ncclAvg) ? ncclSum : op); // Perform reduction on the other device arrays for (int i = 1; i < dataset.numDevices; i++) @@ -51,15 +52,15 @@ namespace CorrectnessTests { switch (dataset.dataType) { - case ncclInt8: resultI1[j] = ReduceOp(op, resultI1[j], arrayI1[j]); break; - case ncclUint8: resultU1[j] = ReduceOp(op, resultU1[j], arrayU1[j]); break; - case ncclInt32: resultI4[j] = ReduceOp(op, resultI4[j], arrayI4[j]); break; - case ncclUint32: resultU4[j] = ReduceOp(op, resultU4[j], arrayU4[j]); break; - case ncclInt64: resultI8[j] = ReduceOp(op, resultI8[j], arrayI8[j]); break; - case ncclUint64: resultU8[j] = ReduceOp(op, resultU8[j], arrayU8[j]); break; - case ncclFloat32: resultF4[j] = ReduceOp(op, resultF4[j], arrayF4[j]); break; - case ncclFloat64: resultF8[j] = ReduceOp(op, resultF8[j], arrayF8[j]); break; - case ncclBfloat16: resultB2[j] = ReduceOp(op, resultB2[j], arrayB2[j]); break; + case ncclInt8: resultI1[j] = ReduceOp(red_op, resultI1[j], arrayI1[j]); break; + case ncclUint8: resultU1[j] = ReduceOp(red_op, resultU1[j], arrayU1[j]); break; + case ncclInt32: resultI4[j] = ReduceOp(red_op, resultI4[j], arrayI4[j]); break; + case ncclUint32: resultU4[j] = ReduceOp(red_op, resultU4[j], arrayU4[j]); break; + case ncclInt64: resultI8[j] = ReduceOp(red_op, resultI8[j], arrayI8[j]); break; + case ncclUint64: resultU8[j] = ReduceOp(red_op, resultU8[j], arrayU8[j]); break; + case ncclFloat32: resultF4[j] = ReduceOp(red_op, resultF4[j], arrayF4[j]); break; + case ncclFloat64: resultF8[j] = ReduceOp(red_op, resultF8[j], arrayF8[j]); break; + case ncclBfloat16: resultB2[j] = ReduceOp(red_op, resultB2[j], arrayB2[j]); break; default: fprintf(stderr, "[ERROR] Unsupported datatype\n"); exit(0); @@ -67,6 +68,9 @@ namespace CorrectnessTests } } + if (op == ncclAvg) + Average(dataset, resultI1); + // Copy results into expected arrays size_t const byteCount = dataset.NumBytes() / dataset.numDevices;