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 = r'''\n#include <torch/extension.h>\ntorch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor w, c10::optional<torch::Tensor> bias_opt,\n int64_t stride, int64_t padding, int64_t dilation, int64_t groups);\n'''\n_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <c10/cuda/CUDAStream.h>\n#include <vector>\n\n__global__ void conv3d_357_c3_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n float* __restrict__ y,\n int N, int OD, int OH, int OW) {\n long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n long long total = (long long)N * 64LL * OD * OH * OW;\n if (idx >= total) return;\n\n int ow = idx % OW;\n long long t = idx / OW;\n int oh = t % OH; t /= OH;\n int od = t % OD; t /= OD;\n int oc = t & 63;\n int n = t >> 6;\n\n int ID = OD + 2;\n int IH = OH + 4;\n int IW = OW + 6;\n const float* xp = x + (((long long)n * 3 * ID * IH * IW) + ((long long)od * IH + oh) * IW + ow);\n const float* wp = w + (long long)oc * 315;\n\n float acc = 0.0f;\n #pragma unroll\n for (int c=0; c<3; ++c) {\n const float* xpc = xp + (long long)c * ID * IH * IW;\n #pragma unroll\n for (int kd=0; kd<3; ++kd) {\n #pragma unroll\n for (int kh=0; kh<5; ++kh) {\n const float* xrow = xpc + ((long long)kd * IH + kh) * IW;\n int woff = ((c*3 + kd)*5 + kh)*7;\n #pragma unroll\n for (int kw=0; kw<7; ++kw) {\n acc = fmaf(__ldg(xrow + kw), __ldg(wp + woff + kw), acc);\n }\n }\n }\n }\n y[idx] = acc;\n}\n\n__global__ void conv3d_generic_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ b,\n float* __restrict__ y,\n int N, int C, int OC,\n int ID, int IH, int IW,\n int KD, int KH, int KW,\n int OD, int OH, int OW,\n int stride, int pad, int dilation,\n int has_bias) {\n long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n long long total = (long long)N * OC * OD * OH * OW;\n if (idx >= total) return;\n int ow = idx % OW;\n long long t = idx / OW;\n int oh = t % OH; t /= OH;\n int od = t % OD; t /= OD;\n int oc = t % OC; int n = t / OC;\n float acc = has_bias ? __ldg(b + oc) : 0.0f;\n for (int c=0; c<C; ++c) {\n for (int kd=0; kd<KD; ++kd) {\n int id = od*stride - pad + kd*dilation;\n if ((unsigned)id >= (unsigned)ID) continue;\n for (int kh=0; kh<KH; ++kh) {\n int ih = oh*stride - pad + kh*dilation;\n if ((unsigned)ih >= (unsigned)IH) continue;\n for (int kw=0; kw<KW; ++kw) {\n int iw = ow*stride - pad + kw*dilation;\n if ((unsigned)iw >= (unsigned)IW) continue;\n float xv = __ldg(x + ((((long long)n*C + c)*ID + id)*IH + ih)*IW + iw);\n float wv = __ldg(w + (((((long long)oc*C + c)*KD + kd)*KH + kh)*KW + kw));\n acc = fmaf(xv, wv, acc);\n }\n }\n }\n }\n y[idx] = acc;\n}\n\ntorch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor w, c10::optional<torch::Tensor> bias_opt,\n int64_t stride, int64_t padding, int64_t dilation, int64_t groups) {\n TORCH_CHECK(x.is_cuda() && w.is_cuda(), \"CUDA tensors expected\");\n TORCH_CHECK(x.scalar_type() == torch::kFloat32 && w.scalar_type() == torch::kFloat32, \"float32 expected\");\n TORCH_CHECK(groups == 1, \"groups=1 only\");\n x = x.contiguous();\n w = w.contiguous();\n int N = x.size(0), C = x.size(1), ID = x.size(2), IH = x.size(3), IW = x.size(4);\n int OC = w.size(0), KD = w.size(2), KH = w.size(3), KW = w.size(4);\n int OD = (ID + 2*padding - dilation*(KD-1) - 1) / stride + 1;\n int OH = (IH + 2*padding - dilation*(KH-1) - 1) / stride + 1;\n int OW = (IW + 2*padding - dilation*(KW-1) - 1) / stride + 1;\n auto y = torch::empty({N, OC, OD, OH, OW}, x.options());\n int threads = 256;\n long long total = (long long)N * OC * OD * OH * OW;\n int blocks = (int)((total + threads - 1) / threads);\n cudaStream_t stream = c10::cuda::getCurrentCUDAStream();\n if (C == 3 && OC == 64 && KD == 3 && KH == 5 && KW == 7 && stride == 1 && padding == 0 && dilation == 1 && !bias_opt.has_value()) {\n conv3d_357_c3_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), y.data_ptr<float>(), N, OD, OH, OW);\n } else {\n const float* bp = nullptr; int has_bias = 0;\n torch::Tensor bc;\n if (bias_opt.has_value()) { bc = bias_opt.value().contiguous(); bp = bc.data_ptr<float>(); has_bias = 1; }\n conv3d_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), bp, y.data_ptr<float>(), N,C,OC,ID,IH,IW,KD,KH,KW,OD,OH,OW,(int)stride,(int)padding,(int)dilation,has_bias);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv3d_357_ext_v2', cpp_sources=_cpp, cuda_sources=_src, functions=['conv3d_forward'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):\n super(ModelNew, self).__init__()\n self.conv3d = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)\n self.stride = stride if isinstance(stride, int) else stride[0]\n self.padding = padding if isinstance(padding, int) else padding[0]\n self.dilation = dilation if isinstance(dilation, int) else dilation[0]\n self.groups = groups\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return _ext.conv3d_forward(x, self.conv3d.weight, self.conv3d.bias, self.stride, self.padding, self.dilation, self.groups)\n"
}