import torch
from torch.utils.cpp_extension import load_inline

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

// Two-pass算法：精度与性能的完美平衡
__global__ void variance_two_pass_kernel(
    const float* __restrict__ x,
    float* __restrict__ variance,
    int batch_size,
    int feature_dim
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int tid = threadIdx.x;
    int base = sample_idx * feature_dim;
    
    // 第一遍：计算均值
    float sum = 0.0f;
    for (int i = tid; i < feature_dim; i += blockDim.x) {
        sum += x[base + i];
    }
    
    // Warp级归约求和
    for (int offset = 16; offset > 0; offset /= 2) {
        sum += __shfl_down_sync(0xffffffff, sum, offset);
    }
    
    float mean = 0.0f;
    if (tid == 0) {
        mean = sum / feature_dim;
    }
    mean = __shfl_sync(0xffffffff, mean, 0);  // 广播均值
    
    // 第二遍：计算方差
    float var_sum = 0.0f;
    for (int i = tid; i < feature_dim; i += blockDim.x) {
        float diff = x[base + i] - mean;
        var_sum += diff * diff;
    }
    
    // Warp级归约求和
    for (int offset = 16; offset > 0; offset /= 2) {
        var_sum += __shfl_down_sync(0xffffffff, var_sum, offset);
    }
    
    if (tid == 0) {
        variance[sample_idx] = var_sum / feature_dim;
    }
}

torch::Tensor variance_cuda(
    torch::Tensor x
) {
    auto x_contig = x.contiguous();
    
    int batch_size = x_contig.size(0);
    int feature_dim = x_contig.size(1);
    
    auto variance = torch::zeros({batch_size}, x_contig.options());
    
    // Two-pass优化：最佳平衡点
    const int block_size = 32;  // warp大小
    
    variance_two_pass_kernel<<<batch_size, block_size>>>(
        x_contig.data_ptr<float>(),
        variance.data_ptr<float>(),
        batch_size,
        feature_dim
    );
    
    return variance;
}
"""

variance_cpp_source = """
torch::Tensor variance_cuda(torch::Tensor x);
"""

# 编译CUDA代码
variance = load_inline(
    name="variance",
    cpp_sources=variance_cpp_source,
    cuda_sources=variance_source,
    functions=["variance_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-std=c++17"
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.variance = variance

    def forward(self, x):
        return self.variance.variance_cuda(x)
