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

Some refactoring for this 1 channel per block kernel

Browse Source
This commit is contained in:
Mohsen Saffari
2025-12-09 09:13:31 +00:00
parent 2b32dd75ee
commit 04d3a0ced0
3 changed files with 101 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -326,21 +326,26 @@ bool run(const ck_tile::ArgParser& arg_parser)
bool save_pass = ck_tile::check_err(save_mean_host, save_mean_ref, "Error: Saved mean incorrect!", 1e-3, 1e-3);
save_pass = save_pass && ck_tile::check_err(save_inv_std_host, save_inv_std_ref, "Error: Saved inv_std incorrect!", 1e-3, 1e-3);
std::cout << "\n=== Saved Statistics ===" << std::endl;
for(ck_tile::index_t c = 0; c < std::min(C, ck_tile::index_t(4)); ++c)
std::cout << "\n=== Saved Statistics (All Channels) ===" << std::endl;
for(ck_tile::index_t c = 0; c < C; ++c)
{
std::cout << "Ch" << std::setw(2) << c
float mean_diff = std::abs(save_mean_ref.mData[c] - save_mean_host.mData[c]);
float inv_std_diff = std::abs(save_inv_std_ref.mData[c] - save_inv_std_host.mData[c]);
std::cout << "Ch" << std::setw(3) << c
<< " mean: Ref=" << std::setw(10) << save_mean_ref.mData[c]
<< " Dev=" << std::setw(10) << save_mean_host.mData[c]
<< " Diff=" << std::setw(10) << mean_diff
<< " | inv_std: Ref=" << std::setw(10) << save_inv_std_ref.mData[c]
<< " Dev=" << std::setw(10) << save_inv_std_host.mData[c] << std::endl;
<< " Dev=" << std::setw(10) << save_inv_std_host.mData[c]
<< " Diff=" << std::setw(10) << inv_std_diff << std::endl;
}
pass = pass && save_pass;
}
if constexpr(kUpdateMovingAverage)
{
if(repeat == 1)
if(repeat == 1 && warmup == 0)
{
running_mean_buf.FromDevice(running_mean_host.mData.data());
running_var_buf.FromDevice(running_var_host.mData.data());
@@ -348,14 +353,19 @@ bool run(const ck_tile::ArgParser& arg_parser)
bool running_pass = ck_tile::check_err(running_mean_host, running_mean_ref, "Error: Running mean incorrect!", 1e-3, 1e-3);
running_pass = running_pass && ck_tile::check_err(running_var_host, running_var_ref, "Error: Running var incorrect!", 1e-3, 1e-3);
std::cout << "\n=== Running Statistics ===" << std::endl;
for(ck_tile::index_t c = 0; c < std::min(C, ck_tile::index_t(4)); ++c)
std::cout << "\n=== Running Statistics (All Channels) ===" << std::endl;
for(ck_tile::index_t c = 0; c < C; ++c)
{
std::cout << "Ch" << std::setw(2) << c
float mean_diff = std::abs(running_mean_ref.mData[c] - running_mean_host.mData[c]);
float var_diff = std::abs(running_var_ref.mData[c] - running_var_host.mData[c]);
std::cout << "Ch" << std::setw(3) << c
<< " mean: Ref=" << std::setw(10) << running_mean_ref.mData[c]
<< " Dev=" << std::setw(10) << running_mean_host.mData[c]
<< " Diff=" << std::setw(10) << mean_diff
<< " | var: Ref=" << std::setw(10) << running_var_ref.mData[c]
<< " Dev=" << std::setw(10) << running_var_host.mData[c] << std::endl;
<< " Dev=" << std::setw(10) << running_var_host.mData[c]
<< " Diff=" << std::setw(10) << var_diff << std::endl;
}
pass = pass && running_pass;
}

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

@@ -1,5 +1,5 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
@@ -9,9 +9,32 @@
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_problem.hpp"
#include "ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp"
/**
* @file batchnorm_fwd_kernel.hpp
* @brief Batch Normalization Forward Pass Kernel
*
* Normalizes inputs per-channel across batch and spatial dimensions using Welford's algorithm.
* Computes: y = gamma * (x - mean) / sqrt(variance + epsilon) + beta
*
* Supports NHWC tensor layout with optional features controlled by compile-time traits:
* - Save mean/inv_std (kSaveMeanInvStd): Stores statistics for backward pass
* - Update running stats (kUpdateMovingAverage): Maintains exponential moving average for inference
*
* **Welford's Algorithm:**
* For each element x_i:
* delta = x_i - mean
* mean = mean + delta / count
* M2 = M2 + delta * (x_i - mean)
* Final: variance = M2 / count
*
* **Running Statistics Update:**
* running = (1 - momentum) * running_old + momentum * batch
*/
namespace ck_tile {
// Host-side arguments for batchnorm forward pass
/// @brief Host arguments for batch normalization forward pass.
/// All tensors use NHWC (channels-last) layout: [N, H, W, C]
struct BatchnormFwdHostArgs
{
const void* p_x; // [N, H, W, C] input tensor (required, NHWC layout)
@@ -33,12 +56,14 @@ struct BatchnormFwdHostArgs
// Note: save/update flags are now in Traits (compile-time), not here (runtime)
};
// BatchnormFwd: Forward pass batch normalization kernel
/// @brief Batch Normalization Forward Pass Kernel
/// @tparam Problem_ Problem specification defining data types, block shape, and traits
template <typename Problem_>
struct BatchnormFwd
{
// Type aliases from Problem
using Problem = remove_cvref_t<Problem_>;
using Pipeline = BatchnormFwdPipeline<Problem>; // Class-level, like layernorm2d
using Pipeline = BatchnormFwdPipeline<Problem>;
using XDataType = typename Problem::XDataType;
using GammaDataType = typename Problem::GammaDataType;
using BetaDataType = typename Problem::BetaDataType;
@@ -47,6 +72,7 @@ struct BatchnormFwd
using MeanVarDataType = typename Problem::MeanVarDataType;
using BlockShape = typename Problem::BlockShape;
// Tile configuration
static constexpr index_t kBlockSize = BlockShape::BlockSize;
static constexpr index_t Block_M = BlockShape::Block_M;
static constexpr index_t Block_N = BlockShape::Block_N;
@@ -56,21 +82,19 @@ struct BatchnormFwd
// Kernel arguments
struct BatchnormFwdKargs
{
const void* p_x;
const void* p_gamma;
const void* p_beta;
void* p_y;
void* p_running_mean;
void* p_running_var;
void* p_save_mean;
void* p_save_inv_std;
const void* p_x; // Input tensor [N,H,W,C]
const void* p_gamma; // Scale parameters [C]
const void* p_beta; // Bias parameters [C]
void* p_y; // Output tensor [N,H,W,C]
void* p_running_mean; // Running mean [C] (optional)
void* p_running_var; // Running variance [C] (optional)
void* p_save_mean; // Saved mean [C] (optional)
void* p_save_inv_std; // Saved 1/sqrt(var+eps) [C] (optional)
float epsilon;
float momentum;
float epsilon; // Numerical stability constant
float momentum; // Exponential moving average factor
index_t N, C, H, W;
// Note: save/update flags now come from Problem::Traits (compile-time)
index_t N, C, H, W; // Batch, channels, height, width
};
using Kargs = BatchnormFwdKargs; // Alias for convenience
@@ -101,7 +125,7 @@ struct BatchnormFwd
return dim3(hargs.C); // One block per channel
}
// Block size (wave32/64 aware like layernorm2d)
// Block size (architecture-aware for wave32/wave64)
CK_TILE_HOST static constexpr auto BlockSize()
{
return is_wave32() ? BlockShape::template GetBlockSize<true>()
@@ -114,6 +138,43 @@ struct BatchnormFwd
return Pipeline::GetSmemSize();
}
// Validate arguments
CK_TILE_HOST static bool IsSupportedArgument(const Hargs& hargs)
{
// Basic validation
if(hargs.N <= 0 || hargs.C <= 0 || hargs.H <= 0 || hargs.W <= 0)
{
return false;
}
// Validate required pointers
if(hargs.p_x == nullptr || hargs.p_y == nullptr ||
hargs.p_gamma == nullptr || hargs.p_beta == nullptr)
{
return false;
}
// Validate optional pointers based on Traits (compile-time)
if constexpr(Problem::Traits::kUpdateMovingAverage)
{
if(hargs.p_running_mean == nullptr || hargs.p_running_var == nullptr)
{
return false;
}
}
if constexpr(Problem::Traits::kSaveMeanInvStd)
{
if(hargs.p_save_mean == nullptr || hargs.p_save_inv_std == nullptr)
{
return false;
}
}
return true;
}
/// @brief Kernel execution - processes one channel per block
CK_TILE_DEVICE void operator()(Kargs kargs) const
{
const index_t c = get_block_id();
@@ -184,41 +245,7 @@ struct BatchnormFwd
smem);
}
// Validate arguments
CK_TILE_HOST static bool IsSupportedArgument(const Hargs& hargs)
{
// Basic validation
if(hargs.N <= 0 || hargs.C <= 0 || hargs.H <= 0 || hargs.W <= 0)
{
return false;
}
// Validate required pointers
if(hargs.p_x == nullptr || hargs.p_y == nullptr ||
hargs.p_gamma == nullptr || hargs.p_beta == nullptr)
{
return false;
}
// Validate optional pointers based on Traits (compile-time)
if constexpr(Problem::Traits::kUpdateMovingAverage)
{
if(hargs.p_running_mean == nullptr || hargs.p_running_var == nullptr)
{
return false;
}
}
if constexpr(Problem::Traits::kSaveMeanInvStd)
{
if(hargs.p_save_mean == nullptr || hargs.p_save_inv_std == nullptr)
{
return false;
}
}
return true;
}
};
} // namespace ck_tile

4
include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp
View File

@@ -11,7 +11,7 @@ namespace ck_tile {
// Defines tile distributions and helper functions
struct BatchnormFwdPipelineDefaultPolicy
{
// Tile distribution for input data (following layernorm2d pattern exactly)
// Tile distribution for input data
template <typename Problem>
CK_TILE_DEVICE static constexpr auto MakeXBlockTileDistribution()
{
@@ -28,7 +28,7 @@ struct BatchnormFwdPipelineDefaultPolicy
sequence<0, 3, 0, 3>>{});
}
// Tile distribution for gamma/beta parameters (following layernorm2d pattern)
// Tile distribution for gamma/beta parameters
template <typename Problem>
CK_TILE_DEVICE static constexpr auto MakeGammaBetaBlockTileDistribution()
{