final submitted kernel (114 lines, 5,155 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp = r'''
#include <torch/extension.h>
torch::Tensor dsconv_forward(torch::Tensor x, torch::Tensor pw_w, torch::Tensor pw_b, torch::Tensor dw_w, torch::Tensor dw_b, torch::Tensor pr_w, torch::Tensor pr_b);
'''
_cuda = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cublas_v2.h>
static inline void ck(cublasStatus_t s) { TORCH_CHECK(s == CUBLAS_STATUS_SUCCESS, "cuBLAS error"); }
__global__ void depthwise_gelu_kernel(const float* __restrict__ pw,
const float* __restrict__ pw_b,
const float* __restrict__ dw_w,
const float* __restrict__ dw_b,
float* __restrict__ act,
int N, int E, int L, int K) {
long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
long long total = (long long)N * E * L;
if (idx >= total) return;
int l = (int)(idx % L);
long long t = idx / L;
int e = (int)(t % E);
int pad = K >> 1;
const float* base = pw + t * (long long)L;
const float* w = dw_w + (long long)e * K;
float sum = dw_b[e];
float bp = pw_b[e];
#pragma unroll
for (int j = 0; j < 9; ++j) {
int ll = l + j - pad;
if (j < K && (unsigned)ll < (unsigned)L) sum = fmaf(w[j], base[ll] + bp, sum);
}
act[idx] = 0.5f * sum * (1.0f + erff(sum * 0.70710678118654752440f));
}
__global__ void add_residual_bias_kernel(float* __restrict__ out,
const float* __restrict__ residual,
const float* __restrict__ bias,
long long total, int C, int L) {
long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
if (idx < total) {
int c = (int)((idx / L) % C);
out[idx] = out[idx] + bias[c] + residual[idx];
}
}
torch::Tensor dsconv_forward(torch::Tensor x,
torch::Tensor pw_w, torch::Tensor pw_b,
torch::Tensor dw_w, torch::Tensor dw_b,
torch::Tensor pr_w, torch::Tensor pr_b) {
c10::cuda::CUDAGuard device_guard(x.device());
int N = (int)x.size(0), C = (int)x.size(1), L = (int)x.size(2);
int E = (int)pw_w.size(0), K = (int)dw_w.size(2);
auto pw = torch::empty({N, E, L}, x.options());
auto act = torch::empty({N, E, L}, x.options());
auto out = torch::empty({N, C, L}, x.options());
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
ck(cublasSetStream(handle, stream));
ck(cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH));
const float alpha = 1.0f, beta = 0.0f;
ck(cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N,
L, E, C, &alpha,
x.data_ptr<float>(), L, (long long)C * L,
pw_w.data_ptr<float>(), C, 0,
&beta, pw.data_ptr<float>(), L, (long long)E * L, N));
long long etotal = (long long)N * E * L;
depthwise_gelu_kernel<<<(etotal + 255) / 256, 256, 0, stream>>>(
pw.data_ptr<float>(), pw_b.data_ptr<float>(), dw_w.data_ptr<float>(), dw_b.data_ptr<float>(),
act.data_ptr<float>(), N, E, L, K);
C10_CUDA_KERNEL_LAUNCH_CHECK();
ck(cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N,
L, C, E, &alpha,
act.data_ptr<float>(), L, (long long)E * L,
pr_w.data_ptr<float>(), E, 0,
&beta, out.data_ptr<float>(), L, (long long)C * L, N));
long long ototal = (long long)N * C * L;
add_residual_bias_kernel<<<(ototal + 255) / 256, 256, 0, stream>>>(
out.data_ptr<float>(), x.data_ptr<float>(), pr_b.data_ptr<float>(), ototal, C, L);
C10_CUDA_KERNEL_LAUNCH_CHECK();
return out;
}
'''
_ext = load_inline(name='dsconv_ext_final1', cpp_sources=_cpp, cuda_sources=_cuda,
functions=['dsconv_forward'], verbose=False,
extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])
class ModelNew(nn.Module):
def __init__(self, channels, expanded_channels, kernel_size):
super(ModelNew, self).__init__()
self.pointwise = nn.Conv1d(channels, expanded_channels, 1)
self.depthwise = nn.Conv1d(expanded_channels, expanded_channels, kernel_size,
padding=kernel_size // 2, groups=expanded_channels)
self.activation = nn.GELU()
self.project_back = nn.Conv1d(expanded_channels, channels, 1)
def forward(self, x):
return _ext.dsconv_forward(x, self.pointwise.weight, self.pointwise.bias,
self.depthwise.weight, self.depthwise.bias,
self.project_back.weight, self.project_back.bias)