From 2b32dd75eef6946104e19ebb9324c58d5ed5c24d Mon Sep 17 00:00:00 2001 From: Mohsen Saffari Date: Fri, 5 Dec 2025 10:54:54 +0000 Subject: [PATCH] Add Saved statistics and running statistics to example to verifykernel calculations --- .../ck_tile/42_batchnorm/batchnorm_simple.cpp | 143 +++++++++++++++--- 1 file changed, 123 insertions(+), 20 deletions(-) diff --git a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp index 2731b43cc9..6aaf002078 100644 --- a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp +++ b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp @@ -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& x, const ck_tile::HostTensor* gamma, const ck_tile::HostTensor* beta, ck_tile::HostTensor& y, + ck_tile::HostTensor* save_mean, + ck_tile::HostTensor* save_inv_std, + ck_tile::HostTensor* running_mean, + ck_tile::HostTensor* 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& x, ComputeDataType inv_std = static_cast(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 beta_host({C}); ck_tile::HostTensor y_host_ref({N, H, W, C}); // NHWC! ck_tile::HostTensor y_host_dev({N, H, W, C}); // NHWC! + + // Allocate buffers for optional features + ck_tile::HostTensor save_mean_host({C}); + ck_tile::HostTensor save_inv_std_host({C}); + ck_tile::HostTensor running_mean_host({C}); + ck_tile::HostTensor running_var_host({C}); + + // Initialize running statistics + ck_tile::FillUniformDistribution{0.0f, 0.0f}(running_mean_host); // Start at 0 + ck_tile::FillUniformDistribution{1.0f, 1.0f}(running_var_host); // Start at 1 // Fill input with random data ck_tile::FillUniformDistribution{-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; - // Define traits (compile-time configuration) - using Traits = ck_tile::BatchnormFwdTraits; // 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; // Define problem with all types using Problem = ck_tile::BatchnormProblem save_mean_ref({C}); + ck_tile::HostTensor save_inv_std_ref({C}); + ck_tile::HostTensor running_mean_ref({C}); + ck_tile::HostTensor 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( - 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; }