import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, adds a bias term, clamps, scales, clamps, and divides.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape, scaling_factor):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.bias = nn.Parameter(torch.randn(bias_shape)) 
        self.scaling_factor = scaling_factor

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x + self.bias
        x = torch.clamp(x, min=0.0, max=1.0)
        x = x * self.scaling_factor
        x = torch.clamp(x, min=0.0, max=1.0)
        x = x / self.scaling_factor
        return x

batch_size = 128
in_channels  = 64  
out_channels = 64  
height = width = 128 
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
bias_shape = (out_channels, 1, 1)
scaling_factor = 2.0

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

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

import torch
torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src=r'''
#include <torch/extension.h>
torch::Tensor deconv_fused_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor cb, torch::Tensor eb);
torch::Tensor deconv_fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor cb, torch::Tensor eb){return deconv_fused_forward_cuda(x,w,cb,eb);}'''
_cuda_src=r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
constexpr int IC=64,OC=64,IH=128,IW=128,OH=256,OW=256,TM=8,TN=16,TK=16;
__global__ __launch_bounds__(128,4)
void deconv2x2_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ cb,const float* __restrict__ eb,float* __restrict__ y,int B){
    __shared__ float As[4][TM][TK]; __shared__ float Bs[9][TK][TN];
    int tx=threadIdx.x,ty=threadIdx.y,row=blockIdx.x*TM+tx,oc=blockIdx.y*TN+ty; int M=B*IH*IW; int ow=row&(IW-1); int t=row>>7; int oh=t&(IH-1); int n=t>>7;
    float bias=0.f; if(row<M&&oc<OC) bias=cb[oc]+eb[oc]; float a00=bias,a01=bias,a10=bias,a11=bias;
    for(int k0=0;k0<IC;k0+=TK){
#pragma unroll
        for(int p=0;p<4;++p){ int dh=p>>1,dw=p&1; int ih=oh+dh,iw=ow+dw; int kk=k0+ty; float av=0.f; if(row<M&&ih<IH&&iw<IW) av=x[((n*IC+kk)*IH+ih)*IW+iw]; As[p][tx][ty]=av; }
        int tid=ty*TM+tx;
#pragma unroll
        for(int idx=tid; idx<9*TK*TN; idx+=TM*TN){ int tmp=idx; int oc2=tmp%TN; tmp/=TN; int kk2=tmp%TK; int q=tmp/TK; int kh=q/3,kw=q-kh*3; int goc=blockIdx.y*TN+oc2; float bv=0.f; if(goc<OC) bv=w[(((k0+kk2)*OC+goc)*3+kh)*3+kw]; Bs[q][kk2][oc2]=bv; }
        __syncthreads();
        if(row<M&&oc<OC){
#pragma unroll
            for(int kk=0;kk<TK;++kk){ float x00=As[0][tx][kk]; float x01=As[1][tx][kk]; float x10=As[2][tx][kk]; float x11=As[3][tx][kk]; a00=fmaf(x00,Bs[4][kk][ty],a00); a01=fmaf(x00,Bs[5][kk][ty],a01); a01=fmaf(x01,Bs[3][kk][ty],a01); a10=fmaf(x00,Bs[7][kk][ty],a10); a10=fmaf(x10,Bs[1][kk][ty],a10); a11=fmaf(x00,Bs[8][kk][ty],a11); a11=fmaf(x01,Bs[6][kk][ty],a11); a11=fmaf(x10,Bs[2][kk][ty],a11); a11=fmaf(x11,Bs[0][kk][ty],a11); }
        }
        __syncthreads(); }
    if(row<M&&oc<OC){ a00=fminf(fmaxf(a00,0.f),0.5f); a01=fminf(fmaxf(a01,0.f),0.5f); a10=fminf(fmaxf(a10,0.f),0.5f); a11=fminf(fmaxf(a11,0.f),0.5f); long base=((long)n*OC+oc)*OH*OW; long off=base+(oh*2)*OW+(ow*2); y[off]=a00; y[off+1]=a01; y[off+OW]=a10; y[off+OW+1]=a11; }}
torch::Tensor deconv_fused_forward_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor cb,torch::Tensor eb){ CHECK_CUDA(x);CHECK_CUDA(w);CHECK_CUDA(cb);CHECK_CUDA(eb);CHECK_CONTIGUOUS(x);CHECK_CONTIGUOUS(w);CHECK_CONTIGUOUS(cb);CHECK_CONTIGUOUS(eb); int B=x.size(0); auto y=torch::empty({B,OC,OH,OW},x.options()); dim3 block(TM,TN); dim3 grid((B*IH*IW+TM-1)/TM,(OC+TN-1)/TN); deconv2x2_kernel<<<grid,block,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),w.data_ptr<float>(),cb.data_ptr<float>(),eb.data_ptr<float>(),y.data_ptr<float>(),B); return y; }'''
_ext=load_inline(name='deconv_fused_ext_v7c',cpp_sources=[_cpp_src],cuda_sources=[_cuda_src],functions=['deconv_fused_forward'],extra_cuda_cflags=['-O3','--use_fast_math'],verbose=False)
class ModelNew(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride,padding,output_padding,bias_shape,scaling_factor):
        super(ModelNew,self).__init__(); self.conv_transpose=nn.ConvTranspose2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding,output_padding=output_padding); self.bias=nn.Parameter(torch.randn(bias_shape)); self.scaling_factor=scaling_factor
    def forward(self,x):
        return _ext.deconv_fused_forward(x.contiguous(),self.conv_transpose.weight.contiguous(),self.conv_transpose.bias.contiguous(),self.bias.contiguous().view(-1))