```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 1 GiB 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.
# Speed tricks vs. the previous version:
#   * weight sums (ksum) are precomputed once in Python and cached
#   * the __threadfence is executed by thread 0 only (the only writer)
#   * 8x-unrolled float4 streaming loads with 32-bit indexing
#   * grid ordered so concurrently-resident blocks sweep memory linearly
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.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).
template <bool VEC>
__global__ void __launch_bounds__(T1, 4)
fused_convt_kernel(const float* __restrict__ x,
                   const float* __restrict__ ksum,   // (Cin, Cout)
                   const float* __restrict__ cbias,  // (Cout) or nullptr
                   const float* __restrict__ bias,   // (Cout)
                   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 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)) {
            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,
                            torch::Tensor counters,
                            double inv_hw, double scale) {
    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;
    const long per_thr = 16;
    long chunks_l = (work + (long)T1 * per_thr - 1) / ((long)T1 * per_thr);
    chunks_l      = std::max<long>(1L, std::min<long>(chunks_l, 65535L));
    const int chunks = (int)chunks_l;

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

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

    dim3 grid((unsigned)chunks, (unsigned)bc, 1);
    if (vec) {
        fused_convt_kernel<true><<<grid, T1, smem, stream>>>(
            x.data_ptr<float>(), ksum.data_ptr<float>(), cb_ptr,
            bias.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>(), ksum.data_ptr<float>(), cb_ptr,
            bias.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 ksum, "
    "torch::Tensor cbias, torch::Tensor bias, torch::Tensor counters, "
    "double inv_hw, double scale);"
)

_fused_op = load_inline(
    name="fused_convt_gap_bias_lse_v4",
    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). 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._counters = 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])

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

                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

                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),
                    ctr,
                    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
```