ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 01:28:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add Saved statistics and running statistics to example to verifykernel calculations

Browse Source
This commit is contained in:
Mohsen Saffari
2025-12-05 10:54:54 +00:00
parent 8d4f766f9c
commit 2b32dd75ee
1 changed files with 123 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -19,7 +19,7 @@ auto create_args(int argc, char* argv[])
.insert("e", "1e-5", "epsilon")
.insert("v", "1", "cpu validation or not")
.insert("prec", "fp16", "precision")
.insert("warmup", "5", "cold iter")
.insert("warmup", "10", "cold iter")
.insert("repeat", "20", "hot iter");
bool result = arg_parser.parse(argc, argv);
@@ -32,6 +32,11 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor<XDataType>& x,
const ck_tile::HostTensor<GammaDataType>* gamma,
const ck_tile::HostTensor<BetaDataType>* beta,
ck_tile::HostTensor<YDataType>& y,
ck_tile::HostTensor<ComputeDataType>* save_mean,
ck_tile::HostTensor<ComputeDataType>* save_inv_std,
ck_tile::HostTensor<ComputeDataType>* running_mean,
ck_tile::HostTensor<ComputeDataType>* running_var,
ComputeDataType momentum,
ck_tile::index_t N,
ck_tile::index_t C,
ck_tile::index_t H,
@@ -95,6 +100,23 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor<XDataType>& x,
ComputeDataType inv_std = static_cast<ComputeDataType>(1.0) /
ck_tile::sqrt(variance + epsilon);
// Save mean and inv_std if requested
if(save_mean != nullptr)
{
save_mean->mData[c] = mean;
}
if(save_inv_std != nullptr)
{
save_inv_std->mData[c] = inv_std;
}
// Update running statistics if requested
if(running_mean != nullptr && running_var != nullptr)
{
running_mean->mData[c] = (1.0f - momentum) * running_mean->mData[c] + momentum * mean;
running_var->mData[c] = (1.0f - momentum) * running_var->mData[c] + momentum * variance;
}
// Normalize all values in this channel with scale and bias
for(ck_tile::index_t n = 0; n < N; ++n)
{
@@ -138,6 +160,16 @@ bool run(const ck_tile::ArgParser& arg_parser)
ck_tile::HostTensor<ComputeDataType> beta_host({C});
ck_tile::HostTensor<YDataType> y_host_ref({N, H, W, C}); // NHWC!
ck_tile::HostTensor<YDataType> y_host_dev({N, H, W, C}); // NHWC!
// Allocate buffers for optional features
ck_tile::HostTensor<ComputeDataType> save_mean_host({C});
ck_tile::HostTensor<ComputeDataType> save_inv_std_host({C});
ck_tile::HostTensor<ComputeDataType> running_mean_host({C});
ck_tile::HostTensor<ComputeDataType> running_var_host({C});
// Initialize running statistics
ck_tile::FillUniformDistribution<ComputeDataType>{0.0f, 0.0f}(running_mean_host); // Start at 0
ck_tile::FillUniformDistribution<ComputeDataType>{1.0f, 1.0f}(running_var_host); // Start at 1
// Fill input with random data
ck_tile::FillUniformDistribution<XDataType>{-5.f, 5.f}(x_host);
@@ -151,17 +183,23 @@ bool run(const ck_tile::ArgParser& arg_parser)
ck_tile::DeviceMem gamma_buf(gamma_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem beta_buf(beta_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem y_buf(y_host_dev.get_element_space_size_in_bytes());
ck_tile::DeviceMem save_mean_buf(save_mean_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem save_inv_std_buf(save_inv_std_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem running_mean_buf(running_mean_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem running_var_buf(running_var_host.get_element_space_size_in_bytes());
x_buf.ToDevice(x_host.data());
gamma_buf.ToDevice(gamma_host.data());
beta_buf.ToDevice(beta_host.data());
running_mean_buf.ToDevice(running_mean_host.data());
running_var_buf.ToDevice(running_var_host.data());
// Define kernel configuration using Generic2dBlockShape
// 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)
using BlockTile = ck_tile::sequence<1, 256>; // Block size: 1 channel, 128 spatial
using ThreadPerBlock = ck_tile::sequence<1, 128>; // 64 threads
using Vector = ck_tile::sequence<1, 1>; // 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
@@ -169,8 +207,11 @@ bool run(const ck_tile::ArgParser& arg_parser)
using Shape = ck_tile::BatchnormShape<BlockTile, ThreadPerBlock, Vector>;
// Define traits (compile-time configuration)
using Traits = ck_tile::BatchnormFwdTraits<false, false>; // No save, no update
// Feature flags - change these to enable/disable testing different features
constexpr bool kSaveMeanInvStd = true; // Set true to test save for backward
constexpr bool kUpdateMovingAverage = true; // Set true to test running stats
using Traits = ck_tile::BatchnormFwdTraits<kSaveMeanInvStd, kUpdateMovingAverage>;
// Define problem with all types
using Problem = ck_tile::BatchnormProblem<XDataType, // input type
@@ -186,17 +227,17 @@ bool run(const ck_tile::ArgParser& arg_parser)
// Prepare host arguments
// Note: save/update behavior is determined by Traits (compile-time), not runtime args
ck_tile::BatchnormFwdHostArgs hargs{
x_buf.GetDeviceBuffer(), // p_x
gamma_buf.GetDeviceBuffer(), // p_gamma
beta_buf.GetDeviceBuffer(), // p_beta
y_buf.GetDeviceBuffer(), // p_y
nullptr, // p_running_mean (not used, Traits::kUpdateMovingAverage=false)
nullptr, // p_running_var
nullptr, // p_save_mean (not used, Traits::kSaveMeanInvStd=false)
nullptr, // p_save_inv_std
epsilon, // epsilon
0.1f, // momentum
N, C, H, W // dimensions
x_buf.GetDeviceBuffer(), // p_x
gamma_buf.GetDeviceBuffer(), // p_gamma
beta_buf.GetDeviceBuffer(), // p_beta
y_buf.GetDeviceBuffer(), // p_y
running_mean_buf.GetDeviceBuffer(), // p_running_mean (now used!)
running_var_buf.GetDeviceBuffer(), // p_running_var
save_mean_buf.GetDeviceBuffer(), // p_save_mean (now used!)
save_inv_std_buf.GetDeviceBuffer(), // p_save_inv_std
epsilon, // epsilon
0.1f, // momentum
N, C, H, W // dimensions
};
// Validate arguments
@@ -230,9 +271,21 @@ bool run(const ck_tile::ArgParser& arg_parser)
if(do_validation)
{
// Compute reference with gamma and beta
// Compute reference (will also save/update statistics)
ck_tile::HostTensor<ComputeDataType> save_mean_ref({C});
ck_tile::HostTensor<ComputeDataType> save_inv_std_ref({C});
ck_tile::HostTensor<ComputeDataType> running_mean_ref({C});
ck_tile::HostTensor<ComputeDataType> running_var_ref({C});
// Copy initial running stats
std::copy(running_mean_host.mData.begin(), running_mean_host.mData.end(), running_mean_ref.mData.begin());
std::copy(running_var_host.mData.begin(), running_var_host.mData.end(), running_var_ref.mData.begin());
reference_batchnorm_fwd<XDataType, YDataType, ComputeDataType, ComputeDataType, ComputeDataType>(
x_host, &gamma_host, &beta_host, y_host_ref, N, C, H, W, epsilon);
x_host, &gamma_host, &beta_host, y_host_ref,
&save_mean_ref, &save_inv_std_ref,
&running_mean_ref, &running_var_ref,
0.1f, N, C, H, W, epsilon);
// Get device result
y_buf.FromDevice(y_host_dev.mData.data());
@@ -261,9 +314,59 @@ bool run(const ck_tile::ArgParser& arg_parser)
}
std::cout << std::endl;
// Check error
// Check output
pass = ck_tile::check_err(y_host_dev, y_host_ref, "Error: Incorrect results!", 1e-2, 1e-2);
// Conditionally verify features based on what's enabled
if constexpr(kSaveMeanInvStd)
{
save_mean_buf.FromDevice(save_mean_host.mData.data());
save_inv_std_buf.FromDevice(save_inv_std_host.mData.data());
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 << "Ch" << std::setw(2) << c
<< " mean: Ref=" << std::setw(10) << save_mean_ref.mData[c]
<< " Dev=" << std::setw(10) << save_mean_host.mData[c]
<< " | inv_std: Ref=" << std::setw(10) << save_inv_std_ref.mData[c]
<< " Dev=" << std::setw(10) << save_inv_std_host.mData[c] << std::endl;
}
pass = pass && save_pass;
}
if constexpr(kUpdateMovingAverage)
{
if(repeat == 1)
{
running_mean_buf.FromDevice(running_mean_host.mData.data());
running_var_buf.FromDevice(running_var_host.mData.data());
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 << "Ch" << std::setw(2) << c
<< " mean: Ref=" << std::setw(10) << running_mean_ref.mData[c]
<< " Dev=" << std::setw(10) << running_mean_host.mData[c]
<< " | var: Ref=" << std::setw(10) << running_var_ref.mData[c]
<< " Dev=" << std::setw(10) << running_var_host.mData[c] << std::endl;
}
pass = pass && running_pass;
}
else
{
std::cout << "\nNOTE: Running statistics validation requires -warmup=0 -repeat=1" << std::endl;
std::cout << "(Multiple iterations accumulate running stats, making validation incorrect)" << std::endl;
}
}
std::cout << std::endl;
std::cout << "Validation: " << (pass ? "PASSED" : "FAILED") << std::endl;
}