call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,\n int64_t stride, int64_t padding, int64_t dilation, int64_t groups);\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x \" must be float32\")\n\n__global__ void conv1d_k3_s1_p0_oc8_kernel(\n const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ b,\n float* __restrict__ y,\n int N, int Cin, int L, int Cout, int Lout, int has_bias)\n{\n constexpr int OC_TILE = 8;\n extern __shared__ float sw[];\n const int oc_base = blockIdx.z * OC_TILE;\n const int oc_eff = min(OC_TILE, Cout - oc_base);\n const int tid = threadIdx.x;\n const int wcount = oc_eff * Cin * 3;\n\n for (int idx = tid; idx < wcount; idx += blockDim.x) {\n int t = idx;\n int k = t % 3; t /= 3;\n int ci = t % Cin; t /= Cin;\n int oj = t;\n sw[idx] = w[((oc_base + oj) * Cin + ci) * 3 + k];\n }\n __syncthreads();\n\n const int l = blockIdx.x * blockDim.x + tid;\n const int n = blockIdx.y;\n if (l >= Lout) return;\n\n float a0=0.f,a1=0.f,a2=0.f,a3=0.f,a4=0.f,a5=0.f,a6=0.f,a7=0.f;\n if (has_bias) {\n if (oc_eff > 0) a0 = b[oc_base+0];\n if (oc_eff > 1) a1 = b[oc_base+1];\n if (oc_eff > 2) a2 = b[oc_base+2];\n if (oc_eff > 3) a3 = b[oc_base+3];\n if (oc_eff > 4) a4 = b[oc_base+4];\n if (oc_eff > 5) a5 = b[oc_base+5];\n if (oc_eff > 6) a6 = b[oc_base+6];\n if (oc_eff > 7) a7 = b[oc_base+7];\n }\n\n const float* xbase = x + ((long long)n * Cin) * L + l;\n for (int ci = 0; ci < Cin; ++ci) {\n const float v0 = xbase[(long long)ci * L + 0];\n const float v1 = xbase[(long long)ci * L + 1];\n const float v2 = xbase[(long long)ci * L + 2];\n const int si = ci * 3;\n if (oc_eff > 0) { const float* ww = sw + 0*Cin*3 + si; a0 = fmaf(v0, ww[0], a0); a0 = fmaf(v1, ww[1], a0); a0 = fmaf(v2, ww[2], a0); }\n if (oc_eff > 1) { const float* ww = sw + 1*Cin*3 + si; a1 = fmaf(v0, ww[0], a1); a1 = fmaf(v1, ww[1], a1); a1 = fmaf(v2, ww[2], a1); }\n if (oc_eff > 2) { const float* ww = sw + 2*Cin*3 + si; a2 = fmaf(v0, ww[0], a2); a2 = fmaf(v1, ww[1], a2); a2 = fmaf(v2, ww[2], a2); }\n if (oc_eff > 3) { const float* ww = sw + 3*Cin*3 + si; a3 = fmaf(v0, ww[0], a3); a3 = fmaf(v1, ww[1], a3); a3 = fmaf(v2, ww[2], a3); }\n if (oc_eff > 4) { const float* ww = sw + 4*Cin*3 + si; a4 = fmaf(v0, ww[0], a4); a4 = fmaf(v1, ww[1], a4); a4 = fmaf(v2, ww[2], a4); }\n if (oc_eff > 5) { const float* ww = sw + 5*Cin*3 + si; a5 = fmaf(v0, ww[0], a5); a5 = fmaf(v1, ww[1], a5); a5 = fmaf(v2, ww[2], a5); }\n if (oc_eff > 6) { const float* ww = sw + 6*Cin*3 + si; a6 = fmaf(v0, ww[0], a6); a6 = fmaf(v1, ww[1], a6); a6 = fmaf(v2, ww[2], a6); }\n if (oc_eff > 7) { const float* ww = sw + 7*Cin*3 + si; a7 = fmaf(v0, ww[0], a7); a7 = fmaf(v1, ww[1], a7); a7 = fmaf(v2, ww[2], a7); }\n }\n\n if (oc_eff > 0) y[((long long)n*Cout + oc_base+0) * Lout + l] = a0;\n if (oc_eff > 1) y[((long long)n*Cout + oc_base+1) * Lout + l] = a1;\n if (oc_eff > 2) y[((long long)n*Cout + oc_base+2) * Lout + l] = a2;\n if (oc_eff > 3) y[((long long)n*Cout + oc_base+3) * Lout + l] = a3;\n if (oc_eff > 4) y[((long long)n*Cout + oc_base+4) * Lout + l] = a4;\n if (oc_eff > 5) y[((long long)n*Cout + oc_base+5) * Lout + l] = a5;\n if (oc_eff > 6) y[((long long)n*Cout + oc_base+6) * Lout + l] = a6;\n if (oc_eff > 7) y[((long long)n*Cout + oc_base+7) * Lout + l] = a7;\n}\n\ntorch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,\n int64_t stride, int64_t padding, int64_t dilation, int64_t groups) {\n CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_FLOAT(x); CHECK_FLOAT(weight);\n CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight);\n TORCH_CHECK(x.dim() == 3 && weight.dim() == 3, \"bad dims\");\n TORCH_CHECK(groups == 1, \"groups other than 1 not supported\");\n TORCH_CHECK(stride == 1 && padding == 0 && dilation == 1, \"only stride=1 padding=0 dilation=1 supported\");\n TORCH_CHECK(weight.size(2) == 3, \"only kernel_size=3 supported\");\n int N = (int)x.size(0), Cin = (int)x.size(1), L = (int)x.size(2);\n int Cout = (int)weight.size(0);\n int Lout = L - 2;\n auto y = torch::empty({N, Cout, Lout}, x.options());\n int has_bias = (bias.defined() && bias.numel() > 0) ? 1 : 0;\n if (has_bias) { CHECK_CUDA(bias); CHECK_FLOAT(bias); CHECK_CONTIGUOUS(bias); }\n const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;\n const int threads = 256;\n dim3 block(threads);\n dim3 grid((Lout + threads - 1) / threads, N, (Cout + 7) / 8);\n size_t shmem = (size_t)8 * Cin * 3 * sizeof(float);\n conv1d_k3_s1_p0_oc8_kernel<<<grid, block, shmem>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(), N, Cin, L, Cout, Lout, has_bias);\n return y;\n}\n'''\n\n_ext = load_inline(name='conv1d_fast_ext_v2', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['conv1d_forward'], extra_cuda_cflags=['-O3'], verbose=False, with_cuda=True)\n\nclass ModelNew(nn.Module):\n 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):\n super(ModelNew, self).__init__()\n self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)\n self.stride = stride\n self.padding = padding\n self.dilation = dilation\n self.groups = groups\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n b = self.conv1d.bias\n if b is None:\n b = torch.empty(0, device=x.device, dtype=x.dtype)\n return _ext.conv1d_forward(x, self.conv1d.weight, b, self.stride, self.padding, self.dilation, self.groups)\n"
}