import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 3D convolution operation with asymmetric input and kernel sizes.

    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 in the form (kernel_size_d, kernel_size_h, kernel_size_w).
        stride (tuple, optional): Stride of the convolution in the form (stride_d, stride_h, stride_w). Defaults to (1, 1, 1).
        padding (tuple, optional): Padding applied to the input in the form (padding_d, padding_h, padding_w). Defaults to (0, 0, 0).
        dilation (tuple, optional): Spacing between kernel elements in the form (dilation_d, dilation_h, dilation_w). Defaults to (1, 1, 1).
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 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: tuple = (1, 1, 1), padding: tuple = (0, 0, 0), dilation: tuple = (1, 1, 1), groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv3d = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 3D convolution.

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

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

# Test code
batch_size = 8
in_channels = 3
out_channels = 64
kernel_size = (3, 5, 7)  # Asymmetric kernel size
depth = 16
height = 128
width = 128

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization

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

cpp_src = r"""
#include <torch/extension.h>
torch::Tensor conv3d_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, int sd,int sh,int sw,int pd,int ph,int pw,int dd,int dh,int dw);
torch::Tensor conv3d_forward_nobias_cuda(torch::Tensor x, torch::Tensor w, int sd,int sh,int sw,int pd,int ph,int pw,int dd,int dh,int dw);
"""

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define TILE_W 32
#define TX 8
#define OCG 8
#define OC_TILE 16
#define MAX_SH 4096

__global__ __launch_bounds__(64,4)
void conv3d_wtile_oc16_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ y,
 int N,int C,int D,int H,int W,int OC,int KD,int KH,int KW,int OD,int OH,int OW,int sd,int sh,int sw,int pd,int ph,int pw,int dd,int dh,int dw,int has_bias,int span){
    __shared__ float smem[MAX_SH];
    int tx=threadIdx.x, ty=threadIdx.y;
    int oy=blockIdx.y;
    int n=oy/(OD*OH);
    int rem=oy-n*(OD*OH);
    int od=rem/OH;
    int oh=rem-od*OH;
    int ow0=blockIdx.x*TILE_W;
    int oc0=blockIdx.z*OC_TILE+ty;
    int oc1=oc0+OCG;
    int tid=ty*TX+tx;
    int total=C*KD*KH*span;
    for(int idx=tid; idx<total; idx+=TX*OCG){
        int t=idx;
        int iw_off=t%span; t/=span;
        int kh=t%KH; t/=KH;
        int kd=t%KD; t/=KD;
        int c=t;
        int id=od*sd-pd+kd*dd;
        int ih=oh*sh-ph+kh*dh;
        int iw=ow0*sw-pw+iw_off;
        float v=0.f;
        if((unsigned)id<(unsigned)D && (unsigned)ih<(unsigned)H && (unsigned)iw<(unsigned)W){
            long long xi=(((long long)n*C+c)*D+id)*H*W+(long long)ih*W+iw;
            v=x[xi];
        }
        smem[idx]=v;
    }
    __syncthreads();

    float init0=(has_bias && oc0<OC)?b[oc0]:0.f;
    float a00=init0,a01=init0,a02=init0,a03=init0;
    float init1=(has_bias && oc1<OC)?b[oc1]:0.f;
    float a10=init1,a11=init1,a12=init1,a13=init1;
    #pragma unroll
    for(int c=0;c<3;++c){
      #pragma unroll
      for(int kd=0;kd<3;++kd){
        #pragma unroll
        for(int kh=0;kh<5;++kh){
          int sbase=(((c*3+kd)*5+kh)*span)+tx;
          int wb0=(((oc0*3+c)*3+kd)*5+kh)*7;
          int wb1=(((oc1*3+c)*3+kd)*5+kh)*7;
          #pragma unroll
          for(int kw=0;kw<7;++kw){
              float xv0=smem[sbase+kw];
              float xv1=smem[sbase+TX+kw];
              float xv2=smem[sbase+2*TX+kw];
              float xv3=smem[sbase+3*TX+kw];
              if(oc0<OC){
                  float ww=w[wb0+kw];
                  a00=fmaf(xv0,ww,a00); a01=fmaf(xv1,ww,a01); a02=fmaf(xv2,ww,a02); a03=fmaf(xv3,ww,a03);
              }
              if(oc1<OC){
                  float ww=w[wb1+kw];
                  a10=fmaf(xv0,ww,a10); a11=fmaf(xv1,ww,a11); a12=fmaf(xv2,ww,a12); a13=fmaf(xv3,ww,a13);
              }
          }
        }
      }
    }
    int ow=ow0+tx;
    if(oc0<OC){
        long long base0=(((long long)n*OC+oc0)*OD+od)*OH*OW+(long long)oh*OW;
        if(ow<OW)y[base0+ow]=a00;
        if(ow+TX<OW)y[base0+ow+TX]=a01;
        if(ow+2*TX<OW)y[base0+ow+2*TX]=a02;
        if(ow+3*TX<OW)y[base0+ow+3*TX]=a03;
    }
    if(oc1<OC){
        long long base1=(((long long)n*OC+oc1)*OD+od)*OH*OW+(long long)oh*OW;
        if(ow<OW)y[base1+ow]=a10;
        if(ow+TX<OW)y[base1+ow+TX]=a11;
        if(ow+2*TX<OW)y[base1+ow+2*TX]=a12;
        if(ow+3*TX<OW)y[base1+ow+3*TX]=a13;
    }
}

static torch::Tensor launch_conv(torch::Tensor x, torch::Tensor w, const float* bptr, bool has_bias, int sd,int sh,int sw,int pd,int ph,int pw,int dd,int dh,int dw){
    x=x.contiguous(); w=w.contiguous();
    int N=x.size(0), C=x.size(1), D=x.size(2), H=x.size(3), W=x.size(4);
    int OC=w.size(0), KD=w.size(2), KH=w.size(3), KW=w.size(4);
    int OD=(D+2*pd-dd*(KD-1)-1)/sd+1;
    int OH=(H+2*ph-dh*(KH-1)-1)/sh+1;
    int OW=(W+2*pw-dw*(KW-1)-1)/sw+1;
    auto y=torch::empty({N,OC,OD,OH,OW},x.options());
    int span=(TILE_W-1)*sw+(KW-1)*dw+1;
    dim3 block(TX,OCG);
    dim3 grid((OW+TILE_W-1)/TILE_W, N*OD*OH, (OC+OC_TILE-1)/OC_TILE);
    conv3d_wtile_oc16_kernel<<<grid,block,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),w.data_ptr<float>(),bptr,y.data_ptr<float>(),N,C,D,H,W,OC,KD,KH,KW,OD,OH,OW,sd,sh,sw,pd,ph,pw,dd,dh,dw,has_bias?1:0,span);
    return y;
}

torch::Tensor conv3d_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, int sd,int sh,int sw,int pd,int ph,int pw,int dd,int dh,int dw){
    return launch_conv(x,w,b.data_ptr<float>(),true,sd,sh,sw,pd,ph,pw,dd,dh,dw);
}

torch::Tensor conv3d_forward_nobias_cuda(torch::Tensor x, torch::Tensor w, int sd,int sh,int sw,int pd,int ph,int pw,int dd,int dh,int dw){
    return launch_conv(x,w,nullptr,false,sd,sh,sw,pd,ph,pw,dd,dh,dw);
}
"""

_ext = load_inline(name="conv3d_asym_final_submit", cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=["conv3d_forward_cuda", "conv3d_forward_nobias_cuda"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: tuple = (1, 1, 1), padding: tuple = (0, 0, 0), dilation: tuple = (1, 1, 1), groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv3d = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.conv3d.bias is None:
            return _ext.conv3d_forward_nobias_cuda(x, self.conv3d.weight,
                                                   self.stride[0], self.stride[1], self.stride[2],
                                                   self.padding[0], self.padding[1], self.padding[2],
                                                   self.dilation[0], self.dilation[1], self.dilation[2])
        return _ext.conv3d_forward_cuda(x, self.conv3d.weight, self.conv3d.bias,
                                        self.stride[0], self.stride[1], self.stride[2],
                                        self.padding[0], self.padding[1], self.padding[2],
                                        self.dilation[0], self.dilation[1], self.dilation[2])