mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-25 20:05:08 -04:00
98 lines
No EOL
3.2 KiB
C++
98 lines
No EOL
3.2 KiB
C++
/*
|
|
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
|