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

# ---------------------------------------------------------------------------
# Kernel set 1: fully fused Conv2d(3x3) + AvgPool(4x4) + sigmoid + per-batch
# global sum. Uses TF32 tensor-core mma (m16n8k8) on channels_last data and
# never materializes the conv output (saves ~2x HBM traffic of the full
# activation tensor). Specialized for: k=3, stride=1, pad=0, Cin<=8,
# Cout<=64, pool=4. Everything else falls back to the kernels below.
# ---------------------------------------------------------------------------
_fused_conv_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));
}

// tile: 8x8 conv outputs per block, 2 warps, warp w owns output rows 4w..4w+3.
// Xs: 10x10 input patch, 8 channels (stride 9 to dodge bank conflicts).
// W2s: rearranged weights [72][68] (k = (kh*3+kw)*8+ic), pre-rounded to tf32.
__global__ void __launch_bounds__(64)
fused_conv3_pool4_sigmoid_sum_kernel(
    const float* __restrict__ x,     // NHWC
    const float* __restrict__ w2,    // [72][64]
    const float* __restrict__ bias,  // [oc]
    float* __restrict__ out,         // [N]
    const int N, const int C_in, const int H, const int W,
    const int OH, const int OW, const int PH, const int PW,
    const int ntx, const int nty)
{
#if __CUDA_ARCH__ >= 800
    __shared__ float Xs[10 * 10 * 9];
    __shared__ float W2s[72 * 68];
    __shared__ float red[2];

    const int b  = blockIdx.z;
    const int tx = blockIdx.x;
    const int ty = blockIdx.y;
    const int th0 = ty * 8;
    const int tw0 = tx * 8;
    const int tid = threadIdx.x;

    // ---- load input patch (10x10x8) ----
    {
        const long xb = (long)b * H * W * C_in;
        #pragma unroll
        for (int it = 0; it < 2; ++it) {
            const int item = tid + it * 64;
            if (item < 200) {
                const int dy = item / 20;
                const int rem = item - dy * 20;
                const int dx = rem >> 1;
                const int q = (rem & 1) * 4;
                const int y = th0 + dy;
                const int xv = tw0 + dx;
                float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
                if (y < H && xv < W) {
                    const float* src = x + xb + ((long)y * W + xv) * C_in + q;
                    if (C_in >= q + 4) {
                        v = *reinterpret_cast<const float4*>(src);
                    } else {
                        float t[4] = {0.f, 0.f, 0.f, 0.f};
                        for (int i = 0; i < C_in - q; ++i) t[i] = src[i];
                        v = make_float4(t[0], t[1], t[2], t[3]);
                    }
                }
                float* dst = &Xs[(dy * 10 + dx) * 9 + q];
                dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w;
            }
        }
        // zero the 9th (pad) channel column
        for (int item = tid; item < 100; item += 64) {
            Xs[item * 9 + 8] = 0.0f;
        }
    }
    // ---- load rearranged weights ----
    for (int item = tid; item < 72 * 64; item += 64) {
        const int k = item >> 6;
        const int n = item & 63;
        W2s[k * 68 + n] = w2[item];
    }
    if (tid < 2) red[tid] = 0.0f;
    __syncthreads();

    const int lane = tid & 31;
    const int w = tid >> 5;      // warp id 0/1 -> window row
    const int gid = lane >> 2;   // 0..7
    const int tig = lane & 3;    // 0..3

    float c[2][8][4];
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt)
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt)
            #pragma unroll
            for (int i = 0; i < 4; ++i) c[mt][nt][i] = 0.0f;

    // rows (local_oh) for the two m-tiles of this warp
    const int oh_mt0 = 4 * w;      // mt 0 -> rows 4w, 4w+1
    const int oh_mt1 = 4 * w + 2;  // mt 1 -> rows 4w+2, 4w+3

    #pragma unroll
    for (int s = 0; s < 9; ++s) {
        const int kh = s / 3;
        const int kw = s - kh * 3;
        const int kbase = s * 8;
        float av[2][4];
        #pragma unroll
        for (int mt = 0; mt < 2; ++mt) {
            const int oh0 = (mt == 0 ? oh_mt0 : oh_mt1) + kh;
            const int row0 = (oh0 * 10 + kw + gid) * 9 + tig;
            const int row1 = row0 + 90;  // next output row (+10*9)
            av[mt][0] = Xs[row0];
            av[mt][1] = Xs[row1];
            av[mt][2] = Xs[row0 + 4];
            av[mt][3] = Xs[row1 + 4];
        }
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            const int bidx0 = (kbase + tig) * 68 + nt * 8 + gid;
            const float bv0 = W2s[bidx0];
            const float bv1 = W2s[bidx0 + 4 * 68];
            #pragma unroll
            for (int mt = 0; mt < 2; ++mt) {
                asm volatile(
                    "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
                    "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
                    : "+f"(c[mt][nt][0]), "+f"(c[mt][nt][1]),
                      "+f"(c[mt][nt][2]), "+f"(c[mt][nt][3])
                    : "r"(__float_as_uint(av[mt][0])), "r"(__float_as_uint(av[mt][1])),
                      "r"(__float_as_uint(av[mt][2])), "r"(__float_as_uint(av[mt][3])),
                      "r"(__float_as_uint(bv0)), "r"(__float_as_uint(bv1)));
            }
        }
    }

    // ---- fused 4x4 pooling + sigmoid + reduction (all within the warp) ----
    const int prow = ty * 2 + w;
    float acc = 0.0f;
    #pragma unroll
    for (int nt = 0; nt < 8; ++nt) {
        float s0 = (c[0][nt][0] + c[0][nt][2]) + (c[1][nt][0] + c[1][nt][2]);
        float s1 = (c[0][nt][1] + c[0][nt][3]) + (c[1][nt][1] + c[1][nt][3]);
        s0 += __shfl_down_sync(0xffffffffu, s0, 8);
        s0 += __shfl_down_sync(0xffffffffu, s0, 4);
        s1 += __shfl_down_sync(0xffffffffu, s1, 8);
        s1 += __shfl_down_sync(0xffffffffu, s1, 4);
        if ((gid & 3) == 0) {
            const int pcol = tx * 2 + (gid >> 2);
            if (prow < PH && pcol < PW) {
                const int ch = nt * 8 + tig * 2;
                acc += sigmoid_f(s0 * 0.0625f + __ldg(bias + ch));
                acc += sigmoid_f(s1 * 0.0625f + __ldg(bias + ch + 1));
            }
        }
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        acc += __shfl_down_sync(0xffffffffu, acc, off);
    if (lane == 0) red[w] = acc;
    __syncthreads();
    if (tid == 0) {
        const float t = red[0] + red[1];
        if (t != 0.0f) atomicAdd(out + b, t);
    }
#endif
}

torch::Tensor fused_conv_pool_sigmoid_sum_cuda(
    torch::Tensor x, torch::Tensor w2, torch::Tensor bias, int64_t pool) {
    const int N = (int)x.size(0);
    const int C_in = (int)x.size(3);   // NHWC
    const int H = (int)x.size(1);
    const int W = (int)x.size(2);
    const int OH = H - 2;
    const int OW = W - 2;
    const int PH = OH / 4;
    const int PW = OW / 4;

    auto out = at::zeros({N}, x.options());
    const int ntx = (OW + 7) / 8;
    const int nty = (OH + 7) / 8;
    dim3 grid(ntx, nty, N);
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_conv3_pool4_sigmoid_sum_kernel<<<grid, 64, 0, stream>>>(
        x.data_ptr<float>(), w2.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), N, C_in, H, W, OH, OW, PH, PW, ntx, nty);
    return out;
}
"""

# ---------------------------------------------------------------------------
# Kernel set 2 (fallback): fused AvgPool + sigmoid + per-batch sum over a conv
# output, for both channels_last and contiguous layouts (general shapes).
# ---------------------------------------------------------------------------
_fused_tail_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);
    }
}

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

__global__ void nchw_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 4D");

    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);
        if (P == 4) {
            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);
        } else {
            nchw_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 fused_conv_pool_sigmoid_sum_cuda(torch::Tensor x, torch::Tensor w2, 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_conv_src + _fused_tail_src,
    functions=["fused_conv_pool_sigmoid_sum_cuda", "pool_sigmoid_sum_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model.

    Fast path (CUDA fp32, 3x3 stride-1 pad-0 conv, Cin<=8, Cout<=64, pool=4,
    8-alignment satisfied): a single custom tensor-core kernel computes the
    whole network. It runs the conv as an implicit-GEMM with TF32
    m16n8k8 mma on channels_last input (weights pre-rearranged and
    pre-rounded to TF32 once at init), then fuses the 4x4 average pool,
    sigmoid and per-batch global sum directly on the mma accumulators.
    The full conv activation tensor is never written to HBM.

    Generic fallback: cuDNN conv in native channels_last layout followed by a
    fused avgpool+sigmoid+sum kernel that reads the conv output exactly once.
    """
    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

        kh, kw = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
        self._fused_ok = (
            kh == 3 and kw == 3
            and self.conv.stride == (1, 1)
            and self.conv.padding == (0, 0)
            and self.conv.dilation == (1, 1)
            and self.conv.groups == 1
            and in_channels <= 8
            and out_channels <= 64
            and pool_kernel_size == 4
        )
        if self._fused_ok:
            with torch.no_grad():
                w = self.conv.weight.detach()
                if in_channels < 8:
                    w = F.pad(w, (0, 0, 0, 0, 0, 8 - in_channels))
                # rearrange to [72][64]: k = (kh*3+kw)*8 + ic
                w2 = w.permute(2, 3, 1, 0).reshape(72, out_channels)
                if out_channels < 64:
                    w2 = F.pad(w2, (0, 64 - out_channels))
                w2 = w2.contiguous()
                if torch.cuda.is_available():
                    w2 = self._round_tf32(w2.cuda())
                else:
                    w2 = w2.cpu()
            self.register_buffer("_w2", w2, persistent=False)
            self.register_buffer(
                "_bias0",
                torch.zeros(out_channels, device=w2.device),
                persistent=False,
            )

    @staticmethod
    def _round_tf32(t):
        # round-to-nearest-even fp32 -> tf32 (as cuDNN does internally)
        u = t.view(torch.int32)
        rounded = (u + 4096 + ((u >> 13) & 1)) & -8192
        return rounded.view(torch.float32).contiguous()

    def _params_ready(self, x):
        return (
            self._w2.device == x.device
            and self.conv.weight.device == x.device
            and self.conv.weight.dtype == x.dtype
        )

    def forward(self, x):
        if (
            self._fused_ok
            and x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 4
            and x.size(0) <= 65535
            and x.size(2) >= 3
            and x.size(3) >= 3
            and (x.size(2) - 2) % 8 == 0
            and (x.size(3) - 2) % 8 == 0
            and torch.cuda.get_device_capability(x.device)[0] >= 8
            and self._params_ready(x)
        ):
            xcl = x.contiguous(memory_format=torch.channels_last)
            bias = self.conv.bias if self.conv.bias is not None else self._bias0
            return _fused_ops.fused_conv_pool_sigmoid_sum_cuda(
                xcl, self._w2, 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)

        # CPU / other dtypes: original op sequence.
        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])
```