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

# Let cuDNN autotune the (fixed-shape) convolution for extra speed.
torch.backends.cudnn.benchmark = True

# Custom fused CUDA kernel: AvgPool2d(K, stride K) -> sigmoid -> sum over (C, H, W).
# This makes a single streaming pass over the convolution output and avoids
# materializing the pooled / sigmoid intermediates (3 kernel launches -> 1).
_fused_epilogue_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BLOCK_SIZE 256

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

// Block-wide reduction of `val`; result valid for all threads after the call.
__device__ __forceinline__ float block_reduce_sum_(float val) {
    __shared__ float sh[BLOCK_SIZE];
    const int tid = threadIdx.x;
    sh[tid] = val;
    __syncthreads();
    #pragma unroll
    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) sh[tid] += sh[tid + s];
        __syncthreads();
    }
    return sh[0];
}

// One thread block handles one (batch, channel) slice of the conv output.
template <int K>
__global__ void pool_sigmoid_sum_kernel(
        const float* __restrict__ x,
        float* __restrict__ out,
        const int C, const int H, const int W,
        const int PH, const int PW) {
    const int b = blockIdx.y;
    const int c = blockIdx.x;
    const float* slice = x + ((int64_t)b * C + c) * (int64_t)H * W;

    const int total = PH * PW;
    const float scale = 1.0f / (float)(K * K);
    float acc = 0.0f;

    for (int idx = threadIdx.x; idx < total; idx += BLOCK_SIZE) {
        const int ph = idx / PW;
        const int pw = idx - ph * PW;
        const float* p = slice + (int64_t)(ph * K) * W + (pw * K);
        float s = 0.0f;
        #pragma unroll
        for (int i = 0; i < K; ++i) {
            #pragma unroll
            for (int j = 0; j < K; ++j) {
                s += p[(int64_t)i * W + j];
            }
        }
        acc += sigmoidf_(s * scale);
    }

    acc = block_reduce_sum_(acc);
    if (threadIdx.x == 0) {
        atomicAdd(&out[b], acc);
    }
}

// Generic fallback for arbitrary pool kernel sizes.
__global__ void pool_sigmoid_sum_kernel_generic(
        const float* __restrict__ x,
        float* __restrict__ out,
        const int C, const int H, const int W,
        const int PH, const int PW, const int K) {
    const int b = blockIdx.y;
    const int c = blockIdx.x;
    const float* slice = x + ((int64_t)b * C + c) * (int64_t)H * W;

    const int total = PH * PW;
    const float scale = 1.0f / (float)(K * K);
    float acc = 0.0f;

    for (int idx = threadIdx.x; idx < total; idx += BLOCK_SIZE) {
        const int ph = idx / PW;
        const int pw = idx - ph * PW;
        const float* p = slice + (int64_t)(ph * K) * W + (pw * K);
        float s = 0.0f;
        for (int i = 0; i < K; ++i) {
            for (int j = 0; j < K; ++j) {
                s += p[(int64_t)i * W + j];
            }
        }
        acc += sigmoidf_(s * scale);
    }

    acc = block_reduce_sum_(acc);
    if (threadIdx.x == 0) {
        atomicAdd(&out[b], acc);
    }
}

torch::Tensor fused_pool_sigmoid_sum(torch::Tensor x, int64_t k) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 4, "x must be a 4D NCHW tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(k >= 1, "pool kernel size must be >= 1");

    const int B = static_cast<int>(x.size(0));
    const int C = static_cast<int>(x.size(1));
    const int H = static_cast<int>(x.size(2));
    const int W = static_cast<int>(x.size(3));
    const int K = static_cast<int>(k);
    const int PH = H / K;
    const int PW = W / K;

    auto out = at::zeros({B}, x.options());
    if (PH == 0 || PW == 0) {
        return out;
    }
    TORCH_CHECK(B <= 65535, "batch size too large for this kernel");

    const dim3 grid(C, B);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xp = x.data_ptr<float>();
    float* op = out.data_ptr<float>();

    switch (K) {
        case 2:
            pool_sigmoid_sum_kernel<2><<<grid, BLOCK_SIZE, 0, stream>>>(xp, op, C, H, W, PH, PW);
            break;
        case 3:
            pool_sigmoid_sum_kernel<3><<<grid, BLOCK_SIZE, 0, stream>>>(xp, op, C, H, W, PH, PW);
            break;
        case 4:
            pool_sigmoid_sum_kernel<4><<<grid, BLOCK_SIZE, 0, stream>>>(xp, op, C, H, W, PH, PW);
            break;
        case 8:
            pool_sigmoid_sum_kernel<8><<<grid, BLOCK_SIZE, 0, stream>>>(xp, op, C, H, W, PH, PW);
            break;
        default:
            pool_sigmoid_sum_kernel_generic<<<grid, BLOCK_SIZE, 0, stream>>>(xp, op, C, H, W, PH, PW, K);
            break;
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_fused_epilogue_cpp_src = (
    "torch::Tensor fused_pool_sigmoid_sum(torch::Tensor x, int64_t k);"
)

# Compile the inline CUDA epilogue.
_fused_epilogue = load_inline(
    name="fused_pool_sigmoid_sum",
    cpp_sources=_fused_epilogue_cpp_src,
    cuda_sources=_fused_epilogue_cuda_src,
    functions=["fused_pool_sigmoid_sum"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    - Conv2d is kept as-is (dispatches to autotuned cuDNN), since it is the
      compute-heavy part and is already highly optimized.
    - AvgPool2d + sigmoid + sum over (C, H, W) are replaced by a single fused
      custom CUDA kernel, eliminating two intermediate tensors and two kernel
      launches, and cutting global-memory traffic roughly in half.
    """
    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 (
            isinstance(self.pool_kernel_size, int)
            and x.is_cuda
            and x.dtype == torch.float32
            and x.is_contiguous()
        ):
            return _fused_epilogue.fused_pool_sigmoid_sum(x, self.pool_kernel_size)
        # Fallback path (CPU tensors / other dtypes) -- identical to reference.
        x = F.avg_pool2d(x, self.pool_kernel_size)
        x = torch.sigmoid(x)
        x = torch.sum(x, dim=[1, 2, 3])
        return x
```