```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 reduction kernel now uses a flat grid over the whole input
#     tensor whose size is an exact multiple of the GPU's resident-block
#     capacity (queried via the occupancy API). This eliminates the wave-
#     quantization tail on any architecture (A100/H100/etc.). Blocks that
#     straddle a channel boundary handle each channel segment separately.
#   * The weight-tap-sum work is fused into the same kernel launch (a few
#     extra blocks at the front of the grid), removing one kernel launch.
#     The per-channel accumulator is zeroed with a cheap async memset.
#   * The final matmul+bias+logsumexp kernel now uses warp-shuffle
#     reductions instead of shared-memory trees (fewer __syncthreads).

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 kernel:
#  * The first `wblocks` blocks compute the per-(in,out)-channel sums of the
#    K convolution taps (tiny; runs concurrently with the streaming blocks).
#  * The remaining blocks stream the whole input tensor as one flat range,
#    split into equal chunks. Each block accumulates float4-vectorized partial
#    sums per channel segment and atomically adds them into Sx[channel].
// Requires: HW % 4 == 0, chunk % 4 == 0, x 16-byte aligned, Sx pre-zeroed.
__global__ void __launch_bounds__(512)
fused_stream_wsum_kernel(const float* __restrict__ x,
                         const float* __restrict__ w,
                         float* __restrict__ Sx,
                         float* __restrict__ Sw,
                         long HW, long total, long chunk,
                         int wtotal, int K, int wblocks) {
    if ((int)blockIdx.x < wblocks) {
        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < wtotal) {
            const float* p = w + (long)idx * K;
            float s = 0.0f;
            for (int i = 0; i < K; ++i) s += p[i];
            Sw[idx] = s;
        }
        return;
    }
    long start = ((long)blockIdx.x - wblocks) * chunk;
    long end = start + chunk;
    if (end > total) end = total;
    const long stride = blockDim.x;
    while (start < end) {
        long channel = start / HW;
        long seg_end = (channel + 1) * HW;
        if (seg_end > end) seg_end = end;
        const float4* p4 = reinterpret_cast<const float4*>(x + start);
        long n4 = (seg_end - start) >> 2;
        float4 acc = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
        long i = threadIdx.x;
        long limit = n4 - 3 * stride;
        for (; i < limit; 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 < n4; 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);
        s = block_reduce(s);
        if (threadIdx.x == 0) atomicAdd(Sx + channel, s);
        start = seg_end;
    }
}

// Fused init for the scalar fallback path: zero Sx AND compute tap sums.
__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;
    }
}

// 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. Uses warp-shuffle
// reductions (block size is always a multiple of 32).
__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 = -3.402823466e38f;
    if (valid) {
        float acc = 0.0f;
        const float* sx = Sx + (long)n * C_in;
        const float* sw = Sw + c;
        #pragma unroll 4
        for (int k = 0; k < C_in; ++k) {
            acc += sx[k] * sw[(long)k * C_out];
        }
        y = acc * inv_hw + conv_bias[c] + bias[c];
    }

    int lane = c & 31;
    int wid = c >> 5;
    int nw = blockDim.x >> 5;
    __shared__ float sh[32];

    // Block reduction: max over channels.
    float m = y;
    for (int off = 16; off > 0; off >>= 1)
        m = fmaxf(m, __shfl_xor_sync(FULL_MASK, m, off));
    if (lane == 0) sh[wid] = m;
    __syncthreads();
    if (wid == 0) {
        m = (lane < nw) ? sh[lane] : -3.402823466e38f;
        for (int off = 16; off > 0; off >>= 1)
            m = fmaxf(m, __shfl_xor_sync(FULL_MASK, m, off));
        if (lane == 0) sh[0] = m;
    }
    __syncthreads();
    m = sh[0];

    // Block reduction: sum of exp(y - max).
    float e = valid ? expf(y - m) : 0.0f;
    for (int off = 16; off > 0; off >>= 1)
        e += __shfl_xor_sync(FULL_MASK, e, off);
    __syncthreads();  // ensure everyone has read sh[0] before reuse
    if (lane == 0) sh[wid] = e;
    __syncthreads();
    if (wid == 0) {
        float s = (lane < nw) ? sh[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            s += __shfl_xor_sync(FULL_MASK, s, off);
        if (lane == 0) out[n] = (m + logf(s)) * 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");
    TORCH_CHECK(x.dim() == 4, "fused_forward expects a 4D 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();
    const float* xptr = x.data_ptr<float>();
    const float* wptr = weight.data_ptr<float>();

    long HW = (long)H * (long)W;
    long NC = (long)N * (long)C_in;
    long total = NC * HW;
    int wtotal = (int)(C_in * C_out);
    int K = (int)(kernel_h * kernel_w);

    bool vec = ((HW & 3) == 0) &&
               ((reinterpret_cast<uintptr_t>(xptr) & 15) == 0);

    if (vec) {
        // Number of blocks that can be resident simultaneously (cached).
        int dev = 0;
        cudaGetDevice(&dev);
        static int wave_cache[64] = {0};
        int wave = (dev >= 0 && dev < 64) ? wave_cache[dev] : 0;
        if (wave == 0) {
            int sms = 0;
            cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
            int nb = 0;
            cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                &nb, fused_stream_wsum_kernel, 512, 0);
            if (nb <= 0) nb = 4;
            if (sms <= 0) sms = 108;
            wave = sms * nb;
            if (dev >= 0 && dev < 64) wave_cache[dev] = wave;
        }
        // Pick a chunk size so the grid is (nearly) an exact multiple of the
        // resident-block capacity: no wave-quantization tail on any GPU.
        const long TARGET = 32768;  // ~128 KB streamed per block
        long denom = (long)wave * TARGET;
        long k = (total + denom - 1) / denom;
        if (k < 1) k = 1;
        long grid_target = (long)wave * k;
        long chunk = (total + grid_target - 1) / grid_target;
        chunk = ((chunk + 3) / 4) * 4;  // keep float4 alignment
        if (chunk > HW) chunk = HW;     // HW % 4 == 0 in this path
        if (chunk < 4) chunk = 4;
        long main_blocks = (total + chunk - 1) / chunk;
        int wblocks = (wtotal + 511) / 512;

        cudaMemsetAsync(Sx, 0, NC * sizeof(float), stream);
        fused_stream_wsum_kernel<<<(int)(main_blocks + wblocks), 512, 0, stream>>>(
            xptr, wptr, Sx, Sw, HW, total, chunk, wtotal, K, wblocks);
    } else {
        int init_threads = 256;
        int NCi = (int)NC;
        int init_max = wtotal > NCi ? wtotal : NCi;
        init_ws_kernel<<<(init_max + init_threads - 1) / init_threads, init_threads, 0, stream>>>(
            wptr, Sw, Sx, wtotal, NCi, K);

        const long CHUNK = 32768;
        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;
        chunks = (HW + chunk_floats - 1) / chunk_floats;
        channel_sum_scalar_kernel<<<(int)(chunks * NC), 512, 0, stream>>>(
            xptr, Sx, HW, chunk_floats, (int)chunks);
    }

    int block = 32;
    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_v4",
    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 two
    fused custom CUDA kernel launches, 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 flat
    grid sized to an exact multiple of the GPU's resident-block capacity (found
    via the occupancy API), eliminating wave-quantization tails on any
    architecture; the tiny weight-reduction is fused into the same launch. The
    final matmul + bias + logsumexp + scale kernel uses warp-shuffle
    reductions. 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)
        self._fast = out_channels <= 1024
        self._fn = fused_module.fused_forward

    def forward(self, x):
        if (x.is_cuda and x.dtype == torch.float32 and self._fast
                and x.dim() == 4):
            ct = self.conv_transpose
            return self._fn(
                x.contiguous(),
                ct.weight.contiguous(),
                ct.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
```