```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, streams ~1 GiB)
#   2) per-(in,out)-channel sum of the weights (tiny)
#   3) small matmul + bias adds + logsumexp + scale (fully fused)
#
# Improvements over the previous version:
#   * The dominant channel-sum kernel now splits each channel into several
#     chunks (8192 blocks instead of 1024). This removes the large wave-
#     quantization tail (1024 blocks ~= 2.4 waves on A100/H100) and pushes the
#     kernel much closer to peak DRAM bandwidth.
#   * Chunks are combined with atomicAdd, and the zero-init of the accumulator
#     buffer is fused into the tiny weight-sum kernel (removes one memset).
#   * Input-sum and weight-sum share one workspace allocation.

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

#define FULL_MASK 0xffffffffu

// Warp + cross-warp block reduction (512 threads = 16 warps). Result valid on thread 0.
__device__ __forceinline__ float block_reduce(float s) {
    __shared__ float sh[16];
    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();
    float r = 0.0f;
    if (wid == 0) {
        r = (lane < 16) ? sh[lane] : 0.0f;
        for (int off = 8; off > 0; off >>= 1)
            r += __shfl_down_sync(FULL_MASK, r, off);
    }
    return r;
}

// Fused init: zero the per-channel input-sum accumulator AND compute the
// per-(in_channel, out_channel) sums of the K convolution taps.
__global__ void init_ws_kernel(const float* __restrict__ w,
                               float* __restrict__ Sw,
                               float* __restrict__ Sx,
                               int total, int NC, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < NC) Sx[idx] = 0.0f;
    if (idx < total) {
        const float* ptr = w + (long)idx * K;
        float s = 0.0f;
        for (int i = 0; i < K; ++i) s += ptr[i];
        Sw[idx] = s;
    }
}

// Each block sums one contiguous chunk of one (batch, channel) plane using
// vectorized float4 streaming loads with 4x unrolling, then atomically
// accumulates the chunk result into the per-channel sum.
__global__ void __launch_bounds__(512)
channel_sum_kernel(const float* __restrict__ x,
                   float* __restrict__ out,
                   long HW, long chunk_floats, int chunks) {
    long channel = blockIdx.x / chunks;
    int cid = blockIdx.x % chunks;
    long start = (long)cid * chunk_floats;
    long len = HW - start;
    if (len > chunk_floats) len = chunk_floats;

    const float* p = x + channel * HW + start;
    long len4 = len >> 2;
    const float4* p4 = reinterpret_cast<const float4*>(p);

    const long stride = blockDim.x;
    const long limit4 = len4 - 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(p4 + i);
        float4 b = __ldcs(p4 + i + stride);
        float4 c = __ldcs(p4 + i + 2 * stride);
        float4 d = __ldcs(p4 + 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 < len4; i += stride) {
        float4 a = __ldcs(p4 + 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);
    for (long j = (len4 << 2) + threadIdx.x; j < len; j += stride) {
        s += p[j];
    }

    s = block_reduce(s);
    if (threadIdx.x == 0) atomicAdd(out + channel, s);
}

// Scalar fallback for the (rare) case where planes are not 16-byte aligned.
__global__ void __launch_bounds__(512)
channel_sum_scalar_kernel(const float* __restrict__ x,
                          float* __restrict__ out,
                          long HW, long chunk_floats, int chunks) {
    long channel = blockIdx.x / chunks;
    int cid = blockIdx.x % chunks;
    long start = (long)cid * chunk_floats;
    long len = HW - start;
    if (len > chunk_floats) len = chunk_floats;

    const float* p = x + channel * HW + start;
    float s = 0.0f;
    for (long j = threadIdx.x; j < len; j += blockDim.x) {
        s += __ldcs(p + j);
    }
    s = block_reduce(s);
    if (threadIdx.x == 0) atomicAdd(out + channel, 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;
        const float* sx = Sx + n * C_in;
        const float* sw = Sw + c;
        #pragma unroll 4
        for (int k = 0; k < C_in; ++k) {
            acc += sx[k] * sw[k * C_out];
        }
        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);

    // Single workspace: Sx (N*C_in) followed by Sw (C_in*C_out).
    auto workspace = torch::empty({N * C_in + C_in * C_out}, x.options());
    float* Sx = workspace.data_ptr<float>();
    float* Sw = Sx + N * C_in;
    auto out = torch::empty({N, 1}, x.options());

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

    int total = (int)(C_in * C_out);
    int NC = (int)(N * C_in);
    int K = (int)(kernel_h * kernel_w);
    int init_threads = 256;
    int init_max = total > NC ? total : NC;
    init_ws_kernel<<<(init_max + init_threads - 1) / init_threads, init_threads, 0, stream>>>(
        weight.data_ptr<float>(), Sw, Sx, total, NC, K);

    long HW = (long)H * (long)W;
    // Split each channel plane into chunks so the grid has many blocks; this
    // removes wave-quantization tails and keeps every SM fed.
    const long CHUNK = 32768;  // floats (128 KB) per block
    long chunks = (HW + CHUNK - 1) / CHUNK;
    if (chunks * NC < 2048) {
        long want = (2048 + NC - 1) / NC;
        long max_chunks = (HW + 2047) / 2048;
        if (want > max_chunks) want = max_chunks;
        if (want > chunks) chunks = want;
        if (chunks < 1) chunks = 1;
    }
    long chunk_floats = (HW + chunks - 1) / chunks;
    chunk_floats = ((chunk_floats + 3) / 4) * 4;  // keep float4 alignment
    chunks = (HW + chunk_floats - 1) / chunk_floats;

    int grid = (int)(chunks * NC);
    const float* xptr = x.data_ptr<float>();
    bool vec = ((HW & 3) == 0) &&
               ((reinterpret_cast<uintptr_t>(xptr) & 15) == 0);
    if (vec) {
        channel_sum_kernel<<<grid, 512, 0, stream>>>(
            xptr, Sx, HW, chunk_floats, (int)chunks);
    } else {
        channel_sum_scalar_kernel<<<grid, 512, 0, stream>>>(
            xptr, Sx, HW, chunk_floats, (int)chunks);
    }

    int block = 1;
    while (block < C_out && block < 1024) block <<= 1;
    matmul_bias_lse_kernel<<<(int)N, block, 0, stream>>>(
        Sx, Sw, 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_v3",
    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 streams the input with vectorized float4 loads over a
    finely chunked grid (atomically combining chunk partial sums) to reach
    near-peak DRAM 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
                and self.conv_transpose.out_channels <= 1024):
            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
```