/*Common utilities for CUDA code.*/#ifndef CUDA_COMMON_H#define CUDA_COMMON_H#include <math.h>#include <string>#include <cuda_runtime.h>#include <nvtx3/nvToolsExt.h>#include <cuda_profiler_api.h>#include <cuda_bf16.h>// ----------------------------------------------------------------------------// Global defines and settings// WarpSize is not a compile time constant// Defining here like this possibly allows the compiler to optimize better#define WARP_SIZE 32U// try to make sure that 2 blocks fit on A100/H100 to maximise latency tolerance// this needs to be defines rather than queried to be used for __launch_bounds__#if __CUDA_ARCH__ == 800 || __CUDA_ARCH__ >= 900#define MAX_1024_THREADS_BLOCKS 2#else#define MAX_1024_THREADS_BLOCKS 1#endif// convenience macro for calculating grid/block dimensions for kernels#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))// ----------------------------------------------------------------------------// Error checking// CUDA error checkingvoid inline cudaCheck(cudaError_t error, const char *file, int line) { if (error != cudaSuccess) { printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, cudaGetErrorString(error)); exit(EXIT_FAILURE); }};#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__))// ----------------------------------------------------------------------------// CUDA Precision settings and definesenum PrecisionMode { PRECISION_FP32, PRECISION_FP16, PRECISION_BF16};// Specific configurations based on the enabled precision#if defined(ENABLE_FP32)typedef float floatX;#define PRECISION_MODE PRECISION_FP32// use fp16 (note: this may require gradient scaler, currently not implemented!)#elif defined(ENABLE_FP16)typedef half floatX;#define PRECISION_MODE PRECISION_FP16#else // Default to bfloat16typedef __nv_bfloat16 floatX;#define PRECISION_MODE PRECISION_BF16#endif// ----------------------------------------------------------------------------// Load and store with streaming cache hints// Older nvcc does not provide __ldcs and __stcs for bfloat16, despite these// actually just being unsigned shorts. We need to be careful here to only define// our own versions if none already exist, otherwise the compiler will complain.// If not, you easily get "no viable overload" (for sm52) and "function already exists" (sm_80)#if defined(ENABLE_BF16) && (__CUDACC_VER_MAJOR__ < 12) && !((__CUDA_ARCH__ >= 800) || !defined(__CUDA_ARCH__))__device__ floatX __ldcs(const floatX* address) { unsigned short bf = __ldcs(reinterpret_cast<const unsigned short*>(address)); return __nv_bfloat16_raw{bf};}__device__ void __stcs(floatX* address, floatX value) { __stcs(reinterpret_cast<unsigned short*>(address), ((__nv_bfloat16_raw)value).x);}#endif// ----------------------------------------------------------------------------// Profiler utilsclass NvtxRange { public: NvtxRange(const char* s) { nvtxRangePush(s); } NvtxRange(const std::string& base_str, int number) { std::string range_string = base_str + " " + std::to_string(number); nvtxRangePush(range_string.c_str()); } ~NvtxRange() { nvtxRangePop(); }};#define NVTX_RANGE_FN() NvtxRange nvtx_range(__FUNCTION__)#endif // CUDA_COMMON_H
/*
Common utilities for CUDA code.
*/
#ifndef CUDA_COMMON_H
#define CUDA_COMMON_H
#include <math.h>
#include <string>
#include <cuda_runtime.h>
#include <nvtx3/nvToolsExt.h>
#include <cuda_profiler_api.h>
#include <cuda_bf16.h>
// ----------------------------------------------------------------------------
// Global defines and settings
// WarpSize is not a compile time constant
// Defining here like this possibly allows the compiler to optimize better
#define WARP_SIZE 32U
// try to make sure that 2 blocks fit on A100/H100 to maximise latency tolerance
// this needs to be defines rather than queried to be used for __launch_bounds__
#if __CUDA_ARCH__ == 800 || __CUDA_ARCH__ >= 900
#define MAX_1024_THREADS_BLOCKS 2
#else
#define MAX_1024_THREADS_BLOCKS 1
#endif
// convenience macro for calculating grid/block dimensions for kernels
#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
// ----------------------------------------------------------------------------
// Error checking
// CUDA error checking
void inline cudaCheck(cudaError_t error, const char *file, int line) {
if (error != cudaSuccess) {
printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, cudaGetErrorString(error));
exit(EXIT_FAILURE);
}
};
#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__))
// ----------------------------------------------------------------------------
// CUDA Precision settings and defines
enum PrecisionMode {
PRECISION_FP32,
PRECISION_FP16,
PRECISION_BF16
};
// Specific configurations based on the enabled precision
#if defined(ENABLE_FP32)
typedef float floatX;
#define PRECISION_MODE PRECISION_FP32
// use fp16 (note: this may require gradient scaler, currently not implemented!)
#elif defined(ENABLE_FP16)
typedef half floatX;
#define PRECISION_MODE PRECISION_FP16
#else // Default to bfloat16
typedef __nv_bfloat16 floatX;
#define PRECISION_MODE PRECISION_BF16
#endif
// ----------------------------------------------------------------------------
// Load and store with streaming cache hints
// Older nvcc does not provide __ldcs and __stcs for bfloat16, despite these
// actually just being unsigned shorts. We need to be careful here to only define
// our own versions if none already exist, otherwise the compiler will complain.
// If not, you easily get "no viable overload" (for sm52) and "function already exists" (sm_80)
#if defined(ENABLE_BF16) && (__CUDACC_VER_MAJOR__ < 12) && !((__CUDA_ARCH__ >= 800) || !defined(__CUDA_ARCH__))
__device__ floatX __ldcs(const floatX* address) {
unsigned short bf = __ldcs(reinterpret_cast<const unsigned short*>(address));
return __nv_bfloat16_raw{bf};
}
__device__ void __stcs(floatX* address, floatX value) {
__stcs(reinterpret_cast<unsigned short*>(address), ((__nv_bfloat16_raw)value).x);
}
#endif
// ----------------------------------------------------------------------------
// Profiler utils
class NvtxRange {
public:
NvtxRange(const char* s) { nvtxRangePush(s); }
NvtxRange(const std::string& base_str, int number) {
std::string range_string = base_str + " " + std::to_string(number);
nvtxRangePush(range_string.c_str());
}
~NvtxRange() { nvtxRangePop(); }
};
#define NVTX_RANGE_FN() NvtxRange nvtx_range(__FUNCTION__)
#endif // CUDA_COMMON_H