import torch
from torch.utils.cpp_extension import load_inline

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

// 向量化加载结构
struct Float4 {
    float x, y, z, w;
    __device__ Float4(float4 v) : x(v.x), y(v.y), z(v.z), w(v.w) {}
    __device__ Float4(float a, float b, float c, float d) : x(a), y(b), z(c), w(d) {}
};

__device__ __forceinline__ Float4 operator+(const Float4& a, const Float4& b) {
    return Float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
}

__device__ __forceinline__ Float4 operator*(const Float4& a, const Float4& b) {
    return Float4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w);
}

// 优化的warp级归约
__device__ __forceinline__ float warp_reduce_sum(float val) {
    #pragma unroll
    for (int mask = 16; mask > 0; mask >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, mask);
    }
    return val;
}

// 高性能GroupNorm内核
__global__ void groupnorm_kernel_optimized(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ y,
    int batch_size, int channels, int height, int width,
    int num_groups, float eps
) {
    int batch_idx = blockIdx.x;
    int group_idx = blockIdx.y;
    
    if (batch_idx >= batch_size || group_idx >= num_groups) return;
    
    int channels_per_group = channels / num_groups;
    int spatial_size = height * width;
    int group_size = channels_per_group * spatial_size;
    
    int tid = threadIdx.x;
    int warp_id = tid / 32;
    int lane_id = tid % 32;
    
    int group_offset = batch_idx * channels * spatial_size + 
                      group_idx * channels_per_group * spatial_size;
    
    // 使用共享内存进行warp级归约
    extern __shared__ float shared_mem[];
    float* warp_sums = shared_mem;
    float* warp_sum_sqs = shared_mem + 32; // 最多32个warp
    
    // 第一阶段：向量化加载和局部计算
    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;
    
    // 向量化处理4个元素
    int vec_size = spatial_size * channels_per_group;
    int vec_aligned_size = (vec_size / 4) * 4;
    
    // 向量化循环
    for (int i = tid * 4; i < vec_aligned_size; i += blockDim.x * 4) {
        if (i + 3 < vec_size) {
            // 向量化加载
            float4 vec_data = *reinterpret_cast<const float4*>(&x[group_offset + i]);
            Float4 val4(vec_data);
            
            // 向量化计算
            local_sum += val4.x + val4.y + val4.z + val4.w;
            local_sum_sq += val4.x * val4.x + val4.y * val4.y + val4.z * val4.z + val4.w * val4.w;
        }
    }
    
    // 处理剩余元素
    for (int i = vec_aligned_size + tid; i < vec_size; i += blockDim.x) {
        float val = x[group_offset + i];
        local_sum += val;
        local_sum_sq += val * val;
    }
    
    // Warp级归约
    float warp_sum = warp_reduce_sum(local_sum);
    float warp_sum_sq = warp_reduce_sum(local_sum_sq);
    
    // 存储warp结果
    if (lane_id == 0) {
        warp_sums[warp_id] = warp_sum;
        warp_sum_sqs[warp_id] = warp_sum_sq;
    }
    __syncthreads();
    
    // 第二阶段：线程块级归约（只由前32个线程执行）
    float final_sum = 0.0f;
    float final_sum_sq = 0.0f;
    
    if (tid < 32) {
        final_sum = (tid < blockDim.x / 32) ? warp_sums[tid] : 0.0f;
        final_sum_sq = (tid < blockDim.x / 32) ? warp_sum_sqs[tid] : 0.0f;
        
        final_sum = warp_reduce_sum(final_sum);
        final_sum_sq = warp_reduce_sum(final_sum_sq);
    }
    
    __shared__ float mean_val;
    __shared__ float inv_std_val;
    
    if (tid == 0) {
        float mean = final_sum / group_size;
        float var = (final_sum_sq / group_size) - (mean * mean);
        var = fmaxf(var, 0.0f);
        mean_val = mean;
        inv_std_val = rsqrtf(var + eps);
    }
    __syncthreads();
    
    // 第三阶段：向量化应用归一化
    for (int c = 0; c < channels_per_group; c++) {
        int channel_idx = group_idx * channels_per_group + c;
        float w = weight[channel_idx];
        float b = bias[channel_idx];
        
        int channel_offset = group_offset + c * spatial_size;
        
        // 向量化应用
        for (int i = tid * 4; i < spatial_size - 3; i += blockDim.x * 4) {
            // 向量化加载
            float4 vec_data = *reinterpret_cast<const float4*>(&x[channel_offset + i]);
            Float4 val4(vec_data);
            
            // 向量化计算
            Float4 norm4(
                (val4.x - mean_val) * inv_std_val,
                (val4.y - mean_val) * inv_std_val,
                (val4.z - mean_val) * inv_std_val,
                (val4.w - mean_val) * inv_std_val
            );
            
            Float4 result4(
                norm4.x * w + b,
                norm4.y * w + b,
                norm4.z * w + b,
                norm4.w * w + b
            );
            
            // 向量化存储
            *reinterpret_cast<float4*>(&y[channel_offset + i]) = make_float4(result4.x, result4.y, result4.z, result4.w);
        }
        
        // 处理剩余元素
        for (int i = (spatial_size / 4) * 4 + tid; i < spatial_size; i += blockDim.x) {
            int idx = channel_offset + i;
            float val = x[idx];
            float norm_val = (val - mean_val) * inv_std_val;
            y[idx] = norm_val * w + b;
        }
    }
}

torch::Tensor groupnorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int num_groups, float eps) {
    auto x_contig = x.contiguous();
    int batch_size = x_contig.size(0);
    int channels = x_contig.size(1);
    int height = x_contig.size(2);
    int width = x_contig.size(3);
    
    auto y = torch::empty_like(x_contig);
    
    // 优化的线程配置
    dim3 blocks(batch_size, num_groups);
    int threads = 256; // 8个warp
    size_t shared_mem = 64 * sizeof(float); // 32个warp * 2个float
    
    groupnorm_kernel_optimized<<<blocks, threads, shared_mem>>>(
        x_contig.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        y.data_ptr<float>(),
        batch_size, channels, height, width, num_groups, eps
    );
    
    return y;
}
"""

groupnorm_cpp_source = """
torch::Tensor groupnorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int num_groups, float eps);
"""

# 编译优化版本
groupnorm = load_inline(
    name="groupnorm_optimized",
    cpp_sources=groupnorm_cpp_source,
    cuda_sources=groupnorm_source,
    functions=["groupnorm_cuda"],
    extra_cuda_cflags=["-O3"],  
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, num_channels, num_groups=8, eps=1e-5, affine=True):
        super(ModelNew, self).__init__()
        self.num_channels = num_channels
        self.num_groups = num_groups
        self.eps = eps
        self.affine = affine
        self.groupnorm = groupnorm  # The module containing the kernel
        
        if affine:
            self.weight = torch.nn.Parameter(torch.ones(num_channels))
            self.bias = torch.nn.Parameter(torch.zeros(num_channels))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

    def forward(self, x):
        if self.affine:
            return self.groupnorm.groupnorm_cuda(x, self.weight, self.bias, self.num_groups, self.eps)
        else:
            weight = torch.ones(self.num_channels, device=x.device)
            bias = torch.zeros(self.num_channels, device=x.device)
            return self.groupnorm.groupnorm_cuda(x, weight, bias, self.num_groups, self.eps)
