ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 11:07:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

correct the batch normalization algorithm to calc reduce over NxHxW

Browse Source
This commit is contained in:
Mohsen Saffari
2025-11-25 16:46:00 +00:00
parent 5fa3ec2b37
commit c7afd540d9
2 changed files with 47 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
example/ck_tile/42_batchnorm/batchnorm_simple.cpp
View File

@@ -36,14 +36,15 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor<XDataType>& 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<XDataType>& x,
sum += ck_tile::type_convert<ComputeDataType>(x.mData[idx]);
}
}
ComputeDataType mean = sum / static_cast<ComputeDataType>(spatial_size);
// Compute variance
ComputeDataType var_sum = 0;
}
ComputeDataType mean = sum / static_cast<ComputeDataType>(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<XDataType>& x,
var_sum += diff * diff;
}
}
ComputeDataType variance = var_sum / static_cast<ComputeDataType>(spatial_size);
// Normalize
ComputeDataType inv_std = static_cast<ComputeDataType>(1.0) /
ck_tile::sqrt(variance + epsilon);
}
ComputeDataType variance = var_sum / static_cast<ComputeDataType>(per_channel_size);
// Normalize all values in this channel
ComputeDataType inv_std = static_cast<ComputeDataType>(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<Problem>;
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))
{

44
include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp
View File

@@ -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<ComputeDataType>(0);
ComputeDataType thread_m2 = type_convert<ComputeDataType>(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<ComputeDataType>(p_x[offset]);
// Online Welford update
@@ -92,14 +89,17 @@ struct BatchnormFwd
BlockWelford<ComputeDataType>::template Run<index_t, kBlockSize>(
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<ComputeDataType>(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<ComputeDataType>(p_x[offset]);
ComputeDataType normalized = (val - block_mean) * inv_std;
p_y[offset] = type_convert<YDataType>(normalized);