callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_sources = r\"\"\"\n#include <torch/extension.h>\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale);\n\"\"\"\n\ncuda_sources = r\"\"\"\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void fused_deconv_pool_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ b,\n float* __restrict__ out,\n int N, int Cin, int Cout, int D, int H, int W,\n float scale) {\n int block = blockIdx.x;\n int co = block % Cout;\n int n = block / Cout;\n int tid = threadIdx.x;\n const int PD = 15;\n const int PH = 31;\n const int PW = 31;\n const int cells = PD * PH * PW;\n float sum = 0.0f;\n float bias = b[co];\n\n for (int idx = tid; idx < cells; idx += blockDim.x) {\n int t = idx;\n int pw = t % PW; t /= PW;\n int ph = t % PH;\n int pd = t / PH;\n\n float v000 = bias, v001 = bias, v010 = bias, v011 = bias;\n float v100 = bias, v101 = bias, v110 = bias, v111 = bias;\n #pragma unroll\n for (int ci = 0; ci < 3; ++ci) {\n const float* xp = x + (((n * Cin + ci) * D + pd) * H + ph) * W + pw;\n float x000 = xp[0];\n float x001 = xp[1];\n float x010 = xp[W];\n float x011 = xp[W + 1];\n float x100 = xp[H * W];\n float x101 = xp[H * W + 1];\n float x110 = xp[H * W + W];\n float x111 = xp[H * W + W + 1];\n const float* wp = w + ((ci * Cout + co) * 27);\n // W(kd,kh,kw) = wp[kd*9 + kh*3 + kw]\n v000 = fmaf(x000, wp[13], v000); // 1,1,1\n\n v001 = fmaf(x000, wp[14], v001); // 1,1,2\n v001 = fmaf(x001, wp[12], v001); // 1,1,0\n\n v010 = fmaf(x000, wp[16], v010); // 1,2,1\n v010 = fmaf(x010, wp[10], v010); // 1,0,1\n\n v011 = fmaf(x000, wp[17], v011); // 1,2,2\n v011 = fmaf(x001, wp[15], v011); // 1,2,0\n v011 = fmaf(x010, wp[11], v011); // 1,0,2\n v011 = fmaf(x011, wp[9], v011); // 1,0,0\n\n v100 = fmaf(x000, wp[22], v100); // 2,1,1\n v100 = fmaf(x100, wp[4], v100); // 0,1,1\n\n v101 = fmaf(x000, wp[23], v101); // 2,1,2\n v101 = fmaf(x001, wp[21], v101); // 2,1,0\n v101 = fmaf(x100, wp[5], v101); // 0,1,2\n v101 = fmaf(x101, wp[3], v101); // 0,1,0\n\n v110 = fmaf(x000, wp[25], v110); // 2,2,1\n v110 = fmaf(x010, wp[19], v110); // 2,0,1\n v110 = fmaf(x100, wp[7], v110); // 0,2,1\n v110 = fmaf(x110, wp[1], v110); // 0,0,1\n\n v111 = fmaf(x000, wp[26], v111); // 2,2,2\n v111 = fmaf(x001, wp[24], v111); // 2,2,0\n v111 = fmaf(x010, wp[20], v111); // 2,0,2\n v111 = fmaf(x011, wp[18], v111); // 2,0,0\n v111 = fmaf(x100, wp[8], v111); // 0,2,2\n v111 = fmaf(x101, wp[6], v111); // 0,2,0\n v111 = fmaf(x110, wp[2], v111); // 0,0,2\n v111 = fmaf(x111, wp[0], v111); // 0,0,0\n }\n float m1 = fmaxf(fmaxf(v000, v001), fmaxf(v010, v011));\n float m2 = fmaxf(fmaxf(v100, v101), fmaxf(v110, v111));\n sum += fmaxf(m1, m2);\n }\n\n __shared__ float smem[256];\n smem[tid] = sum;\n __syncthreads();\n for (int s = blockDim.x >> 1; s > 0; s >>= 1) {\n if (tid < s) smem[tid] += smem[tid + s];\n __syncthreads();\n }\n if (tid == 0) {\n float val = smem[0] * (scale / (float)cells);\n val = fminf(1.0f, fmaxf(0.0f, val));\n out[(n * Cout + co)] = val;\n }\n}\n\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale) {\n TORCH_CHECK(x.is_cuda(), \"x must be CUDA\");\n auto N = (int)x.size(0);\n auto Cin = (int)x.size(1);\n auto D = (int)x.size(2);\n auto H = (int)x.size(3);\n auto W = (int)x.size(4);\n auto Cout = (int)weight.size(1);\n auto out = torch::empty({N, Cout, 1, 1, 1}, x.options());\n int blocks = N * Cout;\n int threads = 256;\n fused_deconv_pool_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(\n x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(),\n N, Cin, Cout, D, H, W, (float)scale);\n return out;\n}\n\"\"\"\n\n_ext = load_inline(name=\"fused_deconv_pool_ext_v1\",\n cpp_sources=cpp_sources,\n cuda_sources=cuda_sources,\n functions=[\"fused_forward\"],\n extra_cuda_cflags=[\"-O3\", \"--use_fast_math\"],\n extra_cflags=[\"-O3\"],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, stride, padding, scale, maxpool_kernel_size):\n super(ModelNew, self).__init__()\n self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)\n self.scale = scale\n self.maxpool_kernel_size = maxpool_kernel_size\n def forward(self, x):\n return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, float(self.scale))\n"
}