AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add bfloat16 support in RCCL

Browse Source 
Preprocessor symbol RCCL_BFLOAT16 is used as feature indicator
This commit is contained in:
Wenkai Du
2019-11-15 10:39:48 -08:00
parent 58a6e535f6
commit 5e109ed400
10 changed files with 370 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/collectives/device/common.h
+9 -6
View File
@@ -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
+12
View File
@@ -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
+80
View File
@@ -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>