final submitted kernel (47 lines, 7,514 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cpp_src = r'''
#include <torch/extension.h>
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor running_mean, torch::Tensor running_var, double eps, int stride, int padding, int kernel_size, bool training, double momentum);
'''
cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
__device__ __forceinline__ float tf32_round(float x){ unsigned int u=__float_as_uint(x); unsigned int r=0x00000fffU+((u>>13)&1U); u+=r; u&=0xffffe000U; return __uint_as_float(u); }
__global__ void compute_eff_kernel(const float* __restrict__ w,float* __restrict__ eff,int IC,int OC){
int idx=blockIdx.x*blockDim.x+threadIdx.x,total=IC*OC*27; if(idx>=total)return; int t=idx; int c=t%3;t/=3; int b=t%3;t/=3; int a=t%3;t/=3; int oc=t%OC,ic=t/OC; float s=0;
int kd0=(a==0)?1:0,kd1=(a==0)?2:((a==1)?2:0); int kh0=(b==0)?1:0,kh1=(b==0)?2:((b==1)?2:0); int kw0=(c==0)?1:0,kw1=(c==0)?2:((c==1)?2:0);
for(int kd=kd0;kd<=kd1;++kd)for(int kh=kh0;kh<=kh1;++kh)for(int kw=kw0;kw<=kw1;++kw)s+=tf32_round(w[(((ic*OC+oc)*3+kd)*3+kh)*3+kw]); eff[idx]=s*(1.f/64.f); }
__device__ __forceinline__ float deconv_val(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ bias,int n,int oc,int od,int oh,int ow,int IC,int D,int H,int W,int OC,int padding,int stride){
float acc=bias[oc]; int ids[2],ihs[2],iws[2],kds[2],khs[2],kws[2]; int nd=0,nh=0,nw=0;
#pragma unroll
for(int kd=0;kd<3;++kd){int t=od+padding-kd;if((t&1)==0){int id=t>>1;if((unsigned)id<(unsigned)D){ids[nd]=id;kds[nd]=kd;nd++;}}}
#pragma unroll
for(int kh=0;kh<3;++kh){int t=oh+padding-kh;if((t&1)==0){int ih=t>>1;if((unsigned)ih<(unsigned)H){ihs[nh]=ih;khs[nh]=kh;nh++;}}}
#pragma unroll
for(int kw=0;kw<3;++kw){int t=ow+padding-kw;if((t&1)==0){int iw=t>>1;if((unsigned)iw<(unsigned)W){iws[nw]=iw;kws[nw]=kw;nw++;}}}
for(int ic=0;ic<IC;++ic)for(int a=0;a<nd;++a)for(int b=0;b<nh;++b)for(int c=0;c<nw;++c){float xv=tf32_round(x[(((int64_t)n*IC+ic)*D+ids[a])*H*W+ihs[b]*W+iws[c]]);float ww=tf32_round(w[(((ic*OC+oc)*3+kds[a])*3+khs[b])*3+kws[c]]);acc=fmaf(xv,ww,acc);} return acc; }
__global__ void stats_kernel(const float* x,const float* w,const float* bias,double* psum,double* psq,int N,int IC,int D,int H,int W,int OC,int CD,int CH,int CW,int padding,int stride,int B){
int oc=blockIdx.y,bid=blockIdx.x,tid=threadIdx.x; int64_t M=(int64_t)N*CD*CH*CW; double sum=0,sq=0; for(int64_t idx=(int64_t)bid*blockDim.x+tid;idx<M;idx+=(int64_t)B*blockDim.x){int ow=idx%CW;int64_t q=idx/CW;int oh=q%CH;q/=CH;int od=q%CD;q/=CD;int n=q;float v=deconv_val(x,w,bias,n,oc,od,oh,ow,IC,D,H,W,OC,padding,stride);sum+=(double)v;sq+=(double)v*(double)v;} __shared__ double shs[256],shq[256]; shs[tid]=sum;shq[tid]=sq;__syncthreads(); for(int off=128;off>0;off>>=1){if(tid<off){shs[tid]+=shs[tid+off];shq[tid]+=shq[tid+off];}__syncthreads();} if(tid==0){psum[oc*B+bid]=shs[0];psq[oc*B+bid]=shq[0];}}
__global__ void reduce_stats_kernel(const double* psum,const double* psq,float* mean,float* var,float* rm,float* rv,int64_t M,int OC,int B,float mom){int oc=blockIdx.x,tid=threadIdx.x; double s=0,q=0; for(int i=tid;i<B;i+=blockDim.x){s+=psum[oc*B+i];q+=psq[oc*B+i];} __shared__ double shs[256],shq[256]; shs[tid]=s;shq[tid]=q;__syncthreads(); for(int off=128;off>0;off>>=1){if(tid<off){shs[tid]+=shs[tid+off];shq[tid]+=shq[tid+off];}__syncthreads();} if(tid==0){double m=shs[0]/(double)M;double v=shq[0]/(double)M-m*m;if(v<0)v=0;mean[oc]=(float)m;var[oc]=(float)v;rm[oc]=rm[oc]*(1.f-mom)+(float)m*mom;rv[oc]=rv[oc]*(1.f-mom)+(float)(v*((double)M/(double)(M-1)))*mom;}}
__global__ void pool_eff_kernel(const float* x,const float* eff,const float* bias,const float* bn_w,const float* bn_b,const float* mean,const float* var,float* y,int N,int IC,int D,int H,int W,int OC,int OD,int OH,int OW,float eps){
int64_t idx=(int64_t)blockIdx.x*blockDim.x+threadIdx.x,total=(int64_t)N*OC*OD*OH*OW; if(idx>=total)return; int ow=idx%OW; int64_t q=idx/OW; int oh=q%OH;q/=OH; int od=q%OD;q/=OD; int oc=q%OC; int n=q/OC; int id0=od<<1,ih0=oh<<1,iw0=ow<<1; float acc=bias[oc];
for(int ic=0;ic<IC;++ic){const float* xp=x+(((int64_t)n*IC+ic)*D+id0)*H*W+ih0*W+iw0; const float* wp=eff+(ic*OC+oc)*27;
#pragma unroll
for(int dz=0;dz<3;++dz)for(int dy=0;dy<3;++dy)for(int dx=0;dx<3;++dx) acc=fmaf(tf32_round(xp[dz*H*W+dy*W+dx]),wp[(dz*3+dy)*3+dx],acc);}
float scale=bn_w[oc]*rsqrtf(var[oc]+eps); y[idx]=(acc-mean[oc])*scale+bn_b[oc]; }
torch::Tensor fused_forward(torch::Tensor x,torch::Tensor w,torch::Tensor bias,torch::Tensor bn_w,torch::Tensor bn_b,torch::Tensor running_mean,torch::Tensor running_var,double eps,int stride,int padding,int kernel_size,bool training,double momentum){
int N=x.size(0),IC=x.size(1),D=x.size(2),H=x.size(3),W=x.size(4),OC=w.size(1); int CD=(D-1)*stride-2*padding+kernel_size,CH=(H-1)*stride-2*padding+kernel_size,CW=(W-1)*stride-2*padding+kernel_size; int OD=(((CD-2)/2+1)-2)/2+1,OH=(((CH-2)/2+1)-2)/2+1,OW=(((CW-2)/2+1)-2)/2+1; auto y=torch::empty({N,OC,OD,OH,OW},x.options()); auto stream=at::cuda::getCurrentCUDAStream(); const float* mean_ptr; const float* var_ptr; torch::Tensor mean_t,var_t,psum,psq;
if(training){int B=256; auto dopts=x.options().dtype(torch::kFloat64); psum=torch::empty({OC,B},dopts); psq=torch::empty({OC,B},dopts); mean_t=torch::empty({OC},x.options()); var_t=torch::empty({OC},x.options()); stats_kernel<<<dim3(B,OC),256,0,stream>>>(x.data_ptr<float>(),w.data_ptr<float>(),bias.data_ptr<float>(),psum.data_ptr<double>(),psq.data_ptr<double>(),N,IC,D,H,W,OC,CD,CH,CW,padding,stride,B); int64_t M=(int64_t)N*CD*CH*CW; reduce_stats_kernel<<<OC,256,0,stream>>>(psum.data_ptr<double>(),psq.data_ptr<double>(),mean_t.data_ptr<float>(),var_t.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),M,OC,B,(float)momentum); mean_ptr=mean_t.data_ptr<float>(); var_ptr=var_t.data_ptr<float>();} else {mean_ptr=running_mean.data_ptr<float>(); var_ptr=running_var.data_ptr<float>();}
auto eff=torch::empty({IC,OC,3,3,3},x.options()); compute_eff_kernel<<<(IC*OC*27+255)/256,256,0,stream>>>(w.data_ptr<float>(),eff.data_ptr<float>(),IC,OC); int64_t total=(int64_t)N*OC*OD*OH*OW; pool_eff_kernel<<<(total+255)/256,256,0,stream>>>(x.data_ptr<float>(),eff.data_ptr<float>(),bias.data_ptr<float>(),bn_w.data_ptr<float>(),bn_b.data_ptr<float>(),mean_ptr,var_ptr,y.data_ptr<float>(),N,IC,D,H,W,OC,OD,OH,OW,(float)eps); return y; }
'''
_ext=load_inline(name='final_kernel_v1_submit',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['fused_forward'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)
class ModelNew(nn.Module):
def __init__(self,in_channels,out_channels,kernel_size,stride,padding,bias_shape):
super().__init__(); self.conv_transpose=nn.ConvTranspose3d(in_channels,out_channels,kernel_size,stride=stride,padding=padding); self.batch_norm=nn.BatchNorm3d(out_channels); self.stride=stride; self.padding=padding; self.kernel_size=kernel_size
def forward(self,x):
return _ext.fused_forward(x.contiguous(),self.conv_transpose.weight,self.conv_transpose.bias,self.batch_norm.weight,self.batch_norm.bias,self.batch_norm.running_mean,self.batch_norm.running_var,float(self.batch_norm.eps),int(self.stride),int(self.padding),int(self.kernel_size),bool(self.training),float(self.batch_norm.momentum if self.batch_norm.momentum is not None else 0.1))