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

# ----------------------------------------------------------------------------
# Fully fused kernel: Conv2d(3x3, stride 1, no pad) + AvgPool + sigmoid +
# per-batch global sum, in ONE kernel. The 4.8GB conv output is never
# materialized in global memory: each block computes an 8x16 tile of conv
# outputs (all out-channels) with TF32 tensor-core WMMA (im2col staged in
# shared memory), then does pool->sigmoid->sum on the tile in shared memory
# and atomically accumulates one partial sum per image.
# ----------------------------------------------------------------------------
_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

#define TILE_H 8
#define TILE_W 16
#define TILE_PX 128
#define KDIM 72
#define A_ELEMS (TILE_PX * KDIM)
#define SMEM_BASE_BYTES (A_ELEMS * 4)

__device__ __forceinline__ float to_tf32(float f) {
#if __CUDA_ARCH__ >= 800
    uint32_t u;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(u) : "f"(f));
    return __uint_as_float(u);
#else
    return f;
#endif
}

__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 rsmem[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) rsmem[wid] = val;
    __syncthreads();
    if (wid == 0) {
        const int nw = (blockDim.x + 31) >> 5;
        float v = (lane < nw) ? rsmem[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);
    }
}

__global__ void __launch_bounds__(256)
fused_conv3_pool_sigmoid_sum_kernel(
    const float* __restrict__ x,     // [N, 8, H, W] NCHW
    const float* __restrict__ w,     // [CO, 8, 3, 3]
    const float* __restrict__ bias,  // [CO] or nullptr
    float* __restrict__ out,         // [N]
    const int H, const int W, const int CO, const int P,
    const int needR, const int needC)
{
#if __CUDA_ARCH__ >= 800
    const int b   = blockIdx.z;
    const int oy0 = blockIdx.y * TILE_H;
    const int ox0 = blockIdx.x * TILE_W;

    extern __shared__ float smem[];
    float* A    = smem;                 // TILE_PX x KDIM
    float* B    = A + A_ELEMS;          // KDIM x CO
    float* Cbuf = A;                    // aliases A after MMA phase

    const int tid = threadIdx.x;
    const float* xb = x + (long)b * KDIM * H * W / 9;  // b*8*H*W

    // ---- build A (im2col), k = ic*9 + kh*3 + kw ----
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        const int u  = tid + (i << 8);
        const int ic = u >> 7;
        const int p  = u & 127;
        const int dy = p >> 4;
        const int dx = p & 15;
        const int iy0 = oy0 + dy;
        const int ix0 = ox0 + dx;
        const float* xp = xb + (long)ic * H * W + (long)iy0 * W + ix0;
        float* Ap = A + p * KDIM + ic * 9;
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            const bool rok = (iy0 + kh) < H;
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                float v = 0.0f;
                if (rok && (ix0 + kw) < W) v = xp[kh * W + kw];
                Ap[kh * 3 + kw] = to_tf32(v);
            }
        }
    }

    // ---- build B: B[k][n] = w[n][k] ----
    const int KB = KDIM * CO;
    for (int idx = tid; idx < KB; idx += 256) {
        const int k = idx / CO;
        const int n = idx - k * CO;
        B[idx] = to_tf32(w[(long)n * KDIM + k]);
    }
    __syncthreads();

    // ---- MMA: 8 warps, warp -> 16-row subtile; loop 4 n-subtiles of 16 ----
    const int wid  = tid >> 5;
    const int nSub = CO >> 4;
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[4];
    #pragma unroll
    for (int l = 0; l < 4; ++l) wmma::fill_fragment(acc[l], 0.0f);

    #pragma unroll
    for (int j = 0; j < 9; ++j) {
        wmma::fragment<wmma::matrix_a, 16, 16, 8,
                       wmma::precision::tf32, wmma::row_major> af;
        wmma::load_matrix_sync(af, A + (wid * 16) * KDIM + j * 8, KDIM);
        for (int l = 0; l < nSub; ++l) {
            wmma::fragment<wmma::matrix_b, 16, 16, 8,
                           wmma::precision::tf32, wmma::row_major> bf;
            wmma::load_matrix_sync(bf, B + (j * 8) * CO + l * 16, CO);
            wmma::mma_sync(acc[l], af, bf, acc[l]);
        }
    }
    __syncthreads();   // A no longer needed; Cbuf may overwrite it

    for (int l = 0; l < nSub; ++l) {
        wmma::store_matrix_sync(Cbuf + (wid * 16) * CO + l * 16,
                                acc[l], CO, wmma::mem_row_major);
    }
    __syncthreads();

    // ---- pool + sigmoid + sum over this tile ----
    const int wy_n  = TILE_H / P;
    const int wx_n  = TILE_W / P;
    const int tasks = wy_n * wx_n * CO;
    const float inv = 1.0f / (float)(P * P);
    const float b0 = bias ? bias[0] : 0.0f;
    float sum = 0.0f;
    for (int t = tid; t < tasks; t += 256) {
        const int wi = t / CO;
        const int n  = t - wi * CO;
        const int wy = wi / wx_n;
        const int wx = wi - wy * wx_n;
        const int gy = oy0 + wy * P;
        const int gx = ox0 + wx * P;
        if (gy < needR && gx < needC) {
            const int py0 = wy * P;
            const int px0 = wx * P;
            float s = 0.0f;
            for (int r = 0; r < P; ++r) {
                const float* crow = Cbuf + ((py0 + r) * TILE_W + px0) * CO + n;
                #pragma unroll 4
                for (int q = 0; q < P; ++q) s += crow[q * CO];
            }
            const float bn = (n == 0) ? b0 : __ldg(bias + n);
            sum += sigmoid_f(s * inv + bn);
        }
    }
    block_sum_reduce(sum, out + b);
#endif
}

torch::Tensor fused_conv_pool_sigmoid_sum_cuda(
    torch::Tensor x, torch::Tensor w, torch::Tensor bias, int64_t pool)
{
    const int N  = (int)x.size(0);
    const int H  = (int)x.size(2);
    const int W  = (int)x.size(3);
    const int CO = (int)w.size(0);
    const int P  = (int)pool;

    const int OHc = H - 2, OWc = W - 2;
    const int needR = (OHc / P) * P;
    const int needC = (OWc / P) * P;

    auto out = at::zeros({N}, x.options());
    if (N == 0 || needR <= 0 || needC <= 0) return out;

    const int tilesR = (needR + TILE_H - 1) / TILE_H;
    const int tilesC = (needC + TILE_W - 1) / TILE_W;
    const int smem_bytes = SMEM_BASE_BYTES + KDIM * CO * 4;

    static int max_smem_set = 0;
    if (max_smem_set < smem_bytes) {
        cudaFuncSetAttribute(fused_conv3_pool_sigmoid_sum_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             160 * 1024);
        max_smem_set = 160 * 1024;
    }

    const float* bias_ptr = (bias.numel() > 0) ? bias.data_ptr<float>() : nullptr;
    dim3 grid(tilesC, tilesR, N);
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_conv3_pool_sigmoid_sum_kernel<<<grid, 256, smem_bytes, stream>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(), H, W, CO, P, needR, needC);
    return out;
}

// ------------------------- epilogue-only kernels ---------------------------
// Used by the cuDNN fallback path (non-fast-path configurations).

__global__ void nhwc_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 b  = blockIdx.y;
    const int oi = blockIdx.x;
    const long WC = (long)W * C;
    const float* xrow = x + ((long)b * H + (long)oi * 4) * WC;

    const int step = blockDim.x;
    const int dc  = step % C;
    const int doj = step / C;
    int c  = (int)(threadIdx.x) % C;
    int oj = (int)(threadIdx.x) / C;

    float acc = 0.0f;
    while (oj < OW) {
        const float* p = xrow + (long)oj * 4 * C + c;
        float s0 = 0.0f, s1 = 0.0f;
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            const float* pr = p + r * WC;
            s0 += pr[0] + pr[2 * C];
            s1 += pr[C] + pr[3 * C];
        }
        acc += sigmoid_f((s0 + s1) * 0.0625f);
        c += dc;
        if (c >= C) { c -= C; ++oj; }
        oj += doj;
    }
    block_sum_reduce(acc, out + b);
}

__global__ void nhwc_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 b  = blockIdx.y;
    const int oi = blockIdx.x;
    const long WC = (long)W * C;
    const float* xrow = x + ((long)b * H + (long)oi * P) * WC;

    const int step = blockDim.x;
    const int dc  = step % C;
    const int doj = step / C;
    int c  = (int)(threadIdx.x) % C;
    int oj = (int)(threadIdx.x) / C;

    float acc = 0.0f;
    while (oj < OW) {
        const float* p = xrow + (long)oj * P * C + c;
        float s = 0.0f;
        for (int r = 0; r < P; ++r) {
            const float* pr = p + r * WC;
            for (int q = 0; q < P; ++q) s += pr[q * C];
        }
        acc += sigmoid_f(s * scale);
        c += dc;
        if (c >= C) { c -= C; ++oj; }
        oj += doj;
    }
    block_sum_reduce(acc, out + b);
}

__global__ void nchw_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 step = blockDim.x;
    const int doi = step / OW;
    const int doj = step - doi * OW;

    float acc = 0.0f;
    for (int c = c0; c < c1; ++c) {
        const float* xc = xb + c * cstride;
        int idx = threadIdx.x;
        int oi = idx / OW;
        int oj = idx - oi * OW;
        while (oi < OH) {
            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);
            oj += doj;
            if (oj >= OW) { oj -= OW; ++oi; }
            oi += doi;
        }
    }
    block_sum_reduce(acc, out + b);
}

torch::Tensor pool_sigmoid_sum_cuda(torch::Tensor x, int64_t pool) {
    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);
    const int P = (int)pool;
    const int OH = H / P;
    const int OW = W / P;

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

    const int block = 256;
    auto stream = at::cuda::getCurrentCUDAStream();

    at::Tensor xc = x;
    bool channels_last = x.is_contiguous(at::MemoryFormat::ChannelsLast);
    if (!channels_last && !x.is_contiguous()) {
        xc = x.contiguous(at::MemoryFormat::ChannelsLast);
        channels_last = true;
    }

    if (channels_last) {
        const dim3 grid(OH, N);
        if (P == 4) {
            nhwc_pool4_sigmoid_sum_kernel<<<grid, block, 0, stream>>>(
                xc.data_ptr<float>(), out.data_ptr<float>(), C, H, W, OH, OW);
        } else {
            nhwc_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));
        }
    } else {
        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);
        nchw_pool4_sigmoid_sum_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), out.data_ptr<float>(),
            C, H, W, OH, OW, c_per_block);
    }
    return out;
}
"""

_fused_cpp_src = (
    "torch::Tensor fused_conv_pool_sigmoid_sum_cuda("
    "    torch::Tensor x, torch::Tensor w, torch::Tensor bias, int64_t pool);\n"
    "torch::Tensor pool_sigmoid_sum_cuda(torch::Tensor x, int64_t pool);"
)

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

# Device capability gate for the tensor-core fast path (needs sm_80+).
_FAST_ARCH = torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8


class ModelNew(nn.Module):
    """
    Optimized model:
      * Fast path (3x3 conv, stride 1, no padding, in_channels=8,
        out_channels multiple of 16 and <= 64, pool in {1,2,4,8,16},
        sm_80+): a single custom kernel computes the conv with TF32
        tensor-core WMMA (im2col staged in shared memory) and fuses
        AvgPool + sigmoid + per-batch global sum into the same pass, so the
        ~4.8GB conv output is never written to or read from global memory,
        and no input layout conversion is needed at all.
      * Otherwise: cuDNN channels-last conv + a fused
        pool+sigmoid+sum epilogue 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
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.benchmark_limit = 0
        torch.backends.cudnn.allow_tf32 = True

    def _fast_ok(self, x):
        c = self.conv
        if not (_FAST_ARCH and x.is_cuda and x.dtype == torch.float32):
            return False
        p = self.pool_kernel_size
        return (
            x.dim() == 4 and x.size(1) == 8 and x.is_contiguous()
            and c.in_channels == 8
            and c.out_channels % 16 == 0 and c.out_channels <= 64
            and c.kernel_size == (3, 3)
            and c.stride == (1, 1)
            and c.padding == (0, 0)
            and c.dilation == (1, 1)
            and c.groups == 1
            and isinstance(p, int) and p >= 1 and (8 % p) == 0 and (16 % p) == 0
            and x.size(2) >= 3 and x.size(3) >= 3
        )

    def forward(self, x):
        if self._fast_ok(x):
            wc = self.conv.weight
            if not wc.is_contiguous():
                wc = wc.contiguous()
            bias = self.conv.bias
            if bias is None:
                bias = x.new_empty(0)
            elif not bias.is_contiguous():
                bias = bias.contiguous()
            return _fused_ops.fused_conv_pool_sigmoid_sum_cuda(
                x, wc, bias, self.pool_kernel_size)
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 4:
            xcl = x.contiguous(memory_format=torch.channels_last)
            wcl = self.conv.weight.contiguous(memory_format=torch.channels_last)
            y = F.conv2d(xcl, wcl, self.conv.bias, self.conv.stride,
                         self.conv.padding, self.conv.dilation, self.conv.groups)
            return _fused_ops.pool_sigmoid_sum_cuda(y, self.pool_kernel_size)
        x = self.conv(x)
        x = F.avg_pool2d(x, self.pool_kernel_size)
        x = torch.sigmoid(x)
        return torch.sum(x, dim=[1, 2, 3])
```