diff --git a/dev/cuda/gelu_forward.cu b/dev/cuda/gelu_forward.cu index 5c6fff5..4b4386e 100644 --- a/dev/cuda/gelu_forward.cu +++ b/dev/cuda/gelu_forward.cu @@ -18,6 +18,65 @@ version 1 is naive port from CPU code to kernel #include #include "common.h" +// 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 +struct alignas(16) Packed128 { + __device__ __forceinline__ Packed128() = default; + __device__ __forceinline__ explicit Packed128(int4 bits) { + static_assert(sizeof(bits) == sizeof(payload), "Size mismatch."); + memcpy(&payload, &bits, sizeof(bits)); + } + + __device__ __forceinline__ ElementType& operator[](int index) { + return 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); + + ElementType payload[size]; +}; + + +// 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))}; +} + +// 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()); +} + // ---------------------------------------------------------------------------- // CPU code reference @@ -44,6 +103,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_cs(inp + i); + 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 +128,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 +144,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);