forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'optimize BatchNorm1d operator' (#17) from xymdaysgone/GPUCodeForces:batchnorm1d into main
This commit is contained in:
commit
e309547055
|
|
@ -8,65 +8,144 @@ batchnorm_source = r"""
|
|||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
|
||||
__global__ void batchnorm_forward_kernel(
|
||||
// 统一的训练 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 feature = blockIdx.x;
|
||||
if (feature >= features) return;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
extern __shared__ float shared[];
|
||||
float* shm_sum = shared;
|
||||
float* shm_sq = shared + blockDim.x;
|
||||
|
||||
float sum = 0.0f;
|
||||
float sum_sq = 0.0f;
|
||||
|
||||
for (int row = tid; row < batch; row += blockDim.x) {
|
||||
float v = x[row * features + feature];
|
||||
sum += v;
|
||||
sum_sq += v * v;
|
||||
}
|
||||
|
||||
shm_sum[tid] = sum;
|
||||
shm_sq[tid] = sum_sq;
|
||||
__syncthreads();
|
||||
|
||||
for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
|
||||
if (tid < offset) {
|
||||
shm_sum[tid] += shm_sum[tid + offset];
|
||||
shm_sq[tid] += shm_sq[tid + offset];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
__shared__ float s_mean;
|
||||
__shared__ float s_inv_std;
|
||||
|
||||
if (tid == 0) {
|
||||
float mean = shm_sum[0] / batch;
|
||||
float var = shm_sq[0] / batch - mean * mean;
|
||||
var = var > 0.f ? var : 0.f;
|
||||
s_mean = mean;
|
||||
s_inv_std = rsqrtf(var + eps);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float mean = s_mean;
|
||||
float inv_std = s_inv_std;
|
||||
float g = gamma[feature];
|
||||
float b = beta[feature];
|
||||
|
||||
for (int row = tid; row < batch; row += blockDim.x) {
|
||||
float v = x[row * features + feature];
|
||||
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[row * features + feature] = norm * g + b;
|
||||
y[idx] = norm * g + b;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -74,7 +153,12 @@ torch::Tensor batchnorm_cuda_forward(
|
|||
torch::Tensor x,
|
||||
torch::Tensor weight,
|
||||
torch::Tensor bias,
|
||||
double eps
|
||||
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");
|
||||
|
|
@ -85,7 +169,7 @@ torch::Tensor batchnorm_cuda_forward(
|
|||
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 between input and weight");
|
||||
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();
|
||||
|
|
@ -96,25 +180,86 @@ torch::Tensor batchnorm_cuda_forward(
|
|||
int features = x_contig.size(1);
|
||||
|
||||
auto y = torch::empty_like(x_contig);
|
||||
|
||||
int threads = 256;
|
||||
if (batch < threads) {
|
||||
threads = 1;
|
||||
while (threads < batch) threads <<= 1;
|
||||
if (threads < 32) threads = 32;
|
||||
}
|
||||
size_t shared_mem = threads * 2 * sizeof(float);
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
batchnorm_forward_kernel<<<features, threads, shared_mem, stream>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
weight_contig.data_ptr<float>(),
|
||||
bias_contig.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
batch,
|
||||
features,
|
||||
static_cast<float>(eps)
|
||||
);
|
||||
|
||||
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;
|
||||
}
|
||||
|
|
@ -125,7 +270,12 @@ torch::Tensor batchnorm_cuda_forward(
|
|||
torch::Tensor x,
|
||||
torch::Tensor weight,
|
||||
torch::Tensor bias,
|
||||
double eps
|
||||
torch::Tensor running_mean,
|
||||
torch::Tensor running_var,
|
||||
bool training,
|
||||
double momentum,
|
||||
double eps,
|
||||
bool track_running_stats
|
||||
);
|
||||
"""
|
||||
|
||||
|
|
@ -140,8 +290,10 @@ batchnorm_cuda = load_inline(
|
|||
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):
|
||||
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].")
|
||||
|
|
@ -151,10 +303,17 @@ class ModelNew(nn.Module):
|
|||
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:
|
||||
|
|
@ -163,6 +322,20 @@ class ModelNew(nn.Module):
|
|||
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)
|
||||
x = batchnorm_cuda.batchnorm_cuda_forward(x, self.bn_weight, self.bn_bias, self.eps)
|
||||
|
||||
# 传递 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)
|
||||
|
|
@ -5,11 +5,16 @@ class BatchNormModel(nn.Module):
|
|||
"""
|
||||
Model that performs matrix multiplication followed by BatchNorm and ReLU activation.
|
||||
"""
|
||||
def __init__(self, weight, num_features=2048, eps=1e-5, momentum=0.1):
|
||||
def __init__(self, weight, num_features=2048, eps=1e-5, momentum=0.1, track_running_stats=True):
|
||||
super(BatchNormModel, self).__init__()
|
||||
self.weight = nn.Parameter(weight)
|
||||
# 设置 track_running_stats=False 使其始终使用当前批次统计量
|
||||
self.bn = nn.BatchNorm1d(num_features, eps=eps, momentum=momentum, track_running_stats=False)
|
||||
# 设置 track_running_stats=True 以跟踪运行时统计量
|
||||
self.bn = nn.BatchNorm1d(
|
||||
num_features,
|
||||
eps=eps,
|
||||
momentum=momentum,
|
||||
track_running_stats=track_running_stats
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -32,8 +32,9 @@ def run_benchmark():
|
|||
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs]
|
||||
|
||||
# 初始化两个模型
|
||||
torch_model = TorchModel(weight.clone(), num_features=output_dim, eps=1e-5).cuda()
|
||||
cuda_model = CudaModel(weight.clone(), bn_weight.clone(), bn_bias.clone(), eps=1e-5).cuda()
|
||||
track_bool = True
|
||||
torch_model = TorchModel(weight.clone(), num_features=output_dim, eps=1e-5, track_running_stats=track_bool).cuda()
|
||||
cuda_model = CudaModel(weight.clone(), bn_weight.clone(), bn_bias.clone(), eps=1e-5, track_running_stats=track_bool).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
|
@ -58,6 +59,7 @@ def run_benchmark():
|
|||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
print(f"track_running_stats = {track_bool}")
|
||||
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
|
||||
|
||||
# Warm up
|
||||
|
|
|
|||
Loading…
Reference in New Issue