final submitted kernel (29 lines, 6,678 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 transition_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor weight, bool training, double eps, double momentum);
'''
cuda_src=r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be CUDA")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
__global__ void partial_stats_kernel(const float* __restrict__ x,double* __restrict__ psum,double* __restrict__ psq,int N,int C,int H,int W,int elems_per_c,int NB){int c=blockIdx.x,b=blockIdx.y,tid=threadIdx.x;double s=0.0,q=0.0;int stride=NB*blockDim.x,start=b*blockDim.x+tid,HW=H*W;for(int idx=start;idx<elems_per_c;idx+=stride){int n=idx/HW;int hw=idx-n*HW;float v=x[((n*C+c)*H*W)+hw];s+=(double)v;q+=(double)v*(double)v;}extern __shared__ double sm[];double* ss=sm;double* qq=sm+blockDim.x;ss[tid]=s;qq[tid]=q;__syncthreads();for(int off=blockDim.x>>1;off>0;off>>=1){if(tid<off){ss[tid]+=ss[tid+off];qq[tid]+=qq[tid+off];}__syncthreads();}if(tid==0){psum[c*NB+b]=ss[0];psq[c*NB+b]=qq[0];}}
__global__ void final_stats_kernel(const double* __restrict__ psum,const double* __restrict__ psq,float* __restrict__ mean,float* __restrict__ invstd,float* __restrict__ running_mean,float* __restrict__ running_var,int C,int M,int NB,double eps,double momentum){int c=blockIdx.x,tid=threadIdx.x;double s=0.0,q=0.0;for(int i=tid;i<NB;i+=blockDim.x){s+=psum[c*NB+i];q+=psq[c*NB+i];}extern __shared__ double sm[];double* ss=sm;double* qq=sm+blockDim.x;ss[tid]=s;qq[tid]=q;__syncthreads();for(int off=blockDim.x>>1;off>0;off>>=1){if(tid<off){ss[tid]+=ss[tid+off];qq[tid]+=qq[tid+off];}__syncthreads();}if(tid==0){double mu=ss[0]/(double)M;double var=qq[0]/(double)M-mu*mu;if(var<0.0)var=0.0;mean[c]=(float)mu;invstd[c]=(float)rsqrt(var+eps);double unb=var*((double)M/(double)(M-1));running_mean[c]=(float)((1.0-momentum)*(double)running_mean[c]+momentum*mu);running_var[c]=(float)((1.0-momentum)*(double)running_var[c]+momentum*unb);}}
__global__ void pool_bn_relu_nhwc_kernel(const float* __restrict__ x,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ invstd,float* __restrict__ tmp,int N,int C,int H,int W){int tx=threadIdx.x,ty=threadIdx.y;int H2=H>>1,W2=W>>1,P=N*H2*W2;int p0=blockIdx.x*32;int p=p0+tx;int c=blockIdx.y*8+ty;__shared__ float sh[8][32];float val=0.0f;if(p<P&&c<C){int t=p;int w2=t%W2;t/=W2;int h2=t%H2;int n=t/H2;int base=((n*C+c)*H+(h2<<1))*W+(w2<<1);float m=mean[c],is=invstd[c],g=gamma[c],be=beta[c];float v0=x[base],v1=x[base+1],v2=x[base+W],v3=x[base+W+1];v0=fmaf((v0-m)*is,g,be);v0=v0>0.0f?v0:0.0f;v1=fmaf((v1-m)*is,g,be);v1=v1>0.0f?v1:0.0f;v2=fmaf((v2-m)*is,g,be);v2=v2>0.0f?v2:0.0f;v3=fmaf((v3-m)*is,g,be);v3=v3>0.0f?v3:0.0f;val=0.25f*(v0+v1+v2+v3);}sh[ty][tx]=val;__syncthreads();int s2=ty*4+(tx>>3);int cc=blockIdx.y*8+(tx&7);int pp=p0+s2;if(pp<P&&cc<C)tmp[pp*C+cc]=sh[tx&7][s2];}
__global__ void conv1x1_nhwc_out2_kernel(const float* __restrict__ tmp,const float* __restrict__ w,float* __restrict__ out,int P,int C,int O){int ox=threadIdx.x;int py=threadIdx.y;int p0=blockIdx.x*8+py;int p1=p0+4;extern __shared__ float sh[];if(ox<C){if(p0<P)sh[py*C+ox]=tmp[p0*C+ox];if(p1<P)sh[(py+4)*C+ox]=tmp[p1*C+ox];}__syncthreads();if(ox<O){float s0=0.0f,s1=0.0f;for(int c=0;c<32;++c){if(c<C){float ww=w[ox*C+c];s0=fmaf(ww,sh[py*C+c],s0);s1=fmaf(ww,sh[(py+4)*C+c],s1);}}if(p0<P)out[p0*O+ox]=s0;if(p1<P)out[p1*O+ox]=s1;}}
__global__ void make_eval_inv_kernel(const float* __restrict__ running_mean,const float* __restrict__ running_var,float* __restrict__ mean,float* __restrict__ invstd,int C,double eps){int c=blockIdx.x*blockDim.x+threadIdx.x;if(c<C){mean[c]=running_mean[c];invstd[c]=(float)rsqrt((double)running_var[c]+eps);}}
torch::Tensor transition_forward(torch::Tensor x,torch::Tensor gamma,torch::Tensor beta,torch::Tensor running_mean,torch::Tensor running_var,torch::Tensor weight,bool training,double eps,double momentum){CHECK_CUDA(x);CHECK_CUDA(gamma);CHECK_CUDA(beta);CHECK_CUDA(weight);CHECK_CONTIGUOUS(x);CHECK_CONTIGUOUS(weight);int N=(int)x.size(0),C=(int)x.size(1),H=(int)x.size(2),W=(int)x.size(3);int O=(int)weight.size(0);int H2=H/2,W2=W/2,P=N*H2*W2;auto out=torch::empty_strided({N,O,H2,W2},{H2*W2*O,1,W2*O,O},x.options());auto mean=torch::empty({C},x.options());auto invstd=torch::empty({C},x.options());cudaStream_t stream=at::cuda::getCurrentCUDAStream();if(training){const int NB=256;auto optsD=torch::TensorOptions().device(x.device()).dtype(torch::kFloat64);auto psum=torch::empty({C,NB},optsD);auto psq=torch::empty({C,NB},optsD);partial_stats_kernel<<<dim3(C,NB),256,2*256*sizeof(double),stream>>>(x.data_ptr<float>(),psum.data_ptr<double>(),psq.data_ptr<double>(),N,C,H,W,N*H*W,NB);final_stats_kernel<<<C,256,2*256*sizeof(double),stream>>>(psum.data_ptr<double>(),psq.data_ptr<double>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),C,N*H*W,NB,eps,momentum);}else{make_eval_inv_kernel<<<(C+255)/256,256,0,stream>>>(running_mean.data_ptr<float>(),running_var.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),C,eps);}auto tmp=torch::empty({P,C},x.options());pool_bn_relu_nhwc_kernel<<<dim3((P+31)/32,(C+7)/8),dim3(32,8),0,stream>>>(x.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),tmp.data_ptr<float>(),N,C,H,W);conv1x1_nhwc_out2_kernel<<<(P+7)/8,dim3(64,4),8*C*sizeof(float),stream>>>(tmp.data_ptr<float>(),weight.data_ptr<float>(),out.data_ptr<float>(),P,C,O);return out;}
'''
_ext=load_inline(name='transition_fused_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['transition_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], with_cuda=True, verbose=False)
class ModelNew(nn.Module):
def __init__(self,num_input_features:int,num_output_features:int):
super(ModelNew,self).__init__(); self.transition=nn.Sequential(nn.BatchNorm2d(num_input_features),nn.ReLU(inplace=True),nn.Conv2d(num_input_features,num_output_features,kernel_size=1,bias=False),nn.AvgPool2d(kernel_size=2,stride=2))
def forward(self,x):
bn=self.transition[0]; conv=self.transition[2]
return _ext.transition_forward(x.contiguous(),bn.weight,bn.bias,bn.running_mean,bn.running_var,conv.weight.contiguous(),self.training,bn.eps,bn.momentum if bn.momentum is not None else 0.1)