```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Fused kernel: sigmoid + sum over dims [1,2,3] in a single pass.
# The input tensor (B, C, H, W) is processed as one contiguous span of
# C*H*W floats per batch element, loaded as float4, block-reduced and
# accumulated with one atomicAdd per block.
_sigmoid_sum_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void sigmoid_sum_kernel(const float4* __restrict__ x,
                                   float* __restrict__ out,
                                   int n4,
                                   int chunk) {
    const long long b = blockIdx.y;
    const float4* __restrict__ base = x + b * (long long)n4;
    const int start = blockIdx.x * chunk;
    const int end = min(start + chunk, n4);

    float sum = 0.f;
    for (int i = start + threadIdx.x; i < end; i += blockDim.x) {
        const float4 v = base[i];
        sum += 1.f / (1.f + __expf(-v.x));
        sum += 1.f / (1.f + __expf(-v.y));
        sum += 1.f / (1.f + __expf(-v.z));
        sum += 1.f / (1.f + __expf(-v.w));
    }

    __shared__ float sh[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, off);
    if (lane == 0) sh[wid] = sum;
    __syncthreads();
    if (wid == 0) {
        const int nwarp = (blockDim.x + 31) >> 5;
        float s = (lane < nwarp) ? sh[lane] : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            s += __shfl_down_sync(0xffffffffu, s, off);
        if (lane == 0) atomicAdd(out + b, s);
    }
}

torch::Tensor sigmoid_sum_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    const int64_t B = x.size(0);
    const int64_t n = x.numel() / B;
    TORCH_CHECK(n % 4 == 0, "per-batch numel must be divisible by 4");
    const int n4 = (int)(n / 4);

    auto out = torch::zeros({B}, x.options());

    const int chunk = 4096;  // float4 elements per block
    const int blocks_per_batch = (n4 + chunk - 1) / chunk;
    dim3 grid(blocks_per_batch, (unsigned)B);
    const int threads = 256;
    sigmoid_sum_kernel<<<grid, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        reinterpret_cast<const float4*>(x.data_ptr<float>()),
        out.data_ptr<float>(), n4, chunk);
    return out;
}
"""

_sigmoid_sum_cpp_src = "torch::Tensor sigmoid_sum_cuda(torch::Tensor x);"

_fused_sigmoid_sum = load_inline(
    name="fused_sigmoid_sum_v1",
    cpp_sources=_sigmoid_sum_cpp_src,
    cuda_sources=_sigmoid_sum_cuda_src,
    functions=["sigmoid_sum_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


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

    Key algorithmic rewrite: conv(kxk, stride 1) followed by average
    pooling (pxp, stride p) is mathematically equivalent to a single
    convolution with a combined (k+p-1)x(k+p-1) kernel at stride p,
    where the combined kernel is the full 2D convolution of the original
    weights with a pxp ones kernel (scaled by 1/p^2). This:
      * cuts conv FLOPs ~4x (16*9 -> 36 mults per output for k=3, p=4),
      * shrinks the intermediate activation from 382x382 to 95x95
        (4.8 GB -> 0.3 GB), eliminating huge memory traffic,
      * lets cuDNN do the heavy lifting on a much smaller problem.

    The remaining sigmoid + global sum is replaced by a single fused
    custom CUDA kernel (one pass, vectorized loads, block reduction).
    The combined weight is cached and only recomputed if the conv
    weight changes in place.
    """

    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)
        if isinstance(pool_kernel_size, (tuple, list)):
            self._ph, self._pw = int(pool_kernel_size[0]), int(pool_kernel_size[1])
        else:
            self._ph = self._pw = int(pool_kernel_size)
        self._stride = (self._ph, self._pw)
        self._wc = None
        self._wc_key = None

    def _combined_weight(self):
        w = self.conv.weight
        key = (w.data_ptr(), w._version, w.device, w.dtype)
        if self._wc is None or self._wc_key != key:
            with torch.no_grad():
                oc, ic, kh, kw = w.shape
                ph, pw = self._ph, self._pw
                ones = torch.ones(1, 1, ph, pw, device=w.device, dtype=w.dtype)
                # full convolution of w with a ph x pw ones kernel via
                # zero-padding + cross-correlation (ones is a palindrome)
                wp = F.pad(w.detach().reshape(oc * ic, 1, kh, kw),
                           (pw - 1, pw - 1, ph - 1, ph - 1))
                wc = F.conv2d(wp, ones).reshape(oc, ic, kh + ph - 1, kw + pw - 1)
                wc = wc * (1.0 / float(ph * pw))
                self._wc = wc.contiguous()
                self._wc_key = key
        return self._wc

    def forward(self, x):
        wc = self._combined_weight()
        z = F.conv2d(x, wc, self.conv.bias, stride=self._stride)
        n_per_batch = z.numel() // z.size(0)
        if z.is_cuda and z.dtype == torch.float32 and n_per_batch % 4 == 0:
            return _fused_sigmoid_sum.sigmoid_sum_cuda(z.contiguous())
        # generic fallback (CPU / other dtypes)
        return torch.sigmoid(z).sum(dim=[1, 2, 3])
```