#ifndef __ADD_RMS_NORM_CUDA_KERNEL_H__
#define __ADD_RMS_NORM_CUDA_KERNEL_H__

#include <cub/block/block_reduce.cuh>

template <unsigned int BLOCK_SIZE, typename Tcompute, typename Tdata, typename Tweight>
__device__ void add_rmsnormBlock(
    Tdata *__restrict__ y,
    Tdata *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const Tdata *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const Tdata *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const Tweight *__restrict__ w,
    size_t nhead,
    size_t dim,
    float epsilon) {
    // Each block takes care of one head in one batch
    // Each thread deals with every block_size element in the row
    size_t batch_idx = blockIdx.x / nhead;
    size_t head_idx = blockIdx.x % nhead;

    auto y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead;
    auto a_ptr = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead;
    auto b_ptr = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead;
    auto w_ptr = w;
    Tdata *residual_out_ptr = residual_out + batch_idx * stride_residual_out_batch + head_idx * stride_residual_out_nhead;

    // Compute add(a, b) and sum of squares in one pass
    Tcompute sum_squared = 0;
    for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) {
        Tcompute sum_val = Tcompute(a_ptr[i]) + Tcompute(b_ptr[i]);
        residual_out_ptr[i] = Tdata(sum_val); // Store add result
        sum_squared += sum_val * sum_val;
    }

    // Block-reduce sum of squares
    using BlockReduce = cub::BlockReduce<Tcompute, BLOCK_SIZE>;
    __shared__ typename BlockReduce::TempStorage temp_storage;
    sum_squared = BlockReduce(temp_storage).Sum(sum_squared);

    // Thread_0 computes RMS=1/sqrt(ss/dim+epsilon) and stores in shared memory
    __shared__ Tcompute rms;
    if (threadIdx.x == 0) {
        rms = Tcompute(rsqrtf(sum_squared / Tcompute(dim) + epsilon));
    }
    __syncthreads();

    // Apply normalization: y = (a + b) * w * rms
    // Reuse stored values from residual_out
    for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) {
        Tcompute sum_val = Tcompute(residual_out_ptr[i]); // Reuse stored value
        y_ptr[i] = Tdata(sum_val * Tcompute(w_ptr[i]) * rms);
    }
}

// dim=4096, block=1024 => 4 elements per thread: full unroll + register-held sums (no 2nd read of residual_out).
template <typename Tcompute, typename Tdata, typename Tweight>
__device__ void add_rmsnormBlock_dim4096_bs1024(
    Tdata *__restrict__ y,
    Tdata *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const Tdata *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const Tdata *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const Tweight *__restrict__ w,
    size_t nhead,
    float epsilon) {
    constexpr unsigned int BS = 1024;
    constexpr size_t DIM = 4096;

    const size_t batch_idx = blockIdx.x / nhead;
    const size_t head_idx = blockIdx.x % nhead;

    Tdata *y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead;
    const Tdata *a_ptr = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead;
    const Tdata *b_ptr = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead;
    const Tweight *w_ptr = w;
    Tdata *residual_out_ptr = residual_out + batch_idx * stride_residual_out_batch + head_idx * stride_residual_out_nhead;

    const unsigned int t = threadIdx.x;

    Tcompute s0 = Tcompute(a_ptr[t]) + Tcompute(b_ptr[t]);
    Tcompute s1 = Tcompute(a_ptr[t + BS]) + Tcompute(b_ptr[t + BS]);
    Tcompute s2 = Tcompute(a_ptr[t + 2 * BS]) + Tcompute(b_ptr[t + 2 * BS]);
    Tcompute s3 = Tcompute(a_ptr[t + 3 * BS]) + Tcompute(b_ptr[t + 3 * BS]);

    residual_out_ptr[t] = Tdata(s0);
    residual_out_ptr[t + BS] = Tdata(s1);
    residual_out_ptr[t + 2 * BS] = Tdata(s2);
    residual_out_ptr[t + 3 * BS] = Tdata(s3);

    Tcompute sum_squared = s0 * s0 + s1 * s1 + s2 * s2 + s3 * s3;

    using BlockReduce = cub::BlockReduce<Tcompute, BS>;
    __shared__ typename BlockReduce::TempStorage temp_storage;
    sum_squared = BlockReduce(temp_storage).Sum(sum_squared);

    __shared__ Tcompute rms;
    if (t == 0) {
        rms = Tcompute(rsqrtf(sum_squared / Tcompute(DIM) + epsilon));
    }
    __syncthreads();

    y_ptr[t] = Tdata(s0 * Tcompute(w_ptr[t]) * rms);
    y_ptr[t + BS] = Tdata(s1 * Tcompute(w_ptr[t + BS]) * rms);
    y_ptr[t + 2 * BS] = Tdata(s2 * Tcompute(w_ptr[t + 2 * BS]) * rms);
    y_ptr[t + 3 * BS] = Tdata(s3 * Tcompute(w_ptr[t + 3 * BS]) * rms);
}

// dim=8192, block=1024 => 8 elements per thread: full unroll + register-held sums (no 2nd read of residual_out).
template <typename Tcompute, typename Tdata, typename Tweight>
__device__ void add_rmsnormBlock_dim8192_bs1024(
    Tdata *__restrict__ y,
    Tdata *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const Tdata *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const Tdata *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const Tweight *__restrict__ w,
    size_t nhead,
    float epsilon) {
    constexpr unsigned int BS = 1024;
    constexpr size_t DIM = 8192;

    const size_t batch_idx = blockIdx.x / nhead;
    const size_t head_idx = blockIdx.x % nhead;

    Tdata *y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead;
    const Tdata *a_ptr = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead;
    const Tdata *b_ptr = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead;
    const Tweight *w_ptr = w;
    Tdata *residual_out_ptr = residual_out + batch_idx * stride_residual_out_batch + head_idx * stride_residual_out_nhead;

    const unsigned int t = threadIdx.x;

    Tcompute s0 = Tcompute(a_ptr[t]) + Tcompute(b_ptr[t]);
    Tcompute s1 = Tcompute(a_ptr[t + BS]) + Tcompute(b_ptr[t + BS]);
    Tcompute s2 = Tcompute(a_ptr[t + 2 * BS]) + Tcompute(b_ptr[t + 2 * BS]);
    Tcompute s3 = Tcompute(a_ptr[t + 3 * BS]) + Tcompute(b_ptr[t + 3 * BS]);
    Tcompute s4 = Tcompute(a_ptr[t + 4 * BS]) + Tcompute(b_ptr[t + 4 * BS]);
    Tcompute s5 = Tcompute(a_ptr[t + 5 * BS]) + Tcompute(b_ptr[t + 5 * BS]);
    Tcompute s6 = Tcompute(a_ptr[t + 6 * BS]) + Tcompute(b_ptr[t + 6 * BS]);
    Tcompute s7 = Tcompute(a_ptr[t + 7 * BS]) + Tcompute(b_ptr[t + 7 * BS]);

    residual_out_ptr[t] = Tdata(s0);
    residual_out_ptr[t + BS] = Tdata(s1);
    residual_out_ptr[t + 2 * BS] = Tdata(s2);
    residual_out_ptr[t + 3 * BS] = Tdata(s3);
    residual_out_ptr[t + 4 * BS] = Tdata(s4);
    residual_out_ptr[t + 5 * BS] = Tdata(s5);
    residual_out_ptr[t + 6 * BS] = Tdata(s6);
    residual_out_ptr[t + 7 * BS] = Tdata(s7);

    Tcompute sum_squared =
        s0 * s0 + s1 * s1 + s2 * s2 + s3 * s3 + s4 * s4 + s5 * s5 + s6 * s6 + s7 * s7;

    using BlockReduce = cub::BlockReduce<Tcompute, BS>;
    __shared__ typename BlockReduce::TempStorage temp_storage;
    sum_squared = BlockReduce(temp_storage).Sum(sum_squared);

    __shared__ Tcompute rms;
    if (t == 0) {
        rms = Tcompute(rsqrtf(sum_squared / Tcompute(DIM) + epsilon));
    }
    __syncthreads();

    y_ptr[t] = Tdata(s0 * Tcompute(w_ptr[t]) * rms);
    y_ptr[t + BS] = Tdata(s1 * Tcompute(w_ptr[t + BS]) * rms);
    y_ptr[t + 2 * BS] = Tdata(s2 * Tcompute(w_ptr[t + 2 * BS]) * rms);
    y_ptr[t + 3 * BS] = Tdata(s3 * Tcompute(w_ptr[t + 3 * BS]) * rms);
    y_ptr[t + 4 * BS] = Tdata(s4 * Tcompute(w_ptr[t + 4 * BS]) * rms);
    y_ptr[t + 5 * BS] = Tdata(s5 * Tcompute(w_ptr[t + 5 * BS]) * rms);
    y_ptr[t + 6 * BS] = Tdata(s6 * Tcompute(w_ptr[t + 6 * BS]) * rms);
    y_ptr[t + 7 * BS] = Tdata(s7 * Tcompute(w_ptr[t + 7 * BS]) * rms);
}

#endif