ECE 408 sem project

kernels_req_3/layernorm.cuh

@@ -0,0 +1,102 @@
+#ifndef __LAYERNORM_KERNEL_CUH__
+#define __LAYERNORM_KERNEL_CUH__
+
+#include <cuda_runtime.h>
+#include <math.h>
+#include <float.h>
+#include "../utils/cuda_utils.cuh"
+
+#define BLOCK_DIM 256
+#define WARP_SIZE 32
+
+//lecture 15 slide ~32 
+__device__ __inline__ float warp_reduce(float val) {
+    float partialSum = val;
+    for(unsigned int stride = WARP_SIZE/2; stride > 0; stride /= 2) {
+        partialSum += __shfl_down_sync(0xffffffff, partialSum, stride);
+    }
+    return partialSum;
+}
+
+__global__ void layernorm_forward_kernel(float* out, float* mean, float* rstd, const float* inp, const float* weight,
+                                         const float* bias, int B, int T, int C) {
+    // Implement this
+    int bt = blockIdx.x;
+    if(bt >= B*T){
+        return;
+    }
+
+    const float* x = inp + bt * C;
+    float* o = out + bt * C;
+
+    // computer mean
+    float partial_sum = 0.0f;
+    for(int c = threadIdx.x ; c < C ; c += BLOCK_DIM){
+        partial_sum += x[c];
+    }
+
+    partial_sum = warp_reduce(partial_sum);
+
+    __shared__ float partial_sum_s[BLOCK_DIM / WARP_SIZE];
+    if(threadIdx.x % WARP_SIZE == 0){
+        partial_sum_s[threadIdx.x / WARP_SIZE] = partial_sum;
+    }
+
+    __syncthreads();
+
+    if(threadIdx.x < WARP_SIZE){
+        partial_sum = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_sum_s[threadIdx.x] : 0.0f;
+        partial_sum = warp_reduce(partial_sum);
+    }
+
+    //thread 0 will have the final sum
+    __shared__ float s_mean;
+    if (threadIdx.x == 0) {
+        s_mean = partial_sum / C;
+    }
+
+    __syncthreads();
+    float m = s_mean;
+
+    //compute variance
+    partial_sum = 0.0f;
+    for (int c = threadIdx.x; c < C; c += BLOCK_DIM) {
+        float diff = x[c] - m;
+        partial_sum += diff * diff;
+    }
+
+    partial_sum = warp_reduce(partial_sum);
+
+    if (threadIdx.x % WARP_SIZE == 0) {
+        partial_sum_s[threadIdx.x / WARP_SIZE] = partial_sum;
+    }
+    __syncthreads();
+
+    if (threadIdx.x < WARP_SIZE) {
+        partial_sum = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_sum_s[threadIdx.x] : 0.0f;
+        partial_sum = warp_reduce(partial_sum);
+    }
+
+    __shared__ float s_rstd;
+    if (threadIdx.x == 0) {
+        s_rstd = rsqrtf(partial_sum / C + 1e-5f);
+    }
+    __syncthreads();
+    float s = s_rstd;
+
+    //normalize
+    for (int c = threadIdx.x; c < C; c += BLOCK_DIM) {
+        float n = s * (x[c] - m);
+        o[c] = n * weight[c] + bias[c];
+    }
+}
+
+// Launch kernel here
+void layernorm_forward(float* out, float* mean, float* rstd, float* inp, float* weight, float* bias,
+                       int B, int T, int C) {
+    // Implement this
+    int num_blocks = B*T;
+    layernorm_forward_kernel<<<num_blocks, BLOCK_DIM>>>(out,mean, rstd, inp, weight, bias,B,T, C);
+}
+
+#endif // __LAYERNORM_KERNEL_CUH__