```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Enable cuDNN autotuning so the (static-shape) convolution picks the fastest algorithm.
torch.backends.cudnn.benchmark = True

_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BLOCK_THREADS 256
#define NUM_WARPS (BLOCK_THREADS / 32)

__device__ __forceinline__ float sigmoid_f(float v) {
    return 1.0f / (1.0f + expf(-v));
}

__device__ __forceinline__ void block_reduce_atomic_add(float val, float* out) {
    __shared__ float sh[NUM_WARPS];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        val += __shfl_down_sync(0xffffffffu, val, off);
    if (lane == 0) sh[warp] = val;
    __syncthreads();
    if (warp == 0) {
        float v = (lane < NUM_WARPS) ? sh[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            v += __shfl_down_sync(0xffffffffu, v, off);
        if (lane == 0) atomicAdd(out, v);
    }
}

// Fused: 4x4 average pool (stride 4) -> sigmoid -> sum over C,H,W.
// One block per (n, c). Grid = (C, N). Reads the conv output exactly once.
template <bool VEC2>
__global__ void fused_pool4_sigmoid_sum_kernel(
        const float* __restrict__ x, float* __restrict__ out,
        int C, int H, int W, int OH, int OW) {
    const int n = blockIdx.y;
    const int c = blockIdx.x;
    const float* slice = x + ((int64_t)n * C + c) * (int64_t)H * W;
    const float scale = 1.0f / 16.0f;
    float local = 0.0f;
    const int total = OH * OW;
    for (int idx = threadIdx.x; idx < total; idx += BLOCK_THREADS) {
        const int py = idx / OW;
        const int px = idx - py * OW;
        const float* r0 = slice + (int64_t)(py * 4) * W + px * 4;
        float s;
        if (VEC2) {
            const float2 v0 = *reinterpret_cast<const float2*>(r0);
            const float2 v1 = *reinterpret_cast<const float2*>(r0 + 2);
            const float2 v2 = *reinterpret_cast<const float2*>(r0 + W);
            const float2 v3 = *reinterpret_cast<const float2*>(r0 + W + 2);
            const float2 v4 = *reinterpret_cast<const float2*>(r0 + 2 * W);
            const float2 v5 = *reinterpret_cast<const float2*>(r0 + 2 * W + 2);
            const float2 v6 = *reinterpret_cast<const float2*>(r0 + 3 * W);
            const float2 v7 = *reinterpret_cast<const float2*>(r0 + 3 * W + 2);
            s = (v0.x + v0.y + v1.x + v1.y)
              + (v2.x + v2.y + v3.x + v3.y)
              + (v4.x + v4.y + v5.x + v5.y)
              + (v6.x + v6.y + v7.x + v7.y);
        } else {
            s = 0.0f;
            #pragma unroll
            for (int r = 0; r < 4; ++r) {
                const float* row = r0 + r * W;
                s += row[0] + row[1] + row[2] + row[3];
            }
        }
        local += sigmoid_f(s * scale);
    }
    block_reduce_atomic_add(local, out + n);
}

// Generic pool-size fallback.
__global__ void fused_pool_generic_sigmoid_sum_kernel(
        const float* __restrict__ x, float* __restrict__ out,
        int C, int H, int W, int OH, int OW, int P) {
    const int n = blockIdx.y;
    const int c = blockIdx.x;
    const float* slice = x + ((int64_t)n * C + c) * (int64_t)H * W;
    const float scale = 1.0f / (float)(P * P);
    float local = 0.0f;
    const int total = OH * OW;
    for (int idx = threadIdx.x; idx < total; idx += BLOCK_THREADS) {
        const int py = idx / OW;
        const int px = idx - py * OW;
        const float* r0 = slice + (int64_t)(py * P) * W + px * P;
        float s = 0.0f;
        for (int r = 0; r < P; ++r) {
            const float* row = r0 + r * W;
            for (int i = 0; i < P; ++i) s += row[i];
        }
        local += sigmoid_f(s * scale);
    }
    block_reduce_atomic_add(local, out + n);
}

torch::Tensor fused_pool_sigmoid_sum(torch::Tensor x, int64_t pool_kernel_size) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(x.dim() == 4, "input must be NCHW");
    auto xc = x.contiguous();
    const int N = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int H = (int)xc.size(2);
    const int W = (int)xc.size(3);
    const int P = (int)pool_kernel_size;
    const int OH = H / P;
    const int OW = W / P;
    auto out = torch::zeros({N}, xc.options());
    if (N == 0 || C == 0 || OH == 0 || OW == 0) return out;
    dim3 grid(C, N);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xp = xc.data_ptr<float>();
    float* op = out.data_ptr<float>();
    if (P == 4) {
        if ((W & 1) == 0) {
            fused_pool4_sigmoid_sum_kernel<true><<<grid, BLOCK_THREADS, 0, stream>>>(xp, op, C, H, W, OH, OW);
        } else {
            fused_pool4_sigmoid_sum_kernel<false><<<grid, BLOCK_THREADS, 0, stream>>>(xp, op, C, H, W, OH, OW);
        }
    } else {
        fused_pool_generic_sigmoid_sum_kernel<<<grid, BLOCK_THREADS, 0, stream>>>(xp, op, C, H, W, OH, OW, P);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_fused_cpp_src = "torch::Tensor fused_pool_sigmoid_sum(torch::Tensor x, int64_t pool_kernel_size);"

_fused_ext = load_inline(
    name="fused_pool_sigmoid_sum_ext",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_pool_sigmoid_sum"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    - The 3x3 convolution stays on the highly-tuned cuDNN path (autotuned via
      cudnn.benchmark), which is compute-dominant for this workload.
    - AvgPool2d + sigmoid + sum(dim=[1,2,3]) are fused into a single custom
      CUDA kernel that reads the convolution output exactly once and never
      materializes the pooled or sigmoid intermediates, eliminating several
      full-tensor global-memory round trips and extra kernel launches.
    """
    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.pool_kernel_size = pool_kernel_size

    def forward(self, x):
        x = self.conv(x)
        return _fused_ext.fused_pool_sigmoid_sum(x, self.pool_kernel_size)
```