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

Add scale and bias to batch normalization

Browse Source
This commit is contained in:
Mohsen Saffari
2025-12-04 16:22:32 +00:00
parent 60145b3cce
commit 8d4f766f9c
3 changed files with 25 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -142,12 +142,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
// Fill input with random data
ck_tile::FillUniformDistribution<XDataType>{-5.f, 5.f}(x_host);
// Set gamma=1.0 and beta=0.0 for testing (identity transform)
for(ck_tile::index_t c = 0; c < C; ++c)
{
gamma_host.mData[c] = static_cast<ComputeDataType>(1.0);
beta_host.mData[c] = static_cast<ComputeDataType>(0.0);
}
// Fill gamma and beta with random values (test scale/bias)
ck_tile::FillUniformDistribution<ComputeDataType>{0.8f, 1.2f}(gamma_host); // Scale around 1
ck_tile::FillUniformDistribution<ComputeDataType>{-0.5f, 0.5f}(beta_host); // Bias around 0
// Allocate device memory
ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes());
@@ -160,11 +157,15 @@ bool run(const ck_tile::ArgParser& arg_parser)
beta_buf.ToDevice(beta_host.data());
// Define kernel configuration using Generic2dBlockShape
// For N=2, H=8, W=8: per-channel elements = 2×8×8 = 128
// Use Repeat_N=2: Block_N = Repeat_N × ThreadPerBlock_N × Vector_N = 2×64×1 = 128 ✓
using BlockTile = ck_tile::sequence<1, 128>; // Block size: 1 channel, 128 spatial
using ThreadPerBlock = ck_tile::sequence<1, 128>; // 64 threads (must be warp size)
using Vector = ck_tile::sequence<1, 1>; // Vector size (start with 1)
// Vector_N controls vectorization: higher = fewer iterations, more elements per thread
// Block_N = ThreadPerBlock_N × Vector_N (must match tile size needed)
using BlockTile = ck_tile::sequence<1, 2048>; // Block size: 1 channel, 128 spatial
using ThreadPerBlock = ck_tile::sequence<1, 1024>; // 64 threads
using Vector = ck_tile::sequence<1, 2>; // Vector_N=2 (try 1,2,4,8)
// With Vector_N=2: 64 threads × 2 elements = 128 elements per tile
// With Vector_N=4: Need ThreadPerBlock=32 for 32×4=128
// Experiment to find optimal!
using Shape = ck_tile::BatchnormShape<BlockTile, ThreadPerBlock, Vector>;

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

@@ -139,28 +139,6 @@ struct BatchnormFwd
return make_tile_window(tmp2_, make_tuple(number<Block_M>{}, number<Block_N>{}), {0, 0});
}();
const auto gamma_window = [&]() {
const GammaDataType* p_gamma = static_cast<const GammaDataType*>(kargs.p_gamma);
const auto tmp_ = make_naive_tensor_view_packed<address_space_enum::global>(
p_gamma + c,
make_tuple(1),
number<1>{});
const auto tmp2_ = pad_tensor_view(tmp_, make_tuple(number<Block_M>{}), sequence<false>{});
return make_tile_window(tmp2_, make_tuple(number<Block_M>{}), {0});
}();
const auto beta_window = [&]() {
const BetaDataType* p_beta = static_cast<const BetaDataType*>(kargs.p_beta);
const auto tmp_ = make_naive_tensor_view_packed<address_space_enum::global>(
p_beta + c,
make_tuple(1),
number<1>{});
const auto tmp2_ = pad_tensor_view(tmp_, make_tuple(number<Block_M>{}), sequence<false>{});
return make_tile_window(tmp2_, make_tuple(number<Block_M>{}), {0});
}();
auto y_window = [&]() {
YDataType* p_y = static_cast<YDataType*>(kargs.p_y);
YDataType* p_y_channel = p_y + c; // Offset by c (NHWC)
@@ -187,10 +165,13 @@ struct BatchnormFwd
MeanVarDataType* p_save_mean = static_cast<MeanVarDataType*>(kargs.p_save_mean);
MeanVarDataType* p_save_inv_std = static_cast<MeanVarDataType*>(kargs.p_save_inv_std);
// Call pipeline with properly distributed tile windows
// Call pipeline with x/y windows and gamma/beta pointers
const GammaDataType* p_gamma = static_cast<const GammaDataType*>(kargs.p_gamma);
const BetaDataType* p_beta = static_cast<const BetaDataType*>(kargs.p_beta);
Pipeline{}(x_window,
gamma_window,
beta_window,
p_gamma,
p_beta,
y_window,
p_running_mean,
p_running_var,

16
include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp
View File

@@ -32,12 +32,10 @@ struct BatchnormFwdPipeline
}
template <typename XWindow,
typename GammaWindow,
typename BetaWindow,
typename YWindow>
CK_TILE_DEVICE void operator()(const XWindow& x_window_,
const GammaWindow& gamma_window_,
const BetaWindow& beta_window_,
const GammaDataType* p_gamma,
const BetaDataType* p_beta,
YWindow& y_window_,
MeanVarDataType* p_running_mean,
MeanVarDataType* p_running_var,
@@ -61,9 +59,9 @@ struct BatchnormFwdPipeline
y_window_,
Policy::template MakeXBlockTileDistribution<Problem>());
// Gamma/beta windows passed in but not used yet (gamma=1, beta=0 in test)
[[maybe_unused]] const auto gamma_window = gamma_window_;
[[maybe_unused]] const auto beta_window = beta_window_;
// Load gamma and beta scalars for this channel
const ComputeDataType gamma_val = type_convert<ComputeDataType>(p_gamma[channel_idx]);
const ComputeDataType beta_val = type_convert<ComputeDataType>(p_beta[channel_idx]);
// Calculate how many tiles needed (like layernorm2d two-pass)
constexpr index_t Block_N = Problem::BlockShape::Block_N;
@@ -125,8 +123,8 @@ struct BatchnormFwdPipeline
sweep_tile(y, [&](auto idx) {
ComputeDataType x_val = type_convert<ComputeDataType>(x[idx]);
// y = (x - mean) / std (no gamma/beta for now)
ComputeDataType normalized = (x_val - block_mean) * inv_std;
// y = gamma * (x - mean) / std + beta
ComputeDataType normalized = gamma_val * ((x_val - block_mean) * inv_std) + beta_val;
y(idx) = type_convert<YDataType>(normalized);
});