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

Merge remote-tracking branch 'rccl/develop' into 2.19.4

Browse Source
This commit is contained in:
BertanDogancay
2024-02-15 13:37:14 -08:00
parent 32cca51894 16d7f372b7
commit 76f83f95ab
9 changed files with 218 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/device/reduce_kernel.h
+1 -29
View File
@@ -239,22 +239,7 @@ struct Apply_Reduce<FuncProd<uint8_t>, /*EltPerPack=*/4> {
SPECIALIZE_REDUCE(FuncMinMax, float, 1, float, fn.isMinNotMax ? fminf(x, y) : fmaxf(x, y))
SPECIALIZE_REDUCE(FuncMinMax, double, 1, double, fn.isMinNotMax ? fmin(x, y) : fmax(x, y))
#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
SPECIALIZE_REDUCE(FuncSum, half, 1, half, __hadd(x, y))
SPECIALIZE_REDUCE(FuncSum, half, 2, half2, __hadd2(x, y))
SPECIALIZE_REDUCE(FuncProd, half, 1, half, __hmul(x, y))
SPECIALIZE_REDUCE(FuncProd, half, 2, half2, __hmul2(x, y))
#else
SPECIALIZE_REDUCE(FuncSum, half, 1, half, __float2half(__half2float(x) + __half2float(y)))
SPECIALIZE_REDUCE(FuncProd, half, 1, half, __float2half(__half2float(x) * __half2float(y)))
#endif
#if __CUDA_ARCH__ >= 800
SPECIALIZE_REDUCE(FuncMinMax, half, 1, half, fn.isMinNotMax ? __hmin(x, y) : __hmax(x, y))
SPECIALIZE_REDUCE(FuncMinMax, half, 2, half2, fn.isMinNotMax ? __hmin2(x, y) : __hmax2(x, y))
#else
SPECIALIZE_REDUCE(FuncMinMax, half, 1, half, __float2half(fn.isMinNotMax ? fminf(__half2float(x), __half2float(y)) : fmaxf(__half2float(x), __half2float(y))))
#endif
SPECIALIZE_REDUCE(FuncMinMax, half, 1, half, fn.isMinNotMax ? __hmin(x, y) : __hmax(x, y))
#if defined(RCCL_BFLOAT16)
#if __CUDA_ARCH__ >= 800
@@ -374,7 +359,6 @@ struct FuncPreMulSum {
template<>
struct FuncPreMulSum<half> {
using EltType = half;
#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
half2 scalar;
__device__ FuncPreMulSum(uint64_t opArg=0) {
union { uint64_t u64; half val; };
@@ -382,14 +366,6 @@ struct FuncPreMulSum<half> {
scalar.x = val;
scalar.y = val;
}
#else
float scalar;
__device__ FuncPreMulSum(uint64_t opArg=0) {
union { uint64_t u64; half val; };
u64 = opArg;
scalar = __half2float(val);
}
#endif
};
#if defined(RCCL_BFLOAT16)
@@ -439,11 +415,7 @@ template<>
struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(half)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half)> a) {
#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
return toPack<half>(__hmul(fromPack<half>(a), fn.scalar.x));
#else
return toPack<half>(__float2half(__half2float(fromPack<half>(a)) * fn.scalar));
#endif
}
};
#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
src/init_nvtx.cc
+2
View File
@@ -11,6 +11,7 @@ static constexpr const nvtxPayloadEnum_t NvtxEnumRedSchema[] = {
// Must be called before the first call to any reduction operation.
void initNvtxRegisteredEnums() {
#ifndef NVTX_NO_IMPL
// Register schemas and strings
constexpr const nvtxPayloadEnumAttr_t eAttr {
.fieldMask = NVTX_PAYLOAD_ENUM_ATTR_ENTRIES | NVTX_PAYLOAD_ENUM_ATTR_NUM_ENTRIES |
@@ -23,4 +24,5 @@ void initNvtxRegisteredEnums() {
};
nvtxPayloadEnumRegister(nvtx3::domain::get<nccl_domain>(), &eAttr);
#endif
}