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>\ntorch::Tensor convt3d_hswish_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add);\n'''\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ __launch_bounds__(256, 2)\nvoid convt3d_hswish_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ conv_bias,\n const float* __restrict__ add,\n float* __restrict__ out,\n long total) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= total) return;\n int ow = idx & 31;\n int oh = (idx >> 5) & 31;\n int od = (idx >> 10) & 31;\n int tmp = idx >> 15;\n int oc = tmp & 63;\n int n = tmp >> 6;\n float acc = conv_bias[oc];\n int id0,kd0,nd,id1,kd1;\n if ((od & 1) == 0) { id0 = od >> 1; kd0 = 1; id1 = 0; kd1 = 0; nd = 1; }\n else { id0 = od >> 1; kd0 = 2; id1 = id0 + 1; kd1 = 0; nd = (id1 < 16) ? 2 : 1; }\n int ih0,kh0,nh,ih1,kh1;\n if ((oh & 1) == 0) { ih0 = oh >> 1; kh0 = 1; ih1 = 0; kh1 = 0; nh = 1; }\n else { ih0 = oh >> 1; kh0 = 2; ih1 = ih0 + 1; kh1 = 0; nh = (ih1 < 16) ? 2 : 1; }\n int iw0,kw0,nw,iw1,kw1;\n if ((ow & 1) == 0) { iw0 = ow >> 1; kw0 = 1; iw1 = 0; kw1 = 0; nw = 1; }\n else { iw0 = ow >> 1; kw0 = 2; iw1 = iw0 + 1; kw1 = 0; nw = (iw1 < 16) ? 2 : 1; }\n\n long nbase = ((long)n * 32) * 4096L;\n for (int ic = 0; ic < 32; ++ic) {\n const float* xic = x + nbase + (long)ic * 4096L;\n const float* wic = w + ((long)ic * 64L + oc) * 27L;\n float xv;\n if (nd == 1 && nh == 1 && nw == 1) {\n xv = xic[(id0 * 16 + ih0) * 16 + iw0];\n acc = fmaf(xv, wic[(kd0 * 3 + kh0) * 3 + kw0], acc);\n } else {\n xv = xic[(id0 * 16 + ih0) * 16 + iw0];\n acc = fmaf(xv, wic[(kd0 * 3 + kh0) * 3 + kw0], acc);\n if (nw == 2) { xv = xic[(id0 * 16 + ih0) * 16 + iw1]; acc = fmaf(xv, wic[(kd0 * 3 + kh0) * 3 + kw1], acc); }\n if (nh == 2) {\n xv = xic[(id0 * 16 + ih1) * 16 + iw0]; acc = fmaf(xv, wic[(kd0 * 3 + kh1) * 3 + kw0], acc);\n if (nw == 2) { xv = xic[(id0 * 16 + ih1) * 16 + iw1]; acc = fmaf(xv, wic[(kd0 * 3 + kh1) * 3 + kw1], acc); }\n }\n if (nd == 2) {\n xv = xic[(id1 * 16 + ih0) * 16 + iw0]; acc = fmaf(xv, wic[(kd1 * 3 + kh0) * 3 + kw0], acc);\n if (nw == 2) { xv = xic[(id1 * 16 + ih0) * 16 + iw1]; acc = fmaf(xv, wic[(kd1 * 3 + kh0) * 3 + kw1], acc); }\n if (nh == 2) {\n xv = xic[(id1 * 16 + ih1) * 16 + iw0]; acc = fmaf(xv, wic[(kd1 * 3 + kh1) * 3 + kw0], acc);\n if (nw == 2) { xv = xic[(id1 * 16 + ih1) * 16 + iw1]; acc = fmaf(xv, wic[(kd1 * 3 + kh1) * 3 + kw1], acc); }\n }\n }\n }\n }\n float y = acc + add[idx];\n float t = fminf(fmaxf(y + 3.0f, 0.0f), 6.0f) * 0.16666666666666666f;\n out[idx] = y * y * t;\n}\n\ntorch::Tensor convt3d_hswish_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add) {\n auto out = torch::empty_like(add);\n long total = add.numel();\n const int threads = 256;\n const dim3 blocks((unsigned int)((total + threads - 1) / threads));\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n convt3d_hswish_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), add.data_ptr<float>(), out.data_ptr<float>(), total);\n return out;\n}\n'''\n_ext = load_inline(name='convt3d_hswish_ext_v2', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['convt3d_hswish_forward'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'])\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape):\n super(ModelNew, self).__init__()\n self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n def forward(self, x, add_input):\n return _ext.convt3d_hswish_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, add_input)\n"
}