From bdac0256a56176c2a638ade921f80a7bb7dab493 Mon Sep 17 00:00:00 2001 From: Wenkai Du Date: Sat, 16 Nov 2019 21:37:20 -0800 Subject: [PATCH] Add bfloat16 all reduce unit test --- test/CorrectnessTest.hpp | 9 +++++++++ test/test_AllReduce.cpp | 3 ++- test/test_AllReduce.hpp | 3 +++ 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/test/CorrectnessTest.hpp b/test/CorrectnessTest.hpp index e7d3c75382..7440bd5ecc 100644 --- a/test/CorrectnessTest.hpp +++ b/test/CorrectnessTest.hpp @@ -11,6 +11,7 @@ #include #include #include "rccl.h" +#include "../include/rccl_bfloat16.h" #define HIP_CALL(x) ASSERT_EQ(x, hipSuccess) #define NCCL_CALL(x) ASSERT_EQ(x, ncclSuccess) @@ -47,6 +48,7 @@ namespace CorrectnessTests case ncclFloat16: return 2; case ncclFloat32: return 4; case ncclFloat64: return 8; + case ncclBfloat16: return 2; default: fprintf(stderr, "[ERROR] Unsupported datatype (%d)\n", dataType); exit(0); @@ -217,6 +219,7 @@ namespace CorrectnessTests uint64_t* arrayU8 = (uint64_t *)arrayI1; float* arrayF4 = (float *)arrayI1; double* arrayF8 = (double *)arrayI1; + rccl_bfloat16* arrayB2 = (rccl_bfloat16 *)arrayI1; // NOTE: Currently half-precision float tests are unsupported due to half being supported // on GPU only and not host @@ -241,6 +244,7 @@ namespace CorrectnessTests case ncclUint64: arrayU8[j] = valueI; break; case ncclFloat32: arrayF4[j] = valueF; break; case ncclFloat64: arrayF8[j] = valueF; break; + case ncclBfloat16: arrayB2[j] = rccl_bfloat16(valueF); break; default: fprintf(stderr, "[ERROR] Unsupported datatype\n"); exit(0); @@ -278,6 +282,7 @@ namespace CorrectnessTests uint64_t* outputU8 = (uint64_t *)outputI1; float* outputF4 = (float *)outputI1; double* outputF8 = (double *)outputI1; + rccl_bfloat16* outputB2 = (rccl_bfloat16 *)outputI1; bool isMatch = true; @@ -295,6 +300,7 @@ namespace CorrectnessTests uint64_t* expectedU8 = (uint64_t *)expectedI1; float* expectedF4 = (float *)expectedI1; double* expectedF8 = (double *)expectedI1; + rccl_bfloat16* expectedB2 = (rccl_bfloat16 *)expectedI1; for (int j = 0; j < dataset.numElements && isMatch; j++) { @@ -308,6 +314,7 @@ namespace CorrectnessTests 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; default: fprintf(stderr, "[ERROR] Unsupported datatype\n"); exit(0); @@ -333,6 +340,8 @@ namespace CorrectnessTests printf("Expected %f. Output %f on device %d[%d]\n", outputF4[j], expectedF4[j], i, j); break; case ncclFloat64: printf("Expected %lf. Output %lf on device %d[%d]\n", outputF8[j], expectedF8[j], i, j); break; + case ncclBfloat16: + printf("Expected %f. Output %f on device %d[%d]\n", (float)outputB2[j], (float)expectedB2[j], i, j); break; default: fprintf(stderr, "[ERROR] Unsupported datatype\n"); exit(0); diff --git a/test/test_AllReduce.cpp b/test/test_AllReduce.cpp index 0fb7474d0e..0fd5eedf91 100644 --- a/test/test_AllReduce.cpp +++ b/test/test_AllReduce.cpp @@ -50,7 +50,8 @@ namespace CorrectnessTests ncclUint64, //ncclFloat16, ncclFloat32, - ncclFloat64), + ncclFloat64, + ncclBfloat16), // Number of elements testing::Values(1024, 1048576), // Number of devices diff --git a/test/test_AllReduce.hpp b/test/test_AllReduce.hpp index d8867cb649..a056be95a7 100644 --- a/test/test_AllReduce.hpp +++ b/test/test_AllReduce.hpp @@ -29,6 +29,7 @@ namespace CorrectnessTests uint64_t* resultU8 = (uint64_t *)resultI1; float* resultF4 = (float *)resultI1; double* resultF8 = (double *)resultI1; + rccl_bfloat16* resultB2 = (rccl_bfloat16 *)resultI1; // Initialize the result with the first device's array memcpy(resultI1, dataset.expected[0], dataset.NumBytes()); @@ -44,6 +45,7 @@ namespace CorrectnessTests uint64_t* arrayU8 = (uint64_t *)arrayI1; float* arrayF4 = (float *)arrayI1; double* arrayF8 = (double *)arrayI1; + rccl_bfloat16* arrayB2 = (rccl_bfloat16 *)arrayI1; for (int j = 0; j < dataset.numElements; j++) { @@ -57,6 +59,7 @@ namespace CorrectnessTests 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; default: fprintf(stderr, "[ERROR] Unsupported datatype\n"); exit(0);