```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

torch.backends.cudnn.benchmark = True

# Single extension:
#   1) fused_conv3x3_c8_pool4_sigmoid_sum: ONE kernel fusing 3x3 conv
#      (fp16 tensor-core mma: 4x m16n8k16 + 1x m16n8k8 => exactly K=72 taps,
#      zero padding waste), 4x4 average pool, sigmoid and global sum.
#      fp32 NCHW input is cast/transposed directly into shared memory and the
#      ~2.4 GB conv output is never materialized. Conflict-free shared-memory
#      layouts for both A (pixel patch, stride 20 px) and B (weights, exact
#      72-half rows). Occupancy pinned at 2 blocks/SM via launch_bounds.
#   2) Backup ops (fused cast+transpose, fused pool+sigmoid+sum) for the cuDNN
#      path used on configurations the mma kernel does not cover.
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

namespace {

#define TILE_PIX 16
#define PATCH_DIM 18
#define PATCH_STRIDE 20
#define WSM_STRIDE 72

__device__ __forceinline__ void mma16816(float* c, const uint32_t* a, const uint32_t* b) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
#endif
}

__device__ __forceinline__ void mma1688(float* c, const uint32_t* a, uint32_t b) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
        "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%0,%1,%2,%3};\n"
        : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
        : "r"(a[0]), "r"(a[1]), "r"(b));
#endif
}

// Each block covers a 16x16 tile of conv output (= 4x4 pooling windows).
// M = the 16 pixels of one 4x4 pooling window, N = 8 output channels,
// K = 72 taps (8 in-ch x 3 x 3, packed as k = (ky*3+kx)*8 + ci), executed as
// 4 k16 mma + 1 k8 mma (no padding). Bias added in fp32 in the epilogue.
__global__ void __launch_bounds__(256, 2)
fused_conv3x3_c8_pool4_sigmoid_sum_kernel(
        const float* __restrict__ x,      // [N,8,H,W] fp32
        const __half* __restrict__ wk,    // [64,72] fp16
        const float* __restrict__ bias,   // [64] fp32
        float* __restrict__ out,          // [N] fp32
        const int H, const int W, const int Hp, const int Wp) {
    __shared__ __half smem_in[PATCH_DIM * PATCH_STRIDE * 8];
    __shared__ __half smem_w[64 * WSM_STRIDE];
    __shared__ float smem_bias[64];
    __shared__ float red[8];

    const int t = threadIdx.x;
    const int n = blockIdx.z;
    const int py0 = blockIdx.y * TILE_PIX;
    const int px0 = blockIdx.x * TILE_PIX;

    // ---- weights [64,72] -> smem (layout identical: 9216B straight copy) ----
    {
        const uint4* wsrc = reinterpret_cast<const uint4*>(wk);
        uint4* wdst = reinterpret_cast<uint4*>(smem_w);
        #pragma unroll
        for (int i = t; i < 576; i += 256) wdst[i] = wsrc[i];
        if (t < 64) smem_bias[t] = bias[t];
    }

    // ---- input patch 18x18 px x 8 ch: fp32 NCHW -> fp16 smem ----
    {
        const int area = H * W;
        const float* xb = x + (size_t)n * 8 * (size_t)area;
        uint32_t* dst = reinterpret_cast<uint32_t*>(smem_in);
        #pragma unroll
        for (int i = t; i < 1296; i += 256) {
            const int q = i & 3;
            const int pix = i >> 2;
            const int dx = pix % PATCH_DIM;
            const int dy = pix / PATCH_DIM;
            const int gy = py0 + dy;
            const int gx = px0 + dx;
            float2 v = make_float2(0.f, 0.f);
            if (gy < H && gx < W) {
                const size_t off = (size_t)gy * W + gx;
                v.x = xb[(size_t)(2 * q) * area + off];
                v.y = xb[(size_t)(2 * q + 1) * area + off];
            }
            __half2 h2 = __floats2half2_rn(v.x, v.y);
            dst[(dy * PATCH_STRIDE + dx) * 4 + q] =
                *reinterpret_cast<uint32_t*>(&h2);
        }
    }
    __syncthreads();

    const int warp = t >> 5;
    const int lane = t & 31;
    const int g = lane >> 2;
    const int tig = lane & 3;

    // window assignment: 2 adjacent windows per warp (same window row)
    const int wslot0 = warp << 1;
    const int w0y = (wslot0 >> 2) << 2;
    const int w0x = (wslot0 & 3) << 2;

    const int gwy = (blockIdx.y << 2) + (wslot0 >> 2);
    const int gwx = (blockIdx.x << 2) + (wslot0 & 3);
    const bool valid0 = (gwy < Hp) && (gwx < Wp);
    const bool valid1 = (gwy < Hp) && (gwx + 1 < Wp);

    float acc[2][8][4];
    #pragma unroll
    for (int w = 0; w < 2; ++w)
        #pragma unroll
        for (int i = 0; i < 8; ++i)
            #pragma unroll
            for (int j = 0; j < 4; ++j) acc[w][i][j] = 0.f;

    const uint32_t* in_u32 = reinterpret_cast<const uint32_t*>(smem_in);
    const uint32_t* w_u32 = reinterpret_cast<const uint32_t*>(smem_w);
    const int ya = w0y + (g >> 2);
    const int yb = ya + 2;
    const int xa0 = w0x + (g & 3);
    const int xa1 = xa0 + 4;

    // per-thread base offsets (in uint32) for both window columns
    const int base0 = ya * (PATCH_STRIDE * 4) + xa0 * 4 + tig;
    const int base1 = ya * (PATCH_STRIDE * 4) + xa1 * 4 + tig;

    // 4 k16 chunks: tap pairs {(0,0),(0,1)}, {(0,2),(1,0)}, {(1,1),(1,2)}, {(2,0),(2,1)}
    const int tap_r0[4] = {0, 0, 1, 2};
    const int tap_q0[4] = {0, 2, 1, 0};
    const int tap_r1[4] = {0, 1, 1, 2};
    const int tap_q1[4] = {1, 0, 2, 1};

    #pragma unroll
    for (int ch = 0; ch < 4; ++ch) {
        const int o0 = (tap_r0[ch] * PATCH_STRIDE + tap_q0[ch]) * 4;
        const int o1 = (tap_r1[ch] * PATCH_STRIDE + tap_q1[ch]) * 4;
        const int dyb = 2 * PATCH_STRIDE * 4;

        uint32_t A[2][4];
        A[0][0] = in_u32[base0 + o0];
        A[0][1] = in_u32[base0 + o0 + dyb];
        A[0][2] = in_u32[base0 + o1];
        A[0][3] = in_u32[base0 + o1 + dyb];
        A[1][0] = in_u32[base1 + o0];
        A[1][1] = in_u32[base1 + o0 + dyb];
        A[1][2] = in_u32[base1 + o1];
        A[1][3] = in_u32[base1 + o1 + dyb];

        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            const uint32_t* wp = w_u32 + (nt * 8 + g) * (WSM_STRIDE / 2) + ch * 8;
            uint32_t B[2];
            B[0] = wp[tig];
            B[1] = wp[tig + 4];
            mma16816(acc[0][nt], A[0], B);
            mma16816(acc[1][nt], A[1], B);
        }
    }

    // k8 chunk: tap (2,2), k = 64..71
    {
        const int o0 = (2 * PATCH_STRIDE + 2) * 4;
        const int dyb = 2 * PATCH_STRIDE * 4;
        uint32_t A8[2][2];
        A8[0][0] = in_u32[base0 + o0];
        A8[0][1] = in_u32[base0 + o0 + dyb];
        A8[1][0] = in_u32[base1 + o0];
        A8[1][1] = in_u32[base1 + o0 + dyb];
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            const uint32_t* wp = w_u32 + (nt * 8 + g) * (WSM_STRIDE / 2) + 32;
            const uint32_t b = wp[tig];
            mma1688(acc[0][nt], A8[0], b);
            mma1688(acc[1][nt], A8[1], b);
        }
    }

    // ---- 4x4 pool (mean) + bias + sigmoid, accumulate per window ----
    float wsum = 0.f;
    #pragma unroll
    for (int w = 0; w < 2; ++w) {
        const bool valid = w ? valid1 : valid0;
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            float* c = acc[w][nt];
            // t0: column (2*tig) pixel-pair total, t1: column 2*tig+1
            float t0 = c[0] + c[2];
            float t1 = c[1] + c[3];
            #pragma unroll
            for (int off = 4; off < 32; off <<= 1) {
                t0 += __shfl_xor_sync(0xffffffffu, t0, off);
                t1 += __shfl_xor_sync(0xffffffffu, t1, off);
            }
            const float p0 = t0 * 0.0625f + smem_bias[nt * 8 + 2 * tig];
            const float p1 = t1 * 0.0625f + smem_bias[nt * 8 + 2 * tig + 1];
            float s = __frcp_rn(1.0f + __expf(-p0)) + __frcp_rn(1.0f + __expf(-p1));
            s = (g == 0) ? s : 0.0f;
            #pragma unroll
            for (int off = 1; off < 4; off <<= 1)
                s += __shfl_xor_sync(0xffffffffu, s, off);
            if (lane == 0 && valid) wsum += s;
        }
    }
    if (lane == 0) red[warp] = wsum;
    __syncthreads();
    if (t == 0) {
        float ssum = 0.f;
        #pragma unroll
        for (int i = 0; i < 8; ++i) ssum += red[i];
        atomicAdd(&out[n], ssum);
    }
}

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

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

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

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 tt = v / cg_count;
        const int wx = tt % Wp;
        const int wy = tt / 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));
}

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 tt = v / Wp;
        const int wy = tt % Hp;
        const int c  = tt / 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 fused_conv3x3_c8_pool4_sigmoid_sum(
        torch::Tensor x, torch::Tensor wk, torch::Tensor bias) {
    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 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(C == 8, "in_channels must be 8");
    TORCH_CHECK(wk.is_cuda() && wk.scalar_type() == torch::kHalf, "wk must be CUDA fp16");
    TORCH_CHECK(wk.numel() == 64 * 72, "wk must be [64,72]");
    TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kFloat, "bias must be CUDA fp32");
    TORCH_CHECK(bias.numel() == 64, "bias must be [64]");

    auto wkc = wk.contiguous();
    auto bc = bias.contiguous();
    const int Ho = H - 2;
    const int Wo = W - 2;
    TORCH_CHECK(Ho > 0 && Wo > 0, "spatial dims too small");
    const int Hp = Ho / 4;
    const int Wp = Wo / 4;

    auto out = torch::zeros({N}, xc.options().dtype(torch::kFloat32));
    if (N == 0 || Hp == 0 || Wp == 0) return out;

    dim3 grid((Wp + 3) / 4, (Hp + 3) / 4, N);
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_conv3x3_c8_pool4_sigmoid_sum_kernel<<<grid, 256, 0, stream>>>(
        xc.data_ptr<float>(),
        reinterpret_cast<const __half*>(wkc.data_ptr<at::Half>()),
        bc.data_ptr<float>(),
        out.data_ptr<float>(),
        H, W, Hp, Wp);
    return out;
}

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 fused_conv3x3_c8_pool4_sigmoid_sum(torch::Tensor x, torch::Tensor wk, torch::Tensor bias);\n"
    "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_v5",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=[
        "fused_conv3x3_c8_pool4_sigmoid_sum",
        "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:
      * Fast path (sm_80+): a single custom tensor-core kernel fuses the 3x3
        conv (fp16 mma: 4x m16n8k16 + 1x m16n8k8 = exactly K=72, no padding),
        the 4x4 average pool, the sigmoid and the global sum. The fp32 NCHW
        input is converted to fp16 directly into shared memory (fusing the
        cast + layout transpose), and the multi-GB conv output is never
        materialized. Shared-memory layouts are bank-conflict-free and
        occupancy is pinned to 2 blocks/SM.
      * Backup path: fused fp32->fp16 cast + NCHW->NHWC transpose, cuDNN fp16
        channels_last conv, and a fused pool+sigmoid+sum kernel.
      * Fallback: original PyTorch ops.
    """
    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._wk = None
        self._bias32 = None
        self._w16 = None
        self._b16 = None
        self._w_key = None

    def _refresh_caches(self):
        w = self.conv.weight
        key = (w.data_ptr(), w._version)
        if self._w_key == key:
            return
        self._w_key = key
        b = self.conv.bias
        co = w.shape[0]
        # mma weight layout: [co, 72] fp16, k = (ky*3+kx)*8 + ci
        self._wk = None
        self._bias32 = None
        try:
            if co == 64 and w.shape[1] == 8 and w.shape[2] == 3 and w.shape[3] == 3:
                wp = w.detach().permute(2, 3, 1, 0).reshape(72, co)
                self._wk = wp.t().half().contiguous()
                if b is not None:
                    self._bias32 = b.detach().float().contiguous()
                else:
                    self._bias32 = torch.zeros(co, dtype=torch.float32, device=w.device)
        except Exception:
            self._wk = None
            self._bias32 = None
        # cuDNN backup weights
        try:
            self._w16 = w.detach().half().contiguous(memory_format=torch.channels_last)
            self._b16 = b.detach().half() if b is not None else None
        except Exception:
            self._w16 = None
            self._b16 = None

    def _mma_capable(self):
        try:
            if not torch.cuda.is_available():
                return False
            major, minor = torch.cuda.get_device_capability()
            return major >= 8
        except Exception:
            return False

    def forward(self, x):
        pool = self.pool_kernel_size
        fast_mma = (
            x.is_cuda
            and x.dtype == torch.float32
            and isinstance(pool, int)
            and pool == 4
            and self.conv.in_channels == 8
            and self.conv.out_channels == 64
            and self.conv.kernel_size == (3, 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 x.size(2) >= 3
            and x.size(3) >= 3
        )
        if fast_mma and self._mma_capable():
            self._refresh_caches()
            if self._wk is not None and self._bias32 is not None:
                return _ext.fused_conv3x3_c8_pool4_sigmoid_sum(
                    x, self._wk, self._bias32)

        fast_v2 = (
            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_v2:
            try:
                self._refresh_caches()
                if self._w16 is not None:
                    x16 = _ext.nchw_to_nhwc_fp16(x)
                    y = F.conv2d(
                        x16,
                        self._w16,
                        self._b16,
                        self.conv.stride,
                        self.conv.padding,
                        self.conv.dilation,
                        self.conv.groups,
                    )
                    if y.is_cont