```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 observation (algorithmic optimization):
#
# For a ConvTranspose2d with padding == 0, every (input element, kernel tap)
# pair contributes exactly once to the output feature map. Therefore the
# spatial sum of the transposed-convolution output is simply:
#
#   sum_{p,q} out[b,o,p,q] = sum_i ( sum_{u,v} W[i,o,u,v] ) * ( sum_{r,s} x[b,i,r,s] )
#
# i.e. a tiny (B,in) @ (in,out) matmul of per-channel sums. Since the next op
# is global average pooling, we never need to materialize the huge
# (B, out, H+k-1, W+k-1) intermediate tensor at all.
#
# Kernel 1 (`channel_sum`) reduces the input to per-channel sums with float64
# accumulation (bandwidth bound: a single pass over the input).
# Kernel 2 (`finalize`) fuses: division by the pooling area, both bias adds,
# the log-sum-exp over channels (online/max-subtracted), the trivial spatial
# sum, and the final multiplication by 10.
# ---------------------------------------------------------------------------

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

#define FULL_MASK 0xffffffffu

__global__ void channel_sum_kernel(const float4* __restrict__ x,
                                   double* __restrict__ out,
                                   const int hw4) {
    const long slice = (long)blockIdx.x;
    const float4* base = x + slice * (long)hw4;

    double acc = 0.0;
    for (int i = threadIdx.x; i < hw4; i += blockDim.x) {
        const float4 v = base[i];
        acc += (double)v.x;
        acc += (double)v.y;
        acc += (double)v.z;
        acc += (double)v.w;
    }

    // block reduction (double precision)
    for (int off = 16; off > 0; off >>= 1)
        acc += __shfl_down_sync(FULL_MASK, acc, off);

    __shared__ double sh[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) sh[warp] = acc;
    __syncthreads();

    if (warp == 0) {
        const int nwarps = (blockDim.x + 31) >> 5;
        double v = (lane < nwarps) ? sh[lane] : 0.0;
        for (int off = 16; off > 0; off >>= 1)
            v += __shfl_down_sync(FULL_MASK, v, off);
        if (lane == 0) out[slice] = v;
    }
}

torch::Tensor channel_sum(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "input must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(x.is_contiguous(), "input must be contiguous");

    const long B = x.size(0);
    const long C = x.size(1);
    const long hw = x.size(2) * x.size(3);
    TORCH_CHECK(hw % 4 == 0, "spatial size must be divisible by 4");

    auto out = torch::empty({B, C},
        torch::TensorOptions().dtype(torch::kFloat64).device(x.device()));

    const int threads = 256;
    const int blocks = (int)(B * C);
    channel_sum_kernel<<<blocks, threads>>>(
        reinterpret_cast<const float4*>(x.data_ptr<float>()),
        out.data_ptr<double>(),
        (int)(hw / 4));
    return out;
}

__global__ void finalize_kernel(const double* __restrict__ prod,
                                const float* __restrict__ cbias,
                                const float* __restrict__ bias,
                                float* __restrict__ out,
                                const int O,
                                const double invN) {
    const int b = blockIdx.x;
    const int o = threadIdx.x;
    const bool valid = (o < O);

    double v = 0.0;
    if (valid)
        v = prod[(long)b * O + o] * invN + (double)cbias[o] + (double)bias[o];

    __shared__ double sh[32];
    const int lane = o & 31;
    const int warp = o >> 5;
    const int nwarps = (blockDim.x + 31) >> 5;

    // max reduction
    double m = valid ? v : -INFINITY;
    for (int off = 16; off > 0; off >>= 1)
        m = fmax(m, __shfl_down_sync(FULL_MASK, m, off));
    if (lane == 0) sh[warp] = m;
    __syncthreads();
    if (warp == 0) {
        double mm = (lane < nwarps) ? sh[lane] : -INFINITY;
        for (int off = 16; off > 0; off >>= 1)
            mm = fmax(mm, __shfl_down_sync(FULL_MASK, mm, off));
        if (lane == 0) sh[0] = mm;
    }
    __syncthreads();
    const double mx = sh[0];
    __syncthreads();

    // sum of exp(v - max)
    double s = valid ? exp(v - mx) : 0.0;
    for (int off = 16; off > 0; off >>= 1)
        s += __shfl_down_sync(FULL_MASK, s, off);
    if (lane == 0) sh[warp] = s;
    __syncthreads();
    if (warp == 0) {
        double ss = (lane < nwarps) ? sh[lane] : 0.0;
        for (int off = 16; off > 0; off >>= 1)
            ss += __shfl_down_sync(FULL_MASK, ss, off);
        if (lane == 0)
            out[b] = (float)(10.0 * (mx + log(ss)));
    }
}

torch::Tensor finalize(torch::Tensor prod,
                       torch::Tensor cbias,
                       torch::Tensor bias,
                       double invN) {
    const long B = prod.size(0);
    const long O = prod.size(1);
    auto out = torch::empty({B, 1},
        torch::TensorOptions().dtype(torch::kFloat32).device(prod.device()));
    int threads = (int)(((O + 31) / 32) * 32);
    if (threads < 32) threads = 32;
    if (threads > 1024) threads = 1024;
    finalize_kernel<<<(int)B, threads>>>(
        prod.data_ptr<double>(),
        cbias.data_ptr<float>(),
        bias.data_ptr<float>(),
        out.data_ptr<float>(),
        (int)O, invN);
    return out;
}
"""

_cpp_source = r"""
torch::Tensor channel_sum(torch::Tensor x);
torch::Tensor finalize(torch::Tensor prod, torch::Tensor cbias, torch::Tensor bias, double invN);
"""

_ext = load_inline(
    name="fused_gap_convtranspose",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["channel_sum", "finalize"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    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.ext = _ext

    def forward(self, x):
        stride = self.conv_transpose.stride
        padding = self.conv_transpose.padding
        dilation = self.conv_transpose.dilation
        groups = self.conv_transpose.groups

        # The algebraic shortcut is exact whenever every kernel-tap
        # contribution lands inside the output feature map (padding == 0)
        # and groups == 1 (true for the given architecture).
        if padding[0] == 0 and padding[1] == 0 and groups == 1 and x.is_cuda:
            B, C, H, W = x.shape
            kh, kw = self.conv_transpose.kernel_size
            h_out = (H - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kh - 1) + 1
            w_out = (W - 1) * stride[1] - 2 * padding[1] + dilation[1] * (kw - 1) + 1
            inv_n = 1.0 / float(h_out * w_out)

            x = x.contiguous()
            xsum = self.ext.channel_sum(x)                                   # (B, in)  f64
            wsum = self.conv_transpose.weight.to(torch.float64).sum(dim=(2, 3))  # (in, out)
            prod = torch.matmul(xsum, wsum)                                  # (B, out) f64
            return self.ext.finalize(prod,
                                     self.conv_transpose.bias,
                                     self.bias.view(-1),
                                     inv_n)

        # Generic fallback (not used for the given architecture).
        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
```