```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 in FP32):
# For ConvTranspose2d with stride=1, padding=0, dilation=1, groups=1 (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 in this version:
#   * The streaming reduction writes per-chunk partial sums with plain stores
#     (no atomicAdd, no accumulator buffer to zero), eliminating the memset op
#     entirely. Only two kernel launches per forward remain.
#   * The tiny weight-tap sums are fused into the front of the streaming
#     kernel's grid (a few extra blocks), as before.
#   * The epilogue kernel first reduces the per-chunk partials and preloads the
#     weight sums into shared memory with fully coalesced accesses, making the
#     matmul + logsumexp epilogue roughly 2x faster than the previous
#     strided-global version.
#   * All tensor .contiguous()/reshape handling moved from Python into C++ to
#     cut per-call CPU overhead; a single workspace allocation holds the
#     partials, the weight sums and the output.

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
#define NEG_INF (-3.402823466e38f)

// 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;
}

// The first `wblocks` blocks of the grid compute the per-(in,out)-channel sums
// of the K convolution taps (tiny; runs concurrently with the streaming blocks).
__device__ __forceinline__ bool do_wsum(const float* __restrict__ w,
                                        float* __restrict__ Sw,
                                        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 true;
    }
    return false;
}

// Each remaining block sums one contiguous chunk of one (batch, channel) plane
// using vectorized float4 streaming loads with 4x unrolling, then writes the
// chunk partial with a plain store (no atomics, no zero-init needed).
// Requires: HW % 4 == 0, chunk_floats % 4 == 0, x 16-byte aligned.
__global__ void __launch_bounds__(512)
stream_sum_kernel(const float* __restrict__ x,
                  const float* __restrict__ w,
                  float* __restrict__ partials,
                  float* __restrict__ Sw,
                  long HW, long chunk_floats, int chunks,
                  int wtotal, int K, int wblocks) {
    if (do_wsum(w, Sw, wtotal, K, wblocks)) return;
    long bid = (long)blockIdx.x - wblocks;
    long channel = bid / chunks;
    int cid = (int)(bid - channel * (long)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);

    s = block_reduce(s);
    if (threadIdx.x == 0) partials[channel * (long)chunks + cid] = s;
}

// Scalar fallback for the (rare) case where planes are not 16-byte aligned.
__global__ void __launch_bounds__(512)
stream_sum_scalar_kernel(const float* __restrict__ x,
                         const float* __restrict__ w,
                         float* __restrict__ partials,
                         float* __restrict__ Sw,
                         long HW, long chunk_floats, int chunks,
                         int wtotal, int K, int wblocks) {
    if (do_wsum(w, Sw, wtotal, K, wblocks)) return;
    long bid = (long)blockIdx.x - wblocks;
    long channel = bid / chunks;
    int cid = (int)(bid - channel * (long)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) partials[channel * (long)chunks + cid] = s;
}

// One block per batch element: reduce the per-chunk partials into per-channel
// input sums, fused matmul (Sx @ Sw), scale, bias adds, log-sum-exp over
// channels and final multiplication by 10. Uses shared-memory preloads with
// fully coalesced global accesses plus warp-shuffle reductions.
__global__ void matmul_lse_kernel(const float* __restrict__ partials,
                                  const float* __restrict__ Sw,
                                  const float* __restrict__ conv_bias,
                                  const float* __restrict__ bias,
                                  float* __restrict__ out,
                                  int C_in, int C_out, int chunks,
                                  float inv_hw, int use_smem) {
    extern __shared__ float sm[];
    int n = blockIdx.x;
    int tid = threadIdx.x;
    bool valid = tid < C_out;
    float y = NEG_INF;

    if (use_smem) {
        float* sx = sm;              // [C_in]
        float* sws = sm + C_in;      // [C_in * C_out]
        for (int k = tid; k < C_in; k += blockDim.x) {
            const float* p = partials + (long)(n * C_in + k) * chunks;
            float s = 0.0f;
            for (int j = 0; j < chunks; ++j) s += p[j];
            sx[k] = s;
        }
        int tot = C_in * C_out;
        for (int i = tid; i < tot; i += blockDim.x) sws[i] = Sw[i];
        __syncthreads();
        if (valid) {
            float acc = 0.0f;
            #pragma unroll 4
            for (int k = 0; k < C_in; ++k) {
                acc += sx[k] * sws[k * C_out + tid];
            }
            y = acc * inv_hw + conv_bias[tid] + bias[tid];
        }
    } else {
        if (valid) {
            float acc = 0.0f;
            for (int k = 0; k < C_in; ++k) {
                const float* p = partials + (long)(n * C_in + k) * chunks;
                float s = 0.0f;
                for (int j = 0; j < chunks; ++j) s += p[j];
                acc += s * Sw[(long)k * C_out + tid];
            }
            y = acc * inv_hw + conv_bias[tid] + bias[tid];
        }
    }

    // Block-wide log-sum-exp over channels (block size is a multiple of 32).
    int lane = tid & 31;
    int wid = tid >> 5;
    int nw = blockDim.x >> 5;
    __shared__ float red[32];

    float m = y;
    for (int off = 16; off > 0; off >>= 1)
        m = fmaxf(m, __shfl_xor_sync(FULL_MASK, m, off));
    if (lane == 0) red[wid] = m;
    __syncthreads();
    if (wid == 0) {
        m = (lane < nw) ? red[lane] : NEG_INF;
        for (int off = 16; off > 0; off >>= 1)
            m = fmaxf(m, __shfl_xor_sync(FULL_MASK, m, off));
        if (lane == 0) red[0] = m;
    }
    __syncthreads();
    m = red[0];
    __syncthreads();  // ensure everyone read red[0] before reuse

    float e = valid ? expf(y - m) : 0.0f;
    for (int off = 16; off > 0; off >>= 1)
        e += __shfl_xor_sync(FULL_MASK, e, off);
    if (lane == 0) red[wid] = e;
    __syncthreads();
    if (wid == 0) {
        float s = (lane < nw) ? red[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");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto cb = conv_bias.contiguous();
    auto b2 = bias.contiguous();

    int64_t N = xc.size(0);
    int64_t C_in = xc.size(1);
    int64_t H = xc.size(2);
    int64_t W = xc.size(3);
    int64_t C_out = wc.size(1);
    TORCH_CHECK(C_out <= 1024, "C_out must be <= 1024");
    TORCH_CHECK(cb.numel() == C_out && b2.numel() == C_out,
                "bias sizes must match C_out");

    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);

    long HW = (long)H * (long)W;
    long NC = (long)N * (long)C_in;

    // 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;

    int64_t wtot = C_in * C_out;
    // Single workspace: partials [NC*chunks] | Sw [C_in*C_out] | out [N].
    auto buf = torch::empty({NC * chunks + wtot + N}, xc.options());
    float* partials = buf.data_ptr<float>();
    float* Sw = partials + NC * chunks;
    float* outptr = Sw + wtot;
    auto out = buf.narrow(0, NC * chunks + wtot, N).view({N, 1});

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xptr = xc.data_ptr<float>();
    const float* wptr = wc.data_ptr<float>();
    int K = (int)(kernel_h * kernel_w);
    int wblocks = (int)((wtot + 511) / 512);
    long grid = NC * chunks + wblocks;

    bool vec = ((HW & 3) == 0) &&
               ((reinterpret_cast<uintptr_t>(xptr) & 15) == 0);
    if (vec) {
        stream_sum_kernel<<<(int)grid, 512, 0, stream>>>(
            xptr, wptr, partials, Sw, HW, chunk_floats, (int)chunks,
            (int)wtot, K, wblocks);
    } else {
        stream_sum_scalar_kernel<<<(int)grid, 512, 0, stream>>>(
            xptr, wptr, partials, Sw, HW, chunk_floats, (int)chunks,
            (int)wtot, K, wblocks);
    }

    int block = 32;
    while (block < C_out && block < 1024) block <<= 1;
    size_t smem = (size_t)C_in * (size_t)(C_out + 1) * sizeof(float);
    int use_smem = (smem <= 46 * 1024) ? 1 : 0;
    if (!use_smem) smem = 0;
    matmul_lse_kernel<<<(int)N, block, smem, stream>>>(
        partials, Sw, cb.data_ptr<float>(), b2.data_ptr<float>(), outptr,
        (int)C_in, (int)C_out, (int)chunks, inv_hw, use_smem);

    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_v6",
    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
    finely chunked grid, writing per-chunk partial sums with plain stores (no
    atomics, no memset). The epilogue kernel reduces the partials and preloads
    the weight sums into shared memory with fully coalesced accesses, then
    fuses the matmul + biases + logsumexp + scale with 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)

        def _is_one(v):
            return v == 1 or tuple(v) == (1, 1)

        def _is_zero(v):
            return v == 0 or tuple(v) == (0, 0)

        ct = self.conv_transpose
        # The sum factorization is exact only for a "full" transposed conv.
        self._fast = (
            out_channels <= 1024
            and ct.groups == 1
            and _is_one(ct.stride)
            and _is_zero(ct.padding)
            and _is_zero(ct.output_padding)
            and _is_one(ct.dilation)
        )
        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,
                ct.weight,
                ct.bias,
                self.bias,
                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
```