```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 per-(batch,channel) spatial sums of the
# input (one memory-bound pass over 1 GiB that never materializes the
# transposed-conv output) plus a tiny fused epilogue (dot product with cached
# kernel-weight sums + biases + numerically stable log-sum-exp + scaling).
#
# Speed tricks vs. the previous version:
#   * TWO kernels instead of one fused kernel: the hot streaming kernel no
#     longer pays a __threadfence + atomicAdd + syncthreads tail in EVERY
#     block (that tail idled each block's threads for ~1-2 us after its loads
#     finished, wasting ~10-15% of DRAM time). The epilogue runs as a tiny
#     second kernel (one block per batch) that overlaps nothing but costs
#     only a few microseconds.
#   * 8x fewer, 4x bigger blocks (128 KB contiguous each): block launch /
#     reduction / store overheads are amortized 4x better and wave
#     quantization at the end of the grid is negligible.
#   * 8x-unrolled float4 __ldcs streaming loads, 32-bit indexing, grid ordered
#     so concurrently-resident blocks sweep memory linearly.
#   * weight sums (ksum) are precomputed once in Python and cached.
_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);
    }
}

// Pure streaming reduction kernel. 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). No fences, no atomics: each block writes one double and exits.
template <bool VEC>
__global__ void __launch_bounds__(T1, 4)
sum_kernel(const float* __restrict__ x,
           double* __restrict__ partials,  // (B*Cin, chunks)
           int hw, int chunks) {
    const int bc = (int)blockIdx.y;
    const double v =
        chunk_sum<VEC>(x + (size_t)bc * (size_t)hw, hw, (int)blockIdx.x, chunks);
    __shared__ double sh[NWARP];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    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 + (int)blockIdx.x] = rr;
    }
}

// Tiny epilogue: one block per batch. Folds the per-channel partial sums,
// computes vv[o] = (sum_i xs[i] * ksum[i,o]) * inv_hw + bias[o] + cbias[o],
// then a numerically stable log-sum-exp over channels and the final scaling.
__global__ void __launch_bounds__(T1)
epilogue_kernel(const double* __restrict__ partials,  // (B*Cin, chunks)
                const float* __restrict__ ksum,       // (Cin, Cout)
                const float* __restrict__ cbias,      // (Cout) or nullptr
                const float* __restrict__ bias,       // (Cout)
                float* __restrict__ out,              // (B)
                int chunks, int Cin, int Cout,
                double inv_hw, double scale) {
    const int b = blockIdx.x;
    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 int lane = t & 31;
    const int warp = t >> 5;

    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();

    // 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)) {
            double val = acc * inv_hw + (double)bias[o];
            if (cbias != nullptr) val += (double)cbias[o];
            vv[o] = val;
        }
    }
    __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 ksum,
                            torch::Tensor cbias, torch::Tensor bias,
                            double inv_hw, double scale) {
    const c10::cuda::CUDAGuard device_guard(x.device());
    const long B    = x.size(0);
    const long Cin  = x.size(1);
    const long hw   = x.size(2) * x.size(3);
    const long Cout = ksum.size(1);
    const long bc   = B * Cin;

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

    const bool vec  = ((hw & 3L) == 0L);
    const long work = vec ? (hw >> 2) : hw;
    // Aim for ~128 KB of contiguous data per block (256 threads x 8 float4
    // loads x 4 unrolled iterations): big enough to amortize all per-block
    // overheads, small enough to keep wave quantization negligible.
    const long per_block = (long)T1 * 8 * 4;
    long chunks_l = (work + per_block - 1) / per_block;
    chunks_l      = std::max<long>(1L, std::min<long>(chunks_l, 64L));
    const int chunks = (int)chunks_l;

    auto partials = torch::empty(
        {bc, (long)chunks}, x.options().dtype(torch::kFloat64));
    auto out = torch::empty({B, 1}, x.options());

    dim3 grid1((unsigned)chunks, (unsigned)bc, 1);
    if (vec) {
        sum_kernel<true><<<grid1, T1, 0, stream>>>(
            x.data_ptr<float>(), partials.data_ptr<double>(), (int)hw, chunks);
    } else {
        sum_kernel<false><<<grid1, T1, 0, stream>>>(
            x.data_ptr<float>(), partials.data_ptr<double>(), (int)hw, chunks);
    }

    const float* cb_ptr =
        (cbias.defined() && cbias.numel() > 0) ? cbias.data_ptr<float>() : nullptr;
    const size_t smem = (size_t)(Cin + Cout + 32) * sizeof(double);
    epilogue_kernel<<<dim3((unsigned)B, 1, 1), T1, smem, stream>>>(
        partials.data_ptr<double>(), ksum.data_ptr<float>(), cb_ptr,
        bias.data_ptr<float>(), out.data_ptr<float>(), chunks, (int)Cin,
        (int)Cout, inv_hw, scale);
    return out;
}
"""

_cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor ksum, "
    "torch::Tensor cbias, torch::Tensor bias, double inv_hw, double scale);"
)

_fused_op = load_inline(
    name="fused_convt_gap_bias_lse_v5",
    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 pair of
    custom CUDA kernels (valid for padding == 0 and groups == 1, which is the
    configuration used by this model): a pure-streaming spatial-sum kernel
    followed by a tiny fused epilogue kernel (dot with cached weight sums,
    bias add, log-sum-exp, scaling). 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._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._ksum = None
        self._ksum_key = None
        self._empty_bias = None

    def forward(self, x):
        if (
            self._static_ok
            and x.dim() == 4
            and x.is_cuda
            and x.dtype == torch.float32
        ):
            ct = self.conv_transpose
            dev = x.device
            B = int(x.shape[0])
            C = int(x.shape[1])
            H = int(x.shape[2])
            W = int(x.shape[3])
            w = ct.weight

            if B * C <= 65535 and H * W < 2147483647 and w.is_cuda:
                # Cache the per-(in,out) kernel-weight sums; they only change
                # when the weights are updated.
                key = (w.data_ptr(), w._version)
                if key != self._ksum_key:
                    self._ksum = w.detach().sum(dim=(2, 3)).contiguous()
                    self._ksum_key = key

                cb = ct.bias
                if cb is None:
                    eb = self._empty_bias
                    if eb is None or eb.device != dev:
                        eb = torch.empty(0, device=dev, dtype=torch.float32)
                        self._empty_bias = eb
                    cb = eb
                else:
                    cb = cb.contiguous()

                # Output spatial size of the transposed convolution.
                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)

                return self._fwd(
                    x.contiguous(),
                    self._ksum,
                    cb,
                    self.bias.contiguous().view(-1),
                    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
```