// Utilities for use in __device__ code#ifndef CUDA_UTILS_CUH#define CUDA_UTILS_CUH#include "cuda_common.h"// ----------------------------------------------------------------------------// Packed128 data structure that forces the compiler to use 128-bit loads/stores// in GPUs that support (the LDG.128 and STS.128 instructions)// This is a bit similar to the use of float4 in the case of 32-bit floats, but// supports arbitrary precision.template<class ElementType>struct alignas(16) Packed128 { Packed128() = default; __device__ explicit Packed128(int4 bits) { static_assert(sizeof(bits) == sizeof(payload), "Size mismatch."); memcpy(&payload, &bits, sizeof(bits)); } __device__ static Packed128 constant(ElementType value) { Packed128 result; for(int k = 0; k < size; ++k) { result.payload[k] = value; } return result; } __device__ static Packed128 zeros() { return constant(0.f); } __device__ static Packed128 ones() { return constant(1.f); } __device__ ElementType& operator[](int index) { return payload[index]; } __device__ const ElementType& operator[](int index) const { return payload[index]; } __device__ int4 get_bits() const { int4 bits; static_assert(sizeof(bits) == sizeof(payload), "Size mismatch."); memcpy(&bits, &payload, sizeof(bits)); return bits; } static constexpr const size_t size = sizeof(int4) / sizeof(ElementType); ElementType payload[size];};// load a Packed128 from an aligned memory addresstemplate<class ElementType>__device__ Packed128<ElementType> load128(const ElementType* address) { return Packed128<ElementType>{*reinterpret_cast<const int4*>(address)};}// load a Packed128 from an aligned memory address with streaming cache hinttemplate<class ElementType>__device__ Packed128<ElementType> load128cs(const ElementType* address) { return Packed128<ElementType>{__ldcs(reinterpret_cast<const int4*>(address))};}// store a Packed128 to an aligned memory addresstemplate<class ElementType>__device__ void store128(ElementType* target, Packed128<ElementType> value) { *reinterpret_cast<int4*>(target) = value.get_bits();}// store a Packed128 to an aligned memory address with streaming cache hinttemplate<class ElementType>__device__ void store128cs(ElementType* target, Packed128<ElementType> value) { __stcs(reinterpret_cast<int4*>(target), value.get_bits());}// store a Packed128 to an aligned memory address while caching in L2 but bypassing L1template<class ElementType>__device__ void store128cg(ElementType* target, Packed128<ElementType> value) { __stcg(reinterpret_cast<int4*>(target), value.get_bits());}// short-form typedefstypedef Packed128<float> f128;typedef Packed128<floatX> x128;// ----------------------------------------------------------------------------// Warp/Block communication primitives// warp-level reduction for summing values__device__ inline float warpReduceSum(float val) { for (int offset = 16; offset > 0; offset /= 2) { val += __shfl_xor_sync(0xFFFFFFFF, val, offset); } return val;}// warp-level reduction for finding the maximum value__device__ inline float warpReduceMax(float val) { for (int offset = 16; offset > 0; offset /= 2) { val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); } return val;}// requires all 32 threads in the warp to be active, but should work for any block size// uses non-dynamic shared memory so every call increases shared memory requirements by 128 bytes// the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end// but if called inside a loop, the shared memory will be implicitly reused, so set final_sync to 1using reduction_func_t = float (*) (float);template<reduction_func_t warp_reduction>__device__ inline float blockReduce(float val, bool final_sync=false, float out_of_bounds=0.0f) { // two reductions of up to 1024 threads: // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) __shared__ float shared_val[WARP_SIZE]; const int lane_id = threadIdx.x % WARP_SIZE; const int warp_id = threadIdx.x / WARP_SIZE; const int num_warps = blockDim.x / WARP_SIZE; float warp_val = warp_reduction(val); if (lane_id == 0) { shared_val[warp_id] = warp_val; } __syncthreads(); warp_val = (lane_id < num_warps) ? shared_val[lane_id] : out_of_bounds; float block_val = warp_reduction(warp_val); if (final_sync) { __syncthreads(); // only needed in loops when effectively reusing shared memory etc. } return block_val;}// ----------------------------------------------------------------------------// Random Number Generation used in Stochastic Rounding// SquirrelNoise5 - Squirrel's Raw Noise utilities (version 5)// This gives us a random number from threadIdx/blockIdx + a single seed for the entire GPU// todo - possibly overkill and we don't need such high quality random numbers? (tbd)// http://eiserloh.net/noise/SquirrelNoise5.hpp__device__ __host__ constexpr unsigned int SquirrelNoise5(int positionX, unsigned int seed){ constexpr unsigned int SQ5_BIT_NOISE1 = 0xd2a80a3f; // 11010010101010000000101000111111 constexpr unsigned int SQ5_BIT_NOISE2 = 0xa884f197; // 10101000100001001111000110010111 constexpr unsigned int SQ5_BIT_NOISE3 = 0x6C736F4B; // 01101100011100110110111101001011 constexpr unsigned int SQ5_BIT_NOISE4 = 0xB79F3ABB; // 10110111100111110011101010111011 constexpr unsigned int SQ5_BIT_NOISE5 = 0x1b56c4f5; // 00011011010101101100010011110101 unsigned int mangledBits = (unsigned int) positionX; mangledBits *= SQ5_BIT_NOISE1; mangledBits += seed; mangledBits ^= (mangledBits >> 9); mangledBits += SQ5_BIT_NOISE2; mangledBits ^= (mangledBits >> 11); mangledBits *= SQ5_BIT_NOISE3; mangledBits ^= (mangledBits >> 13); mangledBits += SQ5_BIT_NOISE4; mangledBits ^= (mangledBits >> 15); mangledBits *= SQ5_BIT_NOISE5; mangledBits ^= (mangledBits >> 17); return mangledBits;}__device__ __host__ constexpr unsigned int Get2dNoiseUint(int indexX, int indexY, unsigned int seed){ constexpr int PRIME_NUMBER = 198491317; // Large prime number with non-boring bits return SquirrelNoise5(indexX + (PRIME_NUMBER * indexY), seed);}// stochastic rounding built on top of Squirel Noise above (with seed updated per step via xorshift)__device__ __forceinline__ void stochastic_rounding(float in, __nv_bfloat16 *out, unsigned int seed) { // todo - is this stochastic rounding *too good*? can we cut any corners? unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed); unsigned int threshold = random & 0xFFFF; unsigned int float_bits = __float_as_uint(in); unsigned int rounded_bits = float_bits & 0x0000FFFF; float_bits = (rounded_bits > threshold) ? (float_bits | 0xFFFF) : (float_bits & ~0xFFFF); *out = __float2bfloat16_rn(__uint_as_float(float_bits));}__device__ __forceinline__ void stochastic_rounding(float in, half *out, unsigned int random) { *out = (float)in; // todo - implement this...}__device__ __forceinline__ void stochastic_rounding(float in, float *out, unsigned int random) { *out = in; // dummy function for when floatX is float (FP32 mode)}#endif
// Utilities for use in __device__ code
#ifndef CUDA_UTILS_CUH
#define CUDA_UTILS_CUH
#include "cuda_common.h"
// ----------------------------------------------------------------------------
// Packed128 data structure that forces the compiler to use 128-bit loads/stores
// in GPUs that support (the LDG.128 and STS.128 instructions)
// This is a bit similar to the use of float4 in the case of 32-bit floats, but
// supports arbitrary precision.
template<class ElementType>
struct alignas(16) Packed128 {
Packed128() = default;
__device__ explicit Packed128(int4 bits) {
static_assert(sizeof(bits) == sizeof(payload), "Size mismatch.");
memcpy(&payload, &bits, sizeof(bits));
}
__device__ static Packed128 constant(ElementType value) {
Packed128 result;
for(int k = 0; k < size; ++k) {
result.payload[k] = value;
}
return result;
}
__device__ static Packed128 zeros() {
return constant(0.f);
}
__device__ static Packed128 ones() {
return constant(1.f);
}
__device__ ElementType& operator[](int index) {
return payload[index];
}
__device__ const ElementType& operator[](int index) const {
return payload[index];
}
__device__ int4 get_bits() const {
int4 bits;
static_assert(sizeof(bits) == sizeof(payload), "Size mismatch.");
memcpy(&bits, &payload, sizeof(bits));
return bits;
}
static constexpr const size_t size = sizeof(int4) / sizeof(ElementType);
ElementType payload[size];
};
// load a Packed128 from an aligned memory address
template<class ElementType>
__device__ Packed128<ElementType> load128(const ElementType* address) {
return Packed128<ElementType>{*reinterpret_cast<const int4*>(address)};
}
// load a Packed128 from an aligned memory address with streaming cache hint
template<class ElementType>
__device__ Packed128<ElementType> load128cs(const ElementType* address) {
return Packed128<ElementType>{__ldcs(reinterpret_cast<const int4*>(address))};
}
// store a Packed128 to an aligned memory address
template<class ElementType>
__device__ void store128(ElementType* target, Packed128<ElementType> value) {
*reinterpret_cast<int4*>(target) = value.get_bits();
}
// store a Packed128 to an aligned memory address with streaming cache hint
template<class ElementType>
__device__ void store128cs(ElementType* target, Packed128<ElementType> value) {
__stcs(reinterpret_cast<int4*>(target), value.get_bits());
}
// store a Packed128 to an aligned memory address while caching in L2 but bypassing L1
template<class ElementType>
__device__ void store128cg(ElementType* target, Packed128<ElementType> value) {
__stcg(reinterpret_cast<int4*>(target), value.get_bits());
}
// short-form typedefs
typedef Packed128<float> f128;
typedef Packed128<floatX> x128;
// ----------------------------------------------------------------------------
// Warp/Block communication primitives
// warp-level reduction for summing values
__device__ inline float warpReduceSum(float val) {
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
}
return val;
}
// warp-level reduction for finding the maximum value
__device__ inline float warpReduceMax(float val) {
for (int offset = 16; offset > 0; offset /= 2) {
val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset));
}
return val;
}
// requires all 32 threads in the warp to be active, but should work for any block size
// uses non-dynamic shared memory so every call increases shared memory requirements by 128 bytes
// the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end
// but if called inside a loop, the shared memory will be implicitly reused, so set final_sync to 1
using reduction_func_t = float (*) (float);
template<reduction_func_t warp_reduction>
__device__ inline float blockReduce(float val, bool final_sync=false, float out_of_bounds=0.0f) {
// two reductions of up to 1024 threads:
// 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle)
__shared__ float shared_val[WARP_SIZE];
const int lane_id = threadIdx.x % WARP_SIZE;
const int warp_id = threadIdx.x / WARP_SIZE;
const int num_warps = blockDim.x / WARP_SIZE;
float warp_val = warp_reduction(val);
if (lane_id == 0) { shared_val[warp_id] = warp_val; }
__syncthreads();
warp_val = (lane_id < num_warps) ? shared_val[lane_id] : out_of_bounds;
float block_val = warp_reduction(warp_val);
if (final_sync) {
__syncthreads(); // only needed in loops when effectively reusing shared memory etc.
}
return block_val;
}
// ----------------------------------------------------------------------------
// Random Number Generation used in Stochastic Rounding
// SquirrelNoise5 - Squirrel's Raw Noise utilities (version 5)
// This gives us a random number from threadIdx/blockIdx + a single seed for the entire GPU
// todo - possibly overkill and we don't need such high quality random numbers? (tbd)
// http://eiserloh.net/noise/SquirrelNoise5.hpp
__device__ __host__ constexpr unsigned int SquirrelNoise5(int positionX, unsigned int seed)
{
constexpr unsigned int SQ5_BIT_NOISE1 = 0xd2a80a3f; // 11010010101010000000101000111111
constexpr unsigned int SQ5_BIT_NOISE2 = 0xa884f197; // 10101000100001001111000110010111
constexpr unsigned int SQ5_BIT_NOISE3 = 0x6C736F4B; // 01101100011100110110111101001011
constexpr unsigned int SQ5_BIT_NOISE4 = 0xB79F3ABB; // 10110111100111110011101010111011
constexpr unsigned int SQ5_BIT_NOISE5 = 0x1b56c4f5; // 00011011010101101100010011110101
unsigned int mangledBits = (unsigned int) positionX;
mangledBits *= SQ5_BIT_NOISE1;
mangledBits += seed;
mangledBits ^= (mangledBits >> 9);
mangledBits += SQ5_BIT_NOISE2;
mangledBits ^= (mangledBits >> 11);
mangledBits *= SQ5_BIT_NOISE3;
mangledBits ^= (mangledBits >> 13);
mangledBits += SQ5_BIT_NOISE4;
mangledBits ^= (mangledBits >> 15);
mangledBits *= SQ5_BIT_NOISE5;
mangledBits ^= (mangledBits >> 17);
return mangledBits;
}
__device__ __host__ constexpr unsigned int Get2dNoiseUint(int indexX, int indexY, unsigned int seed)
{
constexpr int PRIME_NUMBER = 198491317; // Large prime number with non-boring bits
return SquirrelNoise5(indexX + (PRIME_NUMBER * indexY), seed);
}
// stochastic rounding built on top of Squirel Noise above (with seed updated per step via xorshift)
__device__ __forceinline__ void stochastic_rounding(float in, __nv_bfloat16 *out, unsigned int seed) {
// todo - is this stochastic rounding *too good*? can we cut any corners?
unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed);
unsigned int threshold = random & 0xFFFF;
unsigned int float_bits = __float_as_uint(in);
unsigned int rounded_bits = float_bits & 0x0000FFFF;
float_bits = (rounded_bits > threshold) ? (float_bits | 0xFFFF) : (float_bits & ~0xFFFF);
*out = __float2bfloat16_rn(__uint_as_float(float_bits));
}
__device__ __forceinline__ void stochastic_rounding(float in, half *out, unsigned int random) {
*out = (float)in; // todo - implement this...
}
__device__ __forceinline__ void stochastic_rounding(float in, float *out, unsigned int random) {
*out = in; // dummy function for when floatX is float (FP32 mode)
}
#endif