98 lines · 3.2 KB · C
RawDownload
/*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 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 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 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

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t