From 45252d434ca2bddc4e49ab30e018ae2b9057f5db Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 29 Apr 2024 21:53:08 +0300 Subject: [PATCH 1/4] yet another gelu --- dev/cuda/gelu_forward.cu | 129 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/dev/cuda/gelu_forward.cu b/dev/cuda/gelu_forward.cu index 5c6fff5..82cced3 100644 --- a/dev/cuda/gelu_forward.cu +++ b/dev/cuda/gelu_forward.cu @@ -18,6 +18,110 @@ version 1 is naive port from CPU code to kernel #include #include "common.h" +// OK, so this part requires a bit of explanation. Our end goal here is that we get a convenient interface that +// allows us to interact with vectorized loads that read 128 bits of aligned memory in a uniform way, independent +// of the underlying datatype, and the whims of the nvcc compiler. + +// as a first step, we define some constants that indicate which type of memory operation we intend to do. +// these correspond to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#cache-operators +enum class ELoadMode { + CA, CG, CS, LU, CV +}; + +enum class EStoreMode { + WB, CG, CS, WT +}; + + +// in order to enable dispatch _at compile time_, we need to encode the load/store mode into types (unless we want +// to have ugly template syntax at the call sites, like load(address). +// Therefore, we wrap all these into compile-time integers, and provide global objects that can be passed to +// select the correct function overload. +template +using load_mode_t = std::integral_constant; + +template +using store_mode_t = std::integral_constant; + +constexpr load_mode_t LdCA; +constexpr load_mode_t LdCG; +constexpr load_mode_t LdCS; +constexpr load_mode_t LdLU; +constexpr load_mode_t LdCV; + +constexpr store_mode_t StWB; +constexpr store_mode_t StCG; +constexpr store_mode_t StCS; +constexpr store_mode_t StWT; + +// Finally, we define the dispatch mechanism itself. Really just a long if-else chain, except +// that all of this needs to be decided at compile time (hence constexpr if) +template +__device__ T generic_load(const T* address, load_mode_t) { + if constexpr (Mode == ELoadMode::CA) { + return __ldca(address); + } else if constexpr (Mode == ELoadMode::CG) { + return __ldcg(address); + } else if constexpr (Mode == ELoadMode::CS) { + return __ldcs(address); + } else if constexpr (Mode == ELoadMode::LU) { + return __ldlu(address); + } else if constexpr (Mode == ELoadMode::CV) { + return __ldcv(address); + } else { + __builtin_unreachable(); + } +} + +template +__device__ void generic_store(T* address, const T& value, store_mode_t) { + if constexpr (Mode == EStoreMode::WB) { + return __stwb(address, value); + } else if constexpr (Mode == EStoreMode::CG) { + return __stcg(address, value); + } else if constexpr (Mode == EStoreMode::CS) { + return __stcs(address, value); + } else if constexpr (Mode == EStoreMode::WT) { + return __stwt(address, value); + } else { + __builtin_unreachable(); + } +} + +// Finally, we define a wrapper type that contains 128 bits of whatever underlying type we want +// we store the actual data in an int4 vector, and reinterpret its bits, because int4 gets nvcc to +// reliably produce 128-bit instructions +// TODO do we really need this here, or can we get away with the int4 trick just inside the load/store functions +// we allow individual element access with [], and provide a convenience accessor to get the data converted to +// a regular float for mixed-precision operations. +template +struct alignas(16) Packed128 { + __device__ ElementType& operator[](int index) { + return reinterpret_cast(&payload)[index]; + } + __device__ const ElementType& operator[](int index) const { + return reinterpret_cast(&payload)[index]; + } + __device__ float fp32(int index) { + return static_cast(reinterpret_cast(&payload)[index]); + } + static constexpr const size_t size = sizeof(int4) / sizeof(ElementType); + + int4 payload; +}; + +// use this function to load a Packet128 from an aligned memory address +template +__device__ Packed128> load_aligned(ElementType* address, load_mode_t mode = {}) { + return {generic_load(reinterpret_cast(address), mode)}; +} + +// use this function to store a Packet128 to an aligned memory address +template +__device__ void store_aligned(ElementType* target, Packed128 value, store_mode_t mode = {}) { + generic_store(reinterpret_cast(target), value.payload, mode); +} + // ---------------------------------------------------------------------------- // CPU code reference @@ -44,6 +148,22 @@ __global__ void gelu_kernel(float* out, const float* inp, int N) { } } +// elementwise ops are nice and ez +__global__ void gelu_kernel2(float* out, const float* inp, int N) { + using packet_t = Packed128; + int i = (blockIdx.x * blockDim.x + threadIdx.x) * packet_t::size; + if (i < N) { + packet_t packet_out; + packet_t packet_in = load_aligned(inp + i, LdCS); + for(int k = 0; k < packet_in.size; ++k) { + float xi = packet_in[k]; + float cube = 0.044715f * xi * xi * xi; + packet_out[k] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))); + } + store_aligned(out + i, packet_out); + } +} + // ---------------------------------------------------------------------------- // kernel launcher @@ -53,6 +173,12 @@ void gelu_forward1(float* out, const float* inp, int N, const int block_size) { cudaCheck(cudaGetLastError()); } +void gelu_forward2(float* out, const float* inp, int N, const int block_size) { + const int grid_size = ceil_div(N, 4 * block_size); + gelu_kernel2<<>>(out, inp, N); + cudaCheck(cudaGetLastError()); +} + // kernel version dispatch void gelu_forward(int kernel_num, float* out, @@ -63,6 +189,9 @@ void gelu_forward(int kernel_num, case 1: gelu_forward1(out, inp, B * T * C, block_size); break; + case 2: + gelu_forward2(out, inp, B * T * C, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); From 896b644cb9008122e0e5b3e029fbc53302181e29 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 29 Apr 2024 22:19:12 +0300 Subject: [PATCH 2/4] slight simplification --- dev/cuda/gelu_forward.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/cuda/gelu_forward.cu b/dev/cuda/gelu_forward.cu index 82cced3..f2ee658 100644 --- a/dev/cuda/gelu_forward.cu +++ b/dev/cuda/gelu_forward.cu @@ -112,7 +112,7 @@ struct alignas(16) Packed128 { // use this function to load a Packet128 from an aligned memory address template -__device__ Packed128> load_aligned(ElementType* address, load_mode_t mode = {}) { +__device__ Packed128 load_aligned(const ElementType* address, load_mode_t mode = {}) { return {generic_load(reinterpret_cast(address), mode)}; } From 9f07a176e0ddc64715742c2ffd451112e31acd62 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 29 Apr 2024 22:39:49 +0300 Subject: [PATCH 3/4] simplify to only support cs --- dev/cuda/gelu_forward.cu | 124 +++++++++++++-------------------------- 1 file changed, 40 insertions(+), 84 deletions(-) diff --git a/dev/cuda/gelu_forward.cu b/dev/cuda/gelu_forward.cu index f2ee658..db2e1f5 100644 --- a/dev/cuda/gelu_forward.cu +++ b/dev/cuda/gelu_forward.cu @@ -18,108 +18,64 @@ version 1 is naive port from CPU code to kernel #include #include "common.h" -// OK, so this part requires a bit of explanation. Our end goal here is that we get a convenient interface that -// allows us to interact with vectorized loads that read 128 bits of aligned memory in a uniform way, independent -// of the underlying datatype, and the whims of the nvcc compiler. - -// as a first step, we define some constants that indicate which type of memory operation we intend to do. -// these correspond to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#cache-operators -enum class ELoadMode { - CA, CG, CS, LU, CV -}; - -enum class EStoreMode { - WB, CG, CS, WT -}; - - -// in order to enable dispatch _at compile time_, we need to encode the load/store mode into types (unless we want -// to have ugly template syntax at the call sites, like load(address). -// Therefore, we wrap all these into compile-time integers, and provide global objects that can be passed to -// select the correct function overload. -template -using load_mode_t = std::integral_constant; - -template -using store_mode_t = std::integral_constant; - -constexpr load_mode_t LdCA; -constexpr load_mode_t LdCG; -constexpr load_mode_t LdCS; -constexpr load_mode_t LdLU; -constexpr load_mode_t LdCV; - -constexpr store_mode_t StWB; -constexpr store_mode_t StCG; -constexpr store_mode_t StCS; -constexpr store_mode_t StWT; - -// Finally, we define the dispatch mechanism itself. Really just a long if-else chain, except -// that all of this needs to be decided at compile time (hence constexpr if) -template -__device__ T generic_load(const T* address, load_mode_t) { - if constexpr (Mode == ELoadMode::CA) { - return __ldca(address); - } else if constexpr (Mode == ELoadMode::CG) { - return __ldcg(address); - } else if constexpr (Mode == ELoadMode::CS) { - return __ldcs(address); - } else if constexpr (Mode == ELoadMode::LU) { - return __ldlu(address); - } else if constexpr (Mode == ELoadMode::CV) { - return __ldcv(address); - } else { - __builtin_unreachable(); - } -} - -template -__device__ void generic_store(T* address, const T& value, store_mode_t) { - if constexpr (Mode == EStoreMode::WB) { - return __stwb(address, value); - } else if constexpr (Mode == EStoreMode::CG) { - return __stcg(address, value); - } else if constexpr (Mode == EStoreMode::CS) { - return __stcs(address, value); - } else if constexpr (Mode == EStoreMode::WT) { - return __stwt(address, value); - } else { - __builtin_unreachable(); - } -} - // Finally, we define a wrapper type that contains 128 bits of whatever underlying type we want // we store the actual data in an int4 vector, and reinterpret its bits, because int4 gets nvcc to // reliably produce 128-bit instructions // TODO do we really need this here, or can we get away with the int4 trick just inside the load/store functions // we allow individual element access with [], and provide a convenience accessor to get the data converted to // a regular float for mixed-precision operations. + template struct alignas(16) Packed128 { - __device__ ElementType& operator[](int index) { - return reinterpret_cast(&payload)[index]; + __device__ __forceinline__ Packed128() = default; + __device__ __forceinline__ explicit Packed128(int4 bits) { + static_assert(sizeof(bits) == sizeof(payload), "Size mismatch."); + memcpy(&payload, &bits, sizeof(bits)); } - __device__ const ElementType& operator[](int index) const { - return reinterpret_cast(&payload)[index]; + + __device__ __forceinline__ ElementType& operator[](int index) { + return payload[index]; } - __device__ float fp32(int index) { - return static_cast(reinterpret_cast(&payload)[index]); + __device__ __forceinline__ const ElementType& operator[](int index) const { + return payload[index]; } + __device__ __forceinline__ float fp32(int index) { + return static_cast(payload[index]); + } + + __device__ __forceinline__ 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); - int4 payload; + ElementType payload[size]; }; + // use this function to load a Packet128 from an aligned memory address -template -__device__ Packed128 load_aligned(const ElementType* address, load_mode_t mode = {}) { - return {generic_load(reinterpret_cast(address), mode)}; +template +__device__ __forceinline__ Packed128 load_aligned(const ElementType* address) { + return Packed128{*reinterpret_cast(address)}; +} + +template +__device__ __forceinline__ Packed128 load_aligned_cs(const ElementType* address) { + return Packed128{__ldcs(reinterpret_cast(address))}; } // use this function to store a Packet128 to an aligned memory address -template -__device__ void store_aligned(ElementType* target, Packed128 value, store_mode_t mode = {}) { - generic_store(reinterpret_cast(target), value.payload, mode); +template +__device__ __forceinline__ void store_aligned(ElementType* target, Packed128 value) { + *reinterpret_cast(target) = value.get_bits(); +} + +template +__device__ __forceinline__ void store_aligned_cs(ElementType* target, Packed128 value) { + __stcs(reinterpret_cast(target), value.get_bits()); } // ---------------------------------------------------------------------------- @@ -154,7 +110,7 @@ __global__ void gelu_kernel2(float* out, const float* inp, int N) { int i = (blockIdx.x * blockDim.x + threadIdx.x) * packet_t::size; if (i < N) { packet_t packet_out; - packet_t packet_in = load_aligned(inp + i, LdCS); + packet_t packet_in = load_aligned_cs(inp + i); for(int k = 0; k < packet_in.size; ++k) { float xi = packet_in[k]; float cube = 0.044715f * xi * xi * xi; From 4872c5716427533c8752c728cdd49829b328c286 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 29 Apr 2024 22:46:33 +0300 Subject: [PATCH 4/4] comments --- dev/cuda/gelu_forward.cu | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/dev/cuda/gelu_forward.cu b/dev/cuda/gelu_forward.cu index db2e1f5..4b4386e 100644 --- a/dev/cuda/gelu_forward.cu +++ b/dev/cuda/gelu_forward.cu @@ -18,11 +18,8 @@ version 1 is naive port from CPU code to kernel #include #include "common.h" -// Finally, we define a wrapper type that contains 128 bits of whatever underlying type we want -// we store the actual data in an int4 vector, and reinterpret its bits, because int4 gets nvcc to -// reliably produce 128-bit instructions -// TODO do we really need this here, or can we get away with the int4 trick just inside the load/store functions -// we allow individual element access with [], and provide a convenience accessor to get the data converted to +// We define a wrapper type that contains 128 bits of whatever underlying type we want +// We allow individual element access with [], and provide a convenience accessor to get the data converted to // a regular float for mixed-precision operations. template @@ -56,23 +53,25 @@ struct alignas(16) Packed128 { }; -// use this function to load a Packet128 from an aligned memory address +// load a Packet128 from an aligned memory address template __device__ __forceinline__ Packed128 load_aligned(const ElementType* address) { return Packed128{*reinterpret_cast(address)}; } +// load a Packet128 from an aligned memory address with streaming cache hint template __device__ __forceinline__ Packed128 load_aligned_cs(const ElementType* address) { return Packed128{__ldcs(reinterpret_cast(address))}; } -// use this function to store a Packet128 to an aligned memory address +// store a Packet128 to an aligned memory address template __device__ __forceinline__ void store_aligned(ElementType* target, Packed128 value) { *reinterpret_cast(target) = value.get_bits(); } +// store a Packet128 to an aligned memory address with streaming cache hint template __device__ __forceinline__ void store_aligned_cs(ElementType* target, Packed128 value) { __stcs(reinterpret_cast(target), value.get_bits());