```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 algorithmic insight (exact math, not an approximation):
#
# For ConvTranspose2d with stride=1, padding=0, dilation=1, groups=1, every
# input element is scattered into the output exactly once per kernel tap.
# Therefore the spatial sum of the (never-materialized) conv output is:
#
#   sum_{h,w} out[b,o,h,w] = sum_c S[b,c] * Wsum[c,o] + conv_bias[o]*(OH*OW)
#
# where S[b,c] = sum_{h,w} x[b,c,h,w] and Wsum[c,o] = sum_{i,j} w[c,o,i,j].
# The subsequent global average pooling just divides by OH*OW, so the whole
# 16x64x512x512 -> 16x128x514x514 transposed convolution (~0.6 TFLOP, >2GB of
# traffic) collapses into:
#   1) a bandwidth-bound spatial reduction of x (1 GB read), and
#   2) a tiny fused kernel: dot(S, Wsum)/N + bias + log-sum-exp + *10.
# ---------------------------------------------------------------------------

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

#define SSUM_BLOCK 256

// One block per (batch, channel) plane: computes S[b,c] = sum_{h,w} x[b,c,h,w]
// using vectorized float4 loads and a warp/block tree reduction.
__global__ void spatial_sum_kernel(const float4* __restrict__ x,
                                   float* __restrict__ out,
                                   int n4) {
    const float4* base = x + (size_t)blockIdx.x * (size_t)n4;
    float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f;
    const int stride = SSUM_BLOCK;
    int i = threadIdx.x;
    for (; i + 3 * stride < n4; i += 4 * stride) {
        float4 v0 = base[i];
        float4 v1 = base[i + stride];
        float4 v2 = base[i + 2 * stride];
        float4 v3 = base[i + 3 * stride];
        a0 += v0.x + v0.y + v0.z + v0.w;
        a1 += v1.x + v1.y + v1.z + v1.w;
        a2 += v2.x + v2.y + v2.z + v2.w;
        a3 += v3.x + v3.y + v3.z + v3.w;
    }
    for (; i < n4; i += stride) {
        float4 v = base[i];
        a0 += v.x + v.y + v.z + v.w;
    }
    float sum = (a0 + a1) + (a2 + a3);
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, off);
    __shared__ float sh[SSUM_BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) sh[wid] = sum;
    __syncthreads();
    if (wid == 0) {
        sum = (lane < SSUM_BLOCK / 32) ? sh[lane] : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, off);
        if (lane == 0) out[blockIdx.x] = sum;
    }
}

// One block per batch: val[o] = (sum_c S[b,c]*Wsum[c,o]) * inv_n + bias[o]
// then a numerically-stable log-sum-exp over channels, scaled by `scale`.
__global__ void fused_bias_lse_kernel(const float* __restrict__ S,
                                      const float* __restrict__ weight,
                                      const float* __restrict__ bias1,
                                      const float* __restrict__ bias2,
                                      float* __restrict__ out,
                                      int C_in, int C_out, int kk,
                                      float inv_n, float scale, int use_b2) {
    extern __shared__ float shmem[];
    float* shS = shmem;          // C_in floats
    float* shV = shmem + C_in;   // C_out floats
    const int b = blockIdx.x;
    for (int i = threadIdx.x; i < C_in; i += blockDim.x)
        shS[i] = S[(size_t)b * C_in + i];
    __syncthreads();

    for (int o = threadIdx.x; o < C_out; o += blockDim.x) {
        float dot = 0.f;
        for (int c = 0; c < C_in; ++c) {
            const float* wp = weight + ((size_t)c * C_out + o) * kk;
            float ws = 0.f;
            #pragma unroll 4
            for (int k = 0; k < kk; ++k) ws += wp[k];
            dot = fmaf(shS[c], ws, dot);
        }
        const float bb = bias1[o] + (use_b2 ? bias2[o] : 0.f);
        shV[o] = fmaf(dot, inv_n, bb);
    }
    __syncthreads();

    const float neg_inf = __int_as_float(0xff800000);
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    const int nwarp = (blockDim.x + 31) >> 5;
    __shared__ float red[32];

    // max over channels
    float m = neg_inf;
    for (int o = threadIdx.x; o < C_out; o += blockDim.x) m = fmaxf(m, shV[o]);
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, off));
    if (lane == 0) red[wid] = m;
    __syncthreads();
    if (wid == 0) {
        m = (lane < nwarp) ? red[lane] : neg_inf;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, off));
        if (lane == 0) red[0] = m;
    }
    __syncthreads();
    m = red[0];

    // sum of exp(val - max)
    float s = 0.f;
    for (int o = threadIdx.x; o < C_out; o += blockDim.x) s += expf(shV[o] - m);
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) s += __shfl_xor_sync(0xffffffffu, s, off);
    if (lane == 0) red[wid] = s;
    __syncthreads();
    if (wid == 0) {
        s = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) s += __shfl_xor_sync(0xffffffffu, s, off);
        if (lane == 0) out[b] = scale * (logf(s) + m);
    }
}

torch::Tensor fused_gap_lse(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor bias1, torch::Tensor bias2,
                            double inv_n, double scale, int64_t use_b2) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");

    auto x_c = x.contiguous();
    auto w_c = weight.contiguous();
    auto b1_c = bias1.contiguous();
    auto b2_c = bias2.contiguous();

    const int B = (int)x_c.size(0);
    const int C_in = (int)x_c.size(1);
    const int64_t hw = x_c.size(2) * x_c.size(3);
    TORCH_CHECK((hw & 3) == 0, "H*W must be divisible by 4");
    const int C_out = (int)w_c.size(1);
    const int kk = (int)(w_c.size(2) * w_c.size(3));
    TORCH_CHECK(b1_c.numel() >= C_out, "bias1 has wrong size");

    auto S = torch::empty({B, C_in}, x_c.options());
    auto out = torch::empty({B, 1}, x_c.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    spatial_sum_kernel<<<B * C_in, SSUM_BLOCK, 0, stream>>>(
        reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
        S.data_ptr<float>(), (int)(hw >> 2));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    int threads2 = 128;
    while (threads2 < C_out && threads2 < 1024) threads2 <<= 1;
    const size_t shbytes = (size_t)(C_in + C_out) * sizeof(float);
    fused_bias_lse_kernel<<<B, threads2, shbytes, stream>>>(
        S.data_ptr<float>(), w_c.data_ptr<float>(),
        b1_c.data_ptr<float>(), b2_c.data_ptr<float>(),
        out.data_ptr<float>(), C_in, C_out, kk,
        (float)inv_n, (float)scale, (int)use_b2);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

cpp_source = """
torch::Tensor fused_gap_lse(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor bias1, torch::Tensor bias2,
                            double inv_n, double scale, int64_t use_b2);
"""

_ext = load_inline(
    name="fused_gap_lse_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_gap_lse"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Drop-in replacement for Model. Uses the closed-form spatial-sum identity of
    stride-1 transposed convolution so the huge conv output is never computed:
    a custom CUDA reduction kernel computes per-channel sums of the input, and a
    second fused CUDA kernel performs the channel dot-product, bias adds,
    log-sum-exp, and the final scaling. Falls back to the reference PyTorch ops
    for any configuration where the identity does not hold.
    """
    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 x.dim() == 4
                and ct.groups == 1 and ct.stride == (1, 1) and ct.padding == (0, 0)
                and ct.dilation == (1, 1) and ct.output_padding == (0, 0)
                and ((x.shape[2] * x.shape[3]) & 3) == 0):
            kh, kw = ct.kernel_size
            out_h = x.shape[2] + kh - 1  # stride=1, pad=0, dil=1
            out_w = x.shape[3] + kw - 1
            inv_n = 1.0 / float(out_h * out_w)
            bias1 = self.bias.reshape(-1)
            if ct.bias is not None:
                bias2 = ct.bias
                use_b2 = 1
            else:
                bias2 = bias1
                use_b2 = 0
            return _ext.fused_gap_lse(x, ct.weight, bias1, bias2, inv_n, 10.0, use_b2)

        # Generic fallback: identical ops to the reference implementation.
        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
```