call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <c10/cuda/CUDAStream.h>\n\n__device__ __forceinline__ float act2(float y, float m) {\n float a = y >= 0.0f ? y : 0.2f * y;\n float b = a * m;\n return b >= 0.0f ? b : 0.2f * b;\n}\n\n__global__ void fused_deconv_pool_kernel(\n const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ bias,\n const float* __restrict__ mult,\n float* __restrict__ out,\n int N, int Cin, int Din, int Hin, int Win,\n int Cout, int Dp, int Hp, int Wp)\n{\n long long total = (long long)N * Cout * Dp * Hp * Wp;\n long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= total) return;\n\n int pw = idx % Wp;\n long long t = idx / Wp;\n int ph = t % Hp; t /= Hp;\n int pd = t % Dp; t /= Dp;\n int oc = t % Cout; int n = t / Cout;\n\n float b = bias[oc];\n float a000=b, a001=b, a010=b, a011=b, a100=b, a101=b, a110=b, a111=b;\n\n bool hasD1 = (pd + 1 < Din);\n bool hasH1 = (ph + 1 < Hin);\n bool hasW1 = (pw + 1 < Win);\n\n long long base_n = ((long long)n * Cin * Din * Hin * Win);\n int HW = Hin * Win;\n // weight layout ConvTranspose3d: [Cin, Cout, 3,3,3]\n for (int ic = 0; ic < Cin; ++ic) {\n const float* xp = x + base_n + (long long)ic * Din * Hin * Win + (long long)pd * HW + ph * Win + pw;\n const float* wp0 = w + (((long long)ic * Cout + oc) * 27);\n float x000 = __ldg(xp);\n float x001 = hasW1 ? __ldg(xp + 1) : 0.0f;\n float x010 = hasH1 ? __ldg(xp + Win) : 0.0f;\n float x011 = (hasH1 && hasW1) ? __ldg(xp + Win + 1) : 0.0f;\n float x100 = hasD1 ? __ldg(xp + HW) : 0.0f;\n float x101 = (hasD1 && hasW1) ? __ldg(xp + HW + 1) : 0.0f;\n float x110 = (hasD1 && hasH1) ? __ldg(xp + HW + Win) : 0.0f;\n float x111 = (hasD1 && hasH1 && hasW1) ? __ldg(xp + HW + Win + 1) : 0.0f;\n\n // k index macro: kd*9 + kh*3 + kw. even output uses k=1; odd current uses k=2, next uses k=0\n a000 = fmaf(x000, __ldg(wp0 + 13), a000); // 1,1,1\n\n a001 = fmaf(x000, __ldg(wp0 + 14), a001); // 1,1,2\n if (hasW1) a001 = fmaf(x001, __ldg(wp0 + 12), a001); // 1,1,0\n\n a010 = fmaf(x000, __ldg(wp0 + 16), a010); // 1,2,1\n if (hasH1) a010 = fmaf(x010, __ldg(wp0 + 10), a010); //1,0,1\n\n a011 = fmaf(x000, __ldg(wp0 + 17), a011); //1,2,2\n if (hasW1) a011 = fmaf(x001, __ldg(wp0 + 15), a011); //1,2,0\n if (hasH1) a011 = fmaf(x010, __ldg(wp0 + 11), a011); //1,0,2\n if (hasH1 && hasW1) a011 = fmaf(x011, __ldg(wp0 + 9), a011); //1,0,0\n\n a100 = fmaf(x000, __ldg(wp0 + 22), a100); //2,1,1\n if (hasD1) a100 = fmaf(x100, __ldg(wp0 + 4), a100); //0,1,1\n\n a101 = fmaf(x000, __ldg(wp0 + 23), a101); //2,1,2\n if (hasW1) a101 = fmaf(x001, __ldg(wp0 + 21), a101); //2,1,0\n if (hasD1) a101 = fmaf(x100, __ldg(wp0 + 5), a101); //0,1,2\n if (hasD1 && hasW1) a101 = fmaf(x101, __ldg(wp0 + 3), a101); //0,1,0\n\n a110 = fmaf(x000, __ldg(wp0 + 25), a110); //2,2,1\n if (hasH1) a110 = fmaf(x010, __ldg(wp0 + 19), a110); //2,0,1\n if (hasD1) a110 = fmaf(x100, __ldg(wp0 + 7), a110); //0,2,1\n if (hasD1 && hasH1) a110 = fmaf(x110, __ldg(wp0 + 1), a110); //0,0,1\n\n a111 = fmaf(x000, __ldg(wp0 + 26), a111); //2,2,2\n if (hasW1) a111 = fmaf(x001, __ldg(wp0 + 24), a111); //2,2,0\n if (hasH1) a111 = fmaf(x010, __ldg(wp0 + 20), a111); //2,0,2\n if (hasH1 && hasW1) a111 = fmaf(x011, __ldg(wp0 + 18), a111); //2,0,0\n if (hasD1) a111 = fmaf(x100, __ldg(wp0 + 8), a111); //0,2,2\n if (hasD1 && hasW1) a111 = fmaf(x101, __ldg(wp0 + 6), a111); //0,2,0\n if (hasD1 && hasH1) a111 = fmaf(x110, __ldg(wp0 + 2), a111); //0,0,2\n if (hasD1 && hasH1 && hasW1) a111 = fmaf(x111, __ldg(wp0 + 0), a111); //0,0,0\n }\n float m = __ldg(mult + oc); // multiplier shape [Cout,1,1,1]\n float mx = act2(a000, m);\n float v = act2(a001, m); mx = fmaxf(mx, v);\n v = act2(a010, m); mx = fmaxf(mx, v);\n v = act2(a011, m); mx = fmaxf(mx, v);\n v = act2(a100, m); mx = fmaxf(mx, v);\n v = act2(a101, m); mx = fmaxf(mx, v);\n v = act2(a110, m); mx = fmaxf(mx, v);\n v = act2(a111, m); mx = fmaxf(mx, v);\n out[idx] = mx;\n}\n\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mult) {\n int N = x.size(0), Cin = x.size(1), Din = x.size(2), Hin = x.size(3), Win = x.size(4);\n int Cout = weight.size(1);\n int Dp = Din; // for fixed stride=2,pad=1,k=3,outpad=1, then conv out=2*Din and pool /2\n int Hp = Hin;\n int Wp = Win;\n auto out = torch::empty({N, Cout, Dp, Hp, Wp}, x.options());\n long long total = (long long)N*Cout*Dp*Hp*Wp;\n int threads = 256;\n int blocks = (int)((total + threads - 1) / threads);\n fused_deconv_pool_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(\n x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), mult.data_ptr<float>(), out.data_ptr<float>(),\n N, Cin, Din, Hin, Win, Cout, Dp, Hp, Wp);\n return out;\n}\n'''\n\n_ext = load_inline(name='fused_deconv_pool_ext_v1', cpp_sources='', cuda_sources=cuda_src, functions=['fused_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, multiplier_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.multiplier = nn.Parameter(torch.randn(multiplier_shape))\n def forward(self, x):\n return _ext.fused_forward(x.contiguous(), self.conv_transpose.weight, self.conv_transpose.bias, self.multiplier)\n"
}