```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 CUDA kernel: AvgPool2d(pool) -> Sigmoid -> Sum over (C, H, W)
# The convolution (the compute-heavy part) is left to cuDNN, while the three
# memory-bound trailing operators are fused into a single kernel that reads
# the conv output exactly once and produces the per-batch sums directly,
# eliminating all intermediate tensors and kernel launches.
_fused_epilogue_cuda = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

template <int POOL>
__global__ void pool_sigmoid_sum_kernel(const float* __restrict__ x,
                                        float* __restrict__ out,
                                        const int C, const int H, const int W,
                                        const int PH, const int PW) {
    const int b = blockIdx.y;
    const long total = (long)C * PH * PW;
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;

    float local = 0.0f;
    if (idx < total) {
        const int pw = (int)(idx % PW);
        const int ph = (int)((idx / PW) % PH);
        const int c  = (int)(idx / ((long)PW * PH));

        const float* in = x + (((long)b * C + c) * (long)H) * (long)W;
        const int h0 = ph * POOL;
        const int w0 = pw * POOL;

        float s = 0.0f;
        #pragma unroll
        for (int i = 0; i < POOL; ++i) {
            const float* row = in + (long)(h0 + i) * W + w0;
            #pragma unroll
            for (int j = 0; j < POOL; ++j) {
                s += row[j];
            }
        }
        s *= 1.0f / (float)(POOL * POOL);
        local = 1.0f / (1.0f + expf(-s));
    }

    // Warp-level reduction.
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        local += __shfl_down_sync(0xffffffffu, local, off);
    }

    __shared__ float warp_sums[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) warp_sums[wid] = local;
    __syncthreads();

    if (wid == 0) {
        const int num_warps = (blockDim.x + 31) >> 5;
        float v = (lane < num_warps) ? warp_sums[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            v += __shfl_down_sync(0xffffffffu, v, off);
        }
        if (lane == 0 && v != 0.0f) {
            atomicAdd(&out[b], v);
        }
    }
}

torch::Tensor pool_sigmoid_sum_cuda(torch::Tensor x, int64_t pool) {
    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 4D (N, C, H, W)");
    TORCH_CHECK(pool >= 1, "pool size must be >= 1");

    auto xc = x.contiguous();
    const int B = (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 PH = H / (int)pool;
    const int PW = W / (int)pool;

    auto out = at::zeros({B}, xc.options());
    const long total = (long)C * PH * PW;
    if (total == 0 || B == 0) {
        return out;
    }

    const int threads = 256;
    const long blocks = (total + threads - 1) / threads;
    dim3 grid((unsigned int)blocks, (unsigned int)B);

    const float* xp = xc.data_ptr<float>();
    float* op = out.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    switch (pool) {
        case 1: pool_sigmoid_sum_kernel<1><<<grid, threads, 0, stream>>>(xp, op, C, H, W, PH, PW); break;
        case 2: pool_sigmoid_sum_kernel<2><<<grid, threads, 0, stream>>>(xp, op, C, H, W, PH, PW); break;
        case 3: pool_sigmoid_sum_kernel<3><<<grid, threads, 0, stream>>>(xp, op, C, H, W, PH, PW); break;
        case 4: pool_sigmoid_sum_kernel<4><<<grid, threads, 0, stream>>>(xp, op, C, H, W, PH, PW); break;
        case 5: pool_sigmoid_sum_kernel<5><<<grid, threads, 0, stream>>>(xp, op, C, H, W, PH, PW); break;
        case 6: pool_sigmoid_sum_kernel<6><<<grid, threads, 0, stream>>>(xp, op, C, H, W, PH, PW); break;
        case 7: pool_sigmoid_sum_kernel<7><<<grid, threads, 0, stream>>>(xp, op, C, H, W, PH, PW); break;
        case 8: pool_sigmoid_sum_kernel<8><<<grid, threads, 0, stream>>>(xp, op, C, H, W, PH, PW); break;
        default: TORCH_CHECK(false, "unsupported pool_kernel_size");
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_fused_epilogue_cpp = (
    "torch::Tensor pool_sigmoid_sum_cuda(torch::Tensor x, int64_t pool);"
)

_fused_epilogue = load_inline(
    name="fused_pool_sigmoid_sum",
    cpp_sources=_fused_epilogue_cpp,
    cuda_sources=_fused_epilogue_cuda,
    functions=["pool_sigmoid_sum_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuDNN convolution followed by a single fused CUDA kernel
    performing average pooling + sigmoid + full reduction over (C, H, W).
    """
    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
        # Let cuDNN autotune the fastest convolution algorithm for this shape.
        torch.backends.cudnn.benchmark = True

    def forward(self, x):
        x = self.conv(x)
        if x.is_cuda and x.dtype == torch.float32:
            return _fused_epilogue.pool_sigmoid_sum_cuda(x, self.pool_kernel_size)
        # Fallback for non-CUDA inputs.
        x = F.avg_pool2d(x, self.pool_kernel_size)
        x = torch.sigmoid(x)
        return torch.sum(x, dim=[1, 2, 3])
```