```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 AvgPool2d + Sigmoid + Sum(dim=[1, 2, 3]) custom CUDA kernel.
#
# The baseline runs three separate PyTorch operators after the convolution:
#   1) AvgPool2d: reads the whole conv output (N*C*H*W) and writes the pooled map.
#   2) Sigmoid:   reads and writes the pooled map.
#   3) Sum:       reads the pooled map again.
# This fused kernel reads the conv output exactly once, computes the pooled
# average, applies the sigmoid inline, and reduces into a per-batch scalar via
# warp-shuffle block reductions followed by a single atomicAdd per block. This
# eliminates all intermediate tensors and two extra kernel launches.
fused_pool_sigmoid_sum_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

__global__ void fused_pool_sigmoid_sum_kernel(
        const float* __restrict__ in,
        float* __restrict__ out,
        const int C, const int H, const int W,
        const int PH, const int PW, const int pool) {
    const int n   = blockIdx.y;
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = C * PH * PW;

    float val = 0.0f;
    if (tid < total) {
        const int spatial = PH * PW;
        const int c  = tid / spatial;
        const int r  = tid - c * spatial;
        const int py = r / PW;
        const int px = r - py * PW;
        const int y0 = py * pool;
        const int x0 = px * pool;

        const float* base =
            in + (((size_t)n * (size_t)C + (size_t)c) * (size_t)H) * (size_t)W;

        float sum = 0.0f;
        if (pool == 4) {
            #pragma unroll
            for (int dy = 0; dy < 4; ++dy) {
                const float* row = base + (size_t)(y0 + dy) * (size_t)W + (size_t)x0;
                #pragma unroll
                for (int dx = 0; dx < 4; ++dx) {
                    sum += row[dx];
                }
            }
        } else {
            for (int dy = 0; dy < pool; ++dy) {
                const float* row = base + (size_t)(y0 + dy) * (size_t)W + (size_t)x0;
                for (int dx = 0; dx < pool; ++dx) {
                    sum += row[dx];
                }
            }
        }

        const float mean = sum * (1.0f / (float)(pool * pool));
        val = 1.0f / (1.0f + expf(-mean));  // sigmoid
    }

    // Block-wide reduction (all threads participate, inactive ones add 0).
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffffu, val, offset);
    }
    __shared__ float warp_sums[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        warp_sums[wid] = val;
    }
    __syncthreads();
    if (wid == 0) {
        const int num_warps = (blockDim.x + 31) >> 5;
        val = (lane < num_warps) ? warp_sums[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            val += __shfl_down_sync(0xffffffffu, val, offset);
        }
        if (lane == 0) {
            atomicAdd(out + n, val);
        }
    }
}

}  // namespace

torch::Tensor fused_pool_sigmoid_sum(torch::Tensor x, int64_t pool) {
    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.dim() == 4, "x must be a 4D NCHW tensor");

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

    auto out = at::zeros({N}, xc.options());

    const int total = C * PH * PW;
    if (N == 0 || total == 0) {
        return out;
    }

    const int threads = 256;
    const dim3 grid((total + threads - 1) / threads, N);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    fused_pool_sigmoid_sum_kernel<<<grid, threads, 0, stream>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(), C, H, W, PH, PW, p);

    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ",
                cudaGetErrorString(err));

    return out;
}
"""

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

fused_ops = load_inline(
    name="fused_pool_sigmoid_sum_ext",
    cpp_sources=fused_pool_sigmoid_sum_cpp_source,
    cuda_sources=fused_pool_sigmoid_sum_cuda_source,
    functions=["fused_pool_sigmoid_sum"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. The convolution is kept as nn.Conv2d (dispatched
    to cuDNN, autotuned via cudnn.benchmark), while AvgPool2d + Sigmoid + Sum are
    replaced by a single fused custom CUDA kernel.
    """
    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
        self.fused_ops = fused_ops
        # Fixed input shapes: let cuDNN autotune the fastest conv algorithm.
        torch.backends.cudnn.benchmark = True

    def forward(self, x):
        x = self.conv(x)
        if x.is_cuda and x.dtype == torch.float32:
            return self.fused_ops.fused_pool_sigmoid_sum(x, self.pool_kernel_size)
        # Fallback path (CPU tensors / other dtypes) matching the baseline.
        x = F.avg_pool2d(x, self.pool_kernel_size)
        x = torch.sigmoid(x)
        return torch.sum(x, dim=[1, 2, 3])
```