```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

# Key mathematical insight:
# For ConvTranspose2d with padding == 0 and groups == 1, every input pixel's
# "stamp" lies completely inside the output (nothing is clipped at the
# borders). Therefore the global average pooling over the transposed-conv
# output can be computed WITHOUT materializing the (huge) output tensor:
#
#   sum_{y,x} out[b,o,y,x] = sum_i ( sum_{h,w} x[b,i,h,w] )
#                            * ( sum_{kh,kw} W[i,o,kh,kw] )
#   mean_{y,x} out[b,o]     = that_sum / (Ho*Wo) + conv_bias[o]
#
# So the whole model collapses to:
#   1) per-(batch,channel) spatial sum of the input   (custom kernel #1)
#   2) tiny fused kernel: dot with per-channel kernel sums, add biases,
#      log-sum-exp over channels, and scale by 10      (custom kernel #2)
#
# This turns ~600 GFLOP of transposed convolution into one memory-bound
# pass over the input, giving a large speedup while matching the original
# numerics (accumulations done in double for extra accuracy).

_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

// ---------------------------------------------------------------------------
// Kernel 1: per-(batch, channel) sum over the spatial dimensions.
// One thread block per (b, c) pair; float4-vectorized loads; the final
// block reduction is performed in double precision for accuracy.
// ---------------------------------------------------------------------------
__global__ void spatial_sum_kernel(const float* __restrict__ x,
                                   double* __restrict__ out,
                                   long hw) {
    const float* p = x + (long)blockIdx.x * hw;
    const int tid  = threadIdx.x;
    const int nthr = blockDim.x;

    float acc = 0.0f;
    if ((hw & 3L) == 0L) {
        const float4* p4 = reinterpret_cast<const float4*>(p);
        const long n4 = hw >> 2;
        for (long i = tid; i < n4; i += nthr) {
            float4 v = p4[i];
            acc += v.x + v.y + v.z + v.w;
        }
    } else {
        for (long i = tid; i < hw; i += nthr) {
            acc += p[i];
        }
    }

    __shared__ double sh[32];
    double a = (double)acc;
    for (int off = 16; off > 0; off >>= 1)
        a += __shfl_down_sync(0xffffffffu, a, off);
    const int lane = tid & 31;
    const int warp = tid >> 5;
    if (lane == 0) sh[warp] = a;
    __syncthreads();
    const int nwarp = (nthr + 31) >> 5;
    if (warp == 0) {
        double r = (lane < nwarp) ? sh[lane] : 0.0;
        for (int off = 16; off > 0; off >>= 1)
            r += __shfl_down_sync(0xffffffffu, r, off);
        if (lane == 0) out[blockIdx.x] = r;
    }
}

// ---------------------------------------------------------------------------
// Kernel 2 (fully fused): kernel-sum dot product + conv bias + model bias +
// log-sum-exp over channels + final scaling.
// One thread block per batch element.
// ---------------------------------------------------------------------------
__global__ void fused_bias_lse_kernel(const double* __restrict__ xsum,   // (B, Cin)
                                      const float*  __restrict__ weight, // (Cin, Cout, kh, kw)
                                      const float*  __restrict__ cbias,  // (Cout) or nullptr
                                      const float*  __restrict__ bias,   // (Cout)
                                      float*        __restrict__ out,    // (B)
                                      int Cin, int Cout, int khw,
                                      double inv_hw_out, double scale) {
    const int b    = blockIdx.x;
    const int t    = threadIdx.x;
    const int nthr = blockDim.x;

    extern __shared__ double v[];              // Cout doubles
    const double* xs = xsum + (long)b * Cin;

    // v[o] = (sum_i xsum[i] * kernelsum[i,o]) / (Ho*Wo) + conv_bias[o] + bias[o]
    for (int o = t; o < Cout; o += nthr) {
        double acc = 0.0;
        for (int i = 0; i < Cin; ++i) {
            const float* w = weight + ((long)i * Cout + o) * khw;
            float ws = 0.0f;
            for (int k = 0; k < khw; ++k) ws += w[k];
            acc += xs[i] * (double)ws;
        }
        double val = acc * inv_hw_out + (double)bias[o];
        if (cbias != nullptr) val += (double)cbias[o];
        v[o] = val;
    }
    __syncthreads();

    __shared__ double red[32];
    const int lane  = t & 31;
    const int warp  = t >> 5;
    const int nwarp = (nthr + 31) >> 5;

    // block-wide max (for a numerically stable log-sum-exp)
    double m = -1.0e308;
    for (int o = t; o < Cout; o += nthr) m = fmax(m, v[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();

    // block-wide sum of exp(v - mx)
    double s = 0.0;
    for (int o = t; o < Cout; o += nthr) s += exp(v[o] - mx);
    for (int off = 16; off > 0; off >>= 1)
        s += __shfl_down_sync(0xffffffffu, s, off);
    if (lane == 0) red[warp] = s;
    __syncthreads();
    if (warp == 0) {
        double ss = (lane < nwarp) ? red[lane] : 0.0;
        for (int off = 16; off > 0; off >>= 1)
            ss += __shfl_down_sync(0xffffffffu, ss, off);
        if (lane == 0) out[b] = (float)(scale * (mx + log(ss)));
    }
}

// ---------------------------------------------------------------------------
// Host launcher
// ---------------------------------------------------------------------------
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor cbias, torch::Tensor bias,
                            double inv_hw_out, 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);

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto xsum   = torch::empty({B, Cin}, opts_d);

    const int threads1 = 256;
    spatial_sum_kernel<<<(int)(B * Cin), threads1>>>(
        x.data_ptr<float>(), xsum.data_ptr<double>(), hw);

    auto out = torch::empty({B, 1}, x.options());
    const float* cb_ptr =
        (cbias.defined() && cbias.numel() > 0) ? cbias.data_ptr<float>() : nullptr;

    const int threads2 = 128;
    const size_t smem  = (size_t)Cout * sizeof(double);
    fused_bias_lse_kernel<<<(int)B, threads2, smem>>>(
        xsum.data_ptr<double>(), weight.data_ptr<float>(), cb_ptr,
        bias.data_ptr<float>(), out.data_ptr<float>(),
        (int)Cin, (int)Cout, (int)khw, inv_hw_out, scale);

    return out;
}
"""

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

_fused_op = load_inline(
    name="fused_convt_gap_bias_lse",
    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
    fused custom CUDA kernels (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._op = _fused_op

    def forward(self, x):
        ct = self.conv_transpose
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 4
            and ct.groups == 1
            and tuple(ct.padding) == (0, 0)
            and ct.weight.dtype == torch.float32
            and self.bias.numel() == ct.out_channels
        ):
            x = x.contiguous()
            H, W = x.shape[2], x.shape[3]
            kH, kW = ct.kernel_size
            sH, sW = ct.stride
            dH, dW = ct.dilation
            opH, opW = ct.output_padding
            # Output spatial size of the transposed convolution.
            Ho = (H - 1) * sH + dH * (kH - 1) + opH + 1
            Wo = (W - 1) * sW + dW * (kW - 1) + opW + 1
            inv_hw = 1.0 / float(Ho * Wo)

            w = ct.weight.contiguous()
            if ct.bias is not None:
                cb = ct.bias.contiguous()
            else:
                cb = torch.empty(0, device=x.device, dtype=torch.float32)

            return self._op.fused_forward(
                x, w, 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
```