#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <math.h>

#ifndef WARP_SIZE
#define WARP_SIZE 32
#endif

// Optimized GroupNorm forward pass kernel with vectorized loads/stores and loop unrolling.
// Each block processes one (sample, group) pair.
// x: input tensor of shape [N, C, H, W];
// weight, bias: affine parameters with shape [C].
__global__ void group_norm_forward_kernel(const float* __restrict__ x,
                                            const float* __restrict__ weight,
                                            const float* __restrict__ bias,
                                            float* __restrict__ y,
                                            int N, int C, int H, int W,
                                            int G, float eps) {
    // Determine sample index (n) and group index (g) from block index.
    int block_id = blockIdx.x; // Total blocks = N * G.
    int n = block_id / G;
    int g = block_id % G;
    
    int group_channels = C / G;
    int HW = H * W;         // Number of spatial elements per channel.
    int L = group_channels * HW;  // Total number of elements in the group.
    int sample_offset = n * C * HW;
    
    // Determine whether we can use vectorized (float4) loads and stores.
    bool use_vectorized = ((HW & 3) == 0);
    
    float local_sum = 0.0f, local_sum_sq = 0.0f;
    
    // Reduction phase: accumulate sum and sum of squares.
    if (use_vectorized) {
        int num_vec = HW / 4;
        // Loop over channels in the group.
        for (int c = 0; c < group_channels; c++) {
            int c_abs = g * group_channels + c;
            int base = sample_offset + c_abs * HW;
            const float4* p = reinterpret_cast<const float4*>(x + base);
            // Unroll the vector loop if possible.
            for (int vec = threadIdx.x; vec < num_vec; vec += blockDim.x) {
                float4 v = p[vec];
                local_sum += v.x + v.y + v.z + v.w;
                local_sum_sq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
            }
        }
    } else {
        // Fallback to scalar loads.
        for (int c = 0; c < group_channels; c++) {
            int c_abs = g * group_channels + c;
            int base = sample_offset + c_abs * HW;
            for (int hw = threadIdx.x; hw < HW; hw += blockDim.x) {
                float val = x[base + hw];
                local_sum += val;
                local_sum_sq += val * val;
            }
        }
    }
    
    // Warp-level reduction using shuffle.
    unsigned mask = 0xffffffff;
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        local_sum += __shfl_down_sync(mask, local_sum, offset);
        local_sum_sq += __shfl_down_sync(mask, local_sum_sq, offset);
    }
    
    // Shared memory for warp-level partial sums.
    __shared__ float shared_sum[32];    // Assume max 32 warps per block.
    __shared__ float shared_sum_sq[32];
    int warp_id = threadIdx.x / WARP_SIZE;
    int lane = threadIdx.x & (WARP_SIZE - 1);
    if (lane == 0) {
        shared_sum[warp_id] = local_sum;
        shared_sum_sq[warp_id] = local_sum_sq;
    }
    __syncthreads();
    
    float mean = 0.0f, inv_std = 0.0f;
    if (threadIdx.x == 0) {
        int num_warps = (blockDim.x + WARP_SIZE - 1) / WARP_SIZE;
        float total_sum = 0.0f, total_sum_sq = 0.0f;
        for (int i = 0; i < num_warps; i++) {
            total_sum += shared_sum[i];
            total_sum_sq += shared_sum_sq[i];
        }
        mean = total_sum / L;
        float var = total_sum_sq / L - mean * mean;
        inv_std = rsqrtf(var + eps);
        // Broadcast computed mean and inv_std.
        shared_sum[0] = mean;
        shared_sum[1] = inv_std;
    }
    __syncthreads();
    mean = shared_sum[0];
    inv_std = shared_sum[1];
    
    // Normalization phase: apply normalization and affine transformation.
    if (use_vectorized) {
        int num_vec = HW / 4;
        for (int c = 0; c < group_channels; c++) {
            int c_abs = g * group_channels + c;
            int base = sample_offset + c_abs * HW;
            float w_val = weight[c_abs];
            float b_val = bias[c_abs];
            const float4* p_in = reinterpret_cast<const float4*>(x + base);
            float4* p_out = reinterpret_cast<float4*>(y + base);
            for (int vec = threadIdx.x; vec < num_vec; vec += blockDim.x) {
                float4 v = p_in[vec];
                float norm_x = (v.x - mean) * inv_std;
                float norm_y = (v.y - mean) * inv_std;
                float norm_z = (v.z - mean) * inv_std;
                float norm_w = (v.w - mean) * inv_std;
                float4 out;
                out.x = norm_x * w_val + b_val;
                out.y = norm_y * w_val + b_val;
                out.z = norm_z * w_val + b_val;
                out.w = norm_w * w_val + b_val;
                p_out[vec] = out;
            }
        }
    } else {
        // Fallback to scalar normalization.
        for (int c = 0; c < group_channels; c++) {
            int c_abs = g * group_channels + c;
            int base = sample_offset + c_abs * HW;
            float w_val = weight[c_abs];
            float b_val = bias[c_abs];
            for (int hw = threadIdx.x; hw < HW; hw += blockDim.x) {
                int idx = base + hw;
                float val = x[idx];
                float norm = (val - mean) * inv_std;
                y[idx] = norm * w_val + b_val;
            }
        }
    }
}

torch::Tensor group_norm_forward(torch::Tensor x,
                                   torch::Tensor weight,
                                   torch::Tensor bias,
                                   int num_groups,
                                   float eps) {
    c10::cuda::CUDAGuard device_guard(x.device());
    
    // x has shape [N, C, H, W].
    int N = x.size(0);
    int C = x.size(1);
    int H = x.size(2);
    int W = x.size(3);
    TORCH_CHECK(C % num_groups == 0, "C must be divisible by num_groups");
    
    auto y = torch::empty_like(x);
    
    int G = num_groups;
    int total_blocks = N * G;
    int threads = 256; // Tunable: choose a block size that fits your hardware.
    size_t shared_mem = 32 * 2 * sizeof(float); // Shared memory for warp aggregation.
    
    group_norm_forward_kernel<<<total_blocks, threads, shared_mem, c10::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        y.data_ptr<float>(),
        N, C, H, W, G, eps
    );
    
    return y;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("group_norm_forward", &group_norm_forward, "Fused GroupNorm forward (optimized with vectorized loads/stores and loop unrolling)");
}