```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 over ALL candidate engines (cached after warmup).
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.benchmark_limit = 0

# Custom CUDA ops:
#   1) nchw_to_nhwc_fp16: fuses fp32->fp16 cast + NCHW->NHWC transpose in one
#      pass so the convolution can run natively in NHWC (channels_last). This
#      eliminates cuDNN's expensive NHWC->NCHW output re-transpose that
#      dominated the previous version.
#   2) fused_pool_sigmoid_sum_nhwc: non-overlapping AvgPool + sigmoid + global
#      sum, reading the NHWC conv output exactly once with fully coalesced
#      16-byte loads.
#   3) fused_pool_sigmoid_sum_nchw: same fusion for NCHW layout (safety net).
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

namespace {

__device__ __forceinline__ void block_reduce_atomic(float acc, float* out) {
    #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, v);
    }
}

// fp32 NCHW -> fp16 NHWC, one thread per pixel, C == 8 (vectorized 16B store).
__global__ void nchw_to_nhwc_fp16_c8_kernel(
        const float* __restrict__ in, __half* __restrict__ out,
        const size_t HW, const size_t total_pixels) {
    const size_t p = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= total_pixels) return;
    const size_t n = p / HW;
    const size_t s = p - n * HW;
    const float* base = in + n * 8 * HW + s;
    uint4 packed;
    __half2* ph = reinterpret_cast<__half2*>(&packed);
    #pragma unroll
    for (int c = 0; c < 4; ++c) {
        ph[c] = __floats2half2_rn(base[(size_t)(2 * c) * HW],
                                  base[(size_t)(2 * c + 1) * HW]);
    }
    *reinterpret_cast<uint4*>(out + p * 8) = packed;
}

// Generic-C variant (one thread per output element).
__global__ void nchw_to_nhwc_fp16_generic_kernel(
        const float* __restrict__ in, __half* __restrict__ out,
        const int C, const size_t HW, const size_t total) {
    const size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    const int c = (int)(idx % (size_t)C);
    const size_t p = idx / (size_t)C;
    const size_t n = p / HW;
    const size_t s = p - n * HW;
    out[idx] = __float2half_rn(in[(n * (size_t)C + c) * HW + s]);
}

// AvgPool(non-overlapping) + sigmoid + sum over [C,H,W], NHWC input.
// Each thread handles one pooling window for a group of 8 channels using
// 16-byte vectorized loads; consecutive threads read consecutive memory.
template <int POOL>
__global__ void fused_pool_sigmoid_sum_nhwc_kernel(
        const __half* __restrict__ x, float* __restrict__ out,
        const int C, const int H, const int W,
        const int Hp, const int Wp, const int units_per_image) {
    const int n = blockIdx.y;
    const __half* __restrict__ img = x + (size_t)n * ((size_t)H * W * C);
    const int cg_count = C >> 3;

    float acc = 0.0f;
    const int stride = gridDim.x * blockDim.x;
    for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < units_per_image; v += stride) {
        const int cg = v % cg_count;
        const int t  = v / cg_count;
        const int wx = t % Wp;
        const int wy = t / Wp;
        const int c0 = cg << 3;

        const __half* base = img + ((size_t)(wy * POOL) * W + (size_t)(wx * POOL)) * C + c0;
        float sums[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
        #pragma unroll
        for (int r = 0; r < POOL; ++r) {
            const __half* row = base + (size_t)r * W * C;
            #pragma unroll
            for (int q = 0; q < POOL; ++q) {
                const uint4 raw = *reinterpret_cast<const uint4*>(row + (size_t)q * C);
                const __half2* h2 = reinterpret_cast<const __half2*>(&raw);
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    const float2 f = __half22float2(h2[i]);
                    sums[2 * i]     += f.x;
                    sums[2 * i + 1] += f.y;
                }
            }
        }
        constexpr float inv = 1.0f / (float)(POOL * POOL);
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            const float m = sums[i] * inv;
            acc += 1.0f / (1.0f + expf(-m));
        }
    }
    block_reduce_atomic(acc, &out[n]);
}

__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));
}

// AvgPool(non-overlapping) + sigmoid + sum over [C,H,W], NCHW input (backup).
template <typename scalar_t, int POOL>
__global__ void fused_pool_sigmoid_sum_nchw_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 stride = gridDim.x * blockDim.x;
    constexpr float inv = 1.0f / (float)(POOL * POOL);

    for (int v = blockIdx.x * blockDim.x + threadIdx.x; 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_reduce_atomic(acc, &out[n]);
}

} // namespace

torch::Tensor nchw_to_nhwc_fp16(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4, "input must be CUDA NCHW");
    TORCH_CHECK(x.scalar_type() == torch::kFloat, "input must be fp32");
    auto xc = x.contiguous();
    const int64_t N = xc.size(0), C = xc.size(1), H = xc.size(2), W = xc.size(3);

    auto out = torch::empty({N, C, H, W},
        xc.options().dtype(torch::kHalf).memory_format(torch::MemoryFormat::ChannelsLast));

    const size_t HW = (size_t)H * (size_t)W;
    const size_t total_pixels = (size_t)N * HW;
    if (total_pixels == 0) return out;

    auto stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;
    __half* out_ptr = reinterpret_cast<__half*>(out.data_ptr<at::Half>());

    if (C == 8) {
        const size_t blocks = (total_pixels + threads - 1) / threads;
        nchw_to_nhwc_fp16_c8_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), out_ptr, HW, total_pixels);
    } else {
        const size_t total = total_pixels * (size_t)C;
        const size_t blocks = (total + threads - 1) / threads;
        nchw_to_nhwc_fp16_generic_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), out_ptr, (int)C, HW, total);
    }
    return out;
}

torch::Tensor fused_pool_sigmoid_sum_nhwc(torch::Tensor x, int64_t pool) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4, "input must be CUDA NCHW4d");
    TORCH_CHECK(x.scalar_type() == torch::kHalf, "input must be fp16");
    TORCH_CHECK(x.is_contiguous(at::MemoryFormat::ChannelsLast), "input must be channels_last");

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    TORCH_CHECK(C % 8 == 0, "channels must be a multiple of 8");

    const int Hp = H / (int)pool;
    const int Wp = W / (int)pool;
    auto out = torch::zeros({N}, x.options().dtype(torch::kFloat32));
    if (N == 0 || C == 0 || Hp == 0 || Wp == 0) return out;

    const int units = Hp * Wp * (C >> 3);
    const int threads = 256;
    int blocks = (units + threads - 1) / threads;
    if (blocks > 1024) blocks = 1024;
    dim3 grid(blocks, N);
    auto stream = at::cuda::getCurrentCUDAStream();
    const __half* ptr = reinterpret_cast<const __half*>(x.data_ptr<at::Half>());

    if (pool == 4) {
        fused_pool_sigmoid_sum_nhwc_kernel<4><<<grid, threads, 0, stream>>>(
            ptr, out.data_ptr<float>(), C, H, W, Hp, Wp, units);
    } else if (pool == 2) {
        fused_pool_sigmoid_sum_nhwc_kernel<2><<<grid, threads, 0, stream>>>(
            ptr, out.data_ptr<float>(), C, H, W, Hp, Wp, units);
    } else {
        TORCH_CHECK(false, "unsupported pool size");
    }
    return out;
}

torch::Tensor fused_pool_sigmoid_sum_nchw(torch::Tensor x, int64_t pool) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4, "input must be CUDA 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);
    TORCH_CHECK(W % 2 == 0, "width must be even");

    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) {
        const __half* ptr = reinterpret_cast<const __half*>(xc.data_ptr<at::Half>());
        if (pool == 4) {
            fused_pool_sigmoid_sum_nchw_kernel<__half, 4><<<grid, threads, 0, stream>>>(
                ptr, out.data_ptr<float>(), C, H, W, Hp, Wp, num_windows);
        } else if (pool == 2) {
            fused_pool_sigmoid_sum_nchw_kernel<__half, 2><<<grid, threads, 0, stream>>>(
                ptr, 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) {
        const float* ptr = xc.data_ptr<float>();
        if (pool == 4) {
            fused_pool_sigmoid_sum_nchw_kernel<float, 4><<<grid, threads, 0, stream>>>(
                ptr, out.data_ptr<float>(), C, H, W, Hp, Wp, num_windows);
        } else if (pool == 2) {
            fused_pool_sigmoid_sum_nchw_kernel<float, 2><<<grid, threads, 0, stream>>>(
                ptr, 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;
}
"""

_cpp_src = (
    "torch::Tensor nchw_to_nhwc_fp16(torch::Tensor x);\n"
    "torch::Tensor fused_pool_sigmoid_sum_nhwc(torch::Tensor x, int64_t pool);\n"
    "torch::Tensor fused_pool_sigmoid_sum_nchw(torch::Tensor x, int64_t pool);"
)

_ext = load_inline(
    name="conv_pool_sigmoid_fused_v2",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=[
        "nchw_to_nhwc_fp16",
        "fused_pool_sigmoid_sum_nhwc",
        "fused_pool_sigmoid_sum_nchw",
    ],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version:
      * The fp32->fp16 cast and NCHW->NHWC transpose of the input are fused into
        a single custom kernel, so the convolution runs natively in
        channels_last (NHWC) on tensor cores with NO layout re-transposes
        (the NHWC->NCHW output transpose dominated the previous version).
      * fp16 weights in channels_last are cached across calls.
      * AvgPool + sigmoid + global sum are fused into one custom CUDA kernel
        that reads the NHWC conv output exactly once with 16-byte coalesced
        loads and emits one fp32 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
        self._w16 = None
        self._b16 = None
        self._w_key = None

    def _half_weights(self):
        w = self.conv.weight
        key = (w.data_ptr(), w._version)
        if self._w16 is None or self._w_key != key:
            self._w16 = w.detach().half().contiguous(memory_format=torch.channels_last)
            b = self.conv.bias
            self._b16 = b.detach().half() if b is not None else None
            self._w_key = key
        return self._w16, self._b16

    def forward(self, x):
        pool = self.pool_kernel_size
        fast = (
            x.is_cuda
            and x.dtype == torch.float32
            and isinstance(pool, int)
            and pool in (2, 4)
            and self.conv.out_channels % 8 == 0
        )
        if fast:
            try:
                w16, b16 = self._half_weights()
                x16 = _ext.nchw_to_nhwc_fp16(x)
                y = F.conv2d(
                    x16,
                    w16,
                    b16,
                    self.conv.stride,
                    self.conv.padding,
                    self.conv.dilation,
                    self.conv.groups,
                )
                if y.is_contiguous(memory_format=torch.channels_last):
                    return _ext.fused_pool_sigmoid_sum_nhwc(y, pool)
                if y.is_contiguous() and y.shape[-1] % 2 == 0:
                    return _ext.fused_pool_sigmoid_sum_nchw(y, pool)
                # Unexpected layout: cheap torch fallback on the conv output.
                y = self.avg_pool(y.float())
                return torch.sum(torch.sigmoid(y), dim=[1, 2, 3])
            except RuntimeError:
                pass

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