From c7afd540d9d24492d567da8419d83df89cbb5234 Mon Sep 17 00:00:00 2001 From: Mohsen Saffari Date: Tue, 25 Nov 2025 16:46:00 +0000 Subject: [PATCH] correct the batch normalization algorithm to calc reduce over NxHxW --- .../ck_tile/42_batchnorm/batchnorm_simple.cpp | 42 +++++++++++------- .../batchnorm/kernel/batchnorm_fwd_kernel.hpp | 44 +++++++++---------- 2 files changed, 47 insertions(+), 39 deletions(-) diff --git a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp index edf7c05ffe..b81c1f923b 100644 --- a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp +++ b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp @@ -36,14 +36,15 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor& x, ComputeDataType epsilon) { const ck_tile::index_t spatial_size = H * W; + const ck_tile::index_t per_channel_size = N * spatial_size; - // Process each (N, C) combination - for(ck_tile::index_t n = 0; n < N; ++n) + // Process each channel (compute statistics across ALL samples and spatial positions) + for(ck_tile::index_t c = 0; c < C; ++c) { - for(ck_tile::index_t c = 0; c < C; ++c) + // Compute mean across all N samples and H×W positions for this channel + ComputeDataType sum = 0; + for(ck_tile::index_t n = 0; n < N; ++n) { - // Compute mean - ComputeDataType sum = 0; for(ck_tile::index_t h = 0; h < H; ++h) { for(ck_tile::index_t w = 0; w < W; ++w) @@ -52,10 +53,13 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor& x, sum += ck_tile::type_convert(x.mData[idx]); } } - ComputeDataType mean = sum / static_cast(spatial_size); - - // Compute variance - ComputeDataType var_sum = 0; + } + ComputeDataType mean = sum / static_cast(per_channel_size); + + // Compute variance across all N samples and H×W positions for this channel + ComputeDataType var_sum = 0; + for(ck_tile::index_t n = 0; n < N; ++n) + { for(ck_tile::index_t h = 0; h < H; ++h) { for(ck_tile::index_t w = 0; w < W; ++w) @@ -66,12 +70,15 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor& x, var_sum += diff * diff; } } - ComputeDataType variance = var_sum / static_cast(spatial_size); - - // Normalize - ComputeDataType inv_std = static_cast(1.0) / - ck_tile::sqrt(variance + epsilon); - + } + ComputeDataType variance = var_sum / static_cast(per_channel_size); + + // Normalize all values in this channel + ComputeDataType inv_std = static_cast(1.0) / + ck_tile::sqrt(variance + epsilon); + + for(ck_tile::index_t n = 0; n < N; ++n) + { for(ck_tile::index_t h = 0; h < H; ++h) { for(ck_tile::index_t w = 0; w < W; ++w) @@ -131,9 +138,10 @@ bool run(const ck_tile::ArgParser& arg_parser) using Kernel = ck_tile::BatchnormFwd; const ck_tile::index_t kBlockSize = Kernel::BlockSize(); - const ck_tile::index_t kGridSize = N * C; // One block per (N, C) pair + const ck_tile::index_t kGridSize = C; // One block per channel - std::cout << "Kernel config: BlockSize=" << kBlockSize << ", GridSize=" << kGridSize << std::endl; + std::cout << "Kernel config: BlockSize=" << kBlockSize << ", GridSize=" << kGridSize + << " (one block per channel, reducing over N×H×W=" << N*H*W << " elements)" << std::endl; if(!Kernel::IsSupportedArgument(N, C, H, W)) { diff --git a/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp b/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp index 8feedd2c6e..e206eeaa19 100644 --- a/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp +++ b/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp @@ -38,39 +38,36 @@ struct BatchnormFwd ComputeDataType epsilon) const { // For batchnorm: input shape is [N, C, H, W] - // We reduce over H*W (spatial dimensions) for each N*C combination - // Each block handles one or more (N,C) combinations + // We reduce over N*H*W (batch and spatial) for EACH channel + // Each block handles ONE channel const index_t spatial_size = H * W; - const index_t nc_size = N * C; + const index_t per_channel_size = N * spatial_size; // Reduce over N×H×W const index_t thread_id = get_thread_id(); const index_t block_id = get_block_id(); - // For POC: simple mapping - each block handles one (N,C) pair - // More sophisticated tiling will come later - const index_t nc_idx = block_id; + // Each block handles one channel + const index_t c = block_id; - if(nc_idx >= nc_size) + if(c >= C) return; - // Calculate n and c from linear index - const index_t n = nc_idx / C; - const index_t c = nc_idx % C; - - // Calculate base offset for this (N,C) pair - const index_t base_offset = n * C * H * W + c * H * W; - // Thread-local Welford statistics ComputeDataType thread_mean = type_convert(0); ComputeDataType thread_m2 = type_convert(0); index_t thread_count = 0; - // Each thread processes a portion of the spatial dimension - // Simple strided access pattern - for(index_t idx = thread_id; idx < spatial_size; idx += kBlockSize) + // Each thread processes elements across ALL samples (N) and spatial positions (H×W) + // for this channel + for(index_t idx = thread_id; idx < per_channel_size; idx += kBlockSize) { - const index_t offset = base_offset + idx; + // Calculate which sample (n) and spatial position (hw) this is + const index_t n = idx / spatial_size; + const index_t hw = idx % spatial_size; + + // Memory layout: [N, C, H, W] with C-major for H×W + const index_t offset = n * C * H * W + c * H * W + hw; ComputeDataType val = type_convert(p_x[offset]); // Online Welford update @@ -92,14 +89,17 @@ struct BatchnormFwd BlockWelford::template Run( block_mean, block_var, block_count, smem); - // Now all threads have the same mean and variance - // Normalize and write output + // Now all threads have the same mean and variance for this channel + // Normalize and write output for ALL samples ComputeDataType inv_std = type_convert(1) / ck_tile::sqrt(block_var + epsilon); - for(index_t idx = thread_id; idx < spatial_size; idx += kBlockSize) + for(index_t idx = thread_id; idx < per_channel_size; idx += kBlockSize) { - const index_t offset = base_offset + idx; + const index_t n = idx / spatial_size; + const index_t hw = idx % spatial_size; + + const index_t offset = n * C * H * W + c * H * W + hw; ComputeDataType val = type_convert(p_x[offset]); ComputeDataType normalized = (val - block_mean) * inv_std; p_y[offset] = type_convert(normalized);