```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) per-(in,out) sums of the conv_transpose kernel weights (tiny)
#   3) a small fused dot-product + biases + log-sum-exp + scale (tiny)
# which avoids materializing the huge transposed-conv output. Steps 1 and 2
# run together in a single "prep" kernel launch (the tiny weight-sum blocks
# ride along with the big reduction grid), and step 3 is a second small kernel.

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <algorithm>

#define T1 256

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

__device__ __forceinline__ void block_store_partial(double v,
                                                    double* __restrict__ partials,
                                                    long idx) {
    __shared__ double sh[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    v = warp_sum(v);
    if (lane == 0) sh[warp] = v;
    __syncthreads();
    if (warp == 0) {
        double r = (lane < (int)((blockDim.x + 31) >> 5)) ? sh[lane] : 0.0;
        r = warp_sum(r);
        if (lane == 0) partials[idx] = r;
    }
}

// Grid layout: blockIdx.x in [0, bc_count + ksum_blocks), blockIdx.y = chunk.
// Blocks with blockIdx.x >= bc_count (and blockIdx.y == 0) compute the
// per-(in,out) conv_transpose weight sums; all other blocks compute chunked
// per-(batch,channel) spatial sums of the input.
template <bool VEC>
__global__ void prep_kernel(const float* __restrict__ x,
                            double* __restrict__ partials,  // (bc_count, chunks)
                            float*  __restrict__ ksum,      // (Cin*Cout)
                            const float* __restrict__ w,
                            long hw, int chunks,
                            long ksum_total, int khw, long bc_count) {
    if ((long)blockIdx.x >= bc_count) {
        if (blockIdx.y != 0) return;
        const long gstride = (long)(gridDim.x - (unsigned)bc_count) * blockDim.x;
        for (long t = (long)(blockIdx.x - (unsigned)bc_count) * blockDim.x + threadIdx.x;
             t < ksum_total; t += gstride) {
            const float* p = w + t * (long)khw;
            float s = 0.0f;
            int k = 0;
            for (; k + 4 <= khw; k += 4)
                s += (p[k] + p[k + 1]) + (p[k + 2] + p[k + 3]);
            for (; k < khw; ++k) s += p[k];
            ksum[t] = s;
        }
        return;
    }

    const float* row = x + (long)blockIdx.x * hw;
    const int stride = blockDim.x;
    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 * blockIdx.y;
        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;
        }
        block_store_partial((double)((acc.x + acc.y) + (acc.z + acc.w)),
                            partials, (long)blockIdx.x * chunks + blockIdx.y);
    } else {
        const long n = hw;
        const long per = (n + chunks - 1) / chunks;
        const long start = per * blockIdx.y;
        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];
        block_store_partial((double)acc, partials,
                            (long)blockIdx.x * chunks + blockIdx.y);
    }
}

// Fully fused: partial-sum reduction + kernel-sum dot product + conv bias +
// model bias + numerically stable log-sum-exp over channels + final scaling.
// One thread block per batch element.
__global__ void fused_bias_lse_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 Cin, int Cout, int chunks,
                                      double inv_hw, double scale) {
    extern __shared__ double smem[];
    double* xs = smem;        // Cin doubles
    double* v  = smem + Cin;  // Cout doubles

    const int b    = blockIdx.x;
    const int t    = threadIdx.x;
    const int nthr = blockDim.x;

    const double* pb = partials + (long)b * Cin * chunks;
    for (int i = t; i < Cin; i += nthr) {
        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();

    for (int o = t; o < Cout; o += nthr) {
        double acc = 0.0;
        for (int i = 0; i < Cin; ++i)
            acc += xs[i] * (double)ksum[(long)i * Cout + o];
        double val = acc * inv_hw + (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 (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)));
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor cbias, torch::Tensor bias,
                            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();
    auto optsD  = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());

    const long ksum_total  = Cin * Cout;
    const long ksum_blocks =
        std::max<long>(1L, std::min<long>((ksum_total + T1 - 1) / T1, 1024L));

    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}, optsD);
    auto ksum     = torch::empty({ksum_total}, x.options());

    dim3 grid((unsigned)(bc + ksum_blocks), (unsigned)chunks, 1);
    if (vec) {
        prep_kernel<true><<<grid, T1, 0, stream>>>(
            x.data_ptr<float>(), partials.data_ptr<double>(), ksum.data_ptr<float>(),
            weight.data_ptr<float>(), hw, chunks, ksum_total, (int)khw, bc);
    } else {
        prep_kernel<false><<<grid, T1, 0, stream>>>(
            x.data_ptr<float>(), partials.data_ptr<double>(), ksum.data_ptr<float>(),
            weight.data_ptr<float>(), hw, chunks, ksum_total, (int)khw, bc);
    }

    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)(Cin + Cout) * sizeof(double);
    fused_bias_lse_kernel<<<(unsigned)B, threads2, smem, stream>>>(
        partials.data_ptr<double>(), ksum.data_ptr<float>(), cb_ptr,
        bias.data_ptr<float>(), out.data_ptr<float>(),
        (int)Cin, (int)Cout, chunks, inv_hw, scale);
    return out;
}
"""

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

_fused_op = load_inline(
    name="fused_convt_gap_bias_lse_v2",
    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
            and (ct.in_channels + ct.out_channels) * 8 <= 45 * 1024
        ):
            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
```