call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n\ntorch::Tensor depthwise_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, bool has_bias, int64_t stride, int64_t padding);\n\ntorch::Tensor depthwise_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, bool has_bias, int64_t stride, int64_t padding) {\n return depthwise_forward_cuda(x, weight, bias, has_bias, stride, padding);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void depthwise_k9s1p4_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ bias,\n float* __restrict__ y,\n int total, int B, int C, int L, bool has_bias) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride_grid = blockDim.x * gridDim.x;\n for (; idx < total; idx += stride_grid) {\n int l = idx % L;\n int t = idx / L;\n int c = t % C;\n int b = t / C;\n const float* xp = x + ((b * C + c) * L);\n const float* wp = w + c * 9;\n float s = has_bias ? bias[c] : 0.0f;\n // padding=4, output length=input length. Zero outside [0,L).\n int base = l - 4;\n #pragma unroll\n for (int k = 0; k < 9; ++k) {\n int xi = base + k;\n if ((unsigned)xi < (unsigned)L) {\n s = fmaf(xp[xi], wp[k], s);\n }\n }\n y[idx] = s;\n }\n}\n\n__global__ void depthwise_generic_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ bias,\n float* __restrict__ y,\n int total, int B, int C, int Lin, int Lout, int K,\n int stride, int padding, bool has_bias) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int stride_grid = blockDim.x * gridDim.x;\n for (; idx < total; idx += stride_grid) {\n int l = idx % Lout;\n int t = idx / Lout;\n int c = t % C;\n int b = t / C;\n const float* xp = x + ((b * C + c) * Lin);\n const float* wp = w + c * K;\n float s = has_bias ? bias[c] : 0.0f;\n int base = l * stride - padding;\n for (int k = 0; k < K; ++k) {\n int xi = base + k;\n if ((unsigned)xi < (unsigned)Lin) s = fmaf(xp[xi], wp[k], s);\n }\n y[idx] = s;\n }\n}\n\ntorch::Tensor depthwise_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, bool has_bias, int64_t stride64, int64_t padding64) {\n const int B = (int)x.size(0);\n const int C = (int)x.size(1);\n const int Lin = (int)x.size(2);\n const int K = (int)weight.size(2);\n const int stride = (int)stride64;\n const int padding = (int)padding64;\n const int Lout = (Lin + 2 * padding - K) / stride + 1;\n auto y = torch::empty({B, C, Lout}, x.options());\n int total = B * C * Lout;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n if (blocks > 65535) blocks = 65535;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (K == 9 && stride == 1 && padding == 4 && Lout == Lin) {\n depthwise_k9s1p4_kernel<<<blocks, threads, 0, stream>>>(\n x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,\n y.data_ptr<float>(), total, B, C, Lin, has_bias);\n } else {\n depthwise_generic_kernel<<<blocks, threads, 0, stream>>>(\n x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,\n y.data_ptr<float>(), total, B, C, Lin, Lout, K, stride, padding, has_bias);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='dwsep_ext_v1', cpp_sources=_cpp_src, cuda_sources=_cuda_src,\n functions=['depthwise_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):\n super(ModelNew, self).__init__()\n self.depthwise = nn.Conv1d(in_channels, in_channels, kernel_size, stride=stride, padding=padding, groups=in_channels, bias=bias)\n self.pointwise = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=bias)\n self.stride = stride\n self.padding = padding\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n b = self.depthwise.bias\n if b is None:\n b = torch.empty(0, device=x.device, dtype=x.dtype)\n has_b = False\n else:\n has_b = True\n tmp = _ext.depthwise_forward(x.contiguous(), self.depthwise.weight.contiguous(), b, has_b, self.stride, self.padding)\n w = self.pointwise.weight.squeeze(2)\n out = torch.bmm(w.unsqueeze(0).expand(tmp.size(0), -1, -1), tmp)\n if self.pointwise.bias is not None:\n out = out + self.pointwise.bias.view(1, -1, 1)\n return out\n"
}