BertanDogancay
470 строки
14 KiB
C++
470 строки
14 KiB
C++
/*************************************************************************
|
|
* Copyright (c) 2019-2022, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
#ifndef NCCL_BITOPS_H_
|
|
#define NCCL_BITOPS_H_
|
|
|
|
#include <stdint.h>
|
|
#include <string.h>
|
|
|
|
#if !__NVCC__
|
|
#ifndef __host__
|
|
#define __host__
|
|
#endif
|
|
#ifndef __device__
|
|
#define __device__
|
|
#endif
|
|
#endif
|
|
|
|
template<typename Int>
|
|
constexpr static __host__ __device__ Int minval(Int a) { return a; }
|
|
template<typename Int, typename ...More>
|
|
constexpr static __host__ __device__ Int minval(Int a, Int b, More ...more) {
|
|
#if __CUDA_ARCH__
|
|
return minval(min(a, b), more...);
|
|
#else
|
|
return minval(a < b ? a : b, more...);
|
|
#endif
|
|
}
|
|
|
|
template<typename Int>
|
|
constexpr static __host__ __device__ Int maxval(Int a) { return a; }
|
|
template<typename Int, typename ...More>
|
|
constexpr static __host__ __device__ Int maxval(Int a, Int b, More ...more) {
|
|
#if __CUDA_ARCH__
|
|
return maxval(max(a, b), more...);
|
|
#else
|
|
return maxval(a > b ? a : b, more...);
|
|
#endif
|
|
}
|
|
|
|
#define DIVUP(x, y) \
|
|
(((x)+(y)-1)/(y))
|
|
|
|
#define ROUNDUP(x, y) \
|
|
(DIVUP((x), (y))*(y))
|
|
|
|
#define ALIGN_POWER(x, y) \
|
|
((x) > (y) ? ROUNDUP(x, y) : ((y)/((y)/(x))))
|
|
|
|
#define ALIGN_SIZE(size, align) \
|
|
size = ((size + (align) - 1) / (align)) * (align);
|
|
|
|
template<typename X, typename Y, typename Z = decltype(X()+Y())>
|
|
static __host__ __device__ constexpr Z divUp(X x, Y y) {
|
|
return (x+y-1)/y;
|
|
}
|
|
|
|
template<typename X, typename Y, typename Z = decltype(X()+Y())>
|
|
static __host__ __device__ constexpr Z roundUp(X x, Y y) {
|
|
return (x+y-1) - (x+y-1)%y;
|
|
}
|
|
template<typename X, typename Y, typename Z = decltype(X()+Y())>
|
|
static __host__ __device__ constexpr Z roundDown(X x, Y y) {
|
|
return x - x%y;
|
|
}
|
|
|
|
// assumes second argument is a power of 2
|
|
template<typename X, typename Z = decltype(X()+int())>
|
|
static __host__ __device__ constexpr Z alignUp(X x, int a) {
|
|
return (x + a-1) & Z(-a);
|
|
}
|
|
// assumes second argument is a power of 2
|
|
template<typename X, typename Z = decltype(X()+int())>
|
|
static __host__ __device__ constexpr Z alignDown(X x, int a) {
|
|
return x & Z(-a);
|
|
}
|
|
|
|
template<typename Int>
|
|
constexpr __host__ __device__ bool isPow2(Int x) {
|
|
return (x & (x-1)) == 0;
|
|
}
|
|
|
|
template<typename T>
|
|
static __host__ __device__ T add4G(T base, int delta4G) {
|
|
union { T tmp; uint32_t u32[2]; };
|
|
tmp = base;
|
|
u32[1] += delta4G;
|
|
return tmp;
|
|
}
|
|
|
|
template<typename T>
|
|
static __host__ __device__ T incWrap4G(T ptr, uint32_t delta4G, uint32_t lo4G, uint32_t hi4G) {
|
|
union { T tmp; uint32_t u32[2]; };
|
|
tmp = ptr;
|
|
u32[1] += delta4G;
|
|
if (u32[1] >= hi4G) u32[1] -= hi4G-lo4G;
|
|
return tmp;
|
|
}
|
|
|
|
template<typename T>
|
|
static __host__ __device__ T decWrap4G(T ptr, uint32_t delta4G, uint32_t lo4G, uint32_t hi4G) {
|
|
union { T tmp; uint32_t u32[2]; };
|
|
tmp = ptr;
|
|
u32[1] -= delta4G;
|
|
if (u32[1] < lo4G) u32[1] += hi4G-lo4G;
|
|
return tmp;
|
|
}
|
|
|
|
// Produce the reciprocal of x for use in idivByRcp
|
|
constexpr __host__ __device__ uint32_t idivRcp32(uint32_t x) {
|
|
return uint32_t(uint64_t(0x100000000)/x);
|
|
}
|
|
constexpr __host__ __device__ uint64_t idivRcp64(uint64_t x) {
|
|
return uint64_t(-1)/x + isPow2(x);
|
|
}
|
|
|
|
static __host__ __device__ uint32_t mul32hi(uint32_t a, uint32_t b) {
|
|
#if __CUDA_ARCH__
|
|
return __umulhi(a, b);
|
|
#else
|
|
return uint64_t(a)*b >> 32;
|
|
#endif
|
|
}
|
|
static __host__ __device__ uint64_t mul64hi(uint64_t a, uint64_t b) {
|
|
#if __CUDA_ARCH__
|
|
return __umul64hi(a, b);
|
|
#else
|
|
return (uint64_t)(((unsigned __int128)a)*b >> 64);
|
|
#endif
|
|
}
|
|
|
|
// Produce the reciprocal of x*y given their respective reciprocals. This incurs
|
|
// no integer division on device.
|
|
static __host__ __device__ uint32_t imulRcp32(uint32_t x, uint32_t xrcp, uint32_t y, uint32_t yrcp) {
|
|
if (xrcp == 0) return yrcp;
|
|
if (yrcp == 0) return xrcp;
|
|
uint32_t rcp = mul32hi(xrcp, yrcp);
|
|
uint32_t rem = -x*y*rcp;
|
|
if (x*y <= rem) rcp += 1;
|
|
return rcp;
|
|
}
|
|
static __host__ __device__ uint64_t imulRcp64(uint64_t x, uint64_t xrcp, uint64_t y, uint64_t yrcp) {
|
|
if (xrcp == 0) return yrcp;
|
|
if (yrcp == 0) return xrcp;
|
|
uint64_t rcp = mul64hi(xrcp, yrcp);
|
|
uint64_t rem = -x*y*rcp;
|
|
if (x*y <= rem) rcp += 1;
|
|
return rcp;
|
|
}
|
|
|
|
// Fast integer division where divisor has precomputed reciprocal.
|
|
// idivFast(x, y, idivRcp(y)) == x/y
|
|
static __host__ __device__ void idivmodFast32(uint32_t *quo, uint32_t *rem, uint32_t x, uint32_t y, uint32_t yrcp) {
|
|
uint32_t q = x, r = 0;
|
|
if (yrcp != 0) {
|
|
q = mul32hi(x, yrcp);
|
|
r = x - y*q;
|
|
if (r >= y) { q += 1; r -= y; }
|
|
}
|
|
*quo = q;
|
|
*rem = r;
|
|
}
|
|
static __host__ __device__ void idivmodFast64(uint64_t *quo, uint64_t *rem, uint64_t x, uint64_t y, uint64_t yrcp) {
|
|
uint64_t q = x, r = 0;
|
|
if (yrcp != 0) {
|
|
q = mul64hi(x, yrcp);
|
|
r = x - y*q;
|
|
if (r >= y) { q += 1; r -= y; }
|
|
}
|
|
*quo = q;
|
|
*rem = r;
|
|
}
|
|
|
|
static __host__ __device__ uint32_t idivFast32(uint32_t x, uint32_t y, uint32_t yrcp) {
|
|
uint32_t q, r;
|
|
idivmodFast32(&q, &r, x, y, yrcp);
|
|
return q;
|
|
}
|
|
static __host__ __device__ uint32_t idivFast64(uint64_t x, uint64_t y, uint64_t yrcp) {
|
|
uint64_t q, r;
|
|
idivmodFast64(&q, &r, x, y, yrcp);
|
|
return q;
|
|
}
|
|
|
|
static __host__ __device__ uint32_t imodFast32(uint32_t x, uint32_t y, uint32_t yrcp) {
|
|
uint32_t q, r;
|
|
idivmodFast32(&q, &r, x, y, yrcp);
|
|
return r;
|
|
}
|
|
static __host__ __device__ uint32_t imodFast64(uint64_t x, uint64_t y, uint64_t yrcp) {
|
|
uint64_t q, r;
|
|
idivmodFast64(&q, &r, x, y, yrcp);
|
|
return r;
|
|
}
|
|
|
|
template<typename Int>
|
|
static __host__ __device__ int countOneBits(Int x) {
|
|
#if __CUDA_ARCH__
|
|
if (sizeof(Int) <= sizeof(unsigned int)) {
|
|
return __popc((unsigned int)x);
|
|
} else if (sizeof(Int) <= sizeof(unsigned long long)) {
|
|
return __popcll((unsigned long long)x);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(unsigned long long), "Unsupported integer size.");
|
|
return -1;
|
|
}
|
|
#else
|
|
if (sizeof(Int) <= sizeof(unsigned int)) {
|
|
return __builtin_popcount((unsigned int)x);
|
|
} else if (sizeof(Int) <= sizeof(unsigned long)) {
|
|
return __builtin_popcountl((unsigned long)x);
|
|
} else if (sizeof(Int) <= sizeof(unsigned long long)) {
|
|
return __builtin_popcountll((unsigned long long)x);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(unsigned long long), "Unsupported integer size.");
|
|
return -1;
|
|
}
|
|
#endif
|
|
}
|
|
|
|
// Returns index of first one bit or returns -1 if mask is zero.
|
|
template<typename Int>
|
|
static __host__ __device__ int firstOneBit(Int mask) {
|
|
int i;
|
|
#if __CUDA_ARCH__
|
|
if (sizeof(Int) <= sizeof(int)) {
|
|
i = __ffs((int)mask);
|
|
} else if (sizeof(Int) <= sizeof(long long)) {
|
|
i = __ffsll((long long)mask);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(long long), "Unsupported integer size.");
|
|
}
|
|
#else
|
|
if (sizeof(Int) <= sizeof(int)) {
|
|
i = __builtin_ffs((int)mask);
|
|
} else if (sizeof(Int) <= sizeof(long)) {
|
|
i = __builtin_ffsl((long)mask);
|
|
} else if (sizeof(Int) <= sizeof(long long)) {
|
|
i = __builtin_ffsll((long long)mask);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(long long), "Unsupported integer size.");
|
|
}
|
|
#endif
|
|
return i-1;
|
|
}
|
|
|
|
template<typename Int>
|
|
static __host__ __device__ int popFirstOneBit(Int* mask) {
|
|
Int tmp = *mask;
|
|
*mask &= *mask-1;
|
|
return firstOneBit(tmp);
|
|
}
|
|
|
|
template<typename Int>
|
|
static __host__ __device__ int log2Down(Int x) {
|
|
int w, n;
|
|
#if __CUDA_ARCH__
|
|
if (sizeof(Int) <= sizeof(int)) {
|
|
w = 8*sizeof(int);
|
|
n = __clz((int)x);
|
|
} else if (sizeof(Int) <= sizeof(long long)) {
|
|
w = 8*sizeof(long long);
|
|
n = __clzll((long long)x);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(long long), "Unsupported integer size.");
|
|
}
|
|
#else
|
|
if (x == 0) {
|
|
return -1;
|
|
} else if (sizeof(Int) <= sizeof(unsigned int)) {
|
|
w = 8*sizeof(unsigned int);
|
|
n = __builtin_clz((unsigned int)x);
|
|
} else if (sizeof(Int) <= sizeof(unsigned long)) {
|
|
w = 8*sizeof(unsigned long);
|
|
n = __builtin_clzl((unsigned long)x);
|
|
} else if (sizeof(Int) <= sizeof(unsigned long long)) {
|
|
w = 8*sizeof(unsigned long long);
|
|
n = __builtin_clzll((unsigned long long)x);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(unsigned long long), "Unsupported integer size.");
|
|
}
|
|
#endif
|
|
return (w-1)-n;
|
|
}
|
|
|
|
template<typename Int>
|
|
static __host__ __device__ int log2Up(Int x) {
|
|
int w, n;
|
|
if (x != 0) x -= 1;
|
|
#if __CUDA_ARCH__
|
|
if (sizeof(Int) <= sizeof(int)) {
|
|
w = 8*sizeof(int);
|
|
n = __clz((int)x);
|
|
} else if (sizeof(Int) <= sizeof(long long)) {
|
|
w = 8*sizeof(long long);
|
|
n = __clzll((long long)x);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(long long), "Unsupported integer size.");
|
|
}
|
|
#else
|
|
if (x == 0) {
|
|
return 0;
|
|
} else if (sizeof(Int) <= sizeof(unsigned int)) {
|
|
w = 8*sizeof(unsigned int);
|
|
n = __builtin_clz((unsigned int)x);
|
|
} else if (sizeof(Int) <= sizeof(unsigned long)) {
|
|
w = 8*sizeof(unsigned long);
|
|
n = __builtin_clzl((unsigned long)x);
|
|
} else if (sizeof(Int) <= sizeof(unsigned long long)) {
|
|
w = 8*sizeof(unsigned long long);
|
|
n = __builtin_clzll((unsigned long long)x);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(unsigned long long), "Unsupported integer size.");
|
|
}
|
|
#endif
|
|
return w-n;
|
|
}
|
|
|
|
template<typename Int>
|
|
static __host__ __device__ Int pow2Up(Int x) {
|
|
return Int(1)<<log2Up(x);
|
|
}
|
|
|
|
template<typename Int>
|
|
static __host__ __device__ Int pow2Down(Int x) {
|
|
// True, log2Down can return -1, but we don't normally pass 0 as an argument...
|
|
// coverity[negative_shift]
|
|
return Int(1)<<log2Down(x);
|
|
}
|
|
|
|
template<typename UInt, int nSubBits>
|
|
static __host__ __device__ UInt reverseSubBits(UInt x) {
|
|
if (nSubBits >= 16 && 8*sizeof(UInt) == nSubBits) {
|
|
switch (8*sizeof(UInt)) {
|
|
case 16: x = __builtin_bswap16(x); break;
|
|
case 32: x = __builtin_bswap32(x); break;
|
|
case 64: x = __builtin_bswap64(x); break;
|
|
default: static_assert(8*sizeof(UInt) <= 64, "Unsupported integer type.");
|
|
}
|
|
return reverseSubBits<UInt, 8>(x);
|
|
} else if (nSubBits == 1) {
|
|
return x;
|
|
} else {
|
|
UInt m = UInt(-1)/((UInt(1)<<(nSubBits/2))+1);
|
|
x = (x & m)<<(nSubBits/2) | (x & ~m)>>(nSubBits/2);
|
|
return reverseSubBits<UInt, nSubBits/2>(x);
|
|
}
|
|
}
|
|
|
|
template<typename T> struct ncclToUnsigned;
|
|
template<> struct ncclToUnsigned<char> { using type = unsigned char; };
|
|
template<> struct ncclToUnsigned<signed char> { using type = unsigned char; };
|
|
template<> struct ncclToUnsigned<unsigned char> { using type = unsigned char; };
|
|
template<> struct ncclToUnsigned<signed short> { using type = unsigned short; };
|
|
template<> struct ncclToUnsigned<unsigned short> { using type = unsigned short; };
|
|
template<> struct ncclToUnsigned<signed int> { using type = unsigned int; };
|
|
template<> struct ncclToUnsigned<unsigned int> { using type = unsigned int; };
|
|
template<> struct ncclToUnsigned<signed long> { using type = unsigned long; };
|
|
template<> struct ncclToUnsigned<unsigned long> { using type = unsigned long; };
|
|
template<> struct ncclToUnsigned<signed long long> { using type = unsigned long long; };
|
|
template<> struct ncclToUnsigned<unsigned long long> { using type = unsigned long long; };
|
|
|
|
// Reverse the bottom nBits bits of x. The top bits will be overwritten with 0's.
|
|
template<typename Int>
|
|
static __host__ __device__ Int reverseBits(Int x, int nBits) {
|
|
using UInt = typename ncclToUnsigned<Int>::type;
|
|
union { UInt ux; Int sx; };
|
|
sx = x;
|
|
#if __CUDA_ARCH__
|
|
if (sizeof(Int) <= sizeof(unsigned int)) {
|
|
ux = __brev(ux);
|
|
} else if (sizeof(Int) <= sizeof(unsigned long long)) {
|
|
ux = __brevll(ux);
|
|
} else {
|
|
static_assert(sizeof(Int) <= sizeof(unsigned long long), "Unsupported integer type.");
|
|
}
|
|
#else
|
|
ux = reverseSubBits<UInt, 8*sizeof(UInt)>(ux);
|
|
#endif
|
|
ux = nBits==0 ? 0 : ux>>(8*sizeof(UInt)-nBits);
|
|
return sx;
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////////////
|
|
// Custom 8 bit floating point format for approximating 32 bit uints. This format
|
|
// has nearly the full range of uint32_t except it only keeps the top 3 bits
|
|
// beneath the leading 1 bit and thus has a max value of 0xf0000000.
|
|
|
|
static __host__ __device__ uint32_t u32fpEncode(uint32_t x, int bitsPerPow2) {
|
|
int log2x;
|
|
#if __CUDA_ARCH__
|
|
log2x = 31-__clz(x|1);
|
|
#else
|
|
log2x = 31-__builtin_clz(x|1);
|
|
#endif
|
|
uint32_t mantissa = x>>(log2x >= bitsPerPow2 ? log2x-bitsPerPow2 : 0) & ((1u<<bitsPerPow2)-1);
|
|
uint32_t exponent = log2x >= bitsPerPow2 ? log2x-(bitsPerPow2-1) : 0;
|
|
return exponent<<bitsPerPow2 | mantissa;
|
|
}
|
|
|
|
static __host__ __device__ uint32_t u32fpDecode(uint32_t x, int bitsPerPow2) {
|
|
uint32_t exponent = x>>bitsPerPow2;
|
|
uint32_t mantissa = (x & ((1u<<bitsPerPow2)-1)) | (exponent!=0 ? 0x8 : 0);
|
|
if (exponent != 0) exponent -= 1;
|
|
return mantissa<<exponent;
|
|
}
|
|
|
|
constexpr uint32_t u32fp8MaxValue() { return 0xf0000000; }
|
|
|
|
static __host__ __device__ uint8_t u32fp8Encode(uint32_t x) {
|
|
return u32fpEncode(x, 3);
|
|
}
|
|
static __host__ __device__ uint32_t u32fp8Decode(uint8_t x) {
|
|
return u32fpDecode(x, 3);
|
|
}
|
|
|
|
// The hash isn't just a function of the bytes but also where the bytes are split
|
|
// into different calls to eatHash().
|
|
static __host__ __device__ void eatHash(uint64_t acc[2], const void* bytes, size_t size) {
|
|
char const* ptr = (char const*)bytes;
|
|
acc[0] ^= size;
|
|
while (size != 0) {
|
|
// Mix the accumulator bits.
|
|
acc[0] += acc[1];
|
|
acc[1] ^= acc[0];
|
|
acc[0] ^= acc[0] >> 31;
|
|
acc[0] *= 0x9de62bbc8cef3ce3;
|
|
acc[1] ^= acc[1] >> 32;
|
|
acc[1] *= 0x485cd6311b599e79;
|
|
// Read in a chunk of input.
|
|
size_t chunkSize = size < sizeof(uint64_t) ? size : sizeof(uint64_t);
|
|
uint64_t x = 0;
|
|
memcpy(&x, ptr, chunkSize);
|
|
ptr += chunkSize;
|
|
size -= chunkSize;
|
|
// Add to accumulator.
|
|
acc[0] += x;
|
|
}
|
|
}
|
|
|
|
template<typename T>
|
|
static __host__ __device__ void eatHash(uint64_t acc[2], const T* bytes) {
|
|
eatHash(acc, (const void*)bytes, sizeof(T));
|
|
}
|
|
|
|
static __host__ __device__ uint64_t digestHash(uint64_t const acc[2]) {
|
|
uint64_t h = acc[0];
|
|
h ^= h >> 31;
|
|
h *= 0xbac3bd562846de6b;
|
|
h += acc[1];
|
|
h ^= h >> 32;
|
|
h *= 0x995a187a14e7b445;
|
|
return h;
|
|
}
|
|
|
|
static __host__ __device__ uint64_t getHash(const void* bytes, size_t size) {
|
|
uint64_t acc[2] = {1, 1};
|
|
eatHash(acc, bytes, size);
|
|
return digestHash(acc);
|
|
}
|
|
template<typename T>
|
|
static __host__ __device__ uint64_t getHash(const T* bytes) {
|
|
return getHash((const void*)bytes, sizeof(T));
|
|
}
|
|
|
|
#endif
|