import torch
import torch.nn as nn

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

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the square convolution kernel (kernel_size x kernel_size).
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 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: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv3d = nn.Conv3d(in_channels, out_channels, (kernel_size, kernel_size, 1), 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, height, width, depth).

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

# Test code
batch_size = 16
in_channels = 3
out_channels = 64
kernel_size = 3
width = 256
height = 256
depth = 10

def get_inputs():
    x = torch.rand(batch_size, in_channels, height, width, depth)
    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_k31_forward(torch::Tensor x, torch::Tensor weight, c10::optional<torch::Tensor> bias_opt,
                                 int64_t stride, int64_t padding, int64_t dilation, int64_t groups);
'''

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

__global__ void conv3d_k31_default_kernel(const float* __restrict__ x,
                                          const float* __restrict__ w,
                                          const float* __restrict__ b,
                                          float* __restrict__ y,
                                          int S) {
    constexpr int OD=10, OW=254, OH=254, IW=256, IH=256, IC=3;
    constexpr int SPN = OH*OW*OD;
    constexpr int CSTR = IH*IW*OD;
    constexpr int HSTR = IW*OD;
    int lane = threadIdx.x;
    int wy = threadIdx.y;
    int oc = blockIdx.y * 8 + wy;
    int base_s = blockIdx.x * 128 + lane;
    int xb0=0, xb1=0, xb2=0, xb3=0;
    int yo0=0, yo1=0, yo2=0, yo3=0;
    bool full = (base_s + 96) < S;
#pragma unroll
    for (int v=0; v<4; ++v) {
        int s = base_s + v*32;
        if (full || s < S) {
            int t=s;
            int z=t%OD; t/=OD;
            int ow=t%OW; t/=OW;
            int oh=t%OH; t/=OH;
            int n=t;
            int xb = ((n*IC*IH + oh)*IW + ow)*OD + z;
            int pos = s - n*SPN;
            int yy = ((n*64 + oc)*SPN) + pos;
            if (v==0) { xb0=xb; yo0=yy; }
            else if (v==1) { xb1=xb; yo1=yy; }
            else if (v==2) { xb2=xb; yo2=yy; }
            else { xb3=xb; yo3=yy; }
        }
    }
    float acc0=0.f, acc1=0.f, acc2=0.f, acc3=0.f;
    int wk = oc*27;
#pragma unroll
    for (int c=0; c<3; ++c) {
        int co = c*CSTR;
#pragma unroll
        for (int kh=0; kh<3; ++kh) {
            int kho = co + kh*HSTR;
#pragma unroll
            for (int kw=0; kw<3; ++kw) {
                float ww = __ldg(w + wk++);
                int off = kho + kw*OD;
                if (full) {
                    acc0 = fmaf(__ldg(x + xb0 + off), ww, acc0);
                    acc1 = fmaf(__ldg(x + xb1 + off), ww, acc1);
                    acc2 = fmaf(__ldg(x + xb2 + off), ww, acc2);
                    acc3 = fmaf(__ldg(x + xb3 + off), ww, acc3);
                } else {
                    if (base_s < S) acc0 = fmaf(__ldg(x + xb0 + off), ww, acc0);
                    if (base_s + 32 < S) acc1 = fmaf(__ldg(x + xb1 + off), ww, acc1);
                    if (base_s + 64 < S) acc2 = fmaf(__ldg(x + xb2 + off), ww, acc2);
                    if (base_s + 96 < S) acc3 = fmaf(__ldg(x + xb3 + off), ww, acc3);
                }
            }
        }
    }
    if (b != nullptr) { float bb=__ldg(b+oc); acc0+=bb; acc1+=bb; acc2+=bb; acc3+=bb; }
    if (full) { y[yo0]=acc0; y[yo1]=acc1; y[yo2]=acc2; y[yo3]=acc3; }
    else { if (base_s < S) y[yo0]=acc0; if (base_s+32 < S) y[yo1]=acc1; if (base_s+64 < S) y[yo2]=acc2; if (base_s+96 < S) y[yo3]=acc3; }
}

__global__ void conv3d_k31_general_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ y,int N,int IC,int IH,int IW,int ID,int OC,int OH,int OW,int OD,int K,int stride,int pad,int dil,int groups,long long total) {
    long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x; if(idx>=total) return;
    int z=idx%OD; long long q=idx/OD; int ow=q%OW; q/=OW; int oh=q%OH; q/=OH; int oc=q%OC; q/=OC; int n=q;
    int ocpg=OC/groups, icpg=IC/groups, g=oc/ocpg; float acc=b?__ldg(b+oc):0.f; int ih0=oh*stride-pad, iw0=ow*stride-pad, iz=z*stride-pad;
    if((unsigned)iz<(unsigned)ID){ for(int ci=0;ci<icpg;++ci){ int c=g*icpg+ci; for(int kh=0;kh<K;++kh){ int ih=ih0+kh*dil; if((unsigned)ih>=(unsigned)IH) continue; for(int kw=0;kw<K;++kw){ int iw=iw0+kw*dil; if((unsigned)iw>=(unsigned)IW) continue; float xv=__ldg(x+(((n*IC+c)*IH+ih)*IW+iw)*ID+iz); float ww=__ldg(w+(((oc*icpg+ci)*K+kh)*K+kw)); acc=fmaf(xv,ww,acc); } } } }
    y[idx]=acc;
}

torch::Tensor conv3d_k31_forward(torch::Tensor x, torch::Tensor weight, c10::optional<torch::Tensor> bias_opt,int64_t stride,int64_t padding,int64_t dilation,int64_t groups) {
    x=x.contiguous(); weight=weight.contiguous(); const int64_t N=x.size(0),IC=x.size(1),IH=x.size(2),IW=x.size(3),ID=x.size(4); const int64_t OC=weight.size(0),K=weight.size(2);
    int64_t OH=(IH+2*padding-dilation*(K-1)-1)/stride+1; int64_t OW=(IW+2*padding-dilation*(K-1)-1)/stride+1; int64_t OD=(ID+2*padding-1)/stride+1; auto y=torch::empty({N,OC,OH,OW,OD},x.options());
    const float* bptr=nullptr; torch::Tensor bias; if(bias_opt.has_value()&&bias_opt.value().defined()){bias=bias_opt.value().contiguous(); bptr=bias.data_ptr<float>();}
    cudaStream_t stream=at::cuda::getCurrentCUDAStream();
    if(x.scalar_type()==at::kFloat && N==16 && IC==3 && IH==256 && IW==256 && ID==10 && OC==64 && K==3 && stride==1 && padding==0 && dilation==1 && groups==1 && OH==254 && OW==254 && OD==10){ int S=(int)(N*OH*OW*OD); dim3 block(32,8,1); dim3 grid((S+127)/128,8,1); conv3d_k31_default_kernel<<<grid,block,0,stream>>>(x.data_ptr<float>(),weight.data_ptr<float>(),bptr,y.data_ptr<float>(),S); }
    else { long long total=N*OC*OH*OW*OD; int threads=256; int blocks=(total+threads-1)/threads; conv3d_k31_general_kernel<<<blocks,threads,0,stream>>>(x.data_ptr<float>(),weight.data_ptr<float>(),bptr,y.data_ptr<float>(),(int)N,(int)IC,(int)IH,(int)IW,(int)ID,(int)OC,(int)OH,(int)OW,(int)OD,(int)K,(int)stride,(int)padding,(int)dilation,(int)groups,total); }
    return y;
}
'''
_ext=load_inline(name='conv3d_k31_ext_final',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['conv3d_k31_forward'],verbose=False,extra_cuda_cflags=['-O3','--use_fast_math'],extra_cflags=['-O3'])
class ModelNew(nn.Module):
    def __init__(self,in_channels:int,out_channels:int,kernel_size:int,stride:int=1,padding:int=0,dilation:int=1,groups:int=1,bias:bool=False):
        super(ModelNew,self).__init__(); self.conv3d=nn.Conv3d(in_channels,out_channels,(kernel_size,kernel_size,1),stride=stride,padding=padding,dilation=dilation,groups=groups,bias=bias); self.stride=stride; self.padding=padding; self.dilation=dilation; self.groups=groups
    def forward(self,x:torch.Tensor)->torch.Tensor:
        return _ext.conv3d_k31_forward(x,self.conv3d.weight,self.conv3d.bias,self.stride,self.padding,self.dilation,self.groups)