import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, global average pooling, adds a bias, applies log-sum-exp, sum, and multiplication.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.mean(x, dim=(2, 3), keepdim=True)  # Global average pooling
        x = x + self.bias
        x = torch.logsumexp(x, dim=1, keepdim=True)  # Log-sum-exp
        x = torch.sum(x, dim=(2, 3))  # Sum
        x = x * 10.0  # Multiplication
        return x

batch_size = 16
in_channels = 64
out_channels = 128
height = width = 512
kernel_size = 3
bias_shape = (out_channels, 1, 1)

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, bias_shape]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp = r'''
#include <torch/extension.h>

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor extra_bias);
'''

_cuda = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

static __inline__ __device__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

static __inline__ __device__ float warp_reduce_max(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
    return val;
}

__global__ void reduce_input_sums_kernel(const float* __restrict__ x, float* __restrict__ sums,
                                         int B, int C, int HW) {
    int bc = blockIdx.x;
    int tid = threadIdx.x;
    const float* ptr = x + ((long long)bc) * HW;
    float s = 0.0f;
    for (int i = tid; i < HW; i += blockDim.x) {
        s += ptr[i];
    }
    s = warp_reduce_sum(s);
    __shared__ float smem[32];
    if ((tid & 31) == 0) smem[tid >> 5] = s;
    __syncthreads();
    if (tid < 32) {
        float v = (tid < (blockDim.x >> 5)) ? smem[tid] : 0.0f;
        v = warp_reduce_sum(v);
        if (tid == 0) sums[bc] = v;
    }
}

__global__ void final_lse_kernel(const float* __restrict__ sums,
                                 const float* __restrict__ weight,
                                 const float* __restrict__ conv_bias,
                                 const float* __restrict__ extra_bias,
                                 float* __restrict__ out,
                                 int B, int Cin, int Cout, int K, float inv_out_area) {
    int n = blockIdx.x;
    int tid = threadIdx.x;
    float val = -INFINITY;
    if (tid < Cout) {
        float acc = 0.0f;
        for (int ci = 0; ci < Cin; ++ci) {
            const float* wptr = weight + (((ci * Cout + tid) * K) * K);
            float ws = 0.0f;
            #pragma unroll
            for (int kk = 0; kk < 9; ++kk) {
                if (kk < K*K) ws += wptr[kk];
            }
            acc += sums[n * Cin + ci] * ws;
        }
        val = conv_bias[tid] + extra_bias[tid] + acc * inv_out_area;
    }
    float m = warp_reduce_max(val);
    __shared__ float smem_max[32];
    if ((tid & 31) == 0) smem_max[tid >> 5] = m;
    __syncthreads();
    float maxv = -INFINITY;
    if (tid < 32) {
        maxv = (tid < (blockDim.x >> 5)) ? smem_max[tid] : -INFINITY;
        maxv = warp_reduce_max(maxv);
        if (tid == 0) smem_max[0] = maxv;
    }
    __syncthreads();
    maxv = smem_max[0];

    float e = (tid < Cout) ? expf(val - maxv) : 0.0f;
    e = warp_reduce_sum(e);
    __shared__ float smem_sum[32];
    if ((tid & 31) == 0) smem_sum[tid >> 5] = e;
    __syncthreads();
    if (tid < 32) {
        float sv = (tid < (blockDim.x >> 5)) ? smem_sum[tid] : 0.0f;
        sv = warp_reduce_sum(sv);
        if (tid == 0) {
            out[n] = 10.0f * (logf(sv) + maxv);
        }
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor extra_bias) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    x = x.contiguous();
    weight = weight.contiguous();
    conv_bias = conv_bias.contiguous();
    extra_bias = extra_bias.contiguous();
    int B = x.size(0);
    int Cin = x.size(1);
    int H = x.size(2);
    int W = x.size(3);
    int Cout = weight.size(1);
    int K = weight.size(2);
    int HW = H * W;
    int Hout = H + K - 1;
    int Wout = W + K - 1;
    auto opts = x.options();
    auto sums = torch::empty({B, Cin}, opts);
    auto out = torch::empty({B, 1}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    reduce_input_sums_kernel<<<B * Cin, 256, 0, stream>>>(x.data_ptr<float>(), sums.data_ptr<float>(), B, Cin, HW);
    final_lse_kernel<<<B, 128, 0, stream>>>(sums.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), extra_bias.data_ptr<float>(), out.data_ptr<float>(), B, Cin, Cout, K, 1.0f / ((float)Hout * (float)Wout));
    return out;
}
'''

_ext = load_inline(name='fused_convtranspose_gap_lse_ext', cpp_sources=_cpp, cuda_sources=_cuda, functions=['fused_forward'], verbose=False, extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])

class ModelNew(nn.Module):
    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))

    def forward(self, x):
        return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, self.bias)