From 04d3a0ced06f1c6090c796b5006486251eee6dab Mon Sep 17 00:00:00 2001 From: Mohsen Saffari Date: Tue, 9 Dec 2025 09:13:31 +0000 Subject: [PATCH] Some refactoring for this 1 channel per block kernel --- .../ck_tile/42_batchnorm/batchnorm_simple.cpp | 28 ++-- .../batchnorm/kernel/batchnorm_fwd_kernel.hpp | 133 +++++++++++------- .../pipeline/batchnorm_fwd_policy.hpp | 4 +- 3 files changed, 101 insertions(+), 64 deletions(-) diff --git a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp index 6aaf002078..5e00d0e27d 100644 --- a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp +++ b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp @@ -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; } diff --git a/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp b/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp index 9c7a0f9f29..840a5e2f68 100644 --- a/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp +++ b/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp @@ -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 struct BatchnormFwd { + // Type aliases from Problem using Problem = remove_cvref_t; - using Pipeline = BatchnormFwdPipeline; // Class-level, like layernorm2d + using Pipeline = BatchnormFwdPipeline; 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() @@ -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 diff --git a/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp index cd9e7a0150..9f7b9cc64a 100644 --- a/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp +++ b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp @@ -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 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 CK_TILE_DEVICE static constexpr auto MakeGammaBetaBlockTileDistribution() {