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\ntorch::Tensor fused_deconv_ln_gelu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias,\n torch::Tensor ln_weight, torch::Tensor ln_bias,\n double eps, double scaling);\n\"\"\"\n\n_cuda_src = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math.h>\n\n#define CIN 32\n#define COUT 64\n#define DIN 16\n#define HIN 32\n#define WIN 32\n#define KDIM 4\n#define DOUT 32\n#define HOUT 64\n#define WOUT 64\n\n__device__ __forceinline__ float warp_sum(float v) {\n #pragma unroll\n for (int offset = 16; offset > 0; offset >>= 1) {\n v += __shfl_down_sync(0xffffffff, v, offset);\n }\n return v;\n}\n\n__device__ __forceinline__ float conv_one(const float* __restrict__ x,\n const float* __restrict__ w,\n const float* __restrict__ b,\n int n, int co, int od, int oh, int ow) {\n float acc = b[co];\n #pragma unroll\n for (int ci = 0; ci < CIN; ++ci) {\n #pragma unroll\n for (int kd = 0; kd < KDIM; ++kd) {\n int td = od + 1 - kd;\n if ((td & 1) == 0) {\n int id = td >> 1;\n if ((unsigned)id < DIN) {\n #pragma unroll\n for (int kh = 0; kh < KDIM; ++kh) {\n int th = oh + 1 - kh;\n if ((th & 1) == 0) {\n int ih = th >> 1;\n if ((unsigned)ih < HIN) {\n #pragma unroll\n for (int kw = 0; kw < KDIM; ++kw) {\n int tw = ow + 1 - kw;\n if ((tw & 1) == 0) {\n int iw = tw >> 1;\n if ((unsigned)iw < WIN) {\n int xidx = (((n * CIN + ci) * DIN + id) * HIN + ih) * WIN + iw;\n int widx = (((ci * COUT + co) * KDIM + kd) * KDIM + kh) * KDIM + kw;\n acc = fmaf(x[xidx], w[widx], acc);\n }\n }\n }\n }\n }\n }\n }\n }\n }\n }\n return acc;\n}\n\n__global__ __launch_bounds__(256, 1)\nvoid fused_kernel(const float* __restrict__ x,\n const float* __restrict__ weight,\n const float* __restrict__ cbias,\n const float* __restrict__ ln_w,\n const float* __restrict__ ln_b,\n float* __restrict__ out,\n float eps, float scaling, int rows) {\n int lane = threadIdx.x & 31;\n int warp = threadIdx.x >> 5;\n int warps_per_block = blockDim.x >> 5;\n int row = blockIdx.x * warps_per_block + warp;\n if (row >= rows) return;\n\n int tmp = row;\n int oh = tmp % HOUT; tmp /= HOUT;\n int od = tmp % DOUT; tmp /= DOUT;\n int co = tmp % COUT; tmp /= COUT;\n int n = tmp;\n\n float v0 = conv_one(x, weight, cbias, n, co, od, oh, lane);\n float v1 = conv_one(x, weight, cbias, n, co, od, oh, lane + 32);\n\n float s = v0 + v1;\n float sum = warp_sum(s);\n float mean = __shfl_sync(0xffffffff, sum, 0) * (1.0f / 64.0f);\n float d0 = v0 - mean;\n float d1 = v1 - mean;\n float ss = d0*d0 + d1*d1;\n float vsum = warp_sum(ss);\n float var = __shfl_sync(0xffffffff, vsum, 0) * (1.0f / 64.0f);\n float inv = rsqrtf(var + eps);\n\n float z0 = d0 * inv * ln_w[lane] + ln_b[lane];\n float z1 = d1 * inv * ln_w[lane + 32] + ln_b[lane + 32];\n const float inv_sqrt2 = 0.70710678118654752440f;\n float g0 = 0.5f * z0 * (1.0f + erff(z0 * inv_sqrt2)) * scaling;\n float g1 = 0.5f * z1 * (1.0f + erff(z1 * inv_sqrt2)) * scaling;\n int base = row * WOUT;\n out[base + lane] = g0;\n out[base + lane + 32] = g1;\n}\n\ntorch::Tensor fused_deconv_ln_gelu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias,\n torch::Tensor ln_weight, torch::Tensor ln_bias,\n double eps, double scaling) {\n auto out = torch::empty({x.size(0), (int64_t)COUT, (int64_t)DOUT, (int64_t)HOUT, (int64_t)WOUT}, x.options());\n int rows = (int)(x.size(0) * COUT * DOUT * HOUT);\n const int threads = 256;\n const int warps = threads / 32;\n dim3 block(threads);\n dim3 grid((rows + warps - 1) / warps);\n fused_kernel<<<grid, block>>>(x.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(),\n ln_weight.data_ptr<float>(), ln_bias.data_ptr<float>(), out.data_ptr<float>(),\n (float)eps, (float)scaling, rows);\n return out;\n}\n\"\"\"\n\n_ext = load_inline(name=\"fused_deconv_ln_gelu_ext_v1\", cpp_sources=_cpp_src, cuda_sources=_cuda_src,\n functions=[\"fused_deconv_ln_gelu_forward\"], extra_cuda_cflags=[\"-O3\"], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True, eps=1e-5, scaling_factor=1.0):\n super(ModelNew, self).__init__()\n self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)\n self.layer_norm = nn.LayerNorm(out_channels, eps=eps)\n self.scaling_factor = scaling_factor\n self.eps = eps\n def forward(self, x):\n return _ext.fused_deconv_ln_gelu_forward(x.contiguous(), self.conv_transpose.weight.contiguous(), self.conv_transpose.bias.contiguous(), self.layer_norm.weight.contiguous(), self.layer_norm.bias.contiguous(), self.eps, self.scaling_factor)\n"
}