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_src = r"""
#include <torch/extension.h>

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor user_bias);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void reduce_input_sums_kernel(const float* __restrict__ x,
                                         float* __restrict__ sums,
                                         int B, int C, int H, int W) {
    int bc = blockIdx.x;
    int b = bc / C;
    int c = bc - b * C;
    int tid = threadIdx.x;
    int HW = H * W;
    long long base = ((long long)b * C + c) * (long long)HW;

    float acc = 0.0f;
    for (int i = tid; i < HW; i += blockDim.x) {
        acc += x[base + i];
    }

    __shared__ float sh[256];
    sh[tid] = acc;
    __syncthreads();

    for (int s = blockDim.x >> 1; s > 32; s >>= 1) {
        if (tid < s) sh[tid] += sh[tid + s];
        __syncthreads();
    }
    if (tid < 32) {
        volatile float* vsh = sh;
        vsh[tid] += vsh[tid + 32];
        vsh[tid] += vsh[tid + 16];
        vsh[tid] += vsh[tid + 8];
        vsh[tid] += vsh[tid + 4];
        vsh[tid] += vsh[tid + 2];
        vsh[tid] += vsh[tid + 1];
    }
    if (tid == 0) sums[bc] = sh[0];
}

__device__ __forceinline__ float channel_value(int b, int oc,
                                               const float* __restrict__ sums,
                                               const float* __restrict__ weight,
                                               const float* __restrict__ conv_bias,
                                               const float* __restrict__ user_bias,
                                               int IC, int OC, int KH, int KW,
                                               float inv_out_area) {
    float v = conv_bias[oc] + user_bias[oc];
    for (int ic = 0; ic < IC; ++ic) {
        float wsum = 0.0f;
        long long wbase = (((long long)ic * OC + oc) * KH) * KW;
        #pragma unroll
        for (int kh = 0; kh < 8; ++kh) {
            if (kh < KH) {
                for (int kw = 0; kw < KW; ++kw) {
                    wsum += weight[wbase + kh * KW + kw];
                }
            }
        }
        v += sums[(long long)b * IC + ic] * wsum * inv_out_area;
    }
    return v;
}

__global__ void final_lse_kernel(const float* __restrict__ sums,
                                 const float* __restrict__ weight,
                                 const float* __restrict__ conv_bias,
                                 const float* __restrict__ user_bias,
                                 float* __restrict__ out,
                                 int B, int IC, int OC, int H, int W, int KH, int KW) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    float inv_out_area = 1.0f / ((float)(H + KH - 1) * (float)(W + KW - 1));

    float local_max = -CUDART_INF_F;
    for (int oc = tid; oc < OC; oc += blockDim.x) {
        float v = channel_value(b, oc, sums, weight, conv_bias, user_bias, IC, OC, KH, KW, inv_out_area);
        local_max = fmaxf(local_max, v);
    }

    __shared__ float sh[256];
    sh[tid] = local_max;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) sh[tid] = fmaxf(sh[tid], sh[tid + s]);
        __syncthreads();
    }
    float m = sh[0];

    float local_sum = 0.0f;
    for (int oc = tid; oc < OC; oc += blockDim.x) {
        float v = channel_value(b, oc, sums, weight, conv_bias, user_bias, IC, OC, KH, KW, inv_out_area);
        local_sum += expf(v - m);
    }
    sh[tid] = local_sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) sh[tid] += sh[tid + s];
        __syncthreads();
    }
    if (tid == 0) {
        out[b] = 10.0f * (logf(sh[0]) + m);
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor user_bias) {
    const int B = (int)x.size(0);
    const int IC = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int OC = (int)weight.size(1);
    const int KH = (int)weight.size(2);
    const int KW = (int)weight.size(3);

    auto sums = torch::empty({B, IC}, x.options());
    auto out = torch::empty({B, 1}, x.options());

    dim3 block(256);
    reduce_input_sums_kernel<<<B * IC, block>>>(x.data_ptr<float>(), sums.data_ptr<float>(), B, IC, H, W);
    final_lse_kernel<<<B, block>>>(sums.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), user_bias.data_ptr<float>(), out.data_ptr<float>(), B, IC, OC, H, W, KH, KW);
    return out;
}
"""

_ext = load_inline(
    name="fused_deconv_gap_lse_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

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)