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.
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 convolution kernel.
stride (int, optional): Stride of the convolution. Defaults to 1.
padding (int, optional): Padding applied to the input. Defaults to 0.
dilation (int, optional): Spacing between kernel elements. Defaults to 1.
groups (int, optional): Number of blocked connections from input channels to output channels. 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, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
super(Model, self).__init__()
self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, 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 = 32
in_channels = 64
out_channels = 128
kernel_size = 3
length = 131072
def get_inputs():
x = torch.rand(batch_size, in_channels, length)
return [x]
def get_init_inputs():
return [in_channels, out_channels, kernel_size] # Provide in_channels, out_channels, kernel_size for initializationfinal submitted kernel (41 lines, 4,352 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 conv1d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias,int64_t stride,int64_t padding,int64_t dilation,int64_t groups);
'''
_cuda_src=r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void conv1d_k3_opt8_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ y,int L,int has_bias){
constexpr int CIN=64,COUT=128,TILE_L=128,TILE_OC=16,K=3,IN_TILE=130; int Lout=L-2; int tx=threadIdx.x, ty=threadIdx.y, pos0=blockIdx.x*TILE_L, oc0=blockIdx.y*TILE_OC, n=blockIdx.z; extern __shared__ float smem[]; float* sx=smem; float* sw=sx+CIN*IN_TILE; int tid=ty*blockDim.x+tx; int nt=blockDim.x*blockDim.y;
for(int idx=tid;idx<CIN*IN_TILE;idx+=nt){int c=idx/IN_TILE; int t=idx-c*IN_TILE; int ip=pos0+t; sx[idx]=(ip<L)?x[((n*CIN+c)*L)+ip]:0.f;}
for(int idx=tid;idx<TILE_OC*CIN*K;idx+=nt){int rem=idx; int oo=rem/(CIN*K); rem-=oo*(CIN*K); int c=rem/K; int k=rem-c*K; sw[idx]=w[((oc0+oo)*CIN+c)*K+k];}
__syncthreads(); int pos=pos0+tx; if(pos<Lout){
float a[8];
#pragma unroll
for(int m=0;m<8;++m){int oc=oc0+ty+m*2; a[m]=has_bias?bias[oc]:0.f;}
#pragma unroll
for(int c=0;c<CIN;++c){int si=c*IN_TILE+tx; float x0=sx[si],x1=sx[si+1],x2=sx[si+2];
#pragma unroll
for(int m=0;m<8;++m){int wb=(ty+m*2)*(CIN*K)+c*K; a[m]=fmaf(x0,sw[wb],a[m]); a[m]=fmaf(x1,sw[wb+1],a[m]); a[m]=fmaf(x2,sw[wb+2],a[m]);}
}
int base=(n*COUT)*Lout+pos;
#pragma unroll
for(int m=0;m<8;++m){int oc=oc0+ty+m*2; y[base+oc*Lout]=a[m];}
}}
__global__ void conv1d_generic_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 K,int stride,int padding,int dilation,int groups,int Lout,int has_bias){int idx=blockIdx.x*blockDim.x+threadIdx.x,total=N*Cout*Lout;if(idx>=total)return;int pos=idx%Lout,tmp=idx/Lout,oc=tmp%Cout,n=tmp/Cout,cinpg=Cin/groups,coutpg=Cout/groups,g=oc/coutpg,cbase=g*cinpg;float acc=has_bias?bias[oc]:0.f;for(int ci=0;ci<cinpg;++ci){int c=cbase+ci;for(int k=0;k<K;++k){int ip=pos*stride-padding+k*dilation;if((unsigned)ip<(unsigned)L)acc=fmaf(x[(n*Cin+c)*L+ip],w[(oc*cinpg+ci)*K+k],acc);}}y[idx]=acc;}
torch::Tensor conv1d_forward(torch::Tensor x,torch::Tensor w,torch::Tensor bias,int64_t stride,int64_t padding,int64_t dilation,int64_t groups){int N=x.size(0),Cin=x.size(1),L=x.size(2),Cout=w.size(0),K=w.size(2);int Lout=(L+2*(int)padding-(int)dilation*(K-1)-1)/(int)stride+1;auto y=torch::empty({N,Cout,Lout},x.options());int has_bias=bias.numel()!=0;cudaStream_t stream=at::cuda::getCurrentCUDAStream();if(N==32&&Cin==64&&Cout==128&&K==3&&stride==1&&padding==0&&dilation==1&&groups==1&&x.is_contiguous()&&w.is_contiguous()){dim3 block(128,2);dim3 grid((Lout+127)/128,8,N);size_t smem=(64*130+16*64*3)*sizeof(float);conv1d_k3_opt8_kernel<<<grid,block,smem,stream>>>(x.data_ptr<float>(),w.data_ptr<float>(),has_bias?bias.data_ptr<float>():nullptr,y.data_ptr<float>(),L,has_bias);}else{int total=N*Cout*Lout,threads=256;conv1d_generic_kernel<<<(total+threads-1)/threads,threads,0,stream>>>(x.data_ptr<float>(),w.data_ptr<float>(),has_bias?bias.data_ptr<float>():nullptr,y.data_ptr<float>(),N,Cin,L,Cout,K,(int)stride,(int)padding,(int)dilation,(int)groups,Lout,has_bias);}return y;}
'''
_ext=load_inline(name='conv1d_k3_ext_v4',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['conv1d_forward'],extra_cuda_cflags=['-O3','--use_fast_math'],extra_cflags=['-O3'],verbose=False)
class ModelNew(nn.Module):
def __init__(self,in_channels:int,out_channels:int,kernel_size:int,stride:int=1,padding:int=0,dilation:int=1,groups:int=1,bias:bool=False):
super(ModelNew,self).__init__(); self.conv1d=nn.Conv1d(in_channels,out_channels,kernel_size,stride=stride,padding=padding,dilation=dilation,groups=groups,bias=bias); self.stride=stride; self.padding=padding; self.dilation=dilation; self.groups=groups
def forward(self,x:torch.Tensor)->torch.Tensor:
b=self.conv1d.bias
if b is None: b=torch.empty(0,device=x.device,dtype=x.dtype)
return _ext.conv1d_forward(x.contiguous(),self.conv1d.weight.contiguous(),b,self.stride,self.padding,self.dilation,self.groups)