ECE 408 sem project

kernels_op_9/layernorm.cuh

@@ -0,0 +1,51 @@
+#ifndef __LAYERNORM_KERNEL_CUH__
+#define __LAYERNORM_KERNEL_CUH__
+
+#include <cuda_runtime.h>
+#include <math.h>
+#include <float.h>
+#include "../utils/cuda_utils.cuh"
+
+__global__ void layernorm_forward_kernel(float* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd, const float* __restrict__ inp, const float* __restrict__ weight,
+                                         const float* __restrict__ bias, int B, int T, int C) {
+    // Implement this
+    int b = blockDim.x * blockIdx.x + threadIdx.x;
+    int t = blockDim.y * blockIdx.y + threadIdx.y;
+
+    if (b >= B || t >= T) {
+        return;
+    }
+
+    const float* __restrict__ x = inp + b * T * C + t * C;
+    float* __restrict__ o = out + b * T * C + t * C;
+
+    float sum = 0.0f;
+    for (int c = 0; c < C; ++c) {
+        sum += x[c];
+    }
+    float m = sum / C;
+
+    sum = 0.0f;
+    for (int c = 0; c < C; ++c) {
+        float diff = x[c] - m;
+        sum += diff * diff;
+    }
+    float s = rsqrtf(sum / C + 1e-5f);
+    
+    for (int c = 0; c < C; ++c) {
+        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
+    dim3 block(16,16);
+    dim3 grid((B+block.x-1)/block.x, (T+block.y-1)/block.y);
+
+    layernorm_forward_kernel<<<grid, block>>>(out,mean, rstd, inp, weight, bias,B,T, C);
+}
+
+#endif // __LAYERNORM_KERNEL_CUH__