Add bfloat16 support in RCCL

Preprocessor symbol RCCL_BFLOAT16 is used as feature indicator

src/collectives/collectives.h

@@ -39,7 +39,8 @@
   DECL_COLL3(coll, op, u64) \
   DECL_COLL3(coll, op, f16) \
   DECL_COLL3(coll, op, f32) \
-  DECL_COLL3(coll, op, f64)
+  DECL_COLL3(coll, op, f64) \
+  DECL_COLL3(coll, op, b16)
 
 #define DECL_COLL(coll) \
   DECL_COLL2(coll, sum) \

src/collectives/device/common.h

@@ -53,7 +53,8 @@ static inline __device__ void exitIfAbortBarrier(int abort) {
   NCCL_FUNC4(coll, op, u64), \
   NCCL_FUNC4(coll, op, f16), \
   NCCL_FUNC4(coll, op, f32), \
-  NCCL_FUNC4(coll, op, f64)
+  NCCL_FUNC4(coll, op, f64), \
+  NCCL_FUNC4(coll, op, b16)
 #define NCCL_FUNCS3B(coll, op) \
   NCCL_FUNC4(coll, op,  i8), \
   NCCL_FUNC4(coll, op,  i8), \
@@ -63,6 +64,7 @@ static inline __device__ void exitIfAbortBarrier(int abort) {
   NCCL_FUNC4(coll, op,  i8), \
   NCCL_FUNC4(coll, op,  i8), \
   NCCL_FUNC4(coll, op,  i8), \
+  NCCL_FUNC4(coll, op,  i8), \
   NCCL_FUNC4(coll, op,  i8)
 
 // Must be consistent with ncclRedOp_t
@@ -121,20 +123,20 @@ struct Caller<f, f + 1>{
 inline
 __device__
 void NCCL_CALL_FUNCTIONS(struct ncclColl* const c) noexcept {
-  if (c->funcIndex < 144) {
+  if (c->funcIndex < 160) {
     if (c->funcIndex % 4 == 0) ncclBroadcastRing_copy_i8(&c->args);
     else if (c->funcIndex % 4 == 1) ncclBroadcastRingLL_copy_i8(&c->args);
     else if (c->funcIndex % 4 == 2) ncclBroadcastTree_copy_i8(&c->args);
     else ncclBroadcastTreeLL_copy_i8(&c->args);
   }
-  else if (c->funcIndex < 288) Caller<144, 288>::call(c);
-  else if (c->funcIndex < 432) {
+  else if (c->funcIndex < 320) Caller<160, 320>::call(c);
+  else if (c->funcIndex < 480) {
     if (c->funcIndex % 4 == 0) ncclAllGatherRing_copy_i8(&c->args);
     else if (c->funcIndex % 4 == 1) ncclAllGatherRingLL_copy_i8(&c->args);
     else if (c->funcIndex % 4 == 2) ncclAllGatherTree_copy_i8(&c->args);
     else ncclAllGatherTreeLL_copy_i8(&c->args);
   }
-  else Caller<432, 720>::call(c);
+  else Caller<480, 800>::call(c);
 }
 
 static __device__ void load_parallel(void* dst, void* src, size_t size, int tid, uint32_t* abortCount) {
@@ -227,7 +229,8 @@ __global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclColl firstColl) { \
   IMPL_COLL3(coll, op, ncclFunc, u64, uint64_t, ncclColl, ncclOp, ncclUint64) \
   IMPL_COLL3(coll, op, ncclFunc, f16, half,     ncclColl, ncclOp, ncclFloat16) \
   IMPL_COLL3(coll, op, ncclFunc, f32, float,    ncclColl, ncclOp, ncclFloat32) \
-  IMPL_COLL3(coll, op, ncclFunc, f64, double,   ncclColl, ncclOp, ncclFloat64)
+  IMPL_COLL3(coll, op, ncclFunc, f64, double,   ncclColl, ncclOp, ncclFloat64) \
+  IMPL_COLL3(coll, op, ncclFunc, b16, rccl_bfloat16, ncclColl, ncclOp, ncclBfloat16)
 
 #define COLL_UNROLL 2
 

src/collectives/device/common_kernel.h

@@ -241,6 +241,18 @@ template<> inline __device__
 void vStore<half>(volatile half* ptr, const half val) {
   ((half*)ptr)[0] = val;
 }
+
+template<> inline __device__
+rccl_bfloat16 vFetch<rccl_bfloat16>(const volatile rccl_bfloat16* ptr) {
+  rccl_bfloat16 r;
+  r.data = ptr->data;
+  return r;
+}
+
+template<> inline __device__
+void vStore<rccl_bfloat16>(volatile rccl_bfloat16* ptr, const rccl_bfloat16 val) {
+  ptr->data = val.data;
+}
 #endif
 
 typedef ulong2 Pack128;

src/collectives/device/reduce_kernel.h

@@ -134,6 +134,86 @@ struct FuncMin : private FuncBase<T> {
   }
 };
 
+template<>
+struct FuncSum<rccl_bfloat16> {
+  static constexpr auto n = sizeof(PackType) / sizeof(rccl_bfloat16);
+  __device__ PackType operator()(PackType x, PackType y) const
+  {
+    union converter { PackType storage; rccl_bfloat16 vec[n]; };
+    static_assert(sizeof(PackType) == sizeof(converter), "PackType must be the same size of converter.");
+    converter cx, cy, cr;
+    cx.storage = x;
+    cy.storage = y;
+    for (auto i = 0u; i != n; ++i) {
+      cr.vec[i] = cx.vec[i] + cy.vec[i];
+    }
+    return cr.storage;
+  }
+  __device__ rccl_bfloat16 operator()(const rccl_bfloat16 x, const rccl_bfloat16 y) const {
+    return x + y;
+  }
+};
+
+template<>
+struct FuncProd<rccl_bfloat16> {
+  static constexpr auto n = sizeof(PackType) / sizeof(rccl_bfloat16);
+  __device__ PackType operator()(PackType x, PackType y) const
+  {
+    union converter { PackType storage; rccl_bfloat16 vec[n]; };
+    static_assert(sizeof(PackType) == sizeof(converter), "PackType must be the same size of converter.");
+    converter cx, cy, cr;
+    cx.storage = x;
+    cy.storage = y;
+    for (auto i = 0u; i != n; ++i) {
+      cr.vec[i] = cx.vec[i] * cy.vec[i];
+    }
+    return cr.storage;
+  }
+  __device__ rccl_bfloat16 operator()(const rccl_bfloat16 x, const rccl_bfloat16 y) const {
+    return x * y;
+  }
+};
+
+template<>
+struct FuncMax<rccl_bfloat16> {
+  static constexpr auto n = sizeof(PackType) / sizeof(rccl_bfloat16);
+  __device__ PackType operator()(PackType x, PackType y) const
+  {
+    union converter { PackType storage; rccl_bfloat16 vec[n]; };
+    static_assert(sizeof(PackType) == sizeof(converter), "PackType must be the same size of converter.");
+    converter cx, cy, cr;
+    cx.storage = x;
+    cy.storage = y;
+    for (auto i = 0u; i != n; ++i) {
+      cr.vec[i] = cx.vec[i] < cy.vec[i] ? cy.vec[i] : cx.vec[i];
+    }
+    return cr.storage;
+  }
+  __device__ rccl_bfloat16 operator()(const rccl_bfloat16 x, const rccl_bfloat16 y) const {
+    return x < y ? y : x;
+  }
+};
+
+template<>
+struct FuncMin<rccl_bfloat16> {
+  static constexpr auto n = sizeof(PackType) / sizeof(rccl_bfloat16);
+  __device__ PackType operator()(PackType x, PackType y) const
+  {
+    union converter { PackType storage; rccl_bfloat16 vec[n]; };
+    static_assert(sizeof(PackType) == sizeof(converter), "PackType must be the same size of converter.");
+    converter cx, cy, cr;
+    cx.storage = x;
+    cy.storage = y;
+    for (auto i = 0u; i != n; ++i) {
+      cr.vec[i] = cx.vec[i] < cy.vec[i] ? cx.vec[i] : cy.vec[i];
+    }
+    return cr.storage;
+  }
+  __device__ rccl_bfloat16 operator()(const rccl_bfloat16 x, const rccl_bfloat16 y) const {
+    return x < y ? x : y;
+  }
+};
+
 #else
 
 template<typename T>