From 0f36cdd15bc6419d64453f5863744a7b6bee005e Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 16 Jun 2024 10:37:39 +0200 Subject: [PATCH 1/5] Use faster kernel 5 for layernorm forward --- llmc/layernorm.cuh | 84 +++++++++++++++++++++++++++------------------- 1 file changed, 50 insertions(+), 34 deletions(-) diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index a49f0bf..0b79bb3 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -17,50 +17,65 @@ E.g., the layernorms are connected to the residuals so we += in layernorm backwa // ---------------------------------------------------------------------------- // CUDA kernels -__global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd, +__global__ void layernorm_forward_kernel5(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd, const floatX* __restrict__ inp, const floatX* __restrict__ weight, const floatX* __restrict__ bias, int N, int C) { - int lane_id = threadIdx.x % WARP_SIZE; - int warp_id = threadIdx.x / WARP_SIZE; + __shared__ float shared_sum[WARP_SIZE]; // block_size max is 1024 = 32 * 32 warps + __shared__ float shared_sum2[WARP_SIZE]; // warps will be writing into shared memory after warp-reduce + int num_warps = blockDim.x / WARP_SIZE; + int warp_id = threadIdx.x / WARP_SIZE; + int lane_id = threadIdx.x % WARP_SIZE; - int idx = blockIdx.x * num_warps + warp_id; - if(idx >= N) { return; } // guard + int idx = blockIdx.x; // simply one block per row + const floatX* x = inp + idx * C; // the row of input that this group of threads is responsible for - // the row of input that this group of threads is responsible for - const floatX* x = inp + idx * C; - - // mean - float sum = 0.0f; - for (int i = lane_id; i < C; i += WARP_SIZE) { - sum += (float)x[i]; + // thread coarsening through the row, reduce the sum in series + float thread_sum = 0.0; // stores sum(x) + float thread_sum2 = 0.0; // stores sum(x**2) + for (int i = threadIdx.x; i < C; i += blockDim.x) { + float xi = x[i]; + thread_sum += xi; + thread_sum2 += xi * xi; } - sum = warpReduceSum(sum); - float m = sum / C; - if(lane_id == 0 && mean != nullptr) { + + // 2 warp level reductions instead of block level in order to prevent 2 syncthreads + + // warp-level reduction + float warp_sum = warpReduceSum(thread_sum); // sum(x) + float warp_sum2 = warpReduceSum(thread_sum2); // sum(x**2) + // store the warp-level reduction in shared memory + if (lane_id == 0) { + shared_sum[warp_id] = warp_sum; + shared_sum2[warp_id] = warp_sum2; + } + __syncthreads(); + // load results from shared memory to threads, pad with zeros for threads that are out of bounds + warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f; + warp_sum2 = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f; + // now reduce the warp-level reductions + float block_sum = warpReduceSum(warp_sum); // sum(x) + float block_sum2 = warpReduceSum(warp_sum2); // sum(x**2) + + // mean, var, rstd + block_sum /= C; // mean(x) + block_sum2 /= C; // mean(x**2) + float m = block_sum; + float var = block_sum2 - m * m; + float s = rsqrtf(var + 1e-5f); + + if (threadIdx.x == 0 && mean != nullptr) { __stcs(mean + idx, (floatX)m); } - // rstd - sum = 0.0f; - for (int i = lane_id; i < C; i += WARP_SIZE) { - float diff = (float)x[i] - m; - sum += diff * diff; - } - sum = warpReduceSum(sum); - float s = rsqrtf(sum / C + 1e-5f); - if(lane_id == 0 && rstd != nullptr) { + if (threadIdx.x == 0 && rstd != nullptr) { __stcs(rstd + idx, (floatX)s); } - // final normalization and scaling by weight/bias floatX* o = out + idx * C; - for (int c = lane_id; c < C; c += WARP_SIZE) { - // load and store using the .cs "streaming" hint to the compiler, - // indicating that this data will not be reused soon, and can be streamed through the caches - // this allows the threads to get more cache-hits for the (shared) weight and bias parameters - float n = s * ((float)__ldcs(x+c) - m); - __stcs(o+c, (floatX)(n * (float)weight[c] + (float)bias[c])); + for (int i = threadIdx.x; i < C; i += blockDim.x) { + float n = s * ((float)__ldcs(x+i) - m); + __stcs(o+i, (floatX)(n * (float)weight[i] + (float)bias[i])); } } @@ -358,10 +373,11 @@ void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, floatX* inp, const floatX* weight, const floatX* bias, int B, int T, int C, cudaStream_t stream) { NVTX_RANGE_FN(); - const int block_size = 512; + const int block_size = 256; + assert(block_size % WARP_SIZE == 0 && block_size <= 1024); const int N = B * T; - const int grid_size = CEIL_DIV(N * WARP_SIZE, block_size); - layernorm_forward_kernel3<<>>(out, mean, rstd, inp, weight, bias, N, C); + const int grid_size = N; + layernorm_forward_kernel5<<>>(out, mean, rstd, inp, weight, bias, N, C); cudaCheck(cudaGetLastError()); } From c6040eef5fb112e433f447b3a3873e7ace9f14be Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 16 Jun 2024 16:36:50 +0200 Subject: [PATCH 2/5] Even better kernel - modification of what we use for fused residual fwd --- dev/cuda/fused_residual_forward.cu | 18 ++--- dev/cuda/layernorm_forward.cu | 106 +++++++++++++++++++++++++++++ llmc/layernorm.cuh | 96 ++++++++++++++++++++++++-- 3 files changed, 207 insertions(+), 13 deletions(-) diff --git a/dev/cuda/fused_residual_forward.cu b/dev/cuda/fused_residual_forward.cu index b98a67c..9752873 100644 --- a/dev/cuda/fused_residual_forward.cu +++ b/dev/cuda/fused_residual_forward.cu @@ -133,7 +133,7 @@ __global__ void fused_residual_forward2(floatX* residual, floatX* normed, floatX for(int c = 0; c < C; ++c) { float out = (float)inp1[c] + (float)inp2[c]; m += out; - residual[c] = out; + residual[c] = (floatX)out; } m = m / C; @@ -149,11 +149,11 @@ __global__ void fused_residual_forward2(floatX* residual, floatX* normed, floatX for (int c = 0; c < C; c++) { float n = (s * ((float)residual[c] - m)); // normalized output float o = n * (float)weight[c] + (float)bias[c]; // scale and shift it - normed[c] = o; // write + normed[c] = (floatX)o; // write } // cache the mean and rstd for the backward pass later - mean[idx] = m; - rstd[idx] = s; + mean[idx] = (floatX)m; + rstd[idx] = (floatX)s; } // handle one token per warp for coalesced access @@ -232,7 +232,7 @@ __global__ void fused_residual_forward_kernel4(floatX* residual, floatX* normed, const x128 in2 = load128cs(inp2 + c); x128 out; for(int k = 0; k < x128::size; ++k) { - out[k] = (float)in1[k] + (float)in2[k]; + out[k] = (floatX)((float)in1[k] + (float)in2[k]); sum += (float)out[k]; sum_sq += (float)out[k] * (float)out[k]; } @@ -309,7 +309,7 @@ __global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, const x128 in2 = load128cs(inp2 + c); x128 out; for(int k = 0; k < x128::size; ++k) { - out[k] = (float)in1[k] + (float)in2[k]; + out[k] = (floatX)((float)in1[k] + (float)in2[k]); sum += (float)out[k]; } store128cs(residual + c, out); @@ -372,8 +372,8 @@ __global__ void fused_residual_forward_kernel6(floatX* residual, floatX* normed, // weights and biases are shared among all tokens x128* s_weight = reinterpret_cast(params); x128* s_bias = reinterpret_cast(params + C * sizeof(floatX)); - // residual output (input to layernorm) is indpendent for each sub-block indicates by threadIdx.z - x128* s_res = reinterpret_cast(params + (2 + threadIdx.z) * C * sizeof(floatX) ); + // residual output (input to layernorm) is independent for each sub-block indicates by threadIdx.z + x128* s_res = reinterpret_cast(params + (2 + threadIdx.z) * C * sizeof(floatX)); // similarly, each sub-block needs its own reduction buffers float* s_mean = reinterpret_cast(params + (2 + blockDim.z) * C * sizeof(floatX) + threadIdx.z * 32 * sizeof(float)); float* s_var = reinterpret_cast(params + (2 + blockDim.z) * C * sizeof(floatX) + 32 * sizeof(float) * (blockDim.z + threadIdx.z)); @@ -385,10 +385,10 @@ __global__ void fused_residual_forward_kernel6(floatX* residual, floatX* normed, s_weight[c / x128::size] = load128(weight + c); s_bias[c / x128::size] = load128(bias + c); } + // the block-level reductions will cause sync before the first time we read these // => no syncthreads needed here - // loop over all tokens for(int tidx = blockIdx.x * blockDim.z + threadIdx.z; tidx < N; tidx += gridDim.x * blockDim.z) { // adjust pointers to current token diff --git a/dev/cuda/layernorm_forward.cu b/dev/cuda/layernorm_forward.cu index 8631aa1..4e4552b 100644 --- a/dev/cuda/layernorm_forward.cu +++ b/dev/cuda/layernorm_forward.cu @@ -29,6 +29,8 @@ verstion 5 allocates blocks per row instead of warps per row, same alg as 4 othe #include #include "common.h" +#define WARP_SIZE 32 + // ---------------------------------------------------------------------------- // CPU code reference @@ -337,6 +339,82 @@ __global__ void layernorm_forward_kernel5(float* __restrict__ out, float* __rest } } +// Inspired by `fused_residual_forward_kernel5` in fused_residual_forward.cu +__global__ void layernorm_forward_kernel6(float* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd, + const float* __restrict__ inp, const float* __restrict__ weight, + const float* __restrict__ bias, int N, int C) { + assert(blockDim.x == WARP_SIZE); + + // load weights and biases into shared memory + // do this before we allow any threads to exit! + extern __shared__ char params[]; + // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so + // let's keep everything as x128 + x128* s_weight = reinterpret_cast(params); + x128* s_bias = reinterpret_cast(params) + (C / x128::size); + x128* s_in = reinterpret_cast(params) + ((2 + threadIdx.y) * C / x128::size); + + int sidx = (threadIdx.x + WARP_SIZE * threadIdx.y) * x128::size; + for(int i = sidx; i < C; i += blockDim.y * WARP_SIZE * x128::size) { + s_weight[i/x128::size] = load128(weight + i); + s_bias[i/x128::size] = load128(bias + i); + } + __syncthreads(); + + int idx = blockIdx.x * blockDim.y + threadIdx.y; + if(idx >= N) { return; } // guard + + // adjust pointers to current token + inp += idx * C; + out += idx * C; + + const float eps = 1e-5f; + float sum = 0.0f; + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = load128cs(inp + c); + for(int k = 0; k < x128::size; ++k) { + sum += (float)in_data[k]; + } + s_in[c / x128::size] = in_data; + } + + sum = warpReduceSum(sum); + float m = sum / C; + float v = 0.f; + + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = s_in[c / x128::size]; + for(int k = 0; k < x128::size; ++k) { + v += ((float)in_data[k] - m) * ((float)in_data[k] - m); + } + } + + v = warpReduceSum(v) / C; + float s = rsqrtf(v + eps); + + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = s_in[c / x128::size]; + const x128 w = s_weight[c / x128::size]; + const x128 b = s_bias[c / x128::size]; + x128 out_data; + for(int k = 0; k < x128::size; ++k) { + float n = s * ((float)in_data[k] - m); // normalized output + float o = n * (float)w[k] + (float)b[k]; // scale and shift it + out_data[k] = o; + } + + store128cs(out + c, out_data); + } + // cache the mean and rstd for the backward pass later + if(threadIdx.x == 0 && mean != nullptr) { + __stcs(mean + idx, m); + } + // store the rstd, no need to cache it + if(threadIdx.x == 0 && rstd != nullptr) { + __stcs(rstd + idx, s); + } +} + // ---------------------------------------------------------------------------- // kernel launcher @@ -401,6 +479,31 @@ void layernorm_forward5(float* out, float* mean, float* rstd, cudaCheck(cudaGetLastError()); } +void layernorm_forward6(float* out, float* mean, float* rstd, + const float* inp, const float* weight, const float* bias, + int B, int T, int C, + int block_size) { + assert(block_size % 32 == 0); + const int N = B * T; + int block_y = block_size / WARP_SIZE; + const int grid_size = ceil_div(N, block_y); + size_t smem = (2 + block_y) * C * sizeof(float); + + // in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute + // this may fail, in which case we fall back to the smem free implementation. + cudaCheck(cudaGetLastError()); + auto status = cudaFuncSetAttribute(layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + cudaGetLastError(); + if (status == cudaSuccess) { + layernorm_forward_kernel6<<>>(out, mean, rstd, inp, weight, bias, N, C); + } else { + const int grid_size = N; + // fall back to the version without shared memory + layernorm_forward_kernel5<<>>(out, mean, rstd, inp, weight, bias, N, C); + } + cudaCheck(cudaGetLastError()); +} + // kernel version dispatch void layernorm_forward(int kernel_num, float* out, float* mean, float* rstd, @@ -423,6 +526,9 @@ void layernorm_forward(int kernel_num, case 5: layernorm_forward5(out, mean, rstd, inp, weight, bias, B, T, C, block_size); break; + case 6: + layernorm_forward6(out, mean, rstd, inp, weight, bias, B, T, C, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index 0b79bb3..861efd0 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -79,6 +79,81 @@ __global__ void layernorm_forward_kernel5(floatX* __restrict__ out, floatX* __re } } +__global__ void layernorm_forward_kernel6(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd, + const floatX* __restrict__ inp, const floatX* __restrict__ weight, + const floatX* __restrict__ bias, int N, int C) { + assert(blockDim.x == WARP_SIZE); + + // load weights and biases into shared memory + // do this before we allow any threads to exit! + extern __shared__ char* params[]; + // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so + // let's keep everything as x128 + x128* s_weight = reinterpret_cast(params); + x128* s_bias = reinterpret_cast(params) + (C / x128::size); + x128* s_in = reinterpret_cast(params) + ((2 + threadIdx.y) * C / x128::size); + + int sidx = (threadIdx.x + WARP_SIZE * threadIdx.y) * x128::size; + for(int i = sidx; i < C; i += blockDim.y * WARP_SIZE * x128::size) { + s_weight[i/x128::size] = load128(weight + i); + s_bias[i/x128::size] = load128(bias + i); + } + __syncthreads(); + + int idx = blockIdx.x * blockDim.y + threadIdx.y; + if(idx >= N) { return; } // guard + + // adjust pointers to current token + inp += idx * C; + out += idx * C; + + const float eps = 1e-5f; + float sum = 0.0f; + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = load128cs(inp + c); + for(int k = 0; k < x128::size; ++k) { + sum += (float)in_data[k]; + } + s_in[c / x128::size] = in_data; + } + + sum = warpReduceSum(sum); + float m = sum / C; + float v = 0.f; + + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = s_in[c / x128::size]; + for(int k = 0; k < x128::size; ++k) { + v += ((float)in_data[k] - m) * ((float)in_data[k] - m); + } + } + + v = warpReduceSum(v) / C; + float s = rsqrtf(v + eps); + + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = s_in[c / x128::size]; + const x128 w = s_weight[c / x128::size]; + const x128 b = s_bias[c / x128::size]; + x128 out_data; + for(int k = 0; k < x128::size; ++k) { + float n = s * ((float)in_data[k] - m); // normalized output + float o = n * (float)w[k] + (float)b[k]; // scale and shift it + out_data[k] = (floatX)o; + } + + store128cs(out + c, out_data); + } + // cache the mean and rstd for the backward pass later + if(threadIdx.x == 0 && mean != nullptr) { + __stcs(mean + idx, (floatX)m); + } + // store the rstd, no need to cache it + if(threadIdx.x == 0 && rstd != nullptr) { + __stcs(rstd + idx, (floatX)s); + } +} + __global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, floatX* mean, floatX* rstd, const floatX* inp1, const floatX* inp2, const floatX* weight, const floatX* bias, @@ -369,15 +444,28 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with // ---------------------------------------------------------------------------- // kernel launchers +// similar to `fused_residual_forward5` void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, floatX* inp, const floatX* weight, const floatX* bias, int B, int T, int C, cudaStream_t stream) { - NVTX_RANGE_FN(); const int block_size = 256; - assert(block_size % WARP_SIZE == 0 && block_size <= 1024); + int block_y = block_size / WARP_SIZE; const int N = B * T; - const int grid_size = N; - layernorm_forward_kernel5<<>>(out, mean, rstd, inp, weight, bias, N, C); + const int grid_size = CEIL_DIV(N, block_y); + size_t smem = (2 + block_y) * C * sizeof(floatX); + + // in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute + // this may fail, in which case we fall back to the smem free implementation. + cudaCheck(cudaGetLastError()); + auto status = cudaFuncSetAttribute(layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + cudaGetLastError(); + if (status == cudaSuccess) { + layernorm_forward_kernel6<<>>(out, mean, rstd, inp, weight, bias, N, C); + } else { + // fall back to the version without shared memory + const int grid_size = N; + layernorm_forward_kernel5<<>>(out, mean, rstd, inp, weight, bias, N, C); + } cudaCheck(cudaGetLastError()); } From f845efd49a89e514f09354a444dc65143f00ff90 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 16 Jun 2024 18:01:52 +0200 Subject: [PATCH 3/5] Add nvtx range to layernorm forward --- llmc/layernorm.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index 861efd0..978f7ea 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -448,6 +448,7 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, floatX* inp, const floatX* weight, const floatX* bias, int B, int T, int C, cudaStream_t stream) { + NVTX_RANGE_FN(); const int block_size = 256; int block_y = block_size / WARP_SIZE; const int N = B * T; From 7b5ee07e1863adf7a02c264be58fd838a8753141 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 21:53:18 +0200 Subject: [PATCH 4/5] Minor refactor of dataloader --- llmc/dataloader.h | 9 ++++----- llmc/rand.h | 10 +++++----- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/llmc/dataloader.h b/llmc/dataloader.h index d6a5b23..f3b0fca 100644 --- a/llmc/dataloader.h +++ b/llmc/dataloader.h @@ -110,7 +110,8 @@ void prepare_intra_shard_indices_(DataLoader *loader) { free(loader->intra_shard_indices); } loader->intra_shard_indices = (int*)mallocCheck(loader->shard_num_samples * sizeof(int)); - random_permutation_with_init(loader->intra_shard_indices, loader->shard_num_samples, &loader->shuffle_rng, 1); + init_identity_permutation(loader->intra_shard_indices, loader->shard_num_samples); + random_permutation(loader->intra_shard_indices, loader->shard_num_samples, &loader->shuffle_rng); } void dataloader_reset(DataLoader *loader) { @@ -118,7 +119,7 @@ void dataloader_reset(DataLoader *loader) { loader->current_sample_idx = 0; if (loader->should_shuffle) { // shuffle the shards - random_permutation_with_init(loader->shard_indices, loader->glob_result.gl_pathc, &loader->shuffle_rng, 0); + random_permutation(loader->shard_indices, loader->glob_result.gl_pathc, &loader->shuffle_rng); } dataloader_load_shard_(loader, loader->current_shard_idx); @@ -178,9 +179,7 @@ void dataloader_init(DataLoader *loader, manual_seed(&shuffle_rng, 42 + process_rank); loader->shuffle_rng = shuffle_rng; loader->shard_indices = (int*)mallocCheck(loader->glob_result.gl_pathc * sizeof(int)); - for (int i = 0; i < loader->glob_result.gl_pathc; i++) { - loader->shard_indices[i] = i; // start with identity permutation - } + init_identity_permutation(loader->shard_indices, loader->glob_result.gl_pathc); loader->intra_shard_indices = NULL; // dynamically allocated allowing different shard sizes } diff --git a/llmc/rand.h b/llmc/rand.h index 99b10fe..63df35f 100644 --- a/llmc/rand.h +++ b/llmc/rand.h @@ -220,13 +220,13 @@ void normal_(float* data, unsigned int numel, float mean, float std, mt19937_sta } } -void random_permutation_with_init(int* data, int numel, mt19937_state* state, int should_init) { - if (should_init) { - for (int i = 0; i < numel; i++) { - data[i] = i; - } +void init_identity_permutation(int *data, int numel) { + for (int i = 0; i < numel; i++) { + data[i] = i; } +} +void random_permutation(int* data, int numel, mt19937_state* state) { for (int i = numel - 1; i > 0; i--) { // pick an index j in [0, i] with equal probability int j = randint32(state) % (i + 1); From 88d4eff946683c3a96bdf313c402174ee1dfaec5 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 23:05:01 +0200 Subject: [PATCH 5/5] Swap kernel 5 for 3 for fallback --- llmc/layernorm.cuh | 81 +++++++++++++++++++--------------------------- 1 file changed, 33 insertions(+), 48 deletions(-) diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index 978f7ea..b93c347 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -17,65 +17,50 @@ E.g., the layernorms are connected to the residuals so we += in layernorm backwa // ---------------------------------------------------------------------------- // CUDA kernels -__global__ void layernorm_forward_kernel5(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd, +__global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd, const floatX* __restrict__ inp, const floatX* __restrict__ weight, const floatX* __restrict__ bias, int N, int C) { - __shared__ float shared_sum[WARP_SIZE]; // block_size max is 1024 = 32 * 32 warps - __shared__ float shared_sum2[WARP_SIZE]; // warps will be writing into shared memory after warp-reduce - - int num_warps = blockDim.x / WARP_SIZE; - int warp_id = threadIdx.x / WARP_SIZE; int lane_id = threadIdx.x % WARP_SIZE; + int warp_id = threadIdx.x / WARP_SIZE; + int num_warps = blockDim.x / WARP_SIZE; - int idx = blockIdx.x; // simply one block per row - const floatX* x = inp + idx * C; // the row of input that this group of threads is responsible for + int idx = blockIdx.x * num_warps + warp_id; + if(idx >= N) { return; } // guard - // thread coarsening through the row, reduce the sum in series - float thread_sum = 0.0; // stores sum(x) - float thread_sum2 = 0.0; // stores sum(x**2) - for (int i = threadIdx.x; i < C; i += blockDim.x) { - float xi = x[i]; - thread_sum += xi; - thread_sum2 += xi * xi; + // the row of input that this group of threads is responsible for + const floatX* x = inp + idx * C; + + // mean + float sum = 0.0f; + for (int i = lane_id; i < C; i += WARP_SIZE) { + sum += (float)x[i]; } - - // 2 warp level reductions instead of block level in order to prevent 2 syncthreads - - // warp-level reduction - float warp_sum = warpReduceSum(thread_sum); // sum(x) - float warp_sum2 = warpReduceSum(thread_sum2); // sum(x**2) - // store the warp-level reduction in shared memory - if (lane_id == 0) { - shared_sum[warp_id] = warp_sum; - shared_sum2[warp_id] = warp_sum2; - } - __syncthreads(); - // load results from shared memory to threads, pad with zeros for threads that are out of bounds - warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f; - warp_sum2 = (lane_id < num_warps) ? shared_sum2[lane_id] : 0.0f; - // now reduce the warp-level reductions - float block_sum = warpReduceSum(warp_sum); // sum(x) - float block_sum2 = warpReduceSum(warp_sum2); // sum(x**2) - - // mean, var, rstd - block_sum /= C; // mean(x) - block_sum2 /= C; // mean(x**2) - float m = block_sum; - float var = block_sum2 - m * m; - float s = rsqrtf(var + 1e-5f); - - if (threadIdx.x == 0 && mean != nullptr) { + sum = warpReduceSum(sum); + float m = sum / C; + if(lane_id == 0 && mean != nullptr) { __stcs(mean + idx, (floatX)m); } - if (threadIdx.x == 0 && rstd != nullptr) { + // rstd + sum = 0.0f; + for (int i = lane_id; i < C; i += WARP_SIZE) { + float diff = (float)x[i] - m; + sum += diff * diff; + } + sum = warpReduceSum(sum); + float s = rsqrtf(sum / C + 1e-5f); + if(lane_id == 0 && rstd != nullptr) { __stcs(rstd + idx, (floatX)s); } + // final normalization and scaling by weight/bias floatX* o = out + idx * C; - for (int i = threadIdx.x; i < C; i += blockDim.x) { - float n = s * ((float)__ldcs(x+i) - m); - __stcs(o+i, (floatX)(n * (float)weight[i] + (float)bias[i])); + for (int c = lane_id; c < C; c += WARP_SIZE) { + // load and store using the .cs "streaming" hint to the compiler, + // indicating that this data will not be reused soon, and can be streamed through the caches + // this allows the threads to get more cache-hits for the (shared) weight and bias parameters + float n = s * ((float)__ldcs(x+c) - m); + __stcs(o+c, (floatX)(n * (float)weight[c] + (float)bias[c])); } } @@ -464,8 +449,8 @@ void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, layernorm_forward_kernel6<<>>(out, mean, rstd, inp, weight, bias, N, C); } else { // fall back to the version without shared memory - const int grid_size = N; - layernorm_forward_kernel5<<>>(out, mean, rstd, inp, weight, bias, N, C); + const int grid_size_fb = CEIL_DIV(N * WARP_SIZE, block_size); + layernorm_forward_kernel3<<>>(out, mean, rstd, inp, weight, bias, N, C); } cudaCheck(cudaGetLastError()); }