forked from ccf-ai-infra/GPUCodeForces
341 lines
11 KiB
Python
341 lines
11 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from torch.utils.cpp_extension import load_inline
|
||
|
||
batchnorm_source = r"""
|
||
#include <torch/extension.h>
|
||
#include <cuda_runtime.h>
|
||
#include <ATen/cuda/CUDAContext.h>
|
||
#include <c10/cuda/CUDAException.h>
|
||
|
||
// 统一的训练 kernel（计算批次统计量）
|
||
__global__ void batchnorm_forward_train_kernel_optimized(
|
||
const float* __restrict__ x,
|
||
const float* __restrict__ gamma,
|
||
const float* __restrict__ beta,
|
||
float* __restrict__ running_mean,
|
||
float* __restrict__ running_var,
|
||
float* __restrict__ y,
|
||
int batch,
|
||
int features,
|
||
float eps,
|
||
float momentum,
|
||
bool update_stats // 是否更新统计量
|
||
) {
|
||
int feature = blockIdx.x;
|
||
if (feature >= features) return;
|
||
|
||
int tid = threadIdx.x;
|
||
int num_threads = blockDim.x;
|
||
int warp_id = tid / 32;
|
||
int lane_id = tid % 32;
|
||
int num_warps = (num_threads + 31) / 32;
|
||
|
||
const float* x_base = x + feature;
|
||
float* y_base = y + feature;
|
||
|
||
float sum = 0.0f;
|
||
float sum_sq = 0.0f;
|
||
|
||
int row = tid;
|
||
for (; row + num_threads <= batch; row += num_threads) {
|
||
float v = x_base[row * features];
|
||
sum += v;
|
||
sum_sq += v * v;
|
||
}
|
||
|
||
if (row < batch) {
|
||
float v = x_base[row * features];
|
||
sum += v;
|
||
sum_sq += v * v;
|
||
}
|
||
|
||
#pragma unroll
|
||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||
sum += __shfl_down_sync(0xffffffff, sum, offset);
|
||
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
|
||
}
|
||
|
||
__shared__ float shared_sum[32];
|
||
__shared__ float shared_sq[32];
|
||
|
||
if (lane_id == 0) {
|
||
shared_sum[warp_id] = sum;
|
||
shared_sq[warp_id] = sum_sq;
|
||
}
|
||
__syncthreads();
|
||
|
||
if (tid < 32) {
|
||
sum = (tid < num_warps) ? shared_sum[tid] : 0.0f;
|
||
sum_sq = (tid < num_warps) ? shared_sq[tid] : 0.0f;
|
||
|
||
#pragma unroll
|
||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||
sum += __shfl_down_sync(0xffffffff, sum, offset);
|
||
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
|
||
}
|
||
}
|
||
|
||
__shared__ float s_mean;
|
||
__shared__ float s_inv_std;
|
||
__shared__ float s_gamma;
|
||
__shared__ float s_beta;
|
||
|
||
if (tid == 0) {
|
||
float mean = sum / batch;
|
||
float var = (sum_sq / batch) - (mean * mean);
|
||
var = fmaxf(var, 0.0f);
|
||
s_mean = mean;
|
||
s_inv_std = rsqrtf(var + eps);
|
||
s_gamma = gamma[feature];
|
||
s_beta = beta[feature];
|
||
|
||
// 只有需要时才更新 running stats
|
||
if (update_stats) {
|
||
running_mean[feature] = (1.0f - momentum) * running_mean[feature] + momentum * mean;
|
||
float unbiased_var = var * batch / fmaxf(float(batch - 1), 1.0f);
|
||
running_var[feature] = (1.0f - momentum) * running_var[feature] + momentum * unbiased_var;
|
||
}
|
||
}
|
||
__syncthreads();
|
||
|
||
float mean = s_mean;
|
||
float inv_std = s_inv_std;
|
||
float g = s_gamma;
|
||
float b = s_beta;
|
||
|
||
row = tid;
|
||
for (; row + num_threads <= batch; row += num_threads) {
|
||
float v = x_base[row * features];
|
||
float norm = (v - mean) * inv_std;
|
||
y_base[row * features] = norm * g + b;
|
||
}
|
||
|
||
if (row < batch) {
|
||
float v = x_base[row * features];
|
||
float norm = (v - mean) * inv_std;
|
||
y_base[row * features] = norm * g + b;
|
||
}
|
||
}
|
||
|
||
// 推理模式 kernel（使用 running stats）
|
||
__global__ void batchnorm_forward_eval_kernel_optimized(
|
||
const float* __restrict__ x,
|
||
const float* __restrict__ gamma,
|
||
const float* __restrict__ beta,
|
||
const float* __restrict__ running_mean,
|
||
const float* __restrict__ running_var,
|
||
float* __restrict__ y,
|
||
int batch,
|
||
int features,
|
||
float eps
|
||
) {
|
||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||
int total = batch * features;
|
||
int stride = gridDim.x * blockDim.x;
|
||
|
||
for (int idx = tid; idx < total; idx += stride) {
|
||
int feature = idx % features;
|
||
|
||
float mean = running_mean[feature];
|
||
float var = running_var[feature];
|
||
float inv_std = rsqrtf(var + eps);
|
||
float g = gamma[feature];
|
||
float b = beta[feature];
|
||
|
||
float v = x[idx];
|
||
float norm = (v - mean) * inv_std;
|
||
y[idx] = norm * g + b;
|
||
}
|
||
}
|
||
|
||
torch::Tensor batchnorm_cuda_forward(
|
||
torch::Tensor x,
|
||
torch::Tensor weight,
|
||
torch::Tensor bias,
|
||
torch::Tensor running_mean,
|
||
torch::Tensor running_var,
|
||
bool training,
|
||
double momentum,
|
||
double eps,
|
||
bool track_running_stats // 改名：更清晰地表达意图
|
||
) {
|
||
TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
|
||
TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
|
||
TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
|
||
TORCH_CHECK(x.dtype() == torch::kFloat32, "only float32 tensors are supported");
|
||
TORCH_CHECK(weight.dtype() == torch::kFloat32, "weight must be float32");
|
||
TORCH_CHECK(bias.dtype() == torch::kFloat32, "bias must be float32");
|
||
TORCH_CHECK(x.dim() == 2, "input must be 2D [batch, features]");
|
||
TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
|
||
TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
|
||
TORCH_CHECK(x.size(1) == weight.size(0), "feature size mismatch");
|
||
TORCH_CHECK(weight.size(0) == bias.size(0), "weight and bias must have the same length");
|
||
|
||
auto x_contig = x.contiguous();
|
||
auto weight_contig = weight.contiguous();
|
||
auto bias_contig = bias.contiguous();
|
||
|
||
int batch = x_contig.size(0);
|
||
int features = x_contig.size(1);
|
||
|
||
auto y = torch::empty_like(x_contig);
|
||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||
|
||
TORCH_CHECK(running_mean.is_cuda(), "running_mean must be a CUDA tensor");
|
||
TORCH_CHECK(running_var.is_cuda(), "running_var must be a CUDA tensor");
|
||
TORCH_CHECK(running_mean.dim() == 1, "running_mean must be 1D");
|
||
TORCH_CHECK(running_var.dim() == 1, "running_var must be 1D");
|
||
TORCH_CHECK(running_mean.size(0) == features, "running_mean size mismatch");
|
||
TORCH_CHECK(running_var.size(0) == features, "running_var size mismatch");
|
||
|
||
// 关键修改：根据 track_running_stats 决定行为
|
||
// track_running_stats=False: 总是计算批次统计（训练和推理都一样）
|
||
// track_running_stats=True + training: 计算批次统计并更新 running stats
|
||
// track_running_stats=True + eval: 使用 running stats
|
||
|
||
bool use_batch_stats = !track_running_stats || training;
|
||
|
||
if (use_batch_stats) {
|
||
// 使用批次统计量（训练模式 或 track_running_stats=False）
|
||
int threads;
|
||
if (batch <= 16) {
|
||
threads = 32;
|
||
} else if (batch <= 32) {
|
||
threads = 32;
|
||
} else if (batch <= 64) {
|
||
threads = 64;
|
||
} else if (batch <= 128) {
|
||
threads = 128;
|
||
} else if (batch <= 256) {
|
||
threads = 256;
|
||
} else {
|
||
threads = 256;
|
||
}
|
||
|
||
int blocks = features;
|
||
size_t shared_mem = 0;
|
||
|
||
// update_stats = track_running_stats && training
|
||
// track_running_stats=False: 不更新
|
||
// track_running_stats=True + training: 更新
|
||
// track_running_stats=True + eval: 不会走到这里
|
||
bool update_stats = track_running_stats && training;
|
||
|
||
batchnorm_forward_train_kernel_optimized<<<blocks, threads, shared_mem, stream>>>(
|
||
x_contig.data_ptr<float>(),
|
||
weight_contig.data_ptr<float>(),
|
||
bias_contig.data_ptr<float>(),
|
||
running_mean.data_ptr<float>(),
|
||
running_var.data_ptr<float>(),
|
||
y.data_ptr<float>(),
|
||
batch,
|
||
features,
|
||
static_cast<float>(eps),
|
||
static_cast<float>(momentum),
|
||
update_stats
|
||
);
|
||
} else {
|
||
// 使用 running stats（track_running_stats=True + eval 模式）
|
||
int total = batch * features;
|
||
int threads = 256;
|
||
int blocks;
|
||
|
||
if (total <= 4096) {
|
||
blocks = (total + threads - 1) / threads;
|
||
} else {
|
||
blocks = min(1024, (total + threads * 4 - 1) / (threads * 4));
|
||
}
|
||
|
||
batchnorm_forward_eval_kernel_optimized<<<blocks, threads, 0, stream>>>(
|
||
x_contig.data_ptr<float>(),
|
||
weight_contig.data_ptr<float>(),
|
||
bias_contig.data_ptr<float>(),
|
||
running_mean.data_ptr<float>(),
|
||
running_var.data_ptr<float>(),
|
||
y.data_ptr<float>(),
|
||
batch,
|
||
features,
|
||
static_cast<float>(eps)
|
||
);
|
||
}
|
||
|
||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||
return y;
|
||
}
|
||
"""
|
||
|
||
batchnorm_cpp_source = r"""
|
||
torch::Tensor batchnorm_cuda_forward(
|
||
torch::Tensor x,
|
||
torch::Tensor weight,
|
||
torch::Tensor bias,
|
||
torch::Tensor running_mean,
|
||
torch::Tensor running_var,
|
||
bool training,
|
||
double momentum,
|
||
double eps,
|
||
bool track_running_stats
|
||
);
|
||
"""
|
||
|
||
batchnorm_cuda = load_inline(
|
||
name="batchnorm_cuda_ext",
|
||
cpp_sources=batchnorm_cpp_source,
|
||
cuda_sources=batchnorm_source,
|
||
functions=["batchnorm_cuda_forward"],
|
||
verbose=True
|
||
)
|
||
|
||
class ModelNew(nn.Module):
|
||
"""
|
||
Model performing matrix multiplication followed by custom CUDA BatchNorm and ReLU.
|
||
Optimized with Warp-level reduction (Plan 1) and thread configuration (Plan 2).
|
||
"""
|
||
def __init__(self, mat_weight: torch.Tensor, bn_weight: torch.Tensor, bn_bias: torch.Tensor,
|
||
eps: float = 1e-5, momentum: float = 0.1, track_running_stats: bool = True):
|
||
super().__init__()
|
||
if mat_weight.dim() != 2:
|
||
raise ValueError("mat_weight must be a 2D tensor [input_dim, output_dim].")
|
||
if bn_weight.dim() != 1 or bn_bias.dim() != 1:
|
||
raise ValueError("BatchNorm weight and bias must be 1D.")
|
||
if bn_weight.size(0) != mat_weight.size(1):
|
||
raise ValueError("BatchNorm parameter size must match output_dim.")
|
||
if bn_weight.size(0) != bn_bias.size(0):
|
||
raise ValueError("BatchNorm weight and bias must share shape.")
|
||
|
||
self.weight = nn.Parameter(mat_weight.clone())
|
||
self.bn_weight = nn.Parameter(bn_weight.clone())
|
||
self.bn_bias = nn.Parameter(bn_bias.clone())
|
||
self.eps = eps
|
||
self.momentum = momentum
|
||
self.track_running_stats = track_running_stats
|
||
|
||
# 无论 track_running_stats 是什么，都创建 buffer
|
||
self.register_buffer('running_mean', torch.zeros(bn_weight.size(0)))
|
||
self.register_buffer('running_var', torch.ones(bn_weight.size(0)))
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
if not x.is_cuda:
|
||
raise ValueError("Input must be a CUDA tensor.")
|
||
if not self.weight.is_cuda:
|
||
raise ValueError("Model weight must be on CUDA.")
|
||
if not self.bn_weight.is_cuda or not self.bn_bias.is_cuda:
|
||
raise ValueError("BatchNorm parameters must be on CUDA.")
|
||
|
||
x = torch.matmul(x, self.weight)
|
||
|
||
# 传递 track_running_stats 参数到 CUDA kernel
|
||
x = batchnorm_cuda.batchnorm_cuda_forward(
|
||
x,
|
||
self.bn_weight,
|
||
self.bn_bias,
|
||
self.running_mean,
|
||
self.running_var,
|
||
self.training,
|
||
self.momentum,
|
||
self.eps,
|
||
self.track_running_stats
|
||
)
|
||
|
||
return torch.relu(x) |