#ifndef __LAYERNORM_KERNEL_CUH__
#define __LAYERNORM_KERNEL_CUH__

#include <cuda_runtime.h>
#include <math.h>
#include <float.h>
#include "../utils/cuda_utils.cuh"

#define MAX_C 768
__constant__ float d_weight[MAX_C];
__constant__ float d_bias[MAX_C];

__global__ void layernorm_forward_kernel(float* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd, const float* __restrict__ inp, 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 * d_weight[c] + d_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) {
    
    cudaMemcpyToSymbolAsync(d_weight, weight, C * sizeof(float));
    cudaMemcpyToSymbolAsync(d_bias, bias, C * sizeof(float));

    // 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, B,T, C);
}

#endif // __LAYERNORM_KERNEL_CUH__