```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)
#   2) a tiny fused epilogue: kernel-weight sums + dot product + biases +
#      numerically stable log-sum-exp + scaling.
# Both steps now live in a SINGLE kernel launch: the last block of each batch
# (detected with a classic thread-fence/atomic counter) performs the epilogue
# for that batch directly, so no second kernel launch is needed.
_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.
template <bool VEC>
__device__ __forceinline__ double chunk_sum(const float* __restrict__ row,
                                            long hw, int chunk, int chunks) {
    const int stride = T1;
    if (VEC) {
        const long n4 = hw >> 2;
        const float4* p4 = reinterpret_cast<const float4*>(row);
        const long per = (n4 + chunks - 1) / chunks;
        const long start = per * (long)chunk;
        long end = start + per;
        if (end > n4) end = n4;
        float4 acc = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
        long i = start + threadIdx.x;
        for (; i + 3L * stride < end; i += 4L * stride) {
            const float4 a = __ldcs(p4 + i);
            const float4 b = __ldcs(p4 + i + stride);
            const float4 c = __ldcs(p4 + i + 2L * stride);
            const float4 d = __ldcs(p4 + i + 3L * stride);
            acc.x += (a.x + b.x) + (c.x + d.x);
            acc.y += (a.y + b.y) + (c.y + d.y);
            acc.z += (a.z + b.z) + (c.z + d.z);
            acc.w += (a.w + b.w) + (c.w + d.w);
        }
        for (; i < end; i += stride) {
            const float4 a = __ldcs(p4 + i);
            acc.x += a.x; acc.y += a.y; acc.z += a.z; acc.w += a.w;
        }
        return (double)((acc.x + acc.y) + (acc.z + acc.w));
    } else {
        const long n = hw;
        const long per = (n + chunks - 1) / chunks;
        const long start = per * (long)chunk;
        long end = start + per;
        if (end > n) end = n;
        float acc = 0.0f;
        long i = start + threadIdx.x;
        for (; i + 3L * stride < end; i += 4L * stride) {
            acc += (row[i] + row[i + stride]) +
                   (row[i + 2L * stride] + row[i + 3L * stride]);
        }
        for (; i < end; i += stride) acc += row[i];
        return (double)acc;
    }
}

// Grid: blockIdx.x = flattened (batch * Cin), blockIdx.y = chunk index.
// Each block stores one partial spatial sum. The last block of each batch
// fuses the whole epilogue (weight sums, dot product, biases, log-sum-exp).
template <bool VEC>
__global__ void __launch_bounds__(T1, 6)
fused_convt_kernel(const float* __restrict__ x,
                   const float* __restrict__ w,
                   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)
                   long hw, int chunks, int Cin, int Cout, int khw,
                   double inv_hw, double scale) {
    const long bc = (long)blockIdx.x;
    const int b = (int)(bc / (long)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 + bc * hw, hw, (int)blockIdx.y, 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[bc * (long)chunks + (long)blockIdx.y] = rr;
    }

    // ---- 2) last-block-of-batch detection (thread-fence reduction) ----
    __shared__ bool is_last;
    __threadfence();
    __syncthreads();
    if (threadIdx.x == 0) {
        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 + (long)b * (long)Cin * chunks;
    for (int i = t; i < Cin; i += T1) {
        const double* p = pb + (long)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] * sum_k w[i,o,k] ) * inv_hw + bias[o] + cbias[o]
    // Two threads per output channel, each covering half of the inputs.
    const int pairs  = T1 / 2;
    const int oiters = (Cout + pairs - 1) / pairs;
    for (int it = 0; it < oiters; ++it) {
        const int o = it * pairs + (t >> 1);
        const bool active = (o < Cout);
        double acc = 0.0;
        if (active) {
            for (int i = (t & 1); i < Cin; i += 2) {
                const float* wp = w + ((long)i * Cout + o) * (long)khw;
                float ks = 0.0f;
                int k = 0;
                for (; k + 4 <= khw; k += 4)
                    ks += (wp[k] + wp[k + 1]) + (wp[k + 2] + wp[k + 3]);
                for (; k < khw; ++k) ks += wp[k];
                acc += xs[i] * (double)ks;
            }
        }
        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 weight,
                            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 = weight.size(1);
    const long khw  = weight.size(2) * weight.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;
    const long per_thr = 8;
    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)bc, (unsigned)chunks, 1);
    if (vec) {
        fused_convt_kernel<true><<<grid, T1, smem, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), cb_ptr,
            bias.data_ptr<float>(), partials.data_ptr<double>(),
            counters.data_ptr<int>(), out.data_ptr<float>(),
            hw, chunks, (int)Cin, (int)Cout, (int)khw, inv_hw, scale);
    } else {
        fused_convt_kernel<false><<<grid, T1, smem, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), cb_ptr,
            bias.data_ptr<float>(), partials.data_ptr<double>(),
            counters.data_ptr<int>(), out.data_ptr<float>(),
            hw, chunks, (int)Cin, (int)Cout, (int)khw, inv_hw, scale);
    }
    return out;
}
"""

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

_fused_op = load_inline(
    name="fused_convt_gap_bias_lse_v3",
    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
        )
        s = ct.stride
        d = ct.dilation
        k = ct.kernel_size
        op = ct.output_padding
        self._sH, self._sW = s[0], s[1]
        self._dH, self._dW = d[0], d[1]
        self._kH, self._kW = k[0], k[1]
        self._opH, self._opW = op[0], op[1]
        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 = x.shape[0]

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

            H, W = x.shape[2], x.shape[3]
            # 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(),
                ct.weight.contiguous(),
                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
```