import torch
from torch.utils.cpp_extension import load_inline
rmsnorm_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void rmsnorm_kernel(const float* __restrict__ x, float* __restrict__ y, const float* __restrict__ gamma, int dim, float eps) {
    extern __shared__ float sdata[];
    float sum_sq = 0.0f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float val = x[i];
        sum_sq += val * val;
    }
    sdata[threadIdx.x] = sum_sq;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            sdata[threadIdx.x] += sdata[threadIdx.x + stride];
        }
        __syncthreads();
    }
    float mean_sq = sdata[0] / dim;
    __syncthreads();
    float rms = rsqrtf(mean_sq + eps);
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        y[i] = (x[i] * rms) * gamma[i];
    }
}

torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor gamma, float eps) {
    auto y = torch::empty_like(x);
    auto batch_size = x.size(0);
    auto seq_len = x.size(1);
    auto dim = x.size(2);
    const int block_size = 256;
    int num_blocks = batch_size * seq_len;
    int shared_mem_size = sizeof(float) * block_size;
    rmsnorm_kernel<<<num_blocks, block_size, shared_mem_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), gamma.data_ptr<float>(), dim, eps);
    return y;
}
"""

rmsnorm_cpp_source = """
torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor gamma, float eps);
"""

rmsnorm = load_inline(
    name="rmsnorm",
    cpp_sources=rmsnorm_cpp_source,
    cuda_sources=rmsnorm_source,
    functions=["rmsnorm_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, dim, eps=1e-6):
        super(ModelNew, self).__init__()
        self.eps = eps
        self.gamma = torch.nn.Parameter(torch.ones(dim))

    def forward(self, x):
        gamma = self.gamma.to(x.device)
        return rmsnorm.rmsnorm_cuda(x, gamma, self.eps)
