```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 insight (algorithmic change): for a ConvTranspose2d with stride=1 and
# padding=0, summing over all spatial output positions gives:
#   sum_{i,j} out[b,o,i,j] = sum_c (sum_{ki,kj} W[c,o,ki,kj]) * (sum_{i,j} in[b,c,i,j])
#                            + conv_bias[o] * H_out * W_out
# because every spatial shift of the kernel covers the full input exactly once.
# Therefore the entire network (transposed conv -> global avg pool -> bias add ->
# logsumexp -> sum -> scale) can be computed without ever materializing the huge
# (B, 128, 514, 514) conv output. We only need:
#   1) per-(b,c) spatial sums of the input   (memory-bound, reads input once)
#   2) per-(c,o) sums of the 3x3 kernel      (tiny)
#   3) a tiny combine + logsumexp kernel     (tiny)

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

#define BLOCK 256

// Per-(batch, channel) spatial sum of the input feature map.
// One block per (b, c) pair. Vectorized float4 loads with per-lane
// accumulation for both bandwidth and accuracy.
__global__ void spatial_sum_kernel(const float* __restrict__ x,
                                   double* __restrict__ out,
                                   int HW) {
    const long bc = blockIdx.x;
    const float* ptr = x + bc * (long)HW;

    float4 acc4 = make_float4(0.f, 0.f, 0.f, 0.f);
    float acc = 0.f;
    if ((HW & 3) == 0) {
        const float4* p4 = reinterpret_cast<const float4*>(ptr);
        const int n4 = HW >> 2;
        for (int i = threadIdx.x; i < n4; i += blockDim.x) {
            float4 v = p4[i];
            acc4.x += v.x; acc4.y += v.y; acc4.z += v.z; acc4.w += v.w;
        }
    } else {
        for (int i = threadIdx.x; i < HW; i += blockDim.x) {
            acc += ptr[i];
        }
    }
    acc += (acc4.x + acc4.y) + (acc4.z + acc4.w);

    __shared__ float sh[BLOCK];
    sh[threadIdx.x] = acc;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) sh[threadIdx.x] += sh[threadIdx.x + s];
        __syncthreads();
    }
    if (threadIdx.x == 0) out[bc] = (double)sh[0];
}

// Per-(in_channel, out_channel) sum of the KxK transposed-conv kernel.
// ConvTranspose2d weight layout: [C_in, C_out, K, K].
__global__ void weight_sum_kernel(const float* __restrict__ w,
                                  double* __restrict__ wsum,
                                  int C_in, int C_out, int K) {
    const int KK = K * K;
    const int total = C_in * C_out;
    for (int idx = threadIdx.x; idx < total; idx += blockDim.x) {
        const int c = idx / C_out;
        const int o = idx - c * C_out;
        const float* wp = w + (long)idx * KK;
        double a = 0.0;
        for (int t = 0; t < KK; ++t) a += (double)wp[t];
        wsum[(long)o * C_in + c] = a;  // store as [C_out, C_in]
    }
}

// Combine: gap[b,o] = (sum_c wsum[o,c]*insum[b,c]) / HW_out + conv_bias[o]
//          x[b,o]   = gap[b,o] + bias[o]
//          out[b]   = 10 * logsumexp_o(x[b,o])
// One block per batch element.
__global__ void combine_lse_kernel(const double* __restrict__ insum,    // [B, C_in]
                                   const double* __restrict__ wsum,     // [C_out, C_in]
                                   const float* __restrict__ conv_bias, // [C_out]
                                   const float* __restrict__ bias,      // [C_out]
                                   float* __restrict__ out,             // [B]
                                   int C_in, int C_out, double inv_HW) {
    const int b = blockIdx.x;
    extern __shared__ double vals[];  // C_out doubles
    const double* ip = insum + (long)b * C_in;

    for (int o = threadIdx.x; o < C_out; o += blockDim.x) {
        const double* wp = wsum + (long)o * C_in;
        double a = 0.0;
        for (int c = 0; c < C_in; ++c) a += wp[c] * ip[c];
        vals[o] = a * inv_HW + (double)conv_bias[o] + (double)bias[o];
    }
    __syncthreads();

    __shared__ double sh[BLOCK];

    // max for numerically stable logsumexp
    double local_max = -1e300;
    for (int o = threadIdx.x; o < C_out; o += blockDim.x) {
        if (vals[o] > local_max) local_max = vals[o];
    }
    sh[threadIdx.x] = local_max;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) sh[threadIdx.x] = fmax(sh[threadIdx.x], sh[threadIdx.x + s]);
        __syncthreads();
    }
    const double maxv = sh[0];
    __syncthreads();

    double local_sum = 0.0;
    for (int o = threadIdx.x; o < C_out; o += blockDim.x) {
        local_sum += exp(vals[o] - maxv);
    }
    sh[threadIdx.x] = local_sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) sh[threadIdx.x] += sh[threadIdx.x + s];
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        out[b] = (float)(10.0 * (log(sh[0]) + maxv));
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor bias) {
    const int B = x.size(0);
    const int C_in = x.size(1);
    const int H = x.size(2);
    const int W = x.size(3);
    const int C_out = weight.size(1);  // ConvTranspose2d: [C_in, C_out, K, K]
    const int K = weight.size(2);
    const int HW = H * W;
    // stride=1, padding=0, dilation=1, output_padding=0
    const int H_out = H + K - 1;
    const int W_out = W + K - 1;
    const double inv_HW = 1.0 / ((double)H_out * (double)W_out);

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto insum = torch::empty({B, C_in}, opts_d);
    auto wsum = torch::empty({C_out, C_in}, opts_d);
    auto out = torch::empty({B, 1}, x.options());

    auto stream = at::cuda::getCurrentCUDAStream().stream();

    spatial_sum_kernel<<<B * C_in, BLOCK, 0, stream>>>(
        x.data_ptr<float>(), insum.data_ptr<double>(), HW);

    weight_sum_kernel<<<1, BLOCK, 0, stream>>>(
        weight.data_ptr<float>(), wsum.data_ptr<double>(), C_in, C_out, K);

    const size_t shmem = (size_t)C_out * sizeof(double);
    combine_lse_kernel<<<B, BLOCK, shmem, stream>>>(
        insum.data_ptr<double>(), wsum.data_ptr<double>(),
        conv_bias.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), C_in, C_out, inv_HW);

    return out;
}
"""

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

fused_ops = load_inline(
    name="fused_convt_gap_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.

    Uses the linearity of the (stride-1, padding-0) transposed convolution and
    the global average pooling to collapse the whole forward pass into:
      1. a memory-bound spatial reduction of the input (custom CUDA kernel),
      2. a tiny kernel summing each 3x3 conv-transpose kernel (custom CUDA),
      3. a tiny fused kernel computing gap + bias + logsumexp + sum + scale.
    The massive (B, out_channels, H+K-1, W+K-1) conv output is never
    materialized, eliminating both the transposed-conv FLOPs and several GB of
    memory traffic. Parameter names are kept identical to the original Model so
    state dicts are fully compatible.
    """
    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.fused_ops = fused_ops

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32 and self.conv_transpose.bias is not None:
            return self.fused_ops.fused_forward(
                x.contiguous(),
                self.conv_transpose.weight.contiguous(),
                self.conv_transpose.bias.contiguous(),
                self.bias.contiguous(),
            )
        # Fallback (CPU / non-float32): original operator sequence.
        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
```