final submitted kernel (181 lines, 7,622 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cpp_source = r'''
#include <torch/extension.h>
torch::Tensor bn_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
torch::Tensor running_mean, torch::Tensor running_var,
bool training, double momentum, double eps);
'''
cuda_source = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
static __inline__ __device__ float warp_reduce_sum(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_down_sync(0xffffffff, val, offset);
return val;
}
static __inline__ __device__ float block_reduce_sum(float val) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int wid = threadIdx.x >> 5;
val = warp_reduce_sum(val);
if (lane == 0) shared[wid] = val;
__syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
if (wid == 0) val = warp_reduce_sum(val);
return val;
}
__global__ void partial_stats_kernel(const float* __restrict__ x,
float* __restrict__ psum,
float* __restrict__ psq,
int C, int S, int N, int P) {
int p = blockIdx.x;
int c = blockIdx.y;
int tid = threadIdx.x;
int s = p * blockDim.x + tid;
float sum = 0.0f;
float sq = 0.0f;
if (s < S) {
int64_t off = ((int64_t)c) * S + s;
int64_t step = (int64_t)C * S;
#pragma unroll 4
for (int n = 0; n < N; ++n, off += step) {
float v = x[off];
sum += v;
sq = fmaf(v, v, sq);
}
}
sum = block_reduce_sum(sum);
sq = block_reduce_sum(sq);
if (tid == 0) {
psum[c * P + p] = sum;
psq[c * P + p] = sq;
}
}
__global__ void final_stats_kernel(const float* __restrict__ psum,
const float* __restrict__ psq,
float* __restrict__ running_mean,
float* __restrict__ running_var,
float* __restrict__ mean_out,
float* __restrict__ inv_out,
int S, int N, int P,
float momentum, float eps) {
int c = blockIdx.x;
int tid = threadIdx.x;
float sum = 0.0f, sq = 0.0f;
for (int i = tid; i < P; i += blockDim.x) {
sum += psum[c * P + i];
sq += psq[c * P + i];
}
sum = block_reduce_sum(sum);
sq = block_reduce_sum(sq);
if (tid == 0) {
float M = (float)((int64_t)N * S);
float mean = sum / M;
float var = fmaxf(sq / M - mean * mean, 0.0f);
float inv = rsqrtf(var + eps);
mean_out[c] = mean;
inv_out[c] = inv;
float unbiased = var * (M / (M - 1.0f));
running_mean[c] = running_mean[c] * (1.0f - momentum) + momentum * mean;
running_var[c] = running_var[c] * (1.0f - momentum) + momentum * unbiased;
}
}
__global__ void normalize_train_kernel(const float* __restrict__ x,
const float* __restrict__ mean,
const float* __restrict__ inv,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ y,
int C, int S, int N) {
int p = blockIdx.x;
int c = blockIdx.y;
int tid = threadIdx.x;
int s = p * blockDim.x + tid;
if (s >= S) return;
float scale = inv[c] * weight[c];
float b = bias[c] - mean[c] * scale;
int64_t off = ((int64_t)c) * S + s;
int64_t step = (int64_t)C * S;
#pragma unroll 4
for (int n = 0; n < N; ++n, off += step) {
y[off] = fmaf(x[off], scale, b);
}
}
__global__ void normalize_eval_kernel(const float* __restrict__ x,
const float* __restrict__ running_mean,
const float* __restrict__ running_var,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ y,
int64_t total, int C, int S, float eps) {
int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
int64_t stride = (int64_t)gridDim.x * blockDim.x;
for (; idx < total; idx += stride) {
int c = (int)((idx / S) % C);
float scale = rsqrtf(running_var[c] + eps) * weight[c];
float b = bias[c] - running_mean[c] * scale;
y[idx] = fmaf(x[idx], scale, b);
}
}
torch::Tensor bn_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
torch::Tensor running_mean, torch::Tensor running_var,
bool training, double momentum, double eps) {
CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(bias);
TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
TORCH_CHECK(x.dim() == 4, "expected NCHW 4D");
auto y = torch::empty_like(x);
int N = (int)x.size(0);
int C = (int)x.size(1);
int S = (int)(x.size(2) * x.size(3));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const int threads = 256;
if (training) {
int P = (S + threads - 1) / threads;
auto opts = x.options();
auto psum = torch::empty({C, P}, opts);
auto psq = torch::empty({C, P}, opts);
auto mean = torch::empty({C}, opts);
auto inv = torch::empty({C}, opts);
dim3 grid(P, C);
partial_stats_kernel<<<grid, threads, 0, stream>>>(x.data_ptr<float>(), psum.data_ptr<float>(), psq.data_ptr<float>(), C, S, N, P);
final_stats_kernel<<<C, threads, 0, stream>>>(psum.data_ptr<float>(), psq.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), mean.data_ptr<float>(), inv.data_ptr<float>(), S, N, P, (float)momentum, (float)eps);
normalize_train_kernel<<<grid, threads, 0, stream>>>(x.data_ptr<float>(), mean.data_ptr<float>(), inv.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), C, S, N);
} else {
int64_t total = x.numel();
int blocks = 4096;
normalize_eval_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), total, C, S, (float)eps);
}
return y;
}
'''
_bn_ext = load_inline(name='bn2d_custom_ext_final', cpp_sources=cpp_source, cuda_sources=cuda_source,
functions=['bn_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False)
class ModelNew(nn.Module):
def __init__(self, num_features: int):
super(ModelNew, self).__init__()
self.bn = nn.BatchNorm2d(num_features=num_features)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.bn.training:
self.bn.num_batches_tracked.add_(1)
return _bn_ext.bn_forward(x, self.bn.weight, self.bn.bias, self.bn.running_mean, self.bn.running_var,
bool(self.bn.training), float(self.bn.momentum), float(self.bn.eps))