call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\ncpp_source=r'''\n#include <torch/extension.h>\ntorch::Tensor convt3d_forward(torch::Tensor x, torch::Tensor weight, c10::optional<torch::Tensor> bias_opt, int stride, int padding, int output_padding, int groups);\n'''\ncuda_source=r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n__device__ __forceinline__ float tf32_round(float a){ unsigned r; asm(\"cvt.rna.tf32.f32 %0, %1;\" : \"=r\"(r) : \"f\"(a)); return __uint_as_float(r); }\n__global__ void ker8(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ y,int N,int IC,int D,int H,int W,int OC,int OD,int OH,int OW,bool has_bias){\n const int OCT=8; long long tid=(long long)blockIdx.x*blockDim.x+threadIdx.x,total=(long long)N*OD*OH*OW*(OC/OCT); if(tid>=total)return; int oct=tid%(OC/OCT); long long s=tid/(OC/OCT); int ow=s%OW;s/=OW;int oh=s%OH;s/=OH;int od=s%OD;s/=OD;int n=s;int oc0=oct*OCT;\n float a0=has_bias?bias[oc0]:0.f,a1=has_bias?bias[oc0+1]:0.f,a2=has_bias?bias[oc0+2]:0.f,a3=has_bias?bias[oc0+3]:0.f,a4=has_bias?bias[oc0+4]:0.f,a5=has_bias?bias[oc0+5]:0.f,a6=has_bias?bias[oc0+6]:0.f,a7=has_bias?bias[oc0+7]:0.f;\n int kd0=max(0,od-(D-1)),kd1=min(2,od),kh0=max(0,oh-(H-1)),kh1=min(2,oh),kw0=max(0,ow-(W-1)),kw1=min(2,ow); long long xnb=(long long)n*IC*D*H*W;\n for(int ic=0;ic<IC;++ic){ long long xcb=xnb+(long long)ic*D*H*W; long long wib=(long long)ic*OC*27;\n for(int kd=kd0;kd<=kd1;++kd){ int id=od-kd; for(int kh=kh0;kh<=kh1;++kh){ int ih=oh-kh; long long xr=xcb+((long long)id*H+ih)*W; for(int kw=kw0;kw<=kw1;++kw){ float xv=tf32_round(x[xr+ow-kw]); int kidx=(kd*3+kh)*3+kw; long long wb=wib+kidx;\n a0=fmaf(xv,tf32_round(w[wb+(long long)(oc0+0)*27]),a0); a1=fmaf(xv,tf32_round(w[wb+(long long)(oc0+1)*27]),a1); a2=fmaf(xv,tf32_round(w[wb+(long long)(oc0+2)*27]),a2); a3=fmaf(xv,tf32_round(w[wb+(long long)(oc0+3)*27]),a3); a4=fmaf(xv,tf32_round(w[wb+(long long)(oc0+4)*27]),a4); a5=fmaf(xv,tf32_round(w[wb+(long long)(oc0+5)*27]),a5); a6=fmaf(xv,tf32_round(w[wb+(long long)(oc0+6)*27]),a6); a7=fmaf(xv,tf32_round(w[wb+(long long)(oc0+7)*27]),a7);\n }}}\n }\n long long yb=(((long long)n*OC+oc0)*OD+od)*OH*OW+(long long)oh*OW+ow, so=(long long)OD*OH*OW; y[yb]=a0;y[yb+so]=a1;y[yb+2*so]=a2;y[yb+3*so]=a3;y[yb+4*so]=a4;y[yb+5*so]=a5;y[yb+6*so]=a6;y[yb+7*so]=a7;\n}\n__global__ void kergen(const float* x,const float* w,const float* bias,float* y,int N,int IC,int D,int H,int W,int OC,int K,int stride,int padding,int groups,int OD,int OH,int OW,bool has_bias){long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x,total=(long long)N*OC*OD*OH*OW;if(idx>=total)return;int ow=idx%OW;long long t=idx/OW;int oh=t%OH;t/=OH;int od=t%OD;t/=OD;int oc=t%OC;int n=t/OC;int ocpg=OC/groups,icpg=IC/groups,g=oc/ocpg,ocg=oc-g*ocpg;float acc=has_bias?bias[oc]:0.f;for(int icg=0;icg<icpg;++icg){int ic=g*icpg+icg;for(int kd=0;kd<K;++kd){int td=od+padding-kd;if(td%stride)continue;int id=td/stride;if((unsigned)id>=(unsigned)D)continue;for(int kh=0;kh<K;++kh){int th=oh+padding-kh;if(th%stride)continue;int ih=th/stride;if((unsigned)ih>=(unsigned)H)continue;for(int kw=0;kw<K;++kw){int tw=ow+padding-kw;if(tw%stride)continue;int iw=tw/stride;if((unsigned)iw>=(unsigned)W)continue;float xv=tf32_round(x[(((long long)n*IC+ic)*D+id)*H*W+(long long)ih*W+iw]);float wv=tf32_round(w[(((long long)ic*ocpg+ocg)*K+kd)*K*K+(long long)kh*K+kw]);acc=fmaf(xv,wv,acc);}}}}y[idx]=acc;}\ntorch::Tensor convt3d_forward(torch::Tensor x,torch::Tensor w,c10::optional<torch::Tensor> bias_opt,int stride,int padding,int output_padding,int groups){auto bias=bias_opt.has_value()?bias_opt.value():torch::Tensor();bool has_bias=bias.defined()&&bias.numel()>0;int N=x.size(0),IC=x.size(1),D=x.size(2),H=x.size(3),W=x.size(4),K=w.size(2),OC=w.size(1)*groups;int OD=(D-1)*stride-2*padding+K+output_padding,OH=(H-1)*stride-2*padding+K+output_padding,OW=(W-1)*stride-2*padding+K+output_padding;auto y=torch::empty({N,OC,OD,OH,OW},x.options());long long total; if(K==3&&stride==1&&padding==0&&output_padding==0&&groups==1&&OC%8==0){total=(long long)N*OD*OH*OW*(OC/8);ker8<<<(total+255)/256,256>>>(x.data_ptr<float>(),w.data_ptr<float>(),has_bias?bias.data_ptr<float>():nullptr,y.data_ptr<float>(),N,IC,D,H,W,OC,OD,OH,OW,has_bias);}else{total=(long long)N*OC*OD*OH*OW;kergen<<<(total+255)/256,256>>>(x.data_ptr<float>(),w.data_ptr<float>(),has_bias?bias.data_ptr<float>():nullptr,y.data_ptr<float>(),N,IC,D,H,W,OC,K,stride,padding,groups,OD,OH,OW,has_bias);}return y;}\n'''\n_ext=load_inline(name='convt3d_ext_final8tf32',cpp_sources=cpp_source,cuda_sources=cuda_source,functions=['convt3d_forward'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self,in_channels:int,out_channels:int,kernel_size:int,stride:int=1,padding:int=0,output_padding:int=0,groups:int=1,bias:bool=False):\n super().__init__(); self.conv_transpose3d=nn.ConvTranspose3d(in_channels,out_channels,(kernel_size,kernel_size,kernel_size),stride=stride,padding=padding,output_padding=output_padding,groups=groups,bias=bias); self.stride=stride; self.padding=padding; self.output_padding=output_padding; self.groups=groups\n def forward(self,x):\n b=self.conv_transpose3d.bias; return _ext.convt3d_forward(x.contiguous(),self.conv_transpose3d.weight.contiguous(),b.contiguous() if b is not None else None,self.stride,self.padding,self.output_padding,self.groups)\n"
}