Merge remote-tracking branch 'nccl/master' into develop
[ROCm/rccl commit: 08a7be231b]
Αυτή η υποβολή περιλαμβάνεται σε:
@@ -42,18 +42,18 @@ struct IsFloatingPoint<double>: std::true_type {};
|
||||
// 3. Have constructor taking `uint64_t opArg`.
|
||||
|
||||
template<typename T>
|
||||
struct FuncCopy { using EltType = T; __device__ FuncCopy(uint64_t opArg=0) {}; };
|
||||
struct FuncCopy { using EltType = T; __device__ __forceinline__ FuncCopy(uint64_t opArg=0) {}; };
|
||||
template<typename T>
|
||||
struct FuncSum { using EltType = T; __device__ FuncSum(uint64_t opArg=0) {}; };
|
||||
struct FuncSum { using EltType = T; __device__ __forceinline__ FuncSum(uint64_t opArg=0) {}; };
|
||||
template<typename T>
|
||||
struct FuncProd { using EltType = T; __device__ FuncProd(uint64_t opArg=0) {}; };
|
||||
struct FuncProd { using EltType = T; __device__ __forceinline__ FuncProd(uint64_t opArg=0) {}; };
|
||||
|
||||
template<typename T>
|
||||
struct FuncMinMax {
|
||||
using EltType = T;
|
||||
BytePack<sizeof(T)> xormask; // only used by integers
|
||||
bool isMinNotMax; // only used by floats
|
||||
__device__ FuncMinMax(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncMinMax(uint64_t opArg=0) {
|
||||
xormask.native = opArg;
|
||||
isMinNotMax = (opArg&1)==0;
|
||||
}
|
||||
@@ -68,13 +68,13 @@ template<typename T> struct FuncSumPostDiv;
|
||||
template<typename Fn>
|
||||
struct RedOpArg { // default case: no argument
|
||||
static constexpr bool ArgUsed = false;
|
||||
__device__ static uint64_t loadArg(void *ptr) { return 0; }
|
||||
__device__ __forceinline__ static uint64_t loadArg(void *ptr) { return 0; }
|
||||
};
|
||||
|
||||
template<typename T>
|
||||
struct RedOpArg<FuncMinMax<T>> {
|
||||
static constexpr bool ArgUsed = true;
|
||||
__device__ static uint64_t loadArg(void *ptr) {
|
||||
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
|
||||
union { uint64_t u64; T val; };
|
||||
u64 = 0;
|
||||
val = *(T*)ptr;
|
||||
@@ -88,6 +88,11 @@ struct RedOpArg<FuncMinMax<T>> {
|
||||
// of elements. These classes are intended to be specialized for specific
|
||||
// combinations of reduction function and pack size.
|
||||
|
||||
template<typename A, typename B, int EltPerPackA>
|
||||
struct Apply_Cast/*{
|
||||
static BytePack<EltPerPackA*sizeof(B)/sizeof(A)> cast(BytePack<EltPerPackA*sizeof(A)> a);
|
||||
}*/;
|
||||
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_Reduce /*{
|
||||
static BytePack<EltPerPack*sizeof(T)> reduce(
|
||||
@@ -115,16 +120,60 @@ struct Apply_LoadMultimem/*{
|
||||
static BytePack<BytePerPack> load(Fn fn, uintptr_t addr);
|
||||
}*/;
|
||||
|
||||
|
||||
// Helpers for dealing with BytePack<0>'s
|
||||
template<typename A, typename B, int EltPerPack>
|
||||
struct Apply_Cast_MaybeEmpty: Apply_Cast<A, B, EltPerPack> {};
|
||||
template<typename A, typename B>
|
||||
struct Apply_Cast_MaybeEmpty<A, B, /*EltPerPack=*/0> {
|
||||
__device__ constexpr static BytePack<0> cast(BytePack<0> a) { return {}; }
|
||||
};
|
||||
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_Reduce_MaybeEmpty: Apply_Reduce<Fn, EltPerPack> {};
|
||||
template<typename Fn>
|
||||
struct Apply_Reduce_MaybeEmpty<Fn, 0> {
|
||||
__device__ constexpr static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) { return {}; }
|
||||
};
|
||||
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_PreOp_MaybeEmpty: Apply_PreOp<Fn, EltPerPack> {};
|
||||
template<typename Fn>
|
||||
struct Apply_PreOp_MaybeEmpty<Fn, 0> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
__device__ constexpr static BytePack<0> preOp(Fn fn, BytePack<0> a) { return {}; }
|
||||
};
|
||||
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_PostOp_MaybeEmpty: Apply_PostOp<Fn, EltPerPack> {};
|
||||
template<typename Fn>
|
||||
struct Apply_PostOp_MaybeEmpty<Fn, 0> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
__device__ constexpr static BytePack<0> postOp(Fn fn, BytePack<0> a) { return {}; }
|
||||
};
|
||||
|
||||
template<typename Fn, int BytePerPack>
|
||||
struct Apply_LoadMultimem_MaybeEmpty: Apply_LoadMultimem<Fn, BytePerPack> {};
|
||||
template<typename Fn>
|
||||
struct Apply_LoadMultimem_MaybeEmpty<Fn, 0> {
|
||||
__device__ constexpr static BytePack<0> load(Fn fn, uintptr_t addr) { return {}; }
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Public API for calling the trait classes. These take the data elements as a
|
||||
// pack of any type, which could be a BytePack<?> or any integral type (uint64_t,
|
||||
// uint32_t, etc.), and will return a new pack where each element has been
|
||||
// transformed appropriately.
|
||||
|
||||
template<typename A, typename B, typename PackA>
|
||||
__device__ __forceinline__ BytePack<BytePackOf<PackA>::Size*sizeof(B)/sizeof(A)> applyCast(PackA a) {
|
||||
return Apply_Cast_MaybeEmpty<A, B, BytePackOf<PackA>::Size/sizeof(A)>::cast(toPack(a));
|
||||
}
|
||||
|
||||
template<typename Fn, typename Pack>
|
||||
__device__ __forceinline__ Pack applyReduce(Fn fn, Pack a, Pack b) {
|
||||
return fromPack<Pack>(
|
||||
Apply_Reduce<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
Apply_Reduce_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
::reduce(fn, toPack(a), toPack(b))
|
||||
);
|
||||
}
|
||||
@@ -132,7 +181,7 @@ __device__ __forceinline__ Pack applyReduce(Fn fn, Pack a, Pack b) {
|
||||
template<typename Fn, typename Pack>
|
||||
__device__ __forceinline__ Pack applyPreOp(Fn fn, Pack a) {
|
||||
return fromPack<Pack>(
|
||||
Apply_PreOp<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
Apply_PreOp_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
::preOp(fn, toPack(a))
|
||||
);
|
||||
}
|
||||
@@ -140,23 +189,107 @@ __device__ __forceinline__ Pack applyPreOp(Fn fn, Pack a) {
|
||||
template<typename Fn, typename Pack>
|
||||
__device__ __forceinline__ Pack applyPostOp(Fn fn, Pack a) {
|
||||
return fromPack<Pack>(
|
||||
Apply_PostOp<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
Apply_PostOp_MaybeEmpty<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
|
||||
::postOp(fn, toPack(a))
|
||||
);
|
||||
}
|
||||
|
||||
template<typename Fn, int BytePerPack>
|
||||
__device__ __forceinline__ BytePack<BytePerPack> applyLoadMultimem(Fn fn, uintptr_t addr) {
|
||||
return Apply_LoadMultimem<Fn, BytePerPack>::load(fn, addr);
|
||||
return Apply_LoadMultimem_MaybeEmpty<Fn, BytePerPack>::load(fn, addr);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Apply_Cast
|
||||
|
||||
template<typename A, typename B, int EltPerPack>
|
||||
struct Apply_Cast {
|
||||
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(B)> cast(BytePack<EltPerPack*sizeof(A)> a) {
|
||||
BytePack<EltPerPack*sizeof(B)> b;
|
||||
b.half[0] = Apply_Cast<A, B, EltPerPack/2>::cast(a.half[0]);
|
||||
b.half[1] = Apply_Cast<A, B, EltPerPack/2>::cast(a.half[1]);
|
||||
return b;
|
||||
}
|
||||
};
|
||||
|
||||
template<typename A, typename B>
|
||||
struct Apply_Cast<A, B, /*EltPerPack=*/1> {
|
||||
__device__ __forceinline__ static BytePack<sizeof(B)> cast(BytePack<sizeof(A)> a) {
|
||||
return toPack(B(fromPack<A>(a)));
|
||||
}
|
||||
};
|
||||
|
||||
template<>
|
||||
struct Apply_Cast<__half, float, /*EltPerPack=*/1> {
|
||||
__device__ __forceinline__ static BytePack<sizeof(float)> cast(BytePack<sizeof(__half)> a) {
|
||||
return toPack(__half2float(fromPack<__half>(a)));
|
||||
}
|
||||
};
|
||||
template<>
|
||||
struct Apply_Cast<float, __half, /*EltPerPack=*/1> {
|
||||
__device__ __forceinline__ static BytePack<sizeof(__half)> cast(BytePack<sizeof(float)> a) {
|
||||
return toPack(__float2half_rn(fromPack<float>(a)));
|
||||
}
|
||||
};
|
||||
|
||||
template<>
|
||||
struct Apply_Cast<__half, float, /*EltPerPack=*/2> {
|
||||
__device__ __forceinline__ static BytePack<4*2> cast(BytePack<2*2> a) {
|
||||
return toPack(__half22float2(fromPack<__half2>(a)));
|
||||
}
|
||||
};
|
||||
template<>
|
||||
struct Apply_Cast<float, __half, /*EltPerPack=*/2> {
|
||||
__device__ __forceinline__ static BytePack<2*2> cast(BytePack<4*2> a) {
|
||||
return toPack(__float22half2_rn(fromPack<float2>(a)));
|
||||
}
|
||||
};
|
||||
|
||||
#if defined(__CUDA_BF16_TYPES_EXIST__) && (CUDART_RUNTIME >= 12000 || __CUDA_ARCH__ >= 800)
|
||||
template<>
|
||||
struct Apply_Cast<__nv_bfloat16, float, /*EltPerPack=*/2> {
|
||||
__device__ __forceinline__ static BytePack<4*2> cast(BytePack<2*2> a) {
|
||||
return toPack(__bfloat1622float2(fromPack<__nv_bfloat162>(a)));
|
||||
}
|
||||
};
|
||||
template<>
|
||||
struct Apply_Cast<float ,__nv_bfloat16, /*EltPerPack=*/2> {
|
||||
__device__ __forceinline__ static BytePack<2*2> cast(BytePack<4*2> a) {
|
||||
return toPack(__float22bfloat162_rn(fromPack<float2>(a)));
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
#define EASY_CAST(A, B, EltPerPack, VecA, VecB) \
|
||||
template<> \
|
||||
struct Apply_Cast<A, B, EltPerPack> { \
|
||||
__device__ __forceinline__ static BytePack<sizeof(B)*EltPerPack> cast(BytePack<sizeof(A)*EltPerPack> a) { \
|
||||
return toPack(VecB(fromPack<VecA>(a))); \
|
||||
} \
|
||||
}; \
|
||||
template<> \
|
||||
struct Apply_Cast<B, A, EltPerPack> { \
|
||||
__device__ __forceinline__ static BytePack<sizeof(A)*EltPerPack> cast(BytePack<sizeof(B)*EltPerPack> b) { \
|
||||
return toPack(VecA(fromPack<VecB>(b))); \
|
||||
} \
|
||||
};
|
||||
|
||||
#if defined(__CUDA_FP8_TYPES_EXIST__)
|
||||
EASY_CAST(__nv_fp8_e5m2, float, 2, __nv_fp8x2_e5m2, float2)
|
||||
EASY_CAST(__nv_fp8_e5m2, float, 4, __nv_fp8x4_e5m2, float4)
|
||||
|
||||
EASY_CAST(__nv_fp8_e4m3, float, 2, __nv_fp8x2_e4m3, float2)
|
||||
EASY_CAST(__nv_fp8_e4m3, float, 4, __nv_fp8x4_e4m3, float4)
|
||||
#endif
|
||||
#undef EASY_CAST
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Apply_Reduce
|
||||
|
||||
// Nonsensical base case
|
||||
template<typename Fn>
|
||||
struct Apply_Reduce<Fn, /*EltPerPack=*/0> {
|
||||
__device__ static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) {
|
||||
__device__ __forceinline__ static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) {
|
||||
return {};
|
||||
}
|
||||
};
|
||||
@@ -168,7 +301,7 @@ struct Apply_Reduce<Fn, /*EltPerPack=*/0> {
|
||||
template<typename Fn, int EltPerPack>
|
||||
struct Apply_Reduce {
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> reduce(Fn fn, BytePack<Size> a, BytePack<Size> b) {
|
||||
__device__ __forceinline__ static BytePack<Size> reduce(Fn fn, BytePack<Size> a, BytePack<Size> b) {
|
||||
a.half[0] = Apply_Reduce<Fn, EltPerPack/2>::reduce(fn, a.half[0], b.half[0]);
|
||||
a.half[1] = Apply_Reduce<Fn, EltPerPack/2>::reduce(fn, a.half[1], b.half[1]);
|
||||
return a;
|
||||
@@ -178,25 +311,25 @@ struct Apply_Reduce {
|
||||
// Base case definitions (EltPerPack == 1)
|
||||
template<typename T>
|
||||
struct Apply_Reduce<FuncCopy<T>, /*EltPerPack=*/1> {
|
||||
__device__ static BytePack<sizeof(T)> reduce(FuncCopy<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncCopy<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
return a;
|
||||
}
|
||||
};
|
||||
template<typename T>
|
||||
struct Apply_Reduce<FuncSum<T>, /*EltPerPack=*/1> {
|
||||
__device__ static BytePack<sizeof(T)> reduce(FuncSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
return toPack<T>(fromPack<T>(a) + fromPack<T>(b));
|
||||
}
|
||||
};
|
||||
template<typename T>
|
||||
struct Apply_Reduce<FuncProd<T>, /*EltPerPack=*/1> {
|
||||
__device__ static BytePack<sizeof(T)> reduce(FuncProd<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncProd<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
return toPack<T>(fromPack<T>(a) * fromPack<T>(b));
|
||||
}
|
||||
};
|
||||
template<typename T>
|
||||
struct Apply_Reduce<FuncMinMax<T>, /*EltPerPack=*/1> {
|
||||
__device__ static BytePack<sizeof(T)> reduce(FuncMinMax<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> reduce(FuncMinMax<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
|
||||
return (a.native ^ fn.xormask.native) < (b.native ^ fn.xormask.native) ? a : b;
|
||||
}
|
||||
};
|
||||
@@ -204,7 +337,7 @@ struct Apply_Reduce<FuncMinMax<T>, /*EltPerPack=*/1> {
|
||||
// Optimizations for specfic types and element count combinations:
|
||||
template<>
|
||||
struct Apply_Reduce<FuncSum<uint8_t>, /*EltPerPack=*/4> {
|
||||
__device__ static BytePack<4> reduce(FuncSum<uint8_t> fn, BytePack<4> a, BytePack<4> b) {
|
||||
__device__ __forceinline__ static BytePack<4> reduce(FuncSum<uint8_t> fn, BytePack<4> a, BytePack<4> b) {
|
||||
constexpr uint32_t even = 0x00ff00ffu;
|
||||
uint32_t x = (a.native & even) + (b.native & even);
|
||||
uint32_t y = (a.native & ~even) + (b.native & ~even);
|
||||
@@ -240,7 +373,7 @@ struct Apply_Reduce<FuncMinMax<uint8_t>, /*EltPerPack=*/4> {
|
||||
|
||||
// template<>
|
||||
// struct Apply_Reduce<FuncProd<uint8_t>, /*EltPerPack=*/4> {
|
||||
// __device__ static BytePack<4> reduce(FuncProd<uint8_t> fn, BytePack<4> apack, BytePack<4> bpack) {
|
||||
// __device__ __forceinline__ static BytePack<4> reduce(FuncProd<uint8_t> fn, BytePack<4> apack, BytePack<4> bpack) {
|
||||
// uint32_t a = apack.native;
|
||||
// uint32_t b = bpack.native;
|
||||
// uint32_t ab0 = (a*b) & 0xffu;
|
||||
@@ -326,7 +459,7 @@ template<typename Fn, int EltPerPack>
|
||||
struct Apply_PreOp {
|
||||
static constexpr bool IsIdentity = Apply_PreOp<Fn, EltPerPack/2>::IsIdentity;
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
|
||||
__device__ __forceinline__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
|
||||
#if __cpp_if_constexpr
|
||||
if constexpr(!IsIdentity) {
|
||||
#else
|
||||
@@ -346,7 +479,7 @@ template<typename Fn>
|
||||
struct Apply_PreOp<Fn, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
|
||||
__device__ __forceinline__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
|
||||
return a;
|
||||
}
|
||||
};
|
||||
@@ -354,7 +487,7 @@ struct Apply_PreOp<Fn, /*EltPerPack=*/1> {
|
||||
template<typename Fn>
|
||||
struct Apply_PreOp<Fn, /*EltPerPack=*/0> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
__device__ static BytePack<0> preOp(Fn fn, BytePack<0> a) {
|
||||
__device__ __forceinline__ static BytePack<0> preOp(Fn fn, BytePack<0> a) {
|
||||
return {};
|
||||
}
|
||||
};
|
||||
@@ -367,7 +500,7 @@ template<typename Fn, int EltPerPack>
|
||||
struct Apply_PostOp {
|
||||
static constexpr bool IsIdentity = Apply_PostOp<Fn, EltPerPack/2>::IsIdentity;
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
|
||||
__device__ __forceinline__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
|
||||
#if __cpp_if_constexpr
|
||||
if constexpr(!IsIdentity) {
|
||||
#else
|
||||
@@ -387,7 +520,7 @@ template<typename Fn>
|
||||
struct Apply_PostOp<Fn, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
template<int Size>
|
||||
__device__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
|
||||
__device__ __forceinline__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
|
||||
return a;
|
||||
}
|
||||
};
|
||||
@@ -395,7 +528,7 @@ struct Apply_PostOp<Fn, /*EltPerPack=*/1> {
|
||||
template<typename Fn>
|
||||
struct Apply_PostOp<Fn, /*EltPerPack=*/0> {
|
||||
static constexpr bool IsIdentity = true;
|
||||
__device__ static BytePack<0> postOp(Fn fn, BytePack<0> a) {
|
||||
__device__ __forceinline__ static BytePack<0> postOp(Fn fn, BytePack<0> a) {
|
||||
return {};
|
||||
}
|
||||
};
|
||||
@@ -407,7 +540,7 @@ struct Apply_PostOp<Fn, /*EltPerPack=*/0> {
|
||||
template<typename T>
|
||||
struct RedOpArg<FuncPreMulSum<T>> {
|
||||
static constexpr bool ArgUsed = true;
|
||||
__device__ static uint64_t loadArg(void *ptr) {
|
||||
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
|
||||
union { uint64_t u64; T val; };
|
||||
u64 = 0;
|
||||
val = *(T*)ptr;
|
||||
@@ -420,7 +553,7 @@ template<typename T>
|
||||
struct FuncPreMulSum {
|
||||
using EltType = T;
|
||||
T scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; T val; };
|
||||
u64 = opArg;
|
||||
scalar = val;
|
||||
@@ -434,7 +567,7 @@ template<>
|
||||
struct FuncPreMulSum<half> {
|
||||
using EltType = half;
|
||||
half2 scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; __half val; };
|
||||
u64 = opArg;
|
||||
scalar.x = val;
|
||||
@@ -451,7 +584,7 @@ struct FuncPreMulSum<half> {
|
||||
using EltType = hip_bfloat16;
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
__nv_bfloat162 scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; __nv_bfloat16 val; };
|
||||
u64 = opArg;
|
||||
scalar.x = val;
|
||||
@@ -459,7 +592,7 @@ struct FuncPreMulSum<half> {
|
||||
}
|
||||
#else
|
||||
float scalar;
|
||||
__device__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg=0) {
|
||||
union { uint64_t u64; hip_bfloat16 val; };
|
||||
u64 = opArg;
|
||||
scalar = (float)(val);
|
||||
@@ -474,7 +607,7 @@ struct FuncPreMulSum<half> {
|
||||
struct FuncPreMulSum<__nv_fp8_e4m3> {
|
||||
using EltType = __nv_fp8_e4m3;
|
||||
__half2 scalar2;
|
||||
__device__ FuncPreMulSum(uint64_t opArg) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg) {
|
||||
union { uint64_t u64; __nv_fp8_storage_t val; };
|
||||
u64 = opArg;
|
||||
scalar2.x = __half(__nv_cvt_fp8_to_halfraw(val, __NV_E4M3));
|
||||
@@ -486,7 +619,7 @@ struct FuncPreMulSum<half> {
|
||||
struct FuncPreMulSum<__nv_fp8_e5m2> {
|
||||
using EltType = __nv_fp8_e5m2;
|
||||
__half2 scalar2;
|
||||
__device__ FuncPreMulSum(uint64_t opArg) {
|
||||
__device__ __forceinline__ FuncPreMulSum(uint64_t opArg) {
|
||||
union { uint64_t u64; __nv_fp8_storage_t val; };
|
||||
u64 = opArg;
|
||||
scalar2.x = __half(__nv_cvt_fp8_to_halfraw(val, __NV_E5M2));
|
||||
@@ -528,7 +661,7 @@ struct FuncPreMulSum<half> {
|
||||
|
||||
template<typename T, int EltPerPack>
|
||||
struct Apply_Reduce<FuncPreMulSum<T>, EltPerPack> {
|
||||
__device__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
|
||||
// FuncPreMulSum reduce dispatches to FuncSum.
|
||||
return Apply_Reduce<FuncSum<T>, EltPerPack>::reduce(FuncSum<T>(), a, b);
|
||||
}
|
||||
@@ -538,7 +671,7 @@ struct Apply_Reduce<FuncPreMulSum<T>, EltPerPack> {
|
||||
template<typename T>
|
||||
struct Apply_PreOp<FuncPreMulSum<T>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(T)> preOp(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> preOp(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a) {
|
||||
return toPack<T>(fromPack<T>(a) * fn.scalar);
|
||||
}
|
||||
};
|
||||
@@ -549,7 +682,7 @@ struct Apply_PreOp<FuncPreMulSum<T>, /*EltPerPack=*/1> {
|
||||
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) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(half)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half)> a) {
|
||||
return toPack<half>(__hmul(fromPack<half>(a), fn.scalar.x));
|
||||
}
|
||||
};
|
||||
@@ -557,7 +690,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/2> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(half2)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half2)> a) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(half2)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half2)> a) {
|
||||
return toPack<half2>(__hmul2(fromPack<half2>(a), fn.scalar));
|
||||
}
|
||||
};
|
||||
@@ -570,7 +703,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<hip_bfloat16>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(hip_bfloat16)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(hip_bfloat16)> preOp(
|
||||
FuncPreMulSum<hip_bfloat16> fn, BytePack<sizeof(hip_bfloat16)> a
|
||||
) {
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
@@ -584,7 +717,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<hip_bfloat16>, /*EltPerPack=*/2> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_bfloat162)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_bfloat162)> preOp(
|
||||
FuncPreMulSum<__nv_bfloat16> fn, BytePack<sizeof(__nv_bfloat162)> a
|
||||
) {
|
||||
return toPack<__nv_bfloat162>(__hmul2(fromPack<__nv_bfloat162>(a), fn.scalar));
|
||||
@@ -601,7 +734,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e4m3>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_fp8_e4m3)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8_e4m3)> preOp(
|
||||
FuncPreMulSum<__nv_fp8_e4m3> fn, BytePack<sizeof(__nv_fp8_e4m3)> a
|
||||
) {
|
||||
return toPack<__nv_fp8_e4m3>(__nv_fp8_e4m3(__hmul(__half(fromPack<__nv_fp8_e4m3>(a)), fn.scalar2.x)));
|
||||
@@ -610,7 +743,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e4m3>, /*EltPerPack=*/2> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_fp8x2_e4m3)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8x2_e4m3)> preOp(
|
||||
FuncPreMulSum<__nv_fp8_e4m3> fn, BytePack<sizeof(__nv_fp8x2_e4m3)> a
|
||||
) {
|
||||
return toPack<__nv_fp8x2_e4m3>(__nv_fp8x2_e4m3(__hmul2(__half2(fromPack<__nv_fp8x2_e4m3>(a)), fn.scalar2)));
|
||||
@@ -620,7 +753,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e5m2>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_fp8_e5m2)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8_e5m2)> preOp(
|
||||
FuncPreMulSum<__nv_fp8_e5m2> fn, BytePack<sizeof(__nv_fp8_e5m2)> a
|
||||
) {
|
||||
return toPack<__nv_fp8_e5m2>(__nv_fp8_e5m2(__hmul(__half(fromPack<__nv_fp8_e5m2>(a)), fn.scalar2.x)));
|
||||
@@ -629,7 +762,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<>
|
||||
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e5m2>, /*EltPerPack=*/2> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(__nv_fp8x2_e5m2)> preOp(
|
||||
__device__ __forceinline__ static BytePack<sizeof(__nv_fp8x2_e5m2)> preOp(
|
||||
FuncPreMulSum<__nv_fp8_e5m2> fn, BytePack<sizeof(__nv_fp8x2_e5m2)> a
|
||||
) {
|
||||
return toPack<__nv_fp8x2_e5m2>(__nv_fp8x2_e5m2(__hmul2(__half2(fromPack<__nv_fp8x2_e5m2>(a)), fn.scalar2)));
|
||||
@@ -666,7 +799,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
|
||||
template<typename T>
|
||||
struct RedOpArg<FuncSumPostDiv<T>> {
|
||||
static constexpr bool ArgUsed = true;
|
||||
__device__ static uint64_t loadArg(void *ptr) {
|
||||
__device__ __forceinline__ static uint64_t loadArg(void *ptr) {
|
||||
return *(uint64_t*)ptr;
|
||||
}
|
||||
};
|
||||
@@ -709,12 +842,12 @@ struct FuncSumPostDiv {
|
||||
uint32_t divisor:31, isSigned:1;
|
||||
UintType recip;
|
||||
|
||||
__device__ FuncSumPostDiv(uint64_t opArg=0) {
|
||||
__device__ __forceinline__ FuncSumPostDiv(uint64_t opArg=0) {
|
||||
isSigned = opArg & 1;
|
||||
divisor = opArg >> 1;
|
||||
recip = Divider<UintType>::divide(UintType(-1), divisor);
|
||||
}
|
||||
__device__ T divide(T x) {
|
||||
__device__ __forceinline__ T divide(T x) {
|
||||
// x is negative iff we are in signed mode and the top bit is set
|
||||
bool xneg = isSigned && (x & ~(T(-1)>>1));
|
||||
// Compute abs(x):
|
||||
@@ -736,7 +869,7 @@ struct FuncSumPostDiv {
|
||||
template<typename T, int EltPerPack>
|
||||
struct Apply_Reduce<FuncSumPostDiv<T>, EltPerPack>:
|
||||
Apply_Reduce<FuncSum<T>, EltPerPack> {
|
||||
__device__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
|
||||
__device__ __forceinline__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
|
||||
// FuncSumPostDiv reduce dispatches to FuncSum.
|
||||
return Apply_Reduce<FuncSum<T>, EltPerPack>::reduce(FuncSum<T>(), a, b);
|
||||
}
|
||||
@@ -745,7 +878,7 @@ struct Apply_Reduce<FuncSumPostDiv<T>, EltPerPack>:
|
||||
template<typename T>
|
||||
struct Apply_PostOp<FuncSumPostDiv<T>, /*EltPerPack=*/1> {
|
||||
static constexpr bool IsIdentity = false;
|
||||
__device__ static BytePack<sizeof(T)> postOp(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a) {
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> postOp(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a) {
|
||||
return toPack<T>(fn.divide(fromPack<T>(a)));
|
||||
}
|
||||
};
|
||||
@@ -753,120 +886,145 @@ struct Apply_PostOp<FuncSumPostDiv<T>, /*EltPerPack=*/1> {
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Apply_LoadMultimem
|
||||
|
||||
#define SIZEOF_BytePack_field_u16 2
|
||||
#define PTX_REG_BytePack_field_u16 "h"
|
||||
#define RegCode_for_size_1 "r"
|
||||
#define RegCode_for_size_2 "h"
|
||||
#define RegCode_for_size_4 "r"
|
||||
#define RegCode_for_size_8 "l"
|
||||
|
||||
#define SIZEOF_BytePack_field_u32 4
|
||||
#define PTX_REG_BytePack_field_u32 "r"
|
||||
#define RegSize_for_size_1 4
|
||||
#define RegSize_for_size_2 2
|
||||
#define RegSize_for_size_4 4
|
||||
#define RegSize_for_size_8 8
|
||||
|
||||
#define SIZEOF_BytePack_field_u64 8
|
||||
#define PTX_REG_BytePack_field_u64 "l"
|
||||
#define PtxAcc_for_u32
|
||||
#define PtxAcc_for_s32
|
||||
#define PtxAcc_for_s64
|
||||
#define PtxAcc_for_u64
|
||||
#define PtxAcc_for_f32
|
||||
#define PtxAcc_for_f64
|
||||
#if CUDART_VERSION >= 12020
|
||||
#define PtxAcc_for_f16 ".acc::f32"
|
||||
#define PtxAcc_for_bf16 ".acc::f32"
|
||||
#define PtxAcc_for_f16x2 ".acc::f32"
|
||||
#define PtxAcc_for_bf16x2 ".acc::f32"
|
||||
#else
|
||||
#define PtxAcc_for_f16
|
||||
#define PtxAcc_for_bf16
|
||||
#define PtxAcc_for_f16x2
|
||||
#define PtxAcc_for_bf16x2
|
||||
#endif
|
||||
#define PtxAcc_for_e4m3 ".acc::f16"
|
||||
#define PtxAcc_for_e5m2 ".acc::f16"
|
||||
#define PtxAcc_for_e4m3x4 ".acc::f16"
|
||||
#define PtxAcc_for_e5m2x4 ".acc::f16"
|
||||
|
||||
#define DEFINE_Apply_LoadMultimem_sum(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_sum(T, ptx_ty, PackSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, SIZEOF_BytePack_field_##pack_field> { \
|
||||
static constexpr int PackSize = SIZEOF_BytePack_field_##pack_field; \
|
||||
__device__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
BytePack<PackSize> ans; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, PackSize> { \
|
||||
__device__ __forceinline__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
BytePack<RegSize_for_size_##PackSize> reg; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty "." #ptx_ty " %0, [%1];" \
|
||||
: "=" RegCode_for_size_##PackSize(reg.native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
BytePack<PackSize> ans; \
|
||||
ans.native = reg.native; \
|
||||
return ans; \
|
||||
} \
|
||||
};
|
||||
#define DEFINE_Apply_LoadMultimem_minmax(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_minmax(T, ptx_ty, PackSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, SIZEOF_BytePack_field_##pack_field> { \
|
||||
static constexpr int PackSize = SIZEOF_BytePack_field_##pack_field; \
|
||||
__device__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
BytePack<PackSize> ans; \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, PackSize> { \
|
||||
__device__ __forceinline__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
BytePack<RegSize_for_size_##PackSize> reg; \
|
||||
if (fn.isMinNotMax) { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.min." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
|
||||
: "=" RegCode_for_size_##PackSize(reg.native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
} else { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.max." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
|
||||
: "=" RegCode_for_size_##PackSize(reg.native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
} \
|
||||
BytePack<PackSize> ans; \
|
||||
ans.native = reg.native; \
|
||||
return ans; \
|
||||
} \
|
||||
};
|
||||
|
||||
#define DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, VecEltSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, 4*(SIZEOF_BytePack_field_##pack_field)> { \
|
||||
static constexpr int PackSize = 4*(SIZEOF_BytePack_field_##pack_field); \
|
||||
__device__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
BytePack<PackSize> ans; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, 4*(VecEltSize)> { \
|
||||
static constexpr int PackSize = 4*(VecEltSize); \
|
||||
__device__ __forceinline__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
union { BytePack<PackSize> ans; BytePack<VecEltSize> elts[4]; }; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty ".v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
|
||||
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
return ans; \
|
||||
} \
|
||||
};
|
||||
#define DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, VecEltSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, 4*(SIZEOF_BytePack_field_##pack_field)> { \
|
||||
static constexpr int PackSize = 4*(SIZEOF_BytePack_field_##pack_field); \
|
||||
__device__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
BytePack<PackSize> ans; \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, 4*(VecEltSize)> { \
|
||||
static constexpr int PackSize = 4*(VecEltSize); \
|
||||
__device__ __forceinline__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
union { BytePack<PackSize> ans; BytePack<VecEltSize> elts[4]; }; \
|
||||
if (fn.isMinNotMax) { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.min.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
|
||||
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
} else { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.max.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
|
||||
"=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
|
||||
: "=" RegCode_for_size_##VecEltSize(elts[0].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[1].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[2].native), \
|
||||
"=" RegCode_for_size_##VecEltSize(elts[3].native) \
|
||||
: "l"(addr) : "memory"); \
|
||||
} \
|
||||
return ans; \
|
||||
} \
|
||||
};
|
||||
|
||||
#define DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(T, ptx_ty, pack_field) \
|
||||
DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(T, ptx_ty, VecEltSize) \
|
||||
DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, VecEltSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncSum<T>, sizeof(T)> { \
|
||||
__device__ static BytePack<sizeof(T)> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
BytePack<2*sizeof(T)> tmp; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
|
||||
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
|
||||
return tmp.half[(addr/sizeof(T))%2]; \
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> load(FuncSum<T> fn, uintptr_t addr) { \
|
||||
union { BytePack<VecEltSize> tmp; BytePack<sizeof(T)> elts[(VecEltSize)/sizeof(T)]; }; \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.add" PtxAcc_for_##ptx_ty "." #ptx_ty " %0, [%1];" \
|
||||
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
|
||||
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
|
||||
return elts[(addr/sizeof(T))%((VecEltSize)/sizeof(T))]; \
|
||||
} \
|
||||
};
|
||||
#define DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(T, ptx_ty, pack_field) \
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, pack_field) \
|
||||
#define DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(T, ptx_ty, VecEltSize) \
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, VecEltSize) \
|
||||
template<> \
|
||||
struct Apply_LoadMultimem<FuncMinMax<T>, sizeof(T)> { \
|
||||
__device__ static BytePack<sizeof(T)> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
BytePack<2*sizeof(T)> tmp; \
|
||||
__device__ __forceinline__ static BytePack<sizeof(T)> load(FuncMinMax<T> fn, uintptr_t addr) { \
|
||||
union { BytePack<VecEltSize> tmp; BytePack<sizeof(T)> elts[(VecEltSize)/sizeof(T)]; }; \
|
||||
if (fn.isMinNotMax) { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.min." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
|
||||
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
|
||||
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
|
||||
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
|
||||
} else { \
|
||||
asm volatile("multimem.ld_reduce.relaxed.sys.global.max." #ptx_ty " %0, [%1];" \
|
||||
: "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
|
||||
: "l"(addr & -uintptr_t(2*sizeof(T))) : "memory"); \
|
||||
: "=" RegCode_for_size_##VecEltSize(tmp.native) \
|
||||
: "l"(addr & -uintptr_t(VecEltSize)) : "memory"); \
|
||||
} \
|
||||
return tmp.half[(addr/sizeof(T))%2]; \
|
||||
return elts[(addr/sizeof(T))%((VecEltSize)/sizeof(T))]; \
|
||||
} \
|
||||
};
|
||||
|
||||
template<typename Fn, int BytePerPack>
|
||||
struct Apply_LoadMultimem {
|
||||
__device__ static BytePack<BytePerPack> load(Fn fn, uintptr_t addr) {
|
||||
__device__ __forceinline__ static BytePack<BytePerPack> load(Fn fn, uintptr_t addr) {
|
||||
//__trap();
|
||||
return {};
|
||||
}
|
||||
@@ -889,29 +1047,43 @@ struct Apply_LoadMultimem {
|
||||
/*multimem.ld_reduce not supported:*/ 0;
|
||||
};
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(uint32_t, u32, u32)
|
||||
DEFINE_Apply_LoadMultimem_minmax(uint32_t, u32, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum(uint32_t, u32, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax(uint32_t, u32, 4)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(int32_t, s32, u32)
|
||||
DEFINE_Apply_LoadMultimem_minmax(int32_t, s32, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum(int32_t, s32, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax(int32_t, s32, 4)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(uint64_t, u64, u64)
|
||||
DEFINE_Apply_LoadMultimem_minmax(uint64_t, u64, u64)
|
||||
DEFINE_Apply_LoadMultimem_sum(uint64_t, u64, 8)
|
||||
DEFINE_Apply_LoadMultimem_minmax(uint64_t, u64, 8)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(int64_t, u64, u64)
|
||||
DEFINE_Apply_LoadMultimem_minmax(int64_t, s64, u64)
|
||||
DEFINE_Apply_LoadMultimem_sum(int64_t, u64, 8)
|
||||
DEFINE_Apply_LoadMultimem_minmax(int64_t, s64, 8)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(float, f32, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4(float, f32, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum(float, f32, 4)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4(float, f32, 4)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum(double, f64, u64)
|
||||
DEFINE_Apply_LoadMultimem_sum(double, f64, 8)
|
||||
|
||||
DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(half, f16x2, u32)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(half, f16x2, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(half, f16x2, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(half, f16x2, 4)
|
||||
|
||||
#if defined(RCCL_BFLOAT16)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(hip_bfloat16, bf16x2, u32)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(hip_bfloat16, bf16x2, u32)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(hip_bfloat16, bf16x2, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(hip_bfloat16, bf16x2, 4)
|
||||
#endif
|
||||
|
||||
#if defined(RCCL_BFLOAT16)
|
||||
#if NCCL_CUDA_ARCH_FAMILY_SPECIFIC == 1000 || NCCL_CUDA_ARCH_FAMILY_SPECIFIC == 1010 || NCCL_CUDA_ARCH_SPECIFIC == 1200 || NCCL_CUDA_ARCH_SPECIFIC == 1210
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(__nv_fp8_e4m3, e4m3x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(__nv_fp8_e4m3, e4m3x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(__nv_fp8_e5m2, e5m2x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(__nv_fp8_e5m2, e5m2x4, 4)
|
||||
#else
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(rccl_float8, e4m3x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(rccl_float8, e4m3x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_sum_v4_and_xparts(rccl_bfloat8, e5m2x4, 4)
|
||||
DEFINE_Apply_LoadMultimem_minmax_v4_and_xparts(rccl_bfloat8, e5m2x4, 4)
|
||||
#endif
|
||||
#endif
|
||||
#else
|
||||
template<typename Fn>
|
||||
@@ -923,11 +1095,29 @@ struct Apply_LoadMultimem {
|
||||
#undef DEFINE_Apply_LoadMultimem
|
||||
#undef DEFINE_Apply_LoadMultimem_v4
|
||||
#undef DEFINE_Apply_LoadMultimem_v4x2_and_subhalf
|
||||
#undef SIZEOF_BytePack_field_u64
|
||||
#undef PTX_REG_BytePack_field_u64
|
||||
#undef SIZEOF_BytePack_field_u32
|
||||
#undef PTX_REG_BytePack_field_u32
|
||||
#undef SIZEOF_BytePack_field_u16
|
||||
#undef PTX_REG_BytePack_field_u16
|
||||
|
||||
#undef RegCode_for_size_2
|
||||
#undef RegCode_for_size_4
|
||||
#undef RegCode_for_size_8
|
||||
|
||||
#undef RegSize_for_size_1
|
||||
#undef RegSize_for_size_2
|
||||
#undef RegSize_for_size_4
|
||||
#undef RegSize_for_size_8
|
||||
|
||||
#undef PtxAcc_for_u32
|
||||
#undef PtxAcc_for_s32
|
||||
#undef PtxAcc_for_s64
|
||||
#undef PtxAcc_for_u64
|
||||
#undef PtxAcc_for_f32
|
||||
#undef PtxAcc_for_f64
|
||||
#undef PtxAcc_for_f16
|
||||
#undef PtxAcc_for_bf16
|
||||
#undef PtxAcc_for_f16x2
|
||||
#undef PtxAcc_for_bf16x2
|
||||
#undef PtxAcc_for_e4m3
|
||||
#undef PtxAcc_for_e5m2
|
||||
#undef PtxAcc_for_e4m3x4
|
||||
#undef PtxAcc_for_e5m2x4
|
||||
|
||||
#endif // REDUCE_KERNEL_H_
|
||||
|
||||
Αναφορά σε νέο ζήτημα
Block a user