```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 stride=1, padding=0, dilation=1, groups=1, every
# input element contributes exactly kh*kw times to the output (once per kernel
# offset, with no border clipping). Therefore:
#     sum_{i,j} convT(x)[n,co,i,j]
#         = sum_ci ( sum_{h,w} x[n,ci,h,w] ) * ( sum_{r,s} W[ci,co,r,s] )
#           + H_out*W_out*conv_bias[co]
# So the global average pool after the transposed convolution can be computed
# WITHOUT materializing the huge (B, C_out, H+k-1, W+k-1) output tensor:
#     pooled = (S @ W_sum) / (H_out*W_out) + conv_bias
# where S = per-channel spatial sums of x and W_sum = per-(ci,co) kernel sums.
# The remaining ops (bias add, log-sum-exp over channels, sum over the
# singleton spatial dims, *10) are fused into a single tiny kernel.

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

#define SUM_THREADS 256
#define FIN_THREADS 128

// One block per (batch, in_channel): tree-reduce H*W elements -> S[n, ci]
__global__ void channel_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ sums,
                                   int HW) {
    const long idx = blockIdx.x;
    const float* base = x + idx * (long)HW;
    float acc = 0.0f;
    if ((HW & 3) == 0) {
        const float4* base4 = reinterpret_cast<const float4*>(base);
        const int n4 = HW >> 2;
        for (int i = threadIdx.x; i < n4; i += blockDim.x) {
            float4 v = base4[i];
            acc += v.x + v.y + v.z + v.w;
        }
    } else {
        for (int i = threadIdx.x; i < HW; i += blockDim.x) acc += base[i];
    }
    __shared__ float sh[SUM_THREADS];
    sh[threadIdx.x] = acc;
    __syncthreads();
    for (int s = SUM_THREADS >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) sh[threadIdx.x] += sh[threadIdx.x + s];
        __syncthreads();
    }
    if (threadIdx.x == 0) sums[idx] = sh[0];
}

// W_sum[ci, co] = sum over kh*kw kernel taps
__global__ void weight_sum_kernel(const float* __restrict__ w,
                                  float* __restrict__ wsum,
                                  int total, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        const float* base = w + (long)idx * K;
        float acc = 0.0f;
        for (int i = 0; i < K; ++i) acc += base[i];
        wsum[idx] = acc;
    }
}

// One block per batch element: matmul (1xCIN)x(CINxCOUT), scale, add biases,
// numerically-stable log-sum-exp over channels, multiply by 10.
__global__ void finalize_kernel(const float* __restrict__ S,
                                const float* __restrict__ wsum,
                                const float* __restrict__ conv_bias,
                                const float* __restrict__ bias,
                                float* __restrict__ out,
                                int CIN, int COUT, float inv_hw) {
    extern __shared__ float smem[];
    float* shS = smem;          // CIN floats
    float* shZ = smem + CIN;    // COUT floats
    __shared__ float red[FIN_THREADS];

    const int n = blockIdx.x;
    for (int i = threadIdx.x; i < CIN; i += blockDim.x) shS[i] = S[n * CIN + i];
    __syncthreads();

    for (int co = threadIdx.x; co < COUT; co += blockDim.x) {
        float acc = 0.0f;
        for (int ci = 0; ci < CIN; ++ci) {
            acc += shS[ci] * wsum[ci * COUT + co];
        }
        float z = acc * inv_hw;
        if (conv_bias != nullptr) z += conv_bias[co];
        z += bias[co];
        shZ[co] = z;
    }
    __syncthreads();

    // max reduction
    float lmax = -3.402823466e38f;
    for (int co = threadIdx.x; co < COUT; co += blockDim.x) lmax = fmaxf(lmax, shZ[co]);
    red[threadIdx.x] = lmax;
    __syncthreads();
    for (int s = FIN_THREADS >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + s]);
        __syncthreads();
    }
    const float maxv = red[0];
    __syncthreads();

    // sum of exp(z - max)
    float lsum = 0.0f;
    for (int co = threadIdx.x; co < COUT; co += blockDim.x) lsum += expf(shZ[co] - maxv);
    red[threadIdx.x] = lsum;
    __syncthreads();
    for (int s = FIN_THREADS >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) red[threadIdx.x] += red[threadIdx.x + s];
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        out[n] = 10.0f * (maxv + logf(red[0]));
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32,
                "fused_forward expects CUDA float32 input");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();

    const int B    = xc.size(0);
    const int CIN  = xc.size(1);
    const int H    = xc.size(2);
    const int W    = xc.size(3);
    const int COUT = wc.size(1);
    const int KH   = wc.size(2);
    const int KW   = wc.size(3);
    const int HW   = H * W;

    // stride=1, padding=0, dilation=1 output size
    const long H_out = (long)H + KH - 1;
    const long W_out = (long)W + KW - 1;
    const float inv_hw = 1.0f / (float)(H_out * W_out);

    auto S     = torch::empty({B, CIN}, xc.options());
    auto wsum  = torch::empty({CIN, COUT}, xc.options());
    auto out   = torch::empty({B, 1}, xc.options());

    channel_sum_kernel<<<B * CIN, SUM_THREADS>>>(
        xc.data_ptr<float>(), S.data_ptr<float>(), HW);

    const int total = CIN * COUT;
    weight_sum_kernel<<<(total + 255) / 256, 256>>>(
        wc.data_ptr<float>(), wsum.data_ptr<float>(), total, KH * KW);

    torch::Tensor cbc;
    const float* cb_ptr = nullptr;
    if (conv_bias.defined() && conv_bias.numel() > 0) {
        cbc = conv_bias.contiguous();
        cb_ptr = cbc.data_ptr<float>();
    }

    const size_t shmem = (size_t)(CIN + COUT) * sizeof(float);
    finalize_kernel<<<B, FIN_THREADS, shmem>>>(
        S.data_ptr<float>(), wsum.data_ptr<float>(), cb_ptr,
        bc.data_ptr<float>(), out.data_ptr<float>(), CIN, COUT, inv_hw);

    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor bias);"
)

fused_op = load_inline(
    name="fused_convt_gap_logsumexp",
    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 ConvTranspose2d +
    global-average-pool with an exact algebraic shortcut (per-channel input
    sums times per-channel kernel sums), and fuses bias-add + log-sum-exp +
    sum + scale into a single CUDA kernel. Falls back to the original PyTorch
    ops for unsupported configurations.
    """
    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))

    def forward(self, x):
        ct = self.conv_transpose
        if (x.is_cuda and x.dtype == torch.float32
                and ct.stride == (1, 1) and ct.padding == (0, 0)
                and ct.output_padding == (0, 0) and ct.dilation == (1, 1)
                and ct.groups == 1):
            cb = ct.bias if ct.bias is not None else x.new_empty(0)
            return fused_op.fused_forward(x, ct.weight, cb, self.bias)

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