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>\n\ntorch::Tensor conv_gn_forward(torch::Tensor x, torch::Tensor w, torch::Tensor gamma, torch::Tensor beta, int64_t num_groups, int64_t stride, int64_t padding, int64_t dilation, double eps, bool affine);\n\"\"\"\n\ncuda_src = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#define TILE_T 128\n#define BLOCK_THREADS 256\n\n__global__ void zero_kernel(float* s, float* q, int total) {\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total) { s[i] = 0.0f; q[i] = 0.0f; }\n}\n\n__inline__ __device__ float warp_reduce_sum(float v) {\n for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(0xffffffff, v, offset);\n return v;\n}\n\n__inline__ __device__ float block_reduce_sum(float v) {\n static __shared__ float shared[32];\n int lane = threadIdx.x & 31;\n int wid = threadIdx.x >> 5;\n v = warp_reduce_sum(v);\n if (lane == 0) shared[wid] = v;\n __syncthreads();\n v = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;\n if (wid == 0) v = warp_reduce_sum(v);\n return v;\n}\n\n__global__ void conv_stats_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n float* __restrict__ y,\n float* __restrict__ sums,\n float* __restrict__ sumsq,\n int N, int Cin, int T, int Cout, int K,\n int Tout, int groups, int stride, int padding, int dilation) {\n int tile = blockIdx.x;\n int g = blockIdx.y;\n int n = blockIdx.z;\n int cpg = Cout / groups;\n int t0 = tile * TILE_T;\n float lsum = 0.0f;\n float lsq = 0.0f;\n int elems = cpg * TILE_T;\n for (int idx = threadIdx.x; idx < elems; idx += blockDim.x) {\n int lc = idx / TILE_T;\n int tt = idx - lc * TILE_T;\n int t_out = t0 + tt;\n if (t_out < Tout) {\n int co = g * cpg + lc;\n float acc = 0.0f;\n #pragma unroll\n for (int kk = 0; kk < 3; ++kk) {\n if (kk < K) {\n int ti = t_out * stride - padding + kk * dilation;\n if ((unsigned)ti < (unsigned)T) {\n const float* xp = x + (n * Cin * T) + ti;\n const float* wp = w + (co * Cin * K) + kk;\n for (int ci = 0; ci < Cin; ++ci) {\n acc = fmaf(xp[ci * T], wp[ci * K], acc);\n }\n }\n }\n }\n y[(n * Cout + co) * Tout + t_out] = acc;\n lsum += acc;\n lsq = fmaf(acc, acc, lsq);\n }\n }\n float bsum = block_reduce_sum(lsum);\n float bsq = block_reduce_sum(lsq);\n if (threadIdx.x == 0) {\n int si = n * groups + g;\n atomicAdd(sums + si, bsum);\n atomicAdd(sumsq + si, bsq);\n }\n}\n\n__global__ void norm_kernel(float* __restrict__ y,\n const float* __restrict__ sums,\n const float* __restrict__ sumsq,\n const float* __restrict__ gamma,\n const float* __restrict__ beta,\n int N, int Cout, int Tout, int groups, float eps, bool affine) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = N * Cout * Tout;\n int cpg = Cout / groups;\n int elems = cpg * Tout;\n for (int i = idx; i < total; i += gridDim.x * blockDim.x) {\n int tmp = i;\n int t = tmp % Tout; tmp /= Tout;\n int c = tmp % Cout;\n int n = tmp / Cout;\n int g = c / cpg;\n int si = n * groups + g;\n float mean = sums[si] / (float)elems;\n float var = sumsq[si] / (float)elems - mean * mean;\n float v = (y[i] - mean) * rsqrtf(var + eps);\n if (affine) v = v * gamma[c] + beta[c];\n y[i] = v;\n }\n}\n\ntorch::Tensor conv_gn_forward(torch::Tensor x, torch::Tensor w, torch::Tensor gamma, torch::Tensor beta, int64_t num_groups, int64_t stride, int64_t padding, int64_t dilation, double eps, bool affine) {\n int N = (int)x.size(0);\n int Cin = (int)x.size(1);\n int T = (int)x.size(2);\n int Cout = (int)w.size(0);\n int K = (int)w.size(2);\n int Tout = (T + 2*(int)padding - (int)dilation*(K-1) - 1) / (int)stride + 1;\n auto y = torch::empty({N, Cout, Tout}, x.options());\n auto opts = x.options();\n auto sums = torch::empty({N, (int)num_groups}, opts);\n auto sumsq = torch::empty({N, (int)num_groups}, opts);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n int stat_total = N * (int)num_groups;\n zero_kernel<<<(stat_total+255)/256, 256, 0, stream>>>(sums.data_ptr<float>(), sumsq.data_ptr<float>(), stat_total);\n dim3 grid((Tout + TILE_T - 1) / TILE_T, (int)num_groups, N);\n conv_stats_kernel<<<grid, BLOCK_THREADS, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), y.data_ptr<float>(), sums.data_ptr<float>(), sumsq.data_ptr<float>(), N, Cin, T, Cout, K, Tout, (int)num_groups, (int)stride, (int)padding, (int)dilation);\n int total = N * Cout * Tout;\n norm_kernel<<<min((total+255)/256, 4096), 256, 0, stream>>>(y.data_ptr<float>(), sums.data_ptr<float>(), sumsq.data_ptr<float>(), affine ? gamma.data_ptr<float>() : nullptr, affine ? beta.data_ptr<float>() : nullptr, N, Cout, Tout, (int)num_groups, (float)eps, affine);\n return y;\n}\n\"\"\"\n\n_ext = load_inline(name='conv_gn_ext_v1', cpp_sources=[cpp_src], cuda_sources=[cuda_src], functions=['conv_gn_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'], verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels:int, out_channels:int, num_groups:int, kernel_size:int=3, stride:int=1, padding:int=1, dilation:int=1, eps:float=1e-5, affine:bool=True):\n super().__init__()\n self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=False)\n self.norm = nn.GroupNorm(num_groups=num_groups, num_channels=out_channels, eps=eps, affine=affine)\n self.num_groups=num_groups; self.stride=stride; self.padding=padding; self.dilation=dilation; self.eps=eps; self.affine=affine\n def forward(self, x):\n gamma = self.norm.weight if self.affine else torch.empty(0, device=x.device, dtype=x.dtype)\n beta = self.norm.bias if self.affine else torch.empty(0, device=x.device, dtype=x.dtype)\n return _ext.conv_gn_forward(x.contiguous(), self.conv.weight.contiguous(), gamma, beta, self.num_groups, self.stride, self.padding, self.dilation, self.eps, self.affine)\n"
}