import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

rmsnorm_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void rmsnorm_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    float* __restrict__ y,
    int batch,
    int features,
    float eps
) {
    int row = blockIdx.x;
    if (row >= batch) return;
    int tid = threadIdx.x;

    extern __shared__ float sdata[];
    float sum_sq = 0.0f;

    for (int i = tid; i < features; i += blockDim.x) {
        float v = x[row * features + i];
        sum_sq += v * v;
    }
    sdata[tid] = sum_sq;
    __syncthreads();

    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
        if (tid < offset) {
            sdata[tid] += sdata[tid + offset];
        }
        __syncthreads();
    }

    float rms = rsqrtf(sdata[0] / features + eps);
    __syncthreads();  // 保证 rms 可见

    for (int i = tid; i < features; i += blockDim.x) {
        float v = x[row * features + i];
        float w = weight[i];
        y[row * features + i] = v * rms * w;
    }
}

torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, float eps) {
    TORCH_CHECK(x.is_cuda(), "x 必须是 CUDA 张量");
    TORCH_CHECK(weight.is_cuda(), "weight 必须是 CUDA 张量");
    TORCH_CHECK(x.dim() == 2, "当前内核仅支持二维输入张量");
    TORCH_CHECK(weight.dim() == 1, "RMSNorm 权重必须是一维向量");
    TORCH_CHECK(x.size(1) == weight.size(0), "输入最后一维与权重长度不匹配");

    int batch = x.size(0);
    int features = x.size(1);

    auto y = torch::empty_like(x);

    int threads = 256;
    if (features < threads) {
        threads = 1;
        while (threads < features) threads <<= 1;
        if (threads < 32) threads = 32;
    }
    size_t shared = threads * sizeof(float);

    rmsnorm_kernel<<<batch, threads, shared>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        y.data_ptr<float>(),
        batch,
        features,
        eps
    );
    return y;
}
"""

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

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

class ModelNew(nn.Module):
    def __init__(self, weight: torch.Tensor, eps: float = 1e-6):
        super().__init__()
        if weight.dim() != 1:
            raise ValueError("RMSNorm 权重必须是一维向量。")
        self.weight = nn.Parameter(weight.clone())
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return rmsnorm.rmsnorm_cuda(x, self.weight, self.eps)
