diff --git a/S1/10/batchnorm1d_cuda.py b/S1/10/batchnorm1d_cuda.py index 4b671c1..a71ec13 100644 --- a/S1/10/batchnorm1d_cuda.py +++ b/S1/10/batchnorm1d_cuda.py @@ -8,65 +8,144 @@ batchnorm_source = r""" #include #include -__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<<>>( - x_contig.data_ptr(), - weight_contig.data_ptr(), - bias_contig.data_ptr(), - y.data_ptr(), - batch, - features, - static_cast(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<<>>( + x_contig.data_ptr(), + weight_contig.data_ptr(), + bias_contig.data_ptr(), + running_mean.data_ptr(), + running_var.data_ptr(), + y.data_ptr(), + batch, + features, + static_cast(eps), + static_cast(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<<>>( + x_contig.data_ptr(), + weight_contig.data_ptr(), + bias_contig.data_ptr(), + running_mean.data_ptr(), + running_var.data_ptr(), + y.data_ptr(), + batch, + features, + static_cast(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) \ No newline at end of file diff --git a/S1/10/batchnorm1d_torch.py b/S1/10/batchnorm1d_torch.py index 9828be3..2f0e2cc 100644 --- a/S1/10/batchnorm1d_torch.py +++ b/S1/10/batchnorm1d_torch.py @@ -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: """ diff --git a/S1/10/run_code.py b/S1/10/run_code.py index b0349ce..8884d9e 100644 --- a/S1/10/run_code.py +++ b/S1/10/run_code.py @@ -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