reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Performs a transposed 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.
output_padding (int, optional): Additional size added to one side of the output shape. Defaults to 0.
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, output_padding: int = 0, groups: int = 1, bias: bool = False):
super(Model, self).__init__()
self.conv1d_transpose = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, groups=groups, bias=bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Performs the transposed 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_transpose(x)
# Test code
batch_size = 64
in_channels = 128
out_channels = 128
kernel_size = 3
# much larger signal length for heavier workload
length = 65536
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 (81 lines, 4,960 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 convt1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,int64_t stride,int64_t padding,int64_t output_padding,int64_t groups);
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDABlas.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
__global__ void pack_w_k3(const float* __restrict__ w, float* __restrict__ wp, int CI, int CO){
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = 3 * CI * CO;
if(idx >= total) return;
int tap = idx / (CI*CO);
int r = idx - tap*(CI*CO);
int ic = r / CO;
int oc = r - ic*CO;
wp[tap*CI*CO + oc + ic*CO] = w[(ic*CO + oc)*3 + tap];
}
__global__ void convt1d_generic_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ y,int N,int CI,int L,int CO,int K,int stride,int padding,int groups,int Lout,int has_bias){
long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x,total=(long long)N*CO*Lout; if(idx>=total)return; int t=idx%Lout; int oc=(idx/Lout)%CO; int n=idx/((long long)CO*Lout); int co_per_g=CO/groups,ci_per_g=CI/groups,g=oc/co_per_g,ocg=oc-g*co_per_g; float acc=has_bias?b[oc]:0.0f; for(int icg=0;icg<ci_per_g;++icg){int ic=g*ci_per_g+icg; for(int k=0;k<K;++k){int u=t+padding-k; if(u%stride==0){int it=u/stride; if(it>=0&&it<L) acc=fmaf(x[(n*CI+ic)*L+it],w[(ic*co_per_g+ocg)*K+k],acc);}}} y[idx]=acc;
}
torch::Tensor convt1d_forward(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,int64_t stride,int64_t padding,int64_t output_padding,int64_t groups){
x=x.contiguous(); weight=weight.contiguous(); if(bias.numel()>0) bias=bias.contiguous();
int N=x.size(0), CI=x.size(1), L=x.size(2), K=weight.size(2), CO=weight.size(1)*groups;
int Lout=(L-1)*(int)stride-2*(int)padding+K+(int)output_padding;
auto y=torch::empty({N,CO,Lout},x.options());
int has_bias=bias.numel()>0;
cudaStream_t stream=at::cuda::getCurrentCUDAStream();
if(x.scalar_type()==at::kFloat && weight.scalar_type()==at::kFloat && K==3 && stride==1 && padding==0 && output_padding==0 && groups==1 && !has_bias){
auto wp=torch::empty({3,CI,CO}, weight.options());
int total=3*CI*CO;
pack_w_k3<<<(total+255)/256,256,0,stream>>>(weight.data_ptr<float>(), wp.data_ptr<float>(), CI, CO);
cudaMemsetAsync(y.data_ptr<float>(), 0, (size_t)N*CO*Lout*sizeof(float), stream);
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
cublasSetStream(handle, stream);
cublasMath_t old_mode;
cublasGetMathMode(handle, &old_mode);
cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH);
float alpha=1.0f, beta=1.0f;
long long strideA=(long long)CI*L;
long long strideB=0;
long long strideC=(long long)CO*Lout;
for(int tap=0; tap<3; ++tap){
const float* A=x.data_ptr<float>();
const float* B=wp.data_ptr<float>() + (long long)tap*CI*CO;
float* C=y.data_ptr<float>() + tap;
cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_T,
L, CO, CI,
&alpha,
A, L, strideA,
B, CO, strideB,
&beta,
C, Lout, strideC,
N);
}
cublasSetMathMode(handle, old_mode);
} else {
int threads=256; long long total=(long long)N*CO*Lout; const float* bp=has_bias?bias.data_ptr<float>():nullptr;
convt1d_generic_kernel<<<(total+threads-1)/threads,threads,0,stream>>>(x.data_ptr<float>(),weight.data_ptr<float>(),bp,y.data_ptr<float>(),N,CI,L,CO,K,(int)stride,(int)padding,(int)groups,Lout,has_bias);
}
return y;
}
'''
_ext=load_inline(name='convt1d_ext_final',cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=['convt1d_forward'],extra_cuda_cflags=['-O3'],extra_ldflags=['-lcublas'],verbose=False)
class ModelNew(nn.Module):
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):
super().__init__(); self.conv1d_transpose=nn.ConvTranspose1d(in_channels,out_channels,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
def forward(self,x):
bias=self.conv1d_transpose.bias
if bias is None: bias=torch.empty(0,device=x.device,dtype=x.dtype)
return _ext.convt1d_forward(x,self.conv1d_transpose.weight,bias,self.stride,self.padding,self.output_padding,self.groups)