call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor conv_min_softmax_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t dim);\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cfloat>\n\n__global__ void conv3d_min_k3_spatial_kernel(\n const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ b,\n float* __restrict__ tmp,\n int N, int D, int H, int W, int Ho, int Wo, int Do, int total)\n{\n int c = blockIdx.y;\n int pos = blockIdx.x * blockDim.x + threadIdx.x;\n if (pos >= total) return;\n int ow = pos % Wo;\n int t = pos / Wo;\n int oh = t % Ho;\n int n = t / Ho;\n\n const float* wc = w + c * 81;\n float bias = b[c];\n float mn = FLT_MAX;\n\n #pragma unroll 1\n for (int od = 0; od < Do; ++od) {\n float acc = bias;\n #pragma unroll\n for (int ci = 0; ci < 3; ++ci) {\n #pragma unroll\n for (int kz = 0; kz < 3; ++kz) {\n int base = (((n * 3 + ci) * D + (od + kz)) * H + oh) * W + ow;\n int wbase = (ci * 27 + kz * 9);\n acc = fmaf(x[base], wc[wbase], acc);\n acc = fmaf(x[base + 1], wc[wbase + 1], acc);\n acc = fmaf(x[base + 2], wc[wbase + 2], acc);\n acc = fmaf(x[base + W], wc[wbase + 3], acc);\n acc = fmaf(x[base + W+1], wc[wbase + 4], acc);\n acc = fmaf(x[base + W+2], wc[wbase + 5], acc);\n acc = fmaf(x[base + 2*W], wc[wbase + 6], acc);\n acc = fmaf(x[base + 2*W+1], wc[wbase + 7], acc);\n acc = fmaf(x[base + 2*W+2], wc[wbase + 8], acc);\n }\n }\n mn = fminf(mn, acc);\n }\n tmp[(c * N * Ho * Wo) + pos] = mn; // C-major temp for coalesced writes per channel\n}\n\n__global__ void softmax_c24_kernel(const float* __restrict__ tmp, float* __restrict__ y,\n int N, int Cout, int Ho, int Wo, int total) {\n int warp = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;\n int lane = threadIdx.x & 31;\n if (warp >= total) return;\n int pos = warp;\n // tmp is [C][N*Ho*Wo]\n float v = -FLT_MAX;\n if (lane < Cout) v = tmp[lane * total + pos];\n float mx = v;\n #pragma unroll\n for (int off=16; off>0; off>>=1) mx = fmaxf(mx, __shfl_down_sync(0xffffffff, mx, off));\n mx = __shfl_sync(0xffffffff, mx, 0);\n float e = (lane < Cout) ? expf(v - mx) : 0.0f;\n float s = e;\n #pragma unroll\n for (int off=16; off>0; off>>=1) s += __shfl_down_sync(0xffffffff, s, off);\n s = __shfl_sync(0xffffffff, s, 0);\n if (lane < Cout) {\n int ow = pos % Wo;\n int t = pos / Wo;\n int oh = t % Ho;\n int n = t / Ho;\n y[((n * Cout + lane) * Ho + oh) * Wo + ow] = e / s;\n }\n}\n\ntorch::Tensor conv_min_softmax_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t dim) {\n int N = (int)x.size(0), D = (int)x.size(2), H = (int)x.size(3), W = (int)x.size(4);\n int Cout = (int)weight.size(0), K = (int)weight.size(2);\n int Do = D - K + 1, Ho = H - K + 1, Wo = W - K + 1;\n int total = N * Ho * Wo;\n auto tmp = torch::empty({Cout, total}, x.options());\n auto y = torch::empty({N, Cout, Ho, Wo}, x.options());\n dim3 block(256);\n dim3 grid((total + 255) / 256, Cout);\n conv3d_min_k3_spatial_kernel<<<grid, block>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), tmp.data_ptr<float>(), N, D, H, W, Ho, Wo, Do, total);\n softmax_c24_kernel<<<(total*32 + 255)/256, 256>>>(tmp.data_ptr<float>(), y.data_ptr<float>(), N, Cout, Ho, Wo, total);\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_min_softmax_ext_v3', cpp_sources=cpp_src, cuda_sources=cuda_src,\n functions=['conv_min_softmax_forward'], verbose=False,\n extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, dim):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)\n self.dim = dim\n def forward(self, x):\n return _ext.conv_min_softmax_forward(x, self.conv.weight, self.conv.bias, self.dim)\n"
}