reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Performs a standard 1D convolution operation with asymmetric input and a square kernel, potentially dilated and strided.
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.
stride (int, optional): Stride of the convolution. Defaults to 1.
dilation (int, optional): Spacing between kernel elements. 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, dilation: int = 1, bias: bool = False):
super(Model, self).__init__()
self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, bias=bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Performs the 1D convolution.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, length).
Returns:
torch.Tensor: Output tensor of shape (batch_size, out_channels, length_out).
"""
return self.conv1d(x)
# Test code
batch_size = 64
in_channels = 64
out_channels = 128
kernel_size = 3
# longer signal
length = 524280
stride = 3
dilation = 4
def get_inputs():
x = torch.rand(batch_size, in_channels, length)
return [x]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, stride, dilation]final submitted kernel (26 lines, 4,930 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>
void conv1d_forward_cuda(torch::Tensor x, torch::Tensor w, c10::optional<torch::Tensor> b, torch::Tensor y, int stride, int dilation);
torch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor w, c10::optional<torch::Tensor> b, int stride, int dilation){auto Lout=(x.size(2)-dilation*(w.size(2)-1)-1)/stride+1;auto y=torch::empty({x.size(0),w.size(0),Lout},x.options());conv1d_forward_cuda(x.contiguous(),w.contiguous(),b,y,stride,dilation);return y;}
'''
_cuda_src=r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
constexpr int TL=128,OC_TILE=16,THREADS=512;
__device__ __forceinline__ float tf32f(float a){ unsigned r; asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(r) : "f"(a)); return __uint_as_float(r); }
__global__ __launch_bounds__(THREADS,2) void conv1d_k3_s3_d4_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ y,int N,int Cin,int L,int Cout,int Lout,int has_bias){int l0=blockIdx.x*TL,oc0=blockIdx.y*OC_TILE,n=blockIdx.z;int valid_tl=Lout-l0;if(valid_tl>TL)valid_tl=TL;if(valid_tl<=0)return;int span=(valid_tl-1)*3+9;extern __shared__ float sx[];int tid=threadIdx.x;int total=Cin*span;const float* xb=x+((long long)n*Cin)*L+(long long)l0*3;for(int idx=tid;idx<total;idx+=THREADS){int c=idx/span;int p=idx-c*span;sx[idx]=tf32f(xb[(long long)c*L+p]);}__syncthreads();int warp=tid>>5,lane=tid&31,oc=oc0+warp;if(oc>=Cout)return;float acc0=0,acc1=0,acc2=0,acc3=0;int lr0=lane,lr1=lane+32,lr2=lane+64,lr3=lane+96;const float* ww=w+(long long)oc*Cin*3;for(int ic=0;ic<Cin;++ic){const float* sxc=sx+ic*span;float w0=0,w1=0,w2=0;if(lane==0){w0=tf32f(ww[ic*3]);w1=tf32f(ww[ic*3+1]);w2=tf32f(ww[ic*3+2]);}w0=__shfl_sync(0xffffffff,w0,0);w1=__shfl_sync(0xffffffff,w1,0);w2=__shfl_sync(0xffffffff,w2,0);if(lr0<valid_tl){int base=lr0*3;acc0=fmaf(sxc[base],w0,acc0);acc0=fmaf(sxc[base+4],w1,acc0);acc0=fmaf(sxc[base+8],w2,acc0);}if(lr1<valid_tl){int base=lr1*3;acc1=fmaf(sxc[base],w0,acc1);acc1=fmaf(sxc[base+4],w1,acc1);acc1=fmaf(sxc[base+8],w2,acc1);}if(lr2<valid_tl){int base=lr2*3;acc2=fmaf(sxc[base],w0,acc2);acc2=fmaf(sxc[base+4],w1,acc2);acc2=fmaf(sxc[base+8],w2,acc2);}if(lr3<valid_tl){int base=lr3*3;acc3=fmaf(sxc[base],w0,acc3);acc3=fmaf(sxc[base+4],w1,acc3);acc3=fmaf(sxc[base+8],w2,acc3);}}if(has_bias){float bv=bias[oc];acc0+=bv;acc1+=bv;acc2+=bv;acc3+=bv;}float* yb=y+((long long)n*Cout+oc)*Lout+l0;if(lr0<valid_tl)yb[lr0]=acc0;if(lr1<valid_tl)yb[lr1]=acc1;if(lr2<valid_tl)yb[lr2]=acc2;if(lr3<valid_tl)yb[lr3]=acc3;}
__global__ void conv1d_generic_kernel(const float*x,const float*w,const float*bias,float*y,int N,int Cin,int L,int Cout,int K,int Lout,int stride,int dilation,int has_bias){long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x,total=(long long)N*Cout*Lout;for(;idx<total;idx+=(long long)gridDim.x*blockDim.x){int l=idx%Lout;long long t=idx/Lout;int oc=t%Cout;int n=t/Cout;float acc=has_bias?bias[oc]:0;for(int ic=0;ic<Cin;++ic){const float*xp=x+((long long)n*Cin+ic)*L+(long long)l*stride;const float*wp=w+((long long)oc*Cin+ic)*K;for(int k=0;k<K;++k)acc=fmaf(xp[k*dilation],wp[k],acc);}y[idx]=acc;}}
void conv1d_forward_cuda(torch::Tensor x,torch::Tensor w,c10::optional<torch::Tensor> b,torch::Tensor y,int stride,int dilation){int N=x.size(0),Cin=x.size(1),L=x.size(2),Cout=w.size(0),K=w.size(2),Lout=y.size(2);const float*bp=nullptr;int hb=0;if(b.has_value()&&b.value().defined()){bp=b.value().data_ptr<float>();hb=1;}cudaStream_t stream=at::cuda::getCurrentCUDAStream();if(Cin==64&&Cout==128&&K==3&&stride==3&&dilation==4){dim3 block(THREADS),grid((Lout+TL-1)/TL,(Cout+OC_TILE-1)/OC_TILE,N);size_t shmem=(size_t)Cin*((TL-1)*3+9)*sizeof(float);cudaFuncSetAttribute(conv1d_k3_s3_d4_kernel,cudaFuncAttributeMaxDynamicSharedMemorySize,(int)shmem);conv1d_k3_s3_d4_kernel<<<grid,block,shmem,stream>>>(x.data_ptr<float>(),w.data_ptr<float>(),bp,y.data_ptr<float>(),N,Cin,L,Cout,Lout,hb);}else{long long total=(long long)N*Cout*Lout;int blocks=(int)min((long long)65535,(total+255)/256);conv1d_generic_kernel<<<blocks,256,0,stream>>>(x.data_ptr<float>(),w.data_ptr<float>(),bp,y.data_ptr<float>(),N,Cin,L,Cout,K,Lout,stride,dilation,hb);}}
'''
_ext=load_inline(name='conv1d_asym_ext_final',cpp_sources=[_cpp_src],cuda_sources=[_cuda_src],functions=['conv1d_forward'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)
class ModelNew(nn.Module):
def __init__(self,in_channels:int,out_channels:int,kernel_size:int,stride:int=1,dilation:int=1,bias:bool=False):
super().__init__();self.conv1d=nn.Conv1d(in_channels,out_channels,kernel_size,stride=stride,dilation=dilation,bias=bias);self.stride=stride;self.dilation=dilation
def forward(self,x):
return _ext.conv1d_forward(x,self.conv1d.weight,self.conv1d.bias,self.stride,self.dilation)