final submitted kernel (217 lines, 8,493 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
void launch_train_fused(torch::Tensor y, torch::Tensor scale, torch::Tensor gamma, torch::Tensor beta,
torch::Tensor running_mean, torch::Tensor running_var,
torch::Tensor partial, torch::Tensor stats,
double momentum, double eps);
void launch_eval_fused(torch::Tensor y, torch::Tensor scale, torch::Tensor gamma, torch::Tensor beta,
torch::Tensor running_mean, torch::Tensor running_var, double eps);
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
torch::Tensor scale, torch::Tensor gamma, torch::Tensor beta,
torch::Tensor running_mean, torch::Tensor running_var,
bool training, double momentum, double eps) {
// cuBLAS GEMM (same layout as nn.Linear): y = x @ weight.T + bias
auto y = at::addmm(bias, x, weight.t());
if (training) {
const int chunks = 8;
auto partial = at::empty({2 * chunks, y.size(1)}, y.options());
auto stats = at::empty({2, y.size(1)}, y.options());
launch_train_fused(y, scale, gamma, beta, running_mean, running_var, partial, stats, momentum, eps);
} else {
launch_eval_fused(y, scale, gamma, beta, running_mean, running_var, eps);
}
return y;
}
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define CHUNKS 8
__global__ void partial_stats_kernel(const float* __restrict__ y,
const float* __restrict__ scale,
float* __restrict__ partial,
int N, int C) {
int lane = threadIdx.x & 31;
int col_base = blockIdx.x * 32;
int c = col_base + lane;
int chunk = blockIdx.y;
if (c >= C) return;
int start = (N * chunk) / CHUNKS;
int end = (N * (chunk + 1)) / CHUNKS;
float s = scale[c];
float sum = 0.0f;
float sum2 = 0.0f;
for (int n = start; n < end; ++n) {
float z = y[(long long)n * C + c] * s;
sum += z;
sum2 += z * z;
}
partial[(long long)chunk * C + c] = sum;
partial[(long long)(CHUNKS + chunk) * C + c] = sum2;
}
__global__ void finalize_stats_kernel(const float* __restrict__ partial,
float* __restrict__ stats,
float* __restrict__ running_mean,
float* __restrict__ running_var,
int N, int C, float momentum, float eps) {
int lane = threadIdx.x & 31;
int c = blockIdx.x * 32 + lane;
if (c >= C) return;
float sum = 0.0f;
float sum2 = 0.0f;
#pragma unroll
for (int k = 0; k < CHUNKS; ++k) {
sum += partial[(long long)k * C + c];
sum2 += partial[(long long)(CHUNKS + k) * C + c];
}
float invN = 1.0f / (float)N;
float mean = sum * invN;
float var = sum2 * invN - mean * mean;
if (var < 0.0f) var = 0.0f;
float invstd = rsqrtf(var + eps);
stats[c] = mean;
stats[(long long)C + c] = invstd;
// PyTorch BatchNorm updates running_var with the unbiased batch variance.
float unbiased = (N > 1) ? (var * ((float)N / (float)(N - 1))) : var;
running_mean[c] = running_mean[c] * (1.0f - momentum) + mean * momentum;
running_var[c] = running_var[c] * (1.0f - momentum) + unbiased * momentum;
}
__global__ void norm_train_kernel(float* __restrict__ y,
const float* __restrict__ scale,
const float* __restrict__ gamma,
const float* __restrict__ beta,
const float* __restrict__ stats,
int N, int C) {
int lane = threadIdx.x & 31;
int c = blockIdx.x * 32 + lane;
int chunk = blockIdx.y;
if (c >= C) return;
int start = (N * chunk) / CHUNKS;
int end = (N * (chunk + 1)) / CHUNKS;
float s = scale[c];
float mean = stats[c];
float invstd = stats[(long long)C + c];
float g = gamma[c];
float b = beta[c];
for (int n = start; n < end; ++n) {
long long idx = (long long)n * C + c;
float z = y[idx] * s;
y[idx] = (z - mean) * invstd * g + b;
}
}
__global__ void norm_eval_kernel(float* __restrict__ y,
const float* __restrict__ scale,
const float* __restrict__ gamma,
const float* __restrict__ beta,
const float* __restrict__ running_mean,
const float* __restrict__ running_var,
int N, int C, float eps) {
int lane = threadIdx.x & 31;
int c = blockIdx.x * 32 + lane;
int chunk = blockIdx.y;
if (c >= C) return;
int start = (N * chunk) / CHUNKS;
int end = (N * (chunk + 1)) / CHUNKS;
float s = scale[c];
float mean = running_mean[c];
float invstd = rsqrtf(running_var[c] + eps);
float g = gamma[c];
float b = beta[c];
for (int n = start; n < end; ++n) {
long long idx = (long long)n * C + c;
float z = y[idx] * s;
y[idx] = (z - mean) * invstd * g + b;
}
}
void launch_train_fused(torch::Tensor y, torch::Tensor scale, torch::Tensor gamma, torch::Tensor beta,
torch::Tensor running_mean, torch::Tensor running_var,
torch::Tensor partial, torch::Tensor stats,
double momentum, double eps) {
int N = (int)y.size(0);
int C = (int)y.size(1);
dim3 block(32);
dim3 grid((C + 31) / 32, CHUNKS);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
partial_stats_kernel<<<grid, block, 0, stream>>>(
y.data_ptr<float>(), scale.data_ptr<float>(), partial.data_ptr<float>(), N, C);
dim3 grid2((C + 31) / 32);
finalize_stats_kernel<<<grid2, block, 0, stream>>>(
partial.data_ptr<float>(), stats.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
N, C, (float)momentum, (float)eps);
norm_train_kernel<<<grid, block, 0, stream>>>(
y.data_ptr<float>(), scale.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), stats.data_ptr<float>(), N, C);
}
void launch_eval_fused(torch::Tensor y, torch::Tensor scale, torch::Tensor gamma, torch::Tensor beta,
torch::Tensor running_mean, torch::Tensor running_var, double eps) {
int N = (int)y.size(0);
int C = (int)y.size(1);
dim3 block(32);
dim3 grid((C + 31) / 32, CHUNKS);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
norm_eval_kernel<<<grid, block, 0, stream>>>(
y.data_ptr<float>(), scale.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
running_mean.data_ptr<float>(), running_var.data_ptr<float>(), N, C, (float)eps);
}
'''
_ext = load_inline(
name='fused_gemm_scale_bn_inline_ext',
cpp_sources=_cpp_src,
cuda_sources=_cuda_src,
functions=['fused_forward'],
extra_cflags=['-O3'],
extra_cuda_cflags=['-O3'],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1):
super(ModelNew, self).__init__()
self.gemm = nn.Linear(in_features, out_features)
self.scale = nn.Parameter(torch.randn(scale_shape))
self.bn = nn.BatchNorm1d(out_features, eps=eps, momentum=momentum)
def forward(self, x):
if self.training:
# Match BatchNorm1d bookkeeping and cumulative-momentum behavior.
self.bn.num_batches_tracked.add_(1)
if self.bn.momentum is None:
exp_factor = 1.0 / float(self.bn.num_batches_tracked.item())
else:
exp_factor = float(self.bn.momentum)
else:
exp_factor = 0.0
return _ext.fused_forward(
x, self.gemm.weight, self.gemm.bias, self.scale,
self.bn.weight, self.bn.bias, self.bn.running_mean, self.bn.running_var,
bool(self.training), float(exp_factor), float(self.bn.eps)
)