```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] per batch
# element, computed in a single pass over the convolution output (FP32).
_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float sigmoid_f(float v) {
    return 1.0f / (1.0f + expf(-v));
}

__device__ __forceinline__ void block_sum_reduce(float val, float* out) {
    __shared__ float smem[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        val += __shfl_down_sync(0xffffffffu, val, off);
    if (lane == 0) smem[wid] = val;
    __syncthreads();
    if (wid == 0) {
        const int nw = (blockDim.x + 31) >> 5;
        float v = (lane < nw) ? smem[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, v);
    }
}

// Specialized 4x4 average-pool + sigmoid + per-batch sum.
// grid = (channel chunks, batch), block = 256 threads.
__global__ void pool4_sigmoid_sum_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    const int C, const int H, const int W,
    const int OH, const int OW,
    const int c_per_block)
{
    const int b  = blockIdx.y;
    const int c0 = blockIdx.x * c_per_block;
    const int c1 = min(c0 + c_per_block, C);

    const float* xb = x + ((long)b * C) * H * W;
    const long cstride = (long)H * W;
    const int nwin = OH * OW;

    float acc = 0.0f;
    for (int c = c0; c < c1; ++c) {
        const float* xc = xb + c * cstride;
        for (int idx = threadIdx.x; idx < nwin; idx += blockDim.x) {
            const int oi = idx / OW;
            const int oj = idx - oi * OW;
            const float* p = xc + (long)oi * 4 * W + oj * 4;
            float s = 0.0f;
            #pragma unroll
            for (int r = 0; r < 4; ++r) {
                const float* pr = p + r * W;
                s += (pr[0] + pr[1]) + (pr[2] + pr[3]);
            }
            acc += sigmoid_f(s * 0.0625f);
        }
    }
    block_sum_reduce(acc, out + b);
}

// Generic PxP average-pool + sigmoid + per-batch sum (fallback).
__global__ void poolP_sigmoid_sum_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    const int C, const int H, const int W,
    const int OH, const int OW,
    const int P, const float scale,
    const int c_per_block)
{
    const int b  = blockIdx.y;
    const int c0 = blockIdx.x * c_per_block;
    const int c1 = min(c0 + c_per_block, C);

    const float* xb = x + ((long)b * C) * H * W;
    const long cstride = (long)H * W;
    const int nwin = OH * OW;

    float acc = 0.0f;
    for (int c = c0; c < c1; ++c) {
        const float* xc = xb + c * cstride;
        for (int idx = threadIdx.x; idx < nwin; idx += blockDim.x) {
            const int oi = idx / OW;
            const int oj = idx - oi * OW;
            const float* p = xc + (long)oi * P * W + oj * P;
            float s = 0.0f;
            for (int r = 0; r < P; ++r) {
                const float* pr = p + r * W;
                for (int q = 0; q < P; ++q) s += pr[q];
            }
            acc += sigmoid_f(s * scale);
        }
    }
    block_sum_reduce(acc, out + b);
}

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 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 P = (int)pool;
    const int OH = H / P;
    const int OW = W / P;

    auto out = at::zeros({N}, xc.options());
    if (N == 0 || OH <= 0 || OW <= 0) return out;

    // Pick a channel split so we get enough blocks to fill the GPU.
    int chunks = (1024 + N - 1) / N;
    if (chunks < 1) chunks = 1;
    if (chunks > C) chunks = C;
    const int c_per_block = (C + chunks - 1) / chunks;
    chunks = (C + c_per_block - 1) / c_per_block;

    const dim3 grid(chunks, N);
    const int block = 256;
    auto stream = at::cuda::getCurrentCUDAStream();

    if (P == 4) {
        pool4_sigmoid_sum_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), out.data_ptr<float>(),
            C, H, W, OH, OW, c_per_block);
    } else {
        poolP_sigmoid_sum_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), out.data_ptr<float>(),
            C, H, W, OH, OW, P, 1.0f / (float)(P * P), c_per_block);
    }
    return out;
}
"""

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

_fused_pool_sigmoid_sum = load_inline(
    name="fused_pool_sigmoid_sum_v1",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["pool_sigmoid_sum_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: the Conv2d stays on the heavily-optimized cuDNN path
    (with autotuning enabled), while AvgPool + sigmoid + global sum are fused
    into a single custom CUDA kernel that reads the conv output exactly once
    and writes only one float per batch element.
    """
    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
        torch.backends.cudnn.benchmark = True

    def forward(self, x):
        x = self.conv(x)
        if x.is_cuda and x.dtype == torch.float32:
            return _fused_pool_sigmoid_sum.pool_sigmoid_sum_cuda(
                x, self.pool_kernel_size)
        # Fallback (e.g. CPU tensors): original op sequence.
        x = F.avg_pool2d(x, self.pool_kernel_size)
        x = torch.sigmoid(x)
        return torch.sum(x, dim=[1, 2, 3])
```