```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 (unchanged, exact in FP32):
# For ConvTranspose2d with stride=1, padding=0 (a "full" convolution), the global
# sum of each output channel factorizes exactly:
#     sum_{h,w} out[n,c,h,w] = sum_k (sum_{h,w} x[n,k,h,w]) * (sum_{p,q} w[k,c,p,q])
#                              + conv_bias[c] * H_out * W_out
# So the whole network collapses to:
#   1) per-channel sum of the input            (memory-bound reduction)
#   2) per-(in,out)-channel sum of the weights (tiny)
#   3) small matmul + bias adds + logsumexp + scale (fully fused)
#
# Optimization in this version: the dominant channel-sum kernel (which must
# stream the entire ~1 GB input) is rewritten with vectorized float4 streaming
# loads (__ldcs), 4x unrolling for memory-level parallelism, 512 threads/block,
# and warp-shuffle reductions, pushing it close to peak DRAM bandwidth.

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

#define FULL_MASK 0xffffffffu

// One block per (batch, in_channel): sum over the H*W spatial elements.
// Vectorized float4 loads with 4x unrolling and streaming cache hint.
__global__ void channel_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ out,
                                   long HW) {
    long idx = blockIdx.x;
    const float* base = x + idx * HW;
    long HW4 = HW >> 2;
    const float4* ptr = reinterpret_cast<const float4*>(base);

    const long stride = blockDim.x;
    const long limit4 = HW4 - 3 * stride;

    float4 acc = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
    long i = threadIdx.x;
    for (; i < limit4; i += 4 * stride) {
        float4 a = __ldcs(ptr + i);
        float4 b = __ldcs(ptr + i + stride);
        float4 c = __ldcs(ptr + i + 2 * stride);
        float4 d = __ldcs(ptr + i + 3 * 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 < HW4; i += stride) {
        float4 a = __ldcs(ptr + i);
        acc.x += a.x; acc.y += a.y; acc.z += a.z; acc.w += a.w;
    }
    float s = (acc.x + acc.y) + (acc.z + acc.w);
    // Scalar tail in case HW is not divisible by 4.
    for (long j = (HW4 << 2) + threadIdx.x; j < HW; j += stride) {
        s += base[j];
    }

    // Warp + cross-warp reduction via shuffle and a small shared array.
    __shared__ float sh[32];
    for (int off = 16; off > 0; off >>= 1)
        s += __shfl_down_sync(FULL_MASK, s, off);
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) sh[wid] = s;
    __syncthreads();
    if (wid == 0) {
        int nw = (blockDim.x + 31) >> 5;
        s = (lane < nw) ? sh[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            s += __shfl_down_sync(FULL_MASK, s, off);
        if (lane == 0) out[idx] = s;
    }
}

// One thread per (in_channel, out_channel) pair: sum the K spatial taps.
__global__ void weight_sum_kernel(const float* __restrict__ w,
                                  float* __restrict__ out,
                                  int total, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        const float* ptr = w + (long)idx * K;
        float s = 0.0f;
        for (int i = 0; i < K; ++i) s += ptr[i];
        out[idx] = s;
    }
}

// One block per batch element: fused matmul (Sx @ Sw), scale, bias adds,
// log-sum-exp over channels and final multiplication by 10.
__global__ void matmul_bias_lse_kernel(const float* __restrict__ Sx,
                                       const float* __restrict__ Sw,
                                       const float* __restrict__ conv_bias,
                                       const float* __restrict__ bias,
                                       float* __restrict__ out,
                                       int C_in, int C_out,
                                       float inv_hw) {
    int n = blockIdx.x;
    int c = threadIdx.x;
    bool valid = c < C_out;

    float y = 0.0f;
    if (valid) {
        float acc = 0.0f;
        for (int k = 0; k < C_in; ++k) {
            acc += Sx[n * C_in + k] * Sw[k * C_out + c];
        }
        y = acc * inv_hw + conv_bias[c] + bias[c];
    }

    __shared__ float sh[1024];
    // Block reduction: max over channels.
    sh[c] = valid ? y : -3.402823466e38f;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (c < off) sh[c] = fmaxf(sh[c], sh[c + off]);
        __syncthreads();
    }
    float m = sh[0];
    __syncthreads();
    // Block reduction: sum of exp(y - max).
    sh[c] = valid ? expf(y - m) : 0.0f;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (c < off) sh[c] += sh[c + off];
        __syncthreads();
    }
    if (c == 0) {
        out[n] = (m + logf(sh[0])) * 10.0f;
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor bias,
                            int64_t kernel_h, int64_t kernel_w) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32,
                "fused_forward expects a CUDA float32 input");
    int64_t N = x.size(0);
    int64_t C_in = x.size(1);
    int64_t H = x.size(2);
    int64_t W = x.size(3);
    int64_t C_out = weight.size(1);
    TORCH_CHECK(C_out <= 1024, "C_out must be <= 1024");

    int64_t H_out = H + kernel_h - 1;
    int64_t W_out = W + kernel_w - 1;
    float inv_hw = 1.0f / (float)(H_out * W_out);

    auto Sx = torch::empty({N, C_in}, x.options());
    auto Sw = torch::empty({C_in, C_out}, x.options());
    auto out = torch::empty({N, 1}, x.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    long HW = (long)H * (long)W;
    channel_sum_kernel<<<(int)(N * C_in), 512, 0, stream>>>(
        x.data_ptr<float>(), Sx.data_ptr<float>(), HW);

    int total = (int)(C_in * C_out);
    int K = (int)(kernel_h * kernel_w);
    weight_sum_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        weight.data_ptr<float>(), Sw.data_ptr<float>(), total, K);

    int block = 1;
    while (block < C_out && block < 1024) block <<= 1;
    matmul_bias_lse_kernel<<<(int)N, block, 0, stream>>>(
        Sx.data_ptr<float>(), Sw.data_ptr<float>(),
        conv_bias.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), (int)C_in, (int)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, int64_t kernel_h, int64_t kernel_w);"
)

fused_module = load_inline(
    name="fused_convtranspose_gap_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 + bias + logsumexp + sum + scale pipeline with three
    fused custom CUDA kernels, exploiting the exact factorization of the global
    sum of a full (stride-1) transposed convolution. The dominant input
    reduction kernel uses vectorized streaming loads for near-peak bandwidth.
    All math is done in FP32.
    """

    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))
        if isinstance(kernel_size, (tuple, list)):
            self.kernel_h = int(kernel_size[0])
            self.kernel_w = int(kernel_size[1])
        else:
            self.kernel_h = int(kernel_size)
            self.kernel_w = int(kernel_size)

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            return fused_module.fused_forward(
                x.contiguous(),
                self.conv_transpose.weight.contiguous(),
                self.conv_transpose.bias.contiguous(),
                self.bias.reshape(-1).contiguous(),
                self.kernel_h,
                self.kernel_w,
            )
        # Fallback to the reference implementation for non-CUDA / non-FP32 inputs.
        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
```