```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 stays in cuDNN (it is compute bound and already near-optimal),
# while the three memory-bound trailing ops are fused into a single kernel that
# reads the conv output exactly once and produces one scalar per batch element.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256

template <int PK>
__global__ void pool_sigmoid_sum_kernel(const float* __restrict__ x,
                                        float* __restrict__ out,
                                        int C, int H, int W,
                                        int PH, int PW,
                                        int elems_per_batch,
                                        int blocks_per_batch) {
    int b   = blockIdx.x / blocks_per_batch;
    int lid = blockIdx.x - b * blocks_per_batch;
    int e   = lid * THREADS + threadIdx.x;

    float val = 0.f;
    if (e < elems_per_batch) {
        int pw = e % PW;
        int t  = e / PW;
        int ph = t % PH;
        int c  = t / PH;
        const float* base = x + ((((long)b * C + c) * H + (long)ph * PK) * W + (long)pw * PK);
        float s = 0.f;
        #pragma unroll
        for (int i = 0; i < PK; ++i) {
            const float* row = base + (long)i * W;
            #pragma unroll
            for (int j = 0; j < PK; ++j) {
                s += row[j];
            }
        }
        s *= 1.f / (float)(PK * PK);
        val = 1.f / (1.f + __expf(-s));
    }

    // Block-wide reduction (one atomic per block).
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        val += __shfl_down_sync(0xffffffffu, val, off);

    __shared__ float warp_sums[THREADS / 32];
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    if (lane == 0) warp_sums[wid] = val;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < THREADS / 32) ? warp_sums[lane] : 0.f;
        #pragma unroll
        for (int off = THREADS / 64; off > 0; off >>= 1)
            v += __shfl_down_sync(0xffffffffu, v, off);
        if (lane == 0) atomicAdd(out + b, v);
    }
}

__global__ void pool_sigmoid_sum_kernel_generic(const float* __restrict__ x,
                                                float* __restrict__ out,
                                                int C, int H, int W,
                                                int PH, int PW, int pk,
                                                float inv_pk2,
                                                int elems_per_batch,
                                                int blocks_per_batch) {
    int b   = blockIdx.x / blocks_per_batch;
    int lid = blockIdx.x - b * blocks_per_batch;
    int e   = lid * THREADS + threadIdx.x;

    float val = 0.f;
    if (e < elems_per_batch) {
        int pw = e % PW;
        int t  = e / PW;
        int ph = t % PH;
        int c  = t / PH;
        const float* base = x + ((((long)b * C + c) * H + (long)ph * pk) * W + (long)pw * pk);
        float s = 0.f;
        for (int i = 0; i < pk; ++i) {
            const float* row = base + (long)i * W;
            for (int j = 0; j < pk; ++j) {
                s += row[j];
            }
        }
        s *= inv_pk2;
        val = 1.f / (1.f + __expf(-s));
    }

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        val += __shfl_down_sync(0xffffffffu, val, off);

    __shared__ float warp_sums[THREADS / 32];
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    if (lane == 0) warp_sums[wid] = val;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < THREADS / 32) ? warp_sums[lane] : 0.f;
        #pragma unroll
        for (int off = THREADS / 64; off > 0; off >>= 1)
            v += __shfl_down_sync(0xffffffffu, v, off);
        if (lane == 0) atomicAdd(out + b, v);
    }
}

torch::Tensor pool_sigmoid_sum_cuda(torch::Tensor x, int64_t pool_kernel_size) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 4, "input must be NCHW");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");

    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 pk = (int)pool_kernel_size;
    const int PH = H / pk;
    const int PW = W / pk;

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

    const int elems_per_batch = C * PH * PW;
    const int blocks_per_batch = (elems_per_batch + THREADS - 1) / THREADS;
    dim3 grid(B * blocks_per_batch);
    auto stream = at::cuda::getCurrentCUDAStream();

    const float* xp = xc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    switch (pk) {
        case 2:
            pool_sigmoid_sum_kernel<2><<<grid, THREADS, 0, stream>>>(
                xp, op, C, H, W, PH, PW, elems_per_batch, blocks_per_batch);
            break;
        case 3:
            pool_sigmoid_sum_kernel<3><<<grid, THREADS, 0, stream>>>(
                xp, op, C, H, W, PH, PW, elems_per_batch, blocks_per_batch);
            break;
        case 4:
            pool_sigmoid_sum_kernel<4><<<grid, THREADS, 0, stream>>>(
                xp, op, C, H, W, PH, PW, elems_per_batch, blocks_per_batch);
            break;
        case 5:
            pool_sigmoid_sum_kernel<5><<<grid, THREADS, 0, stream>>>(
                xp, op, C, H, W, PH, PW, elems_per_batch, blocks_per_batch);
            break;
        default:
            pool_sigmoid_sum_kernel_generic<<<grid, THREADS, 0, stream>>>(
                xp, op, C, H, W, PH, PW, pk, 1.f / (float)(pk * pk),
                elems_per_batch, blocks_per_batch);
            break;
    }
    return out;
}
"""

fused_cpp_source = "torch::Tensor pool_sigmoid_sum_cuda(torch::Tensor x, int64_t pool_kernel_size);"

fused_pool_sigmoid_sum = load_inline(
    name="fused_pool_sigmoid_sum",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["pool_sigmoid_sum_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuDNN convolution followed by a single fused CUDA kernel
    that performs average pooling, sigmoid and the full spatial/channel sum.
    """
    def __init__(self, in_channels, out_channels, kernel_size, pool_kernel_size):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True
        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)
        x = fused_pool_sigmoid_sum.pool_sigmoid_sum_cuda(x, self.pool_kernel_size)
        return x
```