```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 -> Sigmoid -> Sum over (C, H, W).
# The conv output is read exactly once and only a per-batch scalar is
# produced, eliminating the pooled tensor, the sigmoid tensor, and two
# extra kernel launches / memory round-trips.

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

#define BLOCK_SIZE 256

__device__ __forceinline__ float blockReduceSum(float val) {
    __shared__ float shared[BLOCK_SIZE];
    const int tid = threadIdx.x;
    shared[tid] = val;
    __syncthreads();
    for (int s = BLOCK_SIZE >> 1; s > 0; s >>= 1) {
        if (tid < s) shared[tid] += shared[tid + s];
        __syncthreads();
    }
    return shared[0];
}

// Specialized kernel for pool == 4 using vectorized float2 loads.
__global__ void pool4_sigmoid_sum_kernel(const float* __restrict__ x,
                                         float* __restrict__ out,
                                         const int C, const int H, const int W,
                                         const int H_out, const int W_out) {
    const int b = blockIdx.y;
    const int per_batch = C * H_out * W_out;
    const int stride = gridDim.x * BLOCK_SIZE;
    float acc = 0.f;
    for (int i = blockIdx.x * BLOCK_SIZE + threadIdx.x; i < per_batch; i += stride) {
        const int wo = i % W_out;
        const int ho = (i / W_out) % H_out;
        const int c  = i / (W_out * H_out);
        const float* base = x + ((((long)b * C + c) * H) + (long)ho * 4) * W + (long)wo * 4;
        float s = 0.f;
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            const float2 v0 = *reinterpret_cast<const float2*>(base + r * W);
            const float2 v1 = *reinterpret_cast<const float2*>(base + r * W + 2);
            s += v0.x + v0.y + v1.x + v1.y;
        }
        const float mean = s * 0.0625f;  // divide by 16
        acc += 1.f / (1.f + expf(-mean));
    }
    const float total = blockReduceSum(acc);
    if (threadIdx.x == 0) atomicAdd(out + b, total);
}

// Generic kernel for arbitrary pool sizes.
__global__ void poolN_sigmoid_sum_kernel(const float* __restrict__ x,
                                         float* __restrict__ out,
                                         const int C, const int H, const int W,
                                         const int H_out, const int W_out,
                                         const int P) {
    const int b = blockIdx.y;
    const long per_batch = (long)C * H_out * W_out;
    const long stride = (long)gridDim.x * BLOCK_SIZE;
    float acc = 0.f;
    for (long i = (long)blockIdx.x * BLOCK_SIZE + threadIdx.x; i < per_batch; i += stride) {
        const int wo = (int)(i % W_out);
        const int ho = (int)((i / W_out) % H_out);
        const int c  = (int)(i / ((long)W_out * H_out));
        const float* base = x + ((((long)b * C + c) * H) + (long)ho * P) * W + (long)wo * P;
        float s = 0.f;
        for (int r = 0; r < P; ++r)
            for (int q = 0; q < P; ++q)
                s += base[(long)r * W + q];
        const float mean = s / (float)(P * P);
        acc += 1.f / (1.f + expf(-mean));
    }
    const float total = blockReduceSum(acc);
    if (threadIdx.x == 0) atomicAdd(out + b, total);
}

torch::Tensor fused_pool_sigmoid_sum(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 NCHW");
    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 P  = (int)pool;
    TORCH_CHECK(P > 0, "pool size must be positive");
    const int H_out = H / P;
    const int W_out = W / P;
    auto out = torch::zeros({B}, xc.options());
    if (H_out <= 0 || W_out <= 0) return out;

    const long per_batch = (long)C * H_out * W_out;
    long blocks = (per_batch + BLOCK_SIZE - 1) / BLOCK_SIZE;
    if (blocks > 4194304L) blocks = 4194304L;  // grid-stride loop handles the rest
    dim3 grid((unsigned int)blocks, (unsigned int)B);
    auto stream = at::cuda::getCurrentCUDAStream();

    const float* xp = xc.data_ptr<float>();
    float* op = out.data_ptr<float>();
    const bool vec_ok = (P == 4) && ((W & 1) == 0) &&
                        ((reinterpret_cast<uintptr_t>(xp) & 7) == 0);
    if (vec_ok) {
        pool4_sigmoid_sum_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
            xp, op, C, H, W, H_out, W_out);
    } else {
        poolN_sigmoid_sum_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
            xp, op, C, H, W, H_out, W_out, P);
    }
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_pool_sigmoid_sum(torch::Tensor x, int64_t pool);"
)

# Compile the fused pool+sigmoid+sum CUDA extension
fused_op = load_inline(
    name="fused_pool_sigmoid_sum",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_pool_sigmoid_sum"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Conv2d (cuDNN) followed by a single fused CUDA kernel that performs
    average pooling, sigmoid, and the sum 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

    def forward(self, x):
        x = self.conv(x)
        if x.is_cuda and x.dtype == torch.float32:
            # Fused: AvgPool2d + Sigmoid + Sum(dim=[1,2,3])
            return fused_op.fused_pool_sigmoid_sum(x, self.pool_kernel_size)
        # Fallback path for non-CUDA tensors
        x = F.avg_pool2d(x, self.pool_kernel_size)
        x = torch.sigmoid(x)
        return torch.sum(x, dim=[1, 2, 3])
```