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

// Optimized fused classifier kernel with warp-level reduction and inner loop unrolling.
// For each sample b and each output channel i, we compute:
//   output[b, i] = (1/(H*W)) * Σ_{h=0}^{H-1} Σ_{w=0}^{W-1} max( dot(input[b, :, h, w], weight[i]) + bias[i], 0 )
// Input: [B, Cin, H, W]; Weight: [num_classes, Cin]; Bias: [num_classes]
// Grid layout: gridDim.x = num_classes, gridDim.y = B; each block processes one (b, i) pair.
extern "C" __global__ void fused_classifier_kernel_optimized(
    const float* __restrict__ input,   // [B, Cin, H, W]
    const float* __restrict__ weight,    // [num_classes, Cin]
    const float* __restrict__ bias,      // [num_classes]
    float* __restrict__ output,          // [B, num_classes]
    int Cin,
    int H,
    int W,
    int num_classes
) {
    int b = blockIdx.y;  // sample index
    int i = blockIdx.x;  // output channel index
    int T = H * W;       // total spatial elements

    int tid = threadIdx.x;
    float partial_sum = 0.0f;

    // Process spatial elements in a loop over j (each representing a pixel (h, w)).
    // Unroll the inner loop over channels by factor 4 (Cin is assumed divisible by 4).
    for (int j = tid; j < T; j += blockDim.x) {
        int h = j / W;
        int w = j % W;
        float dot = 0.0f;
        // Compute base index for the spatial location (h, w) for a given sample b.
        // The layout is [B, Cin, H, W] with channel stride = H*W.
        int base = (b * Cin * H * W) + (h * W) + w;
        // Unroll the loop over channels by 4.
        for (int c = 0; c < Cin; c += 4) {
            // Each channel index offset is (H*W)
            int offset0 = c * H * W;
            int offset1 = (c + 1) * H * W;
            int offset2 = (c + 2) * H * W;
            int offset3 = (c + 3) * H * W;
            float in0 = __ldg(&input[ base + offset0 ]);
            float in1 = __ldg(&input[ base + offset1 ]);
            float in2 = __ldg(&input[ base + offset2 ]);
            float in3 = __ldg(&input[ base + offset3 ]);
            int weight_index = i * Cin + c;
            float w0 = __ldg(&weight[ weight_index ]);
            float w1 = __ldg(&weight[ weight_index + 1 ]);
            float w2 = __ldg(&weight[ weight_index + 2 ]);
            float w3 = __ldg(&weight[ weight_index + 3 ]);
            dot += in0 * w0 + in1 * w1 + in2 * w2 + in3 * w3;
        }
        dot += bias[i];
        float relu_val = fmaxf(dot, 0.0f);
        partial_sum += relu_val;
    }
    
    // Allocate shared memory for reduction.
    extern __shared__ float sdata[];
    sdata[tid] = partial_sum;
    __syncthreads();
    
    // Reduce in shared memory until we have at most 32 threads.
    for (unsigned int s = blockDim.x / 2; s >= 32; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // Use warp-level shuffles for final reduction.
    if (tid < 32) {
        float sum = sdata[tid];
        for (int offset = 16; offset > 0; offset /= 2) {
            sum += __shfl_down_sync(0xffffffff, sum, offset);
        }
        if (tid == 0) {
            sdata[0] = sum;
        }
    }
    __syncthreads();
    
    // Thread 0 writes the average to the output.
    if (tid == 0) {
        output[b * num_classes + i] = sdata[0] / ((float)T);
    }
}

torch::Tensor fused_classifier_forward(torch::Tensor input,
                                         torch::Tensor weight,
                                         torch::Tensor bias) {
    c10::cuda::CUDAGuard device_guard(input.device());
    
    int B = input.size(0);
    int Cin = input.size(1);
    int H = input.size(2);
    int W = input.size(3);
    int num_classes = weight.size(0);
    
    auto output = torch::empty({B, num_classes}, input.options());
    
    // Configure grid: gridDim.x = num_classes, gridDim.y = B.
    dim3 gridDim(num_classes, B);
    int threads = 256;
    dim3 blockDim(threads);
    size_t shared_mem = threads * sizeof(float);
    
    fused_classifier_kernel_optimized<<<gridDim, blockDim, shared_mem, c10::cuda::getCurrentCUDAStream()>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        output.data_ptr<float>(),
        Cin, H, W, num_classes
    );
    
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &fused_classifier_forward, "Fused classifier: 1x1 conv + ReLU + Global AvgPool (optimized with warp shuffle and unrolling)");
}