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

cpp_source = """
#include <torch/extension.h>

torch::Tensor gem_pooling_cuda_forward(const torch::Tensor& input, float p, float eps);
"""

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

#define BLOCK_SIZE 128
#define WARP_SIZE 32

struct __align__(16) Float4 {
    float x, y, z, w;
};

// --- Warp Reduce Sum ---
template<typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// Kernel: One Warp per Channel slice
__global__ void gem_pooling_warp_kernel(
    float* __restrict__ output,
    const float* __restrict__ input,
    const int num_slices,  // N * C
    const int hw,          // H * W
    const float p,
    const float eps) 
{
    // 计算当前 Warp 的全局 ID
    int global_warp_id = blockIdx.x * (BLOCK_SIZE / WARP_SIZE) + (threadIdx.x / WARP_SIZE);
    
    if (global_warp_id >= num_slices) return;

    int lane_id = threadIdx.x % WARP_SIZE;

    // 定位输入输出指针
    const float* input_slice = input + global_warp_id * hw;
    
    float sum = 0.0f;

    // --- 向量化循环 (Vectorized Loop) ---
    // 每次处理 128 个元素 (32 lanes * 4 floats)
    int i = lane_id * 4;
    
    while (i < hw) {
        if (i + 4 <= hw) {
            Float4 val4 = *reinterpret_cast<const Float4*>(&input_slice[i]);
            
            // Unrolled computation
            val4.x = fmaxf(val4.x, eps); sum += __powf(val4.x, p);
            val4.y = fmaxf(val4.y, eps); sum += __powf(val4.y, p);
            val4.z = fmaxf(val4.z, eps); sum += __powf(val4.z, p);
            val4.w = fmaxf(val4.w, eps); sum += __powf(val4.w, p);
        } else {
            // Tail handling (scalar fallback) for the last few pixels
            if (i < hw) {
                float val = input_slice[i];
                val = fmaxf(val, eps);
                sum += __powf(val, p);
            }
            if (i + 1 < hw) {
                float val = input_slice[i+1];
                val = fmaxf(val, eps);
                sum += __powf(val, p);
            }
            if (i + 2 < hw) {
                float val = input_slice[i+2];
                val = fmaxf(val, eps);
                sum += __powf(val, p);
            }
            // i+3 必定越界，不用检查
        }
        // Stride: Warp Size * 4
        i += WARP_SIZE * 4;
    }

    // Warp 内规约
    sum = warp_reduce_sum(sum);

    // Lane 0 写入结果
    if (lane_id == 0) {
        float mean = sum / static_cast<float>(hw);
        output[global_warp_id] = __powf(mean, 1.0f / p);
    }
}

torch::Tensor gem_pooling_cuda_forward(const torch::Tensor& input, float p, float eps) {
    TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");
    TORCH_CHECK(input.dim() == 4, "Input must be 4D (N, C, H, W)");

    const int N = input.size(0);
    const int C = input.size(1);
    const int H = input.size(2);
    const int W = input.size(3);
    const int HW = H * W;
    const int num_slices = N * C;

    auto output = torch::empty({N, C}, input.options());

    // 计算 Grid Size
    int warps_per_block = BLOCK_SIZE / WARP_SIZE;
    int grid_size = (num_slices + warps_per_block - 1) / warps_per_block;

    dim3 grid(grid_size);
    dim3 block(BLOCK_SIZE);

    gem_pooling_warp_kernel<<<grid, block>>>(
        output.data_ptr<float>(),
        input.data_ptr<float>(),
        num_slices,
        HW,
        p, eps
    );

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, p=3.0, eps=1e-6):
        super(ModelNew, self).__init__()
        self.p = p
        self.eps = eps
        self.op = load_inline(
            name='gem_pooling_warp_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['gem_pooling_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=['-O3']
        )

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return self.op.gem_pooling_cuda_forward(input_tensor.contiguous(), self.p, self.eps)