```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

# Enable cuDNN autotuning so the convolution picks the fastest algorithm.
torch.backends.cudnn.benchmark = True

# Fused CUDA kernel: AvgPool2d (non-overlapping) + sigmoid + sum over [C, H, W].
# One kernel launch, reads the convolution output exactly once and writes one
# scalar per batch element. Works directly on fp16 or fp32 convolution output,
# which also halves the memory traffic of the (huge) intermediate tensor.
_fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

namespace {

__device__ __forceinline__ float2 load_two(const float* p) {
    return *reinterpret_cast<const float2*>(p);
}

__device__ __forceinline__ float2 load_two(const __half* p) {
    return __half22float2(*reinterpret_cast<const __half2*>(p));
}

template <typename scalar_t, int POOL>
__global__ void fused_pool_sigmoid_sum_kernel(
        const scalar_t* __restrict__ x,
        float* __restrict__ out,
        const int C, const int H, const int W,
        const int Hp, const int Wp, const int num_windows) {
    const int n = blockIdx.y;
    const scalar_t* __restrict__ img = x + (size_t)n * C * H * W;

    float acc = 0.0f;
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int stride = gridDim.x * blockDim.x;
    constexpr float inv = 1.0f / (float)(POOL * POOL);

    for (int v = tid; v < num_windows; v += stride) {
        const int wx = v % Wp;
        const int t  = v / Wp;
        const int wy = t % Hp;
        const int c  = t / Hp;

        const scalar_t* p = img + ((size_t)c * H + (size_t)wy * POOL) * W + (size_t)wx * POOL;
        float s = 0.0f;
        #pragma unroll
        for (int r = 0; r < POOL; ++r) {
            #pragma unroll
            for (int q = 0; q < POOL / 2; ++q) {
                const float2 val = load_two(p + 2 * q);
                s += val.x + val.y;
            }
            p += W;
        }
        const float m = s * inv;
        acc += 1.0f / (1.0f + expf(-m));
    }

    // Block reduction.
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        acc += __shfl_down_sync(0xffffffffu, acc, off);

    __shared__ float warp_sums[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) warp_sums[wid] = acc;
    __syncthreads();
    if (wid == 0) {
        const int nwarps = (blockDim.x + 31) >> 5;
        float v = (lane < nwarps) ? 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) atomicAdd(&out[n], v);
    }
}

} // namespace

torch::Tensor fused_pool_sigmoid_sum(torch::Tensor x, int64_t pool) {
    TORCH_CHECK(x.is_cuda(), "input must be CUDA");
    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 Hp = H / (int)pool;
    const int Wp = W / (int)pool;

    auto out = torch::zeros({N}, xc.options().dtype(torch::kFloat32));
    if (N == 0 || C == 0 || Hp == 0 || Wp == 0) return out;

    const int num_windows = C * Hp * Wp;
    const int threads = 256;
    int blocks = (num_windows + threads - 1) / threads;
    if (blocks > 1024) blocks = 1024;
    dim3 grid(blocks, N);
    auto stream = at::cuda::getCurrentCUDAStream();

    if (xc.scalar_type() == torch::kHalf) {
        if (pool == 4) {
            fused_pool_sigmoid_sum_kernel<__half, 4><<<grid, threads, 0, stream>>>(
                reinterpret_cast<const __half*>(xc.data_ptr<at::Half>()),
                out.data_ptr<float>(), C, H, W, Hp, Wp, num_windows);
        } else if (pool == 2) {
            fused_pool_sigmoid_sum_kernel<__half, 2><<<grid, threads, 0, stream>>>(
                reinterpret_cast<const __half*>(xc.data_ptr<at::Half>()),
                out.data_ptr<float>(), C, H, W, Hp, Wp, num_windows);
        } else {
            TORCH_CHECK(false, "unsupported pool size");
        }
    } else if (xc.scalar_type() == torch::kFloat) {
        if (pool == 4) {
            fused_pool_sigmoid_sum_kernel<float, 4><<<grid, threads, 0, stream>>>(
                xc.data_ptr<float>(), out.data_ptr<float>(), C, H, W, Hp, Wp, num_windows);
        } else if (pool == 2) {
            fused_pool_sigmoid_sum_kernel<float, 2><<<grid, threads, 0, stream>>>(
                xc.data_ptr<float>(), out.data_ptr<float>(), C, H, W, Hp, Wp, num_windows);
        } else {
            TORCH_CHECK(false, "unsupported pool size");
        }
    } else {
        TORCH_CHECK(false, "unsupported dtype");
    }
    return out;
}
"""

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

_fused_ops = load_inline(
    name="fused_pool_sigmoid_sum",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["fused_pool_sigmoid_sum"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized version:
      * Convolution is executed in fp16 on tensor cores (fp32 accumulation inside
        cuDNN), which also halves the size of the huge intermediate tensor.
      * AvgPool + sigmoid + global sum are fused into a single custom CUDA kernel
        that reads the convolution output exactly once and emits one scalar per
        batch element.
    Falls back to the original PyTorch ops for unsupported configurations.
    """
    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.avg_pool = nn.AvgPool2d(pool_kernel_size)
        self.pool_kernel_size = pool_kernel_size

    def forward(self, x):
        use_fused = (
            x.is_cuda
            and x.dtype == torch.float32
            and self.pool_kernel_size in (2, 4)
            and x.shape[-1] % 2 == 0
        )
        if use_fused:
            w = self.conv.weight
            b = self.conv.bias
            conv_out = F.conv2d(
                x.half(),
                w.half(),
                b.half() if b is not None else None,
                self.conv.stride,
                self.conv.padding,
                self.conv.dilation,
                self.conv.groups,
            )
            return _fused_ops.fused_pool_sigmoid_sum(conv_out, self.pool_kernel_size)

        x = self.conv(x)
        x = self.avg_pool(x)
        x = torch.sigmoid(x)
        return torch.sum(x, dim=[1, 2, 3])
```