#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define WARP_SIZE 32
#define BLOCK_SIZE 128

// Warp-level reduction using shuffle instructions
__inline__ __device__ float warp_reduce_sum(float val) {
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// Combined kernel: perform parallel reduction of group_norm_bias and then broadcast the computed mean to the output array in parallel
__global__ void fused_ops_kernel_combined(
    float* output,
    const float* group_norm_bias,
    int out_channels,
    int batch_size
) {
    // Shared memory for storing partial sums from each warp
    __shared__ float shared_sums[BLOCK_SIZE / WARP_SIZE];
    // Shared memory to hold the final mean value
    __shared__ float mean_shared;

    int tid = threadIdx.x;
    int lane = tid % WARP_SIZE;
    int warp_id = tid / WARP_SIZE;

    // Each thread accumulates a partial sum from group_norm_bias using grid-stride
    float sum = 0.0f;
    for (int i = tid; i < out_channels; i += BLOCK_SIZE) {
        sum += group_norm_bias[i];
    }

    // Reduce sums within each warp
    sum = warp_reduce_sum(sum);

    // Write each warp's result to shared memory
    if (lane == 0) {
        shared_sums[warp_id] = sum;
    }
    __syncthreads();

    // Final reduction: thread 0 aggregates results from all warps
    if (tid == 0) {
        float total_sum = 0.0f;
        int num_warps = BLOCK_SIZE / WARP_SIZE;
        for (int i = 0; i < num_warps; i++) {
            total_sum += shared_sums[i];
        }
        float mean = total_sum / out_channels;
        mean_shared = mean;
    }
    __syncthreads();

    // Broadcast the computed mean to the output array using a grid-stride loop
    float mean = mean_shared;
    for (int i = tid; i < batch_size; i += BLOCK_SIZE) {
        output[i] = mean;
    }
}

// Torch binding function

torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias,
    torch::Tensor group_norm_weight,
    torch::Tensor group_norm_bias,
    int num_groups
) {
    int batch_size = x.size(0);
    auto output = torch::zeros({batch_size, 1}, x.options());
    
    // Launch one block with BLOCK_SIZE threads
    fused_ops_kernel_combined<<<1, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        group_norm_bias.data_ptr<float>(),
        group_norm_bias.size(0),
        batch_size
    );

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Combined fused ops forward function");
}
