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 with asymmetric input and square kernel.
Supports padding, striding, and dilation.
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.
padding (int, optional): Padding applied to the input. Defaults to 0.
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, padding: int = 0, dilation: int = 1, bias: bool = False):
super(Model, self).__init__()
self.conv1d_transpose = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, 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 = 16
in_channels = 32
out_channels = 64
kernel_size = 3
# long sequence
length = 131072
stride = 2
padding = 1
dilation = 2
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, padding, dilation]final submitted kernel (155 lines, 6,303 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 conv_t1d_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
int64_t stride, int64_t padding, int64_t dilation, bool has_bias);
torch::Tensor conv_t1d(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
int64_t stride, int64_t padding, int64_t dilation, bool has_bias) {
return conv_t1d_cuda(x, weight, bias, stride, padding, dilation, has_bias);
}
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define BX 128
#define OT 16
__device__ __forceinline__ float rtf32(float x) {
#if __CUDA_ARCH__ >= 800
unsigned int y;
asm("cvt.rna.tf32.f32 %0, %1;" : "=b"(y) : "f"(x));
return __uint_as_float(y);
#else
return x;
#endif
}
__global__ void opt_kernel(const float* __restrict__ x,
const float* __restrict__ w,
const float* __restrict__ bias,
float* __restrict__ out,
int Cin, int Cout, int L, int Lout, bool has_bias) {
int j = blockIdx.x * BX + threadIdx.x;
int ocb = blockIdx.y * OT;
int b = blockIdx.z;
if (j >= L) return;
float acc[OT];
#pragma unroll
for (int o = 0; o < OT; ++o) {
int oc = ocb + o;
acc[o] = (has_bias && oc < Cout) ? bias[oc] : 0.0f;
}
const float* xb = x + (long long)b * Cin * L;
for (int ic = 0; ic < Cin; ++ic) {
const float* xic = xb + (long long)ic * L;
float xv0 = rtf32((j + 1 < L) ? xic[j + 1] : 0.0f);
float xv1 = rtf32(xic[j]);
float xv2 = rtf32((j > 0) ? xic[j - 1] : 0.0f);
const float* wic = w + (long long)ic * Cout * 3 + ocb * 3;
#pragma unroll
for (int o = 0; o < OT; ++o) {
int oc = ocb + o;
if (oc < Cout) {
const float* wp = wic + o * 3;
acc[o] = fmaf(xv0, rtf32(wp[0]), acc[o]);
acc[o] = fmaf(xv1, rtf32(wp[1]), acc[o]);
acc[o] = fmaf(xv2, rtf32(wp[2]), acc[o]);
}
}
}
#pragma unroll
for (int o = 0; o < OT; ++o) {
int oc = ocb + o;
if (oc < Cout) {
float even = has_bias ? bias[oc] : 0.0f;
float* yo = out + ((long long)b * Cout + oc) * Lout;
yo[2 * j] = even;
yo[2 * j + 1] = acc[o];
if (j == L - 1) yo[2 * L] = even;
}
}
}
__global__ void generic_kernel(const float* __restrict__ x,
const float* __restrict__ w,
const float* __restrict__ bias,
float* __restrict__ out,
long long total, int Cin, int Cout, int K, int L, int Lout,
int stride, int padding, int dilation, bool has_bias) {
long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total) return;
int pos = idx % Lout;
long long t = idx / Lout;
int oc = t % Cout;
int b = t / Cout;
float sum = has_bias ? bias[oc] : 0.0f;
for (int ic = 0; ic < Cin; ++ic) {
const float* xic = x + ((long long)b * Cin + ic) * L;
const float* wic = w + ((long long)ic * Cout + oc) * K;
for (int k = 0; k < K; ++k) {
int num = pos + padding - k * dilation;
if (num >= 0 && num % stride == 0) {
int ip = num / stride;
if (ip >= 0 && ip < L) sum = fmaf(xic[ip], wic[k], sum);
}
}
}
out[idx] = sum;
}
torch::Tensor conv_t1d_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
int64_t stride64, int64_t padding64, int64_t dilation64, bool has_bias) {
int B = (int)x.size(0), Cin = (int)x.size(1), L = (int)x.size(2);
int Cout = (int)weight.size(1), K = (int)weight.size(2);
int stride = (int)stride64, padding = (int)padding64, dilation = (int)dilation64;
int Lout = (L - 1) * stride - 2 * padding + dilation * (K - 1) + 1;
auto out = torch::empty({B, Cout, Lout}, x.options());
cudaStream_t s = at::cuda::getCurrentCUDAStream();
if (stride == 2 && padding == 1 && dilation == 2 && K == 3 && Lout == 2 * L + 1) {
dim3 block(BX);
dim3 grid((L + BX - 1) / BX, (Cout + OT - 1) / OT, B);
opt_kernel<<<grid, block, 0, s>>>(x.data_ptr<float>(), weight.data_ptr<float>(),
has_bias ? bias.data_ptr<float>() : nullptr,
out.data_ptr<float>(), Cin, Cout, L, Lout, has_bias);
} else {
long long total = (long long)B * Cout * Lout;
generic_kernel<<<(total + 255) / 256, 256, 0, s>>>(x.data_ptr<float>(), weight.data_ptr<float>(),
has_bias ? bias.data_ptr<float>() : nullptr,
out.data_ptr<float>(), total, Cin, Cout, K, L, Lout,
stride, padding, dilation, has_bias);
}
return out;
}
'''
_ext = load_inline(name='conv_t1d_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src,
functions=['conv_t1d'], 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, padding: int = 0, dilation: int = 1, bias: bool = False):
super(ModelNew, self).__init__()
self.conv1d_transpose = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias)
self.stride = stride
self.padding = padding
self.dilation = dilation
def forward(self, x: torch.Tensor) -> torch.Tensor:
b = self.conv1d_transpose.bias
has_bias = b is not None
if b is None:
b = torch.empty(0, device=x.device, dtype=x.dtype)
return _ext.conv_t1d(x, self.conv1d_transpose.weight, b, self.stride, self.padding, self.dilation, has_bias)