callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,\n int64_t sd, int64_t sh, int64_t sw,\n int64_t pd, int64_t ph, int64_t pw,\n int64_t dd, int64_t dh, int64_t dw,\n int64_t groups);\n'''\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#define BLOCK_W 128\n#define OC_TILE 8\n#define KSIZE 315\n\n__global__ __launch_bounds__(BLOCK_W, 2)\nvoid conv3d_3x5x7_c3_kernel(const float* __restrict__ x, const float* __restrict__ w,\n const float* __restrict__ b, float* __restrict__ y,\n int N, int D, int H, int W, int OC, int DO, int HO, int WO, int has_bias) {\n __shared__ float sh_in[BLOCK_W];\n __shared__ float sh_w[OC_TILE * KSIZE];\n int tx = threadIdx.x;\n int row = blockIdx.x;\n int oc_base = blockIdx.y * OC_TILE;\n int tmp = row;\n int oh = tmp % HO; tmp /= HO;\n int od = tmp % DO; tmp /= DO;\n int n = tmp;\n for (int idx = tx; idx < OC_TILE * KSIZE; idx += BLOCK_W) {\n int to = idx / KSIZE;\n int k = idx - to * KSIZE;\n int oc = oc_base + to;\n sh_w[idx] = (oc < OC) ? w[oc * KSIZE + k] : 0.0f;\n }\n __syncthreads();\n float acc0=0.f, acc1=0.f, acc2=0.f, acc3=0.f, acc4=0.f, acc5=0.f, acc6=0.f, acc7=0.f;\n int ow = tx;\n bool valid = (ow < WO);\n int kidx = 0;\n#pragma unroll\n for (int ic=0; ic<3; ++ic) {\n#pragma unroll\n for (int kd=0; kd<3; ++kd) {\n#pragma unroll\n for (int kh=0; kh<5; ++kh) {\n long long in_off = (((((long long)n*3 + ic)*D + (od+kd))*H + (oh+kh))*W + tx);\n sh_in[tx] = x[in_off];\n __syncthreads();\n if (valid) {\n#pragma unroll\n for (int kw=0; kw<7; ++kw) {\n float xv = sh_in[ow + kw];\n int kk = kidx + kw;\n acc0 = fmaf(xv, sh_w[0*KSIZE + kk], acc0);\n acc1 = fmaf(xv, sh_w[1*KSIZE + kk], acc1);\n acc2 = fmaf(xv, sh_w[2*KSIZE + kk], acc2);\n acc3 = fmaf(xv, sh_w[3*KSIZE + kk], acc3);\n acc4 = fmaf(xv, sh_w[4*KSIZE + kk], acc4);\n acc5 = fmaf(xv, sh_w[5*KSIZE + kk], acc5);\n acc6 = fmaf(xv, sh_w[6*KSIZE + kk], acc6);\n acc7 = fmaf(xv, sh_w[7*KSIZE + kk], acc7);\n }\n }\n kidx += 7;\n __syncthreads();\n }\n }\n }\n if (valid) {\n long long base = (((((long long)n * OC) * DO + od) * HO + oh) * WO + ow);\n long long ch_stride = (long long)DO * HO * WO;\n if (has_bias) {\n if (oc_base+0<OC) acc0 += b[oc_base+0]; if (oc_base+1<OC) acc1 += b[oc_base+1];\n if (oc_base+2<OC) acc2 += b[oc_base+2]; if (oc_base+3<OC) acc3 += b[oc_base+3];\n if (oc_base+4<OC) acc4 += b[oc_base+4]; if (oc_base+5<OC) acc5 += b[oc_base+5];\n if (oc_base+6<OC) acc6 += b[oc_base+6]; if (oc_base+7<OC) acc7 += b[oc_base+7];\n }\n if (oc_base+0<OC) y[base + 0*ch_stride] = acc0;\n if (oc_base+1<OC) y[base + 1*ch_stride] = acc1;\n if (oc_base+2<OC) y[base + 2*ch_stride] = acc2;\n if (oc_base+3<OC) y[base + 3*ch_stride] = acc3;\n if (oc_base+4<OC) y[base + 4*ch_stride] = acc4;\n if (oc_base+5<OC) y[base + 5*ch_stride] = acc5;\n if (oc_base+6<OC) y[base + 6*ch_stride] = acc6;\n if (oc_base+7<OC) y[base + 7*ch_stride] = acc7;\n }\n}\n\n__global__ void conv3d_generic_kernel(const float* __restrict__ x, const float* __restrict__ w,\n const float* __restrict__ b, float* __restrict__ y, int N,int C,int D,int H,int W,int OC,\n int KD,int KH,int KW,int DO,int HO,int WO,int sd,int sh,int sw,int pd,int ph,int pw,\n int dd,int dh,int dw,int groups,int has_bias,long long total) {\n long long idx = blockIdx.x * (long long)blockDim.x + threadIdx.x;\n if (idx >= total) return;\n int ow=idx%WO; idx/=WO; int oh=idx%HO; idx/=HO; int od=idx%DO; idx/=DO; int oc=idx%OC; idx/=OC; int n=idx;\n int ocpg=OC/groups, icpg=C/groups, g=oc/ocpg;\n float acc = has_bias ? b[oc] : 0.0f;\n for (int ic0=0; ic0<icpg; ++ic0) { int ic=g*icpg+ic0;\n for (int kd=0; kd<KD; ++kd) { int iz=od*sd-pd+kd*dd; if ((unsigned)iz >= (unsigned)D) continue;\n for (int kh=0; kh<KH; ++kh) { int iy=oh*sh-ph+kh*dh; if ((unsigned)iy >= (unsigned)H) continue;\n for (int kw=0; kw<KW; ++kw) { int ix=ow*sw-pw+kw*dw; if ((unsigned)ix >= (unsigned)W) continue;\n float xv = x[(((((long long)n*C+ic)*D+iz)*H+iy)*W+ix)];\n float wv = w[(((((long long)oc*icpg+ic0)*KD+kd)*KH+kh)*KW+kw)];\n acc = fmaf(xv, wv, acc);\n }}}}\n y[(((((long long)n*OC+oc)*DO+od)*HO+oh)*WO+ow)] = acc;\n}\n\ntorch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,\n int64_t sd,int64_t sh,int64_t sw,int64_t pd,int64_t ph,int64_t pw,\n int64_t dd,int64_t dh,int64_t dw,int64_t groups) {\n TORCH_CHECK(x.is_cuda() && weight.is_cuda(), \"CUDA only\");\n x = x.contiguous(); weight = weight.contiguous();\n if (bias.defined() && bias.numel() > 0) bias = bias.contiguous();\n int N=x.size(0), C=x.size(1), D=x.size(2), H=x.size(3), W=x.size(4);\n int OC=weight.size(0), KD=weight.size(2), KH=weight.size(3), KW=weight.size(4);\n int DO=(D+2*(int)pd-(int)dd*(KD-1)-1)/(int)sd+1;\n int HO=(H+2*(int)ph-(int)dh*(KH-1)-1)/(int)sh+1;\n int WO=(W+2*(int)pw-(int)dw*(KW-1)-1)/(int)sw+1;\n auto y = torch::empty({N,OC,DO,HO,WO}, x.options());\n int has_bias = (bias.defined() && bias.numel() > 0) ? 1 : 0;\n const float* bp = has_bias ? bias.data_ptr<float>() : nullptr;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (C==3 && OC==64 && KD==3 && KH==5 && KW==7 && sd==1 && sh==1 && sw==1 &&\n pd==0 && ph==0 && pw==0 && dd==1 && dh==1 && dw==1 && groups==1 && WO<=BLOCK_W) {\n dim3 block(BLOCK_W); dim3 grid((unsigned)(N*DO*HO), (unsigned)((OC+OC_TILE-1)/OC_TILE));\n conv3d_3x5x7_c3_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp,\n y.data_ptr<float>(), N,D,H,W,OC,DO,HO,WO,has_bias);\n } else {\n long long total=(long long)N*OC*DO*HO*WO; int threads=256; int blocks=(int)((total+threads-1)/threads);\n conv3d_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(),\n N,C,D,H,W,OC,KD,KH,KW,DO,HO,WO,(int)sd,(int)sh,(int)sw,(int)pd,(int)ph,(int)pw,(int)dd,(int)dh,(int)dw,(int)groups,has_bias,total);\n }\n C10_CUDA_KERNEL_LAUNCH_CHECK();\n return y;\n}\n'''\n_ext = load_inline(name='conv3d_3x5x7_c3_ext3', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['conv3d_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'], verbose=False)\nclass ModelNew(nn.Module):\n 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):\n super(ModelNew,self).__init__()\n self.conv3d=nn.Conv3d(in_channels,out_channels,kernel_size,stride=stride,padding=padding,dilation=dilation,groups=groups,bias=bias)\n self.stride=stride; self.padding=padding; self.dilation=dilation; self.groups=groups\n def forward(self,x:torch.Tensor)->torch.Tensor:\n b=self.conv3d.bias if self.conv3d.bias is not None else torch.empty(0,device=x.device,dtype=x.dtype)\n return _ext.conv3d_forward(x,self.conv3d.weight,b,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],self.groups)\n"
}