```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
from torch.utils.cpp_extension import load_inline

# Mathematical shortcut (valid for ConvTranspose2d with padding == 0 and
# groups == 1): every input pixel's kernel "stamp" lies fully inside the
# transposed-convolution output, so the global average of its output is
#   mean(out)[b,o] = ( sum_i sum(x[b,i]) * sum(W[i,o]) ) / (Ho*Wo) + cbias[o]
# The whole model therefore reduces to:
#   1) per-(batch,channel) spatial sums of the input (one memory-bound pass
#      over the input that never materializes the transposed-conv output)
#   2) a tiny fused epilogue: dot product with cached kernel-weight sums +
#      biases + numerically stable log-sum-exp + scaling.
# Both steps run in ONE kernel launch: the last block of each batch (detected
# with a thread-fence/atomic counter) performs the epilogue for that batch,
# overlapping with the still-running blocks of the other batches.
#
# Speed tricks in this version:
#   * all per-call parameters (kernel-weight sums, conv bias, extra bias) are
#     packed ONCE into a single cached tensor -> fewer pybind/tensor args and
#     no per-call branches for the optional conv bias
#   * partials + counters buffers are cached across calls (counter self-resets
#     inside the kernel), so only the tiny output tensor is allocated per call
#   * inv_hw and the packed params are cached in Python, keyed on versions
#   * ~256 KB contiguous per block (tuned against SM count) with 8x-unrolled
#     float4 __ldcs streaming loads and 32-bit indexing
#   * grid ordered so concurrently-resident blocks sweep memory linearly
#   * __threadfence executed only by the single writer thread of each block
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <algorithm>

#define T1 256
#define NWARP (T1 / 32)

__device__ __forceinline__ double warp_sum(double v) {
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, off);
    return v;
}

// Sum of one chunk of a contiguous row of hw floats (32-bit indexing,
// 8x-unrolled streaming loads).
template <bool VEC>
__device__ __forceinline__ double chunk_sum(const float* __restrict__ row,
                                            int hw, int chunk, int chunks) {
    if (VEC) {
        const int n4 = hw >> 2;
        const int per = (n4 + chunks - 1) / chunks;
        const int start = per * chunk;
        int end = start + per;
        if (end > n4) end = n4;
        const float4* p4 = reinterpret_cast<const float4*>(row) + start;
        const int cnt = end - start;
        float4 aA = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
        float4 aB = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
        int i = threadIdx.x;
        for (; i + 7 * T1 < cnt; i += 8 * T1) {
            const float4 v0 = __ldcs(p4 + i);
            const float4 v1 = __ldcs(p4 + i + T1);
            const float4 v2 = __ldcs(p4 + i + 2 * T1);
            const float4 v3 = __ldcs(p4 + i + 3 * T1);
            const float4 v4 = __ldcs(p4 + i + 4 * T1);
            const float4 v5 = __ldcs(p4 + i + 5 * T1);
            const float4 v6 = __ldcs(p4 + i + 6 * T1);
            const float4 v7 = __ldcs(p4 + i + 7 * T1);
            aA.x += (v0.x + v1.x) + (v2.x + v3.x);
            aA.y += (v0.y + v1.y) + (v2.y + v3.y);
            aA.z += (v0.z + v1.z) + (v2.z + v3.z);
            aA.w += (v0.w + v1.w) + (v2.w + v3.w);
            aB.x += (v4.x + v5.x) + (v6.x + v7.x);
            aB.y += (v4.y + v5.y) + (v6.y + v7.y);
            aB.z += (v4.z + v5.z) + (v6.z + v7.z);
            aB.w += (v4.w + v5.w) + (v6.w + v7.w);
        }
        for (; i + 3 * T1 < cnt; i += 4 * T1) {
            const float4 v0 = __ldcs(p4 + i);
            const float4 v1 = __ldcs(p4 + i + T1);
            const float4 v2 = __ldcs(p4 + i + 2 * T1);
            const float4 v3 = __ldcs(p4 + i + 3 * T1);
            aA.x += (v0.x + v1.x) + (v2.x + v3.x);
            aA.y += (v0.y + v1.y) + (v2.y + v3.y);
            aA.z += (v0.z + v1.z) + (v2.z + v3.z);
            aA.w += (v0.w + v1.w) + (v2.w + v3.w);
        }
        for (; i < cnt; i += T1) {
            const float4 v0 = __ldcs(p4 + i);
            aA.x += v0.x; aA.y += v0.y; aA.z += v0.z; aA.w += v0.w;
        }
        return (double)(((aA.x + aB.x) + (aA.y + aB.y)) +
                        ((aA.z + aB.z) + (aA.w + aB.w)));
    } else {
        const int per = (hw + chunks - 1) / chunks;
        const int start = per * chunk;
        int end = start + per;
        if (end > hw) end = hw;
        const float* p = row + start;
        const int cnt = end - start;
        float aA = 0.0f, aB = 0.0f;
        int i = threadIdx.x;
        for (; i + 7 * T1 < cnt; i += 8 * T1) {
            aA += (p[i] + p[i + T1]) + (p[i + 2 * T1] + p[i + 3 * T1]);
            aB += (p[i + 4 * T1] + p[i + 5 * T1]) + (p[i + 6 * T1] + p[i + 7 * T1]);
        }
        for (; i + 3 * T1 < cnt; i += 4 * T1)
            aA += (p[i] + p[i + T1]) + (p[i + 2 * T1] + p[i + 3 * T1]);
        for (; i < cnt; i += T1) aA += p[i];
        return (double)(aA + aB);
    }
}

// Grid: blockIdx.x = chunk index, blockIdx.y = flattened (batch * Cin).
// Consecutive blockIdx.x covers adjacent memory, so the resident blocks sweep
// the tensor linearly (good DRAM locality). The last block of each batch
// fuses the whole epilogue (dot with cached weight sums, biases, log-sum-exp).
// packed layout: [ ksum (Cin*Cout) | cbias (Cout) | bias (Cout) ].
template <bool VEC>
__global__ void __launch_bounds__(T1, 4)
fused_convt_kernel(const float* __restrict__ x,
                   const float* __restrict__ packed,
                   double* __restrict__ partials,    // (B*Cin, chunks)
                   int*   __restrict__ counters,     // (B), zeroed on entry/exit
                   float* __restrict__ out,          // (B)
                   int hw, int chunks, int Cin, int Cout,
                   double inv_hw, double scale) {
    const float* __restrict__ ksum  = packed;
    const float* __restrict__ cbias = packed + (size_t)Cin * Cout;
    const float* __restrict__ bias  = cbias + Cout;

    const int bc    = (int)blockIdx.y;
    const int chunk = (int)blockIdx.x;
    const int b     = bc / Cin;
    const int lane  = threadIdx.x & 31;
    const int warp  = threadIdx.x >> 5;

    // ---- 1) block-wide partial spatial sum of x[b, c] chunk ----
    const double v = chunk_sum<VEC>(x + (size_t)bc * (size_t)hw, hw, chunk, chunks);
    __shared__ double sh[NWARP];
    double r = warp_sum(v);
    if (lane == 0) sh[warp] = r;
    __syncthreads();
    if (warp == 0) {
        double rr = (lane < NWARP) ? sh[lane] : 0.0;
        rr = warp_sum(rr);
        if (lane == 0) partials[(size_t)bc * chunks + chunk] = rr;
    }

    // ---- 2) last-block-of-batch detection (fence only in the writer) ----
    __shared__ bool is_last;
    if (threadIdx.x == 0) {
        __threadfence();
        const int old = atomicAdd(&counters[b], 1);
        is_last = (old == Cin * chunks - 1);
        if (is_last) atomicExch(&counters[b], 0);  // reset for the next call
    }
    __syncthreads();
    if (!is_last) return;
    __threadfence();

    // ---- 3) fused epilogue for batch b (single block) ----
    extern __shared__ double smem[];
    double* xs  = smem;        // Cin doubles
    double* vv  = smem + Cin;  // Cout doubles
    double* red = vv + Cout;   // 32 doubles

    const int t = threadIdx.x;
    const double* pb = partials + (size_t)b * (size_t)Cin * chunks;
    for (int i = t; i < Cin; i += T1) {
        const double* p = pb + (size_t)i * chunks;
        double s = 0.0;
        for (int c = 0; c < chunks; ++c) s += p[c];
        xs[i] = s;
    }
    __syncthreads();

    // vv[o] = ( sum_i xs[i] * ksum[i,o] ) * inv_hw + bias[o] + cbias[o]
    // Two threads per output channel, each covering half of the inputs.
    const int pairs = T1 / 2;
    for (int o0 = 0; o0 < Cout; o0 += pairs) {
        const int o = o0 + (t >> 1);
        const bool active = (o < Cout);
        double acc = 0.0;
        if (active) {
            const float* kp = ksum + o;
            for (int i = (t & 1); i < Cin; i += 2)
                acc += xs[i] * (double)kp[(size_t)i * Cout];
        }
        acc += __shfl_xor_sync(0xffffffffu, acc, 1);
        if (active && ((t & 1) == 0))
            vv[o] = acc * inv_hw + (double)bias[o] + (double)cbias[o];
    }
    __syncthreads();

    // ---- numerically stable log-sum-exp over channels ----
    double m = -1.0e308;
    for (int o = t; o < Cout; o += T1) m = fmax(m, vv[o]);
    for (int off = 16; off > 0; off >>= 1)
        m = fmax(m, __shfl_down_sync(0xffffffffu, m, off));
    if (lane == 0) red[warp] = m;
    __syncthreads();
    if (warp == 0) {
        double mm = (lane < NWARP) ? red[lane] : -1.0e308;
        for (int off = 16; off > 0; off >>= 1)
            mm = fmax(mm, __shfl_down_sync(0xffffffffu, mm, off));
        if (lane == 0) red[0] = mm;
    }
    __syncthreads();
    const double mx = red[0];
    __syncthreads();

    double s = 0.0;
    for (int o = t; o < Cout; o += T1) s += exp(vv[o] - mx);
    s = warp_sum(s);
    if (lane == 0) red[warp] = s;
    __syncthreads();
    if (warp == 0) {
        double ss = (lane < NWARP) ? red[lane] : 0.0;
        ss = warp_sum(ss);
        if (lane == 0) out[b] = (float)(scale * (mx + log(ss)));
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor packed,
                            torch::Tensor partials, torch::Tensor counters,
                            int64_t Cin, int64_t Cout,
                            double inv_hw, double scale) {
    const c10::cuda::CUDAGuard device_guard(x.device());
    const long B  = x.size(0);
    const long hw = x.size(2) * x.size(3);
    const long bc = B * Cin;

    auto stream = at::cuda::getCurrentCUDAStream();

    const bool vec  = ((hw & 3L) == 0L);
    const long work = vec ? (hw >> 2) : hw;

    // ~256 KB of contiguous data per block; big enough to amortize all
    // per-block overheads, small enough that there are still plenty of blocks
    // to saturate every SM. If that would leave the GPU under-filled (small
    // batch*channels), split rows further until ~8 blocks per SM exist.
    const int sm_count =
        at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const long unit      = (long)T1 * 8;      // one 8x-unrolled iteration
    const long per_block = unit * 8;          // 256 KB per block (vec case)
    long chunks_l = (work + per_block - 1) / per_block;
    const long want = (long)sm_count * 8;
    if (chunks_l * bc < want) {
        long c2 = (want + bc - 1) / bc;
        if (c2 > chunks_l)
            chunks_l = std::min(c2, (work + unit - 1) / unit);
    }
    chunks_l      = std::max<long>(1L, std::min<long>(chunks_l, 64L));
    const int chunks = (int)chunks_l;

    auto out = torch::empty({B, 1}, x.options());

    const size_t smem = (size_t)(Cin + Cout + 32) * sizeof(double);
    dim3 grid((unsigned)chunks, (unsigned)bc, 1);
    if (vec) {
        fused_convt_kernel<true><<<grid, T1, smem, stream>>>(
            x.data_ptr<float>(), packed.data_ptr<float>(),
            partials.data_ptr<double>(), counters.data_ptr<int>(),
            out.data_ptr<float>(), (int)hw, chunks, (int)Cin, (int)Cout,
            inv_hw, scale);
    } else {
        fused_convt_kernel<false><<<grid, T1, smem, stream>>>(
            x.data_ptr<float>(), packed.data_ptr<float>(),
            partials.data_ptr<double>(), counters.data_ptr<int>(),
            out.data_ptr<float>(), (int)hw, chunks, (int)Cin, (int)Cout,
            inv_hw, scale);
    }
    return out;
}
"""

_cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor packed, "
    "torch::Tensor partials, torch::Tensor counters, int64_t Cin, "
    "int64_t Cout, double inv_hw, double scale);"
)

_fused_op = load_inline(
    name="fused_convt_gap_bias_lse_v6",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Replaces the expensive transposed convolution
    + global-average-pooling chain with a mathematically equivalent single
    fused custom CUDA kernel (valid for padding == 0 and groups == 1, which is
    the configuration used by this model). Per-call state (packed weight sums,
    conv bias and extra bias, partials buffer, batch counters, inv_hw) is
    cached so the steady-state host path is minimal. Falls back to the
    original PyTorch ops for any other configuration.
    """

    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self._fwd = _fused_op.fused_forward

        ct = self.conv_transpose
        self._static_ok = (
            ct.groups == 1
            and tuple(ct.padding) == (0, 0)
            and ct.weight.dtype == torch.float32
            and self.bias.dtype == torch.float32
            and self.bias.numel() == ct.out_channels
            and (ct.in_channels + ct.out_channels + 32) * 8 <= 32 * 1024
        )
        self._Cin = ct.in_channels
        self._Cout = ct.out_channels
        self._sH, self._sW = ct.stride
        self._dH, self._dW = ct.dilation
        self._kH, self._kW = ct.kernel_size
        self._opH, self._opW = ct.output_padding
        self._packed = None
        self._packed_key = None
        self._partials = None
        self._counters = None
        self._inv_hw = {}

    def forward(self, x):
        if (
            self._static_ok
            and x.dim() == 4
            and x.is_cuda
            and x.dtype == torch.float32
        ):
            B, C, H, W = x.shape
            if B * C <= 65535 and H * W < 2147483647 and C == self._Cin:
                ct = self.conv_transpose
                w = ct.weight
                if w.is_cuda and w.dtype == torch.float32:
                    dev = x.device
                    cb = ct.bias

                    # Pack kernel-weight sums + conv bias + extra bias once;
                    # rebuild only when one of the parameters changes.
                    key = (
                        w.data_ptr(), w._version,
                        self.bias.data_ptr(), self.bias._version,
                        0 if cb is None else cb.data_ptr(),
                        0 if cb is None else cb._version,
                    )
                    if key != self._packed_key:
                        ksum = w.detach().sum(dim=(2, 3)).reshape(-1)
                        bf = self.bias.detach().reshape(-1)
                        if cb is None:
                            cbf = torch.zeros_like(bf)
                        else:
                            cbf = cb.detach().reshape(-1).to(torch.float32)
                        self._packed = torch.cat((ksum, cbf, bf)).contiguous()
                        self._packed_key = key

                    # Cached scratch buffers (partials content is fully
                    # rewritten every call; counters self-reset to zero).
                    need_p = B * C * 64
                    p = self._partials
                    if p is None or p.device != dev or p.numel() < need_p:
                        p = torch.empty(need_p, device=dev, dtype=torch.float64)
                        self._partials = p
                    ctr = self._counters
                    if ctr is None or ctr.device != dev or ctr.numel() < B:
                        ctr = torch.zeros(max(B, 64), device=dev, dtype=torch.int32)
                        self._counters = ctr

                    inv_hw = self._inv_hw.get((H, W))
                    if inv_hw is None:
                        Ho = (H - 1) * self._sH + self._dH * (self._kH - 1) + self._opH + 1
                        Wo = (W - 1) * self._sW + self._dW * (self._kW - 1) + self._opW + 1
                        inv_hw = 1.0 / float(Ho * Wo)
                        self._inv_hw[(H, W)] = inv_hw

                    if not x.is_contiguous():
                        x = x.contiguous()
                    return self._fwd(
                        x, self._packed, p, ctr, C, self._Cout, inv_hw, 10.0
                    )

        # Generic fallback: original computation.
        x = self.conv_transpose(x)
        x = torch.mean(x, dim=(2, 3), keepdim=True)
        x = x + self.bias
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = torch.sum(x, dim=(2, 3))
        x = x * 10.0
        return x
```