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>\n#include <vector>\n\ntorch::Tensor maxpool3d_forward(torch::Tensor x, int64_t kernel_size, c10::optional<int64_t> stride_opt, int64_t padding, int64_t dilation, bool return_indices, bool ceil_mode);\n\"\"\"\n\n_cuda_src = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n#include <limits>\n\n__global__ void maxpool3d_k3s2p1d3_128_kernel(const float* __restrict__ x, float* __restrict__ y, long total) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= total) return;\n\n // output shape: N*C x 62 x 62 x 62, input spatial 128^3\n int ow = idx % 62;\n long t = idx / 62;\n int oh = t % 62;\n t /= 62;\n int od = t % 62;\n long nc = t / 62;\n\n int iw0 = ow * 2 - 1;\n int ih0 = oh * 2 - 1;\n int id0 = od * 2 - 1;\n const long nc_base = nc * 2097152L; // 128*128*128\n float m = -CUDART_INF_F;\n\n if (ow > 0 && oh > 0 && od > 0) {\n // all 27 taps are in bounds for this problem\n #pragma unroll\n for (int kd = 0; kd < 3; ++kd) {\n int id = id0 + kd * 3;\n long dbase = nc_base + (long)id * 16384L;\n #pragma unroll\n for (int kh = 0; kh < 3; ++kh) {\n int ih = ih0 + kh * 3;\n long hbase = dbase + (long)ih * 128L + iw0;\n float v0 = x[hbase];\n float v1 = x[hbase + 3];\n float v2 = x[hbase + 6];\n m = fmaxf(m, v0);\n m = fmaxf(m, v1);\n m = fmaxf(m, v2);\n }\n }\n } else {\n #pragma unroll\n for (int kd = 0; kd < 3; ++kd) {\n int id = id0 + kd * 3;\n if ((unsigned)id < 128u) {\n #pragma unroll\n for (int kh = 0; kh < 3; ++kh) {\n int ih = ih0 + kh * 3;\n if ((unsigned)ih < 128u) {\n long hbase = nc_base + (long)id * 16384L + (long)ih * 128L;\n #pragma unroll\n for (int kw = 0; kw < 3; ++kw) {\n int iw = iw0 + kw * 3;\n if ((unsigned)iw < 128u) {\n m = fmaxf(m, x[hbase + iw]);\n }\n }\n }\n }\n }\n }\n }\n y[idx] = m;\n}\n\n__global__ void maxpool3d_generic_kernel(const float* __restrict__ x, float* __restrict__ y,\n long total, int N, int C, int D, int H, int W, int OD, int OH, int OW,\n int K, int S, int P, int Dil) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= total) return;\n int ow = idx % OW;\n long t = idx / OW;\n int oh = t % OH; t /= OH;\n int od = t % OD; t /= OD;\n long nc = t;\n int id0 = od*S - P, ih0 = oh*S - P, iw0 = ow*S - P;\n long nc_base = nc * (long)D*H*W;\n float m = -CUDART_INF_F;\n for (int kd=0; kd<K; ++kd) {\n int id = id0 + kd*Dil;\n if ((unsigned)id < (unsigned)D) {\n for (int kh=0; kh<K; ++kh) {\n int ih = ih0 + kh*Dil;\n if ((unsigned)ih < (unsigned)H) {\n long hbase = nc_base + ((long)id*H + ih)*W;\n for (int kw=0; kw<K; ++kw) {\n int iw = iw0 + kw*Dil;\n if ((unsigned)iw < (unsigned)W) m = fmaxf(m, x[hbase + iw]);\n }\n }\n }\n }\n }\n y[idx] = m;\n}\n\ntorch::Tensor maxpool3d_forward(torch::Tensor x, int64_t kernel_size, c10::optional<int64_t> stride_opt, int64_t padding, int64_t dilation, bool return_indices, bool ceil_mode) {\n int64_t N = x.size(0), C = x.size(1), D = x.size(2), H = x.size(3), W = x.size(4);\n int64_t K = kernel_size;\n int64_t S = stride_opt.has_value() ? stride_opt.value() : K;\n int64_t P = padding;\n int64_t Dil = dilation;\n int64_t OD, OH, OW;\n if (ceil_mode) {\n OD = (D + 2*P - Dil*(K-1) - 1 + S - 1) / S + 1;\n OH = (H + 2*P - Dil*(K-1) - 1 + S - 1) / S + 1;\n OW = (W + 2*P - Dil*(K-1) - 1 + S - 1) / S + 1;\n } else {\n OD = (D + 2*P - Dil*(K-1) - 1) / S + 1;\n OH = (H + 2*P - Dil*(K-1) - 1) / S + 1;\n OW = (W + 2*P - Dil*(K-1) - 1) / S + 1;\n }\n auto y = torch::empty({N, C, OD, OH, OW}, x.options());\n long total = N*C*OD*OH*OW;\n const int threads = 256;\n dim3 block(threads);\n dim3 grid((total + threads - 1) / threads);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (K == 3 && S == 2 && P == 1 && Dil == 3 && !ceil_mode && D == 128 && H == 128 && W == 128 && OD == 62 && OH == 62 && OW == 62 && x.scalar_type() == at::kFloat) {\n maxpool3d_k3s2p1d3_128_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), total);\n } else {\n maxpool3d_generic_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), total, (int)N, (int)C, (int)D, (int)H, (int)W, (int)OD, (int)OH, (int)OW, (int)K, (int)S, (int)P, (int)Dil);\n }\n return y;\n}\n\"\"\"\n\n_ext = load_inline(name=\"maxpool3d_ext_v1\", cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=[\"maxpool3d_forward\"], extra_cuda_cflags=[\"-O3\", \"--use_fast_math\"], extra_cflags=[\"-O3\"], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, kernel_size: int, stride: int = None, padding: int = 0, dilation: int = 1, return_indices: bool = False, ceil_mode: bool = False):\n super().__init__()\n self.kernel_size = int(kernel_size)\n self.stride = stride\n self.padding = int(padding)\n self.dilation = int(dilation)\n self.return_indices = bool(return_indices)\n self.ceil_mode = bool(ceil_mode)\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return _ext.maxpool3d_forward(x, self.kernel_size, self.stride, self.padding, self.dilation, self.return_indices, self.ceil_mode)\n"
}