Add bfloat16 all reduce unit test

test/CorrectnessTest.hpp

@@ -11,6 +11,7 @@
 #include <vector>
 #include <gtest/gtest.h>
 #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);

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

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);