Merge pull request 'optimize BatchNorm1d operator' (#17) from xymdaysgone/GPUCodeForces:batchnorm1d into main

S1/10/batchnorm1d_cuda.py

@@ -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 eps,
+    float momentum,
+    bool update_stats  // 是否更新统计量
 ) {
     int feature = blockIdx.x;
     if (feature >= features) return;
-    int tid = threadIdx.x;
     
-    extern __shared__ float shared[];
+    int tid = threadIdx.x;
-    float* shm_sum = shared;
+    int num_threads = blockDim.x;
-    float* shm_sq = shared + 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;
     
-    for (int row = tid; row < batch; row += blockDim.x) {
+    int row = tid;
-        float v = x[row * features + feature];
+    for (; row + num_threads <= batch; row += num_threads) {
+        float v = x_base[row * features];
         sum += v;
         sum_sq += v * v;
     }
     
-    shm_sum[tid] = sum;
+    if (row < batch) {
-    shm_sq[tid] = sum_sq;
+        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();
     
-    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
+    if (tid < 32) {
-        if (tid < offset) {
+        sum = (tid < num_warps) ? shared_sum[tid] : 0.0f;
-            shm_sum[tid] += shm_sum[tid + offset];
+        sum_sq = (tid < num_warps) ? shared_sq[tid] : 0.0f;
-            shm_sq[tid] += shm_sq[tid + offset];
+        
+        #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);
         }
-        __syncthreads();
     }
     
     __shared__ float s_mean;
     __shared__ float s_inv_std;
+    __shared__ float s_gamma;
+    __shared__ float s_beta;
     
     if (tid == 0) {
-        float mean = shm_sum[0] / batch;
+        float mean = sum / batch;
-        float var = shm_sq[0] / batch - mean * mean;
+        float var = (sum_sq / batch) - (mean * mean);
-        var = var > 0.f ? var : 0.f;
+        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 = gamma[feature];
+    float g = s_gamma;
-    float b = beta[feature];
+    float b = s_beta;
     
-    for (int row = tid; row < batch; row += blockDim.x) {
+    row = tid;
-        float v = x[row * features + feature];
+    for (; row + num_threads <= batch; row += num_threads) {
+        float v = x_base[row * features];
         float norm = (v - mean) * inv_std;
-        y[row * features + feature] = norm * g + b;
+        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;
     }
 }
 
@@ -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>(),
+    TORCH_CHECK(running_mean.is_cuda(), "running_mean must be a CUDA tensor");
-        weight_contig.data_ptr<float>(),
+    TORCH_CHECK(running_var.is_cuda(), "running_var must be a CUDA tensor");
-        bias_contig.data_ptr<float>(),
+    TORCH_CHECK(running_mean.dim() == 1, "running_mean must be 1D");
-        y.data_ptr<float>(),
+    TORCH_CHECK(running_var.dim() == 1, "running_var must be 1D");
-        batch,
+    TORCH_CHECK(running_mean.size(0) == features, "running_mean size mismatch");
-        features,
+    TORCH_CHECK(running_var.size(0) == features, "running_var size mismatch");
-        static_cast<float>(eps)
+
-    );
+    // 关键修改：根据 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)

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 使其始终使用当前批次统计量
+        # 设置 track_running_stats=True 以跟踪运行时统计量
-        self.bn = nn.BatchNorm1d(num_features, eps=eps, momentum=momentum, track_running_stats=False)
+        self.bn = nn.BatchNorm1d(
+            num_features, 
+            eps=eps, 
+            momentum=momentum, 
+            track_running_stats=track_running_stats
+        )
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """

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()
+    track_bool = True
-    cuda_model = CudaModel(weight.clone(), bn_weight.clone(), bn_bias.clone(), eps=1e-5).cuda()
+    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