import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 2D convolution operation with square input and asymmetric kernel, with dilation and padding.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Size of the convolution kernel (height, width). 
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (tuple, optional): Padding applied to the input (top/bottom, left/right). Defaults to (0, 0).
        dilation (tuple, optional): Spacing between kernel elements (height, width). Defaults to (1, 1).
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: int = 1, padding: tuple = (0, 0), dilation: tuple = (1, 1), bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 2D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv2d(x)

# Test code
batch_size = 8
in_channels = 32
out_channels = 64
kernel_size = (5, 9)
width = 512
height = 512
stride = 1
padding = (2, 4)
dilation = (2, 3)

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

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

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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor conv2d_forward_cuda(torch::Tensor x, torch::Tensor w, int stride, int pad_h, int pad_w, int dil_h, int dil_w);
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w, int stride, int pad_h, int pad_w, int dil_h, int dil_w) { return conv2d_forward_cuda(x,w,stride,pad_h,pad_w,dil_h,dil_w); }
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define VEC 4
__global__ __launch_bounds__(128,2) void conv2d_5x9_oc2_kernel(const float* __restrict__ x,const float* __restrict__ w,float* __restrict__ y,int N,int H,int W){
    const int tid=threadIdx.x, ow0=tid*VEC, oh=blockIdx.y, z=blockIdx.z; const int n=z>>5; const int oc0=(z&31)<<1; const int Hout=H-4, Wout=W-16; if(ow0>=Wout) return;
    float a00=0.f,a01=0.f,a02=0.f,a03=0.f,a10=0.f,a11=0.f,a12=0.f,a13=0.f;
    if(oh>=2 && oh < H-6 && ow0>=4 && ow0+23 < W){
      #pragma unroll
      for(int c=0;c<32;++c){
        #pragma unroll
        for(int kh=0;kh<5;++kh){int ih=oh+kh*2-2; const float* xp=x+((n*32+c)*H+ih)*W; const float* wp0=w+(((oc0*32+c)*5+kh)*9); const float* wp1=wp0+32*5*9;
          #pragma unroll
          for(int kw=0;kw<9;++kw){int iw=ow0+kw*3-4; float x0=__ldg(xp+iw),x1=__ldg(xp+iw+1),x2=__ldg(xp+iw+2),x3=__ldg(xp+iw+3); float wt0=__ldg(wp0+kw),wt1=__ldg(wp1+kw); a00=fmaf(x0,wt0,a00);a01=fmaf(x1,wt0,a01);a02=fmaf(x2,wt0,a02);a03=fmaf(x3,wt0,a03);a10=fmaf(x0,wt1,a10);a11=fmaf(x1,wt1,a11);a12=fmaf(x2,wt1,a12);a13=fmaf(x3,wt1,a13);}
        }
      }
    } else {
      #pragma unroll
      for(int c=0;c<32;++c){
        #pragma unroll
        for(int kh=0;kh<5;++kh){int ih=oh+kh*2-2; if((unsigned)ih<(unsigned)H){const float* xp=x+((n*32+c)*H+ih)*W; const float* wp0=w+(((oc0*32+c)*5+kh)*9); const float* wp1=wp0+32*5*9;
          #pragma unroll
          for(int kw=0;kw<9;++kw){int iw=ow0+kw*3-4; float wt0=__ldg(wp0+kw),wt1=__ldg(wp1+kw); if((unsigned)iw<(unsigned)W){float xv=__ldg(xp+iw);a00=fmaf(xv,wt0,a00);a10=fmaf(xv,wt1,a10);} if((unsigned)(iw+1)<(unsigned)W&&ow0+1<Wout){float xv=__ldg(xp+iw+1);a01=fmaf(xv,wt0,a01);a11=fmaf(xv,wt1,a11);} if((unsigned)(iw+2)<(unsigned)W&&ow0+2<Wout){float xv=__ldg(xp+iw+2);a02=fmaf(xv,wt0,a02);a12=fmaf(xv,wt1,a12);} if((unsigned)(iw+3)<(unsigned)W&&ow0+3<Wout){float xv=__ldg(xp+iw+3);a03=fmaf(xv,wt0,a03);a13=fmaf(xv,wt1,a13);}}
        }}
      }
    }
    float* yp0=y+((n*64+oc0)*Hout+oh)*Wout+ow0; float* yp1=yp0+Hout*Wout; yp0[0]=a00;if(ow0+1<Wout)yp0[1]=a01;if(ow0+2<Wout)yp0[2]=a02;if(ow0+3<Wout)yp0[3]=a03; yp1[0]=a10;if(ow0+1<Wout)yp1[1]=a11;if(ow0+2<Wout)yp1[2]=a12;if(ow0+3<Wout)yp1[3]=a13;
}
__global__ void conv2d_generic_kernel(const float* __restrict__ x,const float* __restrict__ w,float* __restrict__ y,int N,int C,int H,int W,int O,int KH,int KW,int Hout,int Wout,int stride,int pad_h,int pad_w,int dil_h,int dil_w){int idx=blockIdx.x*blockDim.x+threadIdx.x,total=N*O*Hout*Wout;if(idx>=total)return;int ow=idx%Wout;int t=idx/Wout;int oh=t%Hout;t/=Hout;int oc=t%O;int n=t/O;float acc=0.f;for(int c=0;c<C;++c)for(int kh=0;kh<KH;++kh){int ih=oh*stride+kh*dil_h-pad_h;if((unsigned)ih>=(unsigned)H)continue;for(int kw=0;kw<KW;++kw){int iw=ow*stride+kw*dil_w-pad_w;if((unsigned)iw<(unsigned)W)acc=fmaf(x[((n*C+c)*H+ih)*W+iw],w[((oc*C+c)*KH+kh)*KW+kw],acc);}}y[idx]=acc;}
torch::Tensor conv2d_forward_cuda(torch::Tensor x,torch::Tensor w,int stride,int pad_h,int pad_w,int dil_h,int dil_w){int N=x.size(0),C=x.size(1),H=x.size(2),W=x.size(3),O=w.size(0),KH=w.size(2),KW=w.size(3);int Hout=(H+2*pad_h-dil_h*(KH-1)-1)/stride+1,Wout=(W+2*pad_w-dil_w*(KW-1)-1)/stride+1;auto y=torch::empty({N,O,Hout,Wout},x.options());if(C==32&&O==64&&KH==5&&KW==9&&stride==1&&pad_h==2&&pad_w==4&&dil_h==2&&dil_w==3){dim3 block(128),grid(1,Hout,N*32);conv2d_5x9_oc2_kernel<<<grid,block>>>(x.data_ptr<float>(),w.data_ptr<float>(),y.data_ptr<float>(),N,H,W);}else{int total=N*O*Hout*Wout,threads=256,blocks=(total+threads-1)/threads;conv2d_generic_kernel<<<blocks,threads>>>(x.data_ptr<float>(),w.data_ptr<float>(),y.data_ptr<float>(),N,C,H,W,O,KH,KW,Hout,Wout,stride,pad_h,pad_w,dil_h,dil_w);}return y;}
'''
_ext=load_inline(name='conv2d_dil_ext_final',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['conv2d_forward'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)
class ModelNew(nn.Module):
    def __init__(self,in_channels:int,out_channels:int,kernel_size:tuple,stride:int=1,padding:tuple=(0,0),dilation:tuple=(1,1),bias:bool=False):
        super().__init__(); self.conv2d=nn.Conv2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding,dilation=dilation,bias=bias); self.stride=stride; self.padding=padding; self.dilation=dilation
    def forward(self,x:torch.Tensor)->torch.Tensor:
        return _ext.conv2d_forward(x.contiguous(),self.conv2d.weight.contiguous(),int(self.stride),int(self.padding[0]),int(self.padding[1]),int(self.dilation[0]),int(self.dilation[1]))