import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, LogSumExp, HardSwish, subtraction, clamp operations.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.bias = nn.Parameter(torch.randn(1, 1, 1, 1)) 

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = x * torch.sigmoid(x + 3) / 6
        x = x - self.bias
        x = torch.clamp(x, min=-1, max=1)
        return x

batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
bias_shape = (1, 1, 1, 1)  

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

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

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

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

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor scalar_bias);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor scalar_bias) {
    return fused_forward_cuda(x, weight, conv_bias, scalar_bias);
}
'''

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

__global__ __launch_bounds__(256, 2) void fused_deconv_lse_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ cb,
    const float* __restrict__ sb,
    float* __restrict__ out,
    int N, int D, int H, int W, int Do, int Ho, int Wo)
{
    int ow = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y * blockDim.y + threadIdx.y; // row = ((n*Do + od)*Ho + oh)
    if (ow >= Wo) return;
    int rows = N * Do * Ho;
    if (row >= rows) return;

    int oh = row % Ho;
    int tmp = row / Ho;
    int od = tmp % Do;
    int n = tmp / Do;

    float a0=cb[0], a1=cb[1], a2=cb[2], a3=cb[3];
    float a4=cb[4], a5=cb[5], a6=cb[6], a7=cb[7];
    float a8=cb[8], a9=cb[9], a10=cb[10], a11=cb[11];
    float a12=cb[12], a13=cb[13], a14=cb[14], a15=cb[15];

    int pd = (od + 1) & 1;
    int ph = (oh + 1) & 1;
    int pw = (ow + 1) & 1;

    #pragma unroll
    for (int ic=0; ic<3; ++ic) {
        for (int kd=pd; kd<3; kd+=2) {
            int id = (od + 1 - kd) >> 1;
            if ((unsigned)id >= (unsigned)D) continue;
            for (int kh=ph; kh<3; kh+=2) {
                int ih = (oh + 1 - kh) >> 1;
                if ((unsigned)ih >= (unsigned)H) continue;
                for (int kw=pw; kw<3; kw+=2) {
                    int iw = (ow + 1 - kw) >> 1;
                    if ((unsigned)iw >= (unsigned)W) continue;
                    float xv = x[((((n*3 + ic)*D + id)*H + ih)*W + iw)];
                    int base = ic*432 + kd*9 + kh*3 + kw;
                    a0=fmaf(xv,w[base+0*27],a0);    a1=fmaf(xv,w[base+1*27],a1);
                    a2=fmaf(xv,w[base+2*27],a2);    a3=fmaf(xv,w[base+3*27],a3);
                    a4=fmaf(xv,w[base+4*27],a4);    a5=fmaf(xv,w[base+5*27],a5);
                    a6=fmaf(xv,w[base+6*27],a6);    a7=fmaf(xv,w[base+7*27],a7);
                    a8=fmaf(xv,w[base+8*27],a8);    a9=fmaf(xv,w[base+9*27],a9);
                    a10=fmaf(xv,w[base+10*27],a10); a11=fmaf(xv,w[base+11*27],a11);
                    a12=fmaf(xv,w[base+12*27],a12); a13=fmaf(xv,w[base+13*27],a13);
                    a14=fmaf(xv,w[base+14*27],a14); a15=fmaf(xv,w[base+15*27],a15);
                }
            }
        }
    }

    float m = fmaxf(fmaxf(fmaxf(a0,a1),fmaxf(a2,a3)), fmaxf(fmaxf(a4,a5),fmaxf(a6,a7)));
    float m2 = fmaxf(fmaxf(fmaxf(a8,a9),fmaxf(a10,a11)), fmaxf(fmaxf(a12,a13),fmaxf(a14,a15)));
    m = fmaxf(m, m2);
    float s = __expf(a0-m)+__expf(a1-m)+__expf(a2-m)+__expf(a3-m)
            + __expf(a4-m)+__expf(a5-m)+__expf(a6-m)+__expf(a7-m)
            + __expf(a8-m)+__expf(a9-m)+__expf(a10-m)+__expf(a11-m)
            + __expf(a12-m)+__expf(a13-m)+__expf(a14-m)+__expf(a15-m);
    float v = m + __logf(s);
    v = v * (1.0f / (1.0f + __expf(-(v + 3.0f)))) * 0.16666666666666666f;
    v -= sb[0];
    v = fminf(1.0f, fmaxf(-1.0f, v));
    out[((long)row * Wo + ow)] = v;
}

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor scalar_bias) {
    int N = x.size(0);
    int D = x.size(2);
    int H = x.size(3);
    int W = x.size(4);
    int Do = 2 * D - 1;
    int Ho = 2 * H - 1;
    int Wo = 2 * W - 1;
    auto out = torch::empty({N, 1, Do, Ho, Wo}, x.options());
    dim3 block(32, 8, 1);
    dim3 grid((Wo + block.x - 1) / block.x, (N * Do * Ho + block.y - 1) / block.y, 1);
    fused_deconv_lse_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), scalar_bias.data_ptr<float>(), out.data_ptr<float>(),
        N, D, H, W, Do, Ho, Wo);
    return out;
}
'''

_ext = load_inline(name='fused_deconv_lse_ext_v2', cpp_sources=[cpp_source], cuda_sources=[cuda_source], functions=['fused_forward'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.bias = nn.Parameter(torch.randn(1, 1, 1, 1))
    def forward(self, x):
        return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, self.bias)