import torch
import torch.nn as nn

class Model(nn.Module):
    """
    This model performs a convolution, average pooling, applies sigmoid, and sums the result.
    """
    def __init__(self, in_channels, out_channels, kernel_size, pool_kernel_size):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.avg_pool = nn.AvgPool2d(pool_kernel_size)

    def forward(self, x):
        x = self.conv(x)
        x = self.avg_pool(x)
        x = torch.sigmoid(x)
        x = torch.sum(x, dim=[1,2,3]) # Sum over all spatial dimensions
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 384, 384
kernel_size = 3
pool_kernel_size = 4

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, pool_kernel_size]

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


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define FULL_MASK 0xffffffffu

__inline__ __device__ float warp_sum(float v) {
    v += __shfl_down_sync(FULL_MASK, v, 16);
    v += __shfl_down_sync(FULL_MASK, v, 8);
    v += __shfl_down_sync(FULL_MASK, v, 4);
    v += __shfl_down_sync(FULL_MASK, v, 2);
    v += __shfl_down_sync(FULL_MASK, v, 1);
    return v;
}

// Build transposed effective 6x6 filters: effT[(c*36+r*6+s)*64 + oc]
// for avgpool4(conv3(x)).  This makes per-warp weight loads coalesced over oc.
__global__ void make_eff_kernel(const float* __restrict__ w, float* __restrict__ effT) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = 64 * 8 * 36;
    if (idx >= total) return;

    int oc = idx & 63;           // fastest varying: output channel
    int q = idx >> 6;
    int s = q % 6;
    int r = (q / 6) % 6;
    int c = q / 36;

    float acc = 0.0f;
    #pragma unroll
    for (int ky = 0; ky < 3; ++ky) {
        int dy = r - ky;
        if ((unsigned)dy < 4u) {
            #pragma unroll
            for (int kx = 0; kx < 3; ++kx) {
                int dx = s - kx;
                if ((unsigned)dx < 4u) {
                    acc += w[(((oc * 8 + c) * 3 + ky) * 3 + kx)];
                }
            }
        }
    }
    effT[idx] = acc * 0.0625f;
}

// One block handles several spatial pooled locations for one batch item.
// Each pair of warps computes all 64 output channels for one spatial location:
// lanes are output channels, and lane 0 broadcasts each input value.  This avoids
// reloading the same input patch 64 times and keeps effT loads coalesced.
__global__ void fused_kernel(const float* __restrict__ x,
                             const float* __restrict__ effT,
                             const float* __restrict__ bias,
                             float* __restrict__ out,
                             int H, int W, int PH, int PW,
                             int spatial, int blocks_per_n) {
    int n = blockIdx.y;
    int bx = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;          // 0..7 for 256 threads
    int pair = warp >> 1;         // 0..3 spatial locations per iteration
    int half = warp & 1;          // 0: oc 0..31, 1: oc 32..63
    int oc = (half << 5) + lane;

    const float* xN = x + ((long long)n * 8 * H * W);
    float local = 0.0f;

    for (int pos = bx * 4 + pair; pos < spatial; pos += blocks_per_n * 4) {
        int ph = pos / PW;
        int pw = pos - ph * PW;
        int base_h = ph << 2;
        int base_w = pw << 2;

        float acc = bias[oc];
        int q = 0;
        #pragma unroll
        for (int c = 0; c < 8; ++c) {
            const float* xC = xN + ((long long)c * H * W);
            #pragma unroll
            for (int r = 0; r < 6; ++r) {
                const float* row = xC + ((base_h + r) * W + base_w);
                #pragma unroll
                for (int s = 0; s < 6; ++s, ++q) {
                    float xv = 0.0f;
                    if (lane == 0) xv = __ldg(row + s);
                    xv = __shfl_sync(FULL_MASK, xv, 0);
                    float wt = __ldg(effT + q * 64 + oc);
                    acc = fmaf(wt, xv, acc);
                }
            }
        }

        float sig = 1.0f / (1.0f + __expf(-acc));
        float ws = warp_sum(sig);
        if (lane == 0) local += ws;
    }

    __shared__ float smem[8];
    float wsum = warp_sum(local);
    if (lane == 0) smem[warp] = wsum;
    __syncthreads();

    if (warp == 0) {
        float v = (lane < 8) ? smem[lane] : 0.0f;
        v = warp_sum(v);
        if (lane == 0) atomicAdd(out + n, v);
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    const int N = (int)x.size(0);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    // Specialized to Conv2d(kernel=3,stride=1,padding=0) then AvgPool2d(kernel=stride=4)
    const int PH = ((H - 2) - 4) / 4 + 1;
    const int PW = ((W - 2) - 4) / 4 + 1;
    const int spatial = PH * PW;

    auto out = torch::empty({N}, x.options());
    auto effT = torch::empty({8 * 36 * 64}, x.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    make_eff_kernel<<<(64 * 8 * 36 + 255) / 256, 256, 0, stream>>>(
        weight.data_ptr<float>(), effT.data_ptr<float>());

    cudaMemsetAsync(out.data_ptr<float>(), 0, (size_t)N * sizeof(float), stream);

    const int threads = 256;
    const int blocks_per_n = 256;
    dim3 grid(blocks_per_n, N, 1);
    fused_kernel<<<grid, threads, 0, stream>>>(
        x.data_ptr<float>(), effT.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(),
        H, W, PH, PW, spatial, blocks_per_n);

    return out;
}
"""

_ext = load_inline(
    name="fused_conv_avgpool_sigmoid_sum_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, pool_kernel_size):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.avg_pool = nn.AvgPool2d(pool_kernel_size)

    def forward(self, x):
        return _ext.fused_forward(x, self.conv.weight, self.conv.bias)