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/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>