import torch
from torch.utils.cpp_extension import load_inline

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

// RMSNorm + Residual融合kernel
__global__ void rmsnorm_residual_kernel(
    const float* __restrict__ x,
    const float* __restrict__ residual,
    const float* __restrict__ weight,
    float* __restrict__ y,
    int batch_size,
    int seq_len,
    int dim,
    float eps
) {
    // 每个block处理一个序列位置
    int seq_idx = blockIdx.x;
    int batch_idx = blockIdx.y;
    
    if (seq_idx >= seq_len || batch_idx >= batch_size) return;
    
    extern __shared__ float sdata[];
    
    // 计算当前序列位置的起始索引
    int base_idx = batch_idx * seq_len * dim + seq_idx * dim;
    
    // 第一步：计算平方和
    float sum_sq = 0.0f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float val = x[base_idx + 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();
    }
    
    // 第三步：计算RMS
    float mean_sq = sdata[0] / dim;
    float rms = rsqrtf(mean_sq + eps);
    __syncthreads();
    
    // 第四步：应用RMSNorm并加上residual
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        int idx = base_idx + i;
        // RMSNorm(x) + residual
        y[idx] = (x[idx] * rms * weight[i]) + residual[idx];
    }
}

torch::Tensor rmsnorm_residual_cuda(
    torch::Tensor x,
    torch::Tensor residual,
    torch::Tensor weight,
    float eps
) {
    auto batch_size = x.size(0);
    auto seq_len = x.size(1);
    auto dim = x.size(2);
    
    auto y = torch::empty_like(x);
    
    const int block_size = 256;
    dim3 grid(seq_len, batch_size);
    int shared_mem_size = sizeof(float) * block_size;
    
    rmsnorm_residual_kernel<<<grid, block_size, shared_mem_size>>>(
        x.data_ptr<float>(),
        residual.data_ptr<float>(),
        weight.data_ptr<float>(),
        y.data_ptr<float>(),
        batch_size,
        seq_len,
        dim,
        eps
    );
    
    return y;
}
"""

rmsnorm_residual_cpp_source = """
torch::Tensor rmsnorm_residual_cuda(
    torch::Tensor x,
    torch::Tensor residual,
    torch::Tensor weight,
    float eps
);
"""

# 编译CUDA扩展
rmsnorm_residual = load_inline(
    name="rmsnorm_residual_fused",
    cpp_sources=rmsnorm_residual_cpp_source,
    cuda_sources=rmsnorm_residual_source,
    functions=["rmsnorm_residual_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-std=c++17"
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    """
    RMSNorm + Residual融合模型
    """
    def __init__(self, dim, eps=1e-6):
        super(ModelNew, self).__init__()
        self.eps = eps
        self.weight = torch.nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
        """
        融合实现：RMSNorm + Residual
        
        Args:
            x (torch.Tensor): 输入张量 shape (B, L, D)
            residual (torch.Tensor): 残差张量 shape (B, L, D)
        
        Returns:
            torch.Tensor: RMSNorm(x) + residual
        """
        return rmsnorm_residual.rmsnorm_residual_cuda(x, residual, self.weight, self.eps)
