From 3194b653f78b2ab2acb202a5933af13fcef11df8 Mon Sep 17 00:00:00 2001 From: Mohsen Saffari Date: Fri, 28 Nov 2025 18:29:10 +0000 Subject: [PATCH] add tiling, pipeline, follow layernorm2d --- .../ck_tile/42_batchnorm/batchnorm_simple.cpp | 22 +- include/ck_tile/ops/batchnorm.hpp | 2 + .../batchnorm/kernel/batchnorm_fwd_kernel.hpp | 194 ++++++++---------- .../pipeline/batchnorm_fwd_pipeline.hpp | 177 ++++++++++++++++ .../pipeline/batchnorm_fwd_policy.hpp | 59 ++++++ .../batchnorm/pipeline/batchnorm_shape.hpp | 44 +--- 6 files changed, 347 insertions(+), 151 deletions(-) create mode 100644 include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp create mode 100644 include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp diff --git a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp index e38ed2a336..d65618e2bc 100644 --- a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp +++ b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp @@ -139,8 +139,13 @@ bool run(const ck_tile::ArgParser& arg_parser) // Fill input with random data ck_tile::FillUniformDistribution{-5.f, 5.f}(x_host); - ck_tile::FillUniformDistribution{0.8f, 1.2f}(gamma_host); // Scale around 1.0 - ck_tile::FillUniformDistribution{-0.5f, 0.5f}(beta_host); // Bias around 0.0 + + // 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(1.0); + beta_host.mData[c] = static_cast(0.0); + } // Allocate device memory ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes()); @@ -152,13 +157,14 @@ bool run(const ck_tile::ArgParser& arg_parser) gamma_buf.ToDevice(gamma_host.data()); beta_buf.ToDevice(beta_host.data()); - // Define kernel configuration - using BlockWarps = ck_tile::sequence<4, 1>; - using BlockTile = ck_tile::sequence<1, 256>; // Simplified for POC - using WarpTile = ck_tile::sequence<1, 256>; - using Vector = ck_tile::sequence<1, 1>; + // 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) - using Shape = ck_tile::BatchnormShape; + using Shape = ck_tile::BatchnormShape; // Define traits (compile-time configuration) using Traits = ck_tile::BatchnormFwdTraits; // No save, no update diff --git a/include/ck_tile/ops/batchnorm.hpp b/include/ck_tile/ops/batchnorm.hpp index 64a00b336b..2afcf8322b 100644 --- a/include/ck_tile/ops/batchnorm.hpp +++ b/include/ck_tile/ops/batchnorm.hpp @@ -5,6 +5,8 @@ #include "ck_tile/ops/batchnorm/block/block_welford.hpp" #include "ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp" +#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp" +#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp" #include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_traits.hpp" #include "ck_tile/ops/batchnorm/pipeline/batchnorm_problem.hpp" #include "ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp" 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 02caacc839..40b27d9720 100644 --- a/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp +++ b/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp @@ -5,6 +5,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/batchnorm/block/block_welford.hpp" +#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp" #include "ck_tile/ops/batchnorm/pipeline/batchnorm_problem.hpp" #include "ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp" @@ -38,11 +39,14 @@ struct BatchnormFwd { using Problem = remove_cvref_t; using XDataType = typename Problem::XDataType; + using GammaDataType = typename Problem::GammaDataType; + using BetaDataType = typename Problem::BetaDataType; using ComputeDataType = typename Problem::ComputeDataType; using YDataType = typename Problem::YDataType; + using MeanVarDataType = typename Problem::MeanVarDataType; using BlockShape = typename Problem::BlockShape; - static constexpr index_t kBlockSize = BlockShape::kBlockSize; + static constexpr index_t kBlockSize = BlockShape::BlockSize; // Kernel arguments struct BatchnormFwdKargs @@ -98,130 +102,110 @@ struct BatchnormFwd return kBlockSize; } + // Shared memory size (must be constexpr for __shared__ allocation) + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + return BatchnormFwdPipeline::GetSmemSize(); + } + CK_TILE_DEVICE void operator()(Kargs kargs) const { - // Cast pointers to typed pointers - const XDataType* p_x = static_cast(kargs.p_x); - YDataType* p_y = static_cast(kargs.p_y); + using Pipeline = BatchnormFwdPipeline; const index_t N = kargs.N; const index_t C = kargs.C; const index_t H = kargs.H; const index_t W = kargs.W; - const ComputeDataType epsilon = static_cast(kargs.epsilon); - - const index_t spatial_size = H * W; - const index_t per_channel_size = N * spatial_size; // Reduce over N×H×W - const index_t thread_id = get_thread_id(); const index_t block_id = get_block_id(); - - // Each block handles one channel - const index_t c = block_id; + const index_t c = block_id; // Channel index if(c >= C) return; - // Thread-local Welford statistics - ComputeDataType thread_mean = type_convert(0); - ComputeDataType thread_m2 = type_convert(0); - index_t thread_count = 0; + const index_t spatial_size = H * W; + const index_t per_channel_size = N * spatial_size; - // Each thread processes elements across ALL samples (N) and spatial positions (H×W) - // for this channel - for(index_t idx = thread_id; idx < per_channel_size; idx += kBlockSize) - { - // Calculate which sample (n) and spatial position (hw) this is - const index_t n = idx / spatial_size; - const index_t hw = idx % spatial_size; + // Use block dimensions from BlockShape (like layernorm2d) + static constexpr index_t Block_M = BlockShape::Block_M; + static constexpr index_t Block_N = BlockShape::Block_N; + + // Create tensor views WITHOUT distributions (will be applied in pipeline) + const auto x_window = [&]() { + const XDataType* p_x = static_cast(kargs.p_x); + const auto tmp_ = make_naive_tensor_view( + p_x + c * spatial_size, + make_tuple(N, spatial_size), + make_tuple(C * spatial_size, 1), + number<1>{}, + number<1>{}); - // Memory layout: [N, C, H, W] with C-major for H×W - const index_t offset = n * C * H * W + c * H * W + hw; - ComputeDataType val = type_convert(p_x[offset]); + const auto tmp2_ = pad_tensor_view( + tmp_, make_tuple(number{}, number{}), sequence{}); - // Online Welford update - thread_count++; - ComputeDataType delta = val - thread_mean; - thread_mean += delta / type_convert(thread_count); - ComputeDataType delta2 = val - thread_mean; - thread_m2 += delta * delta2; - } + return make_tile_window(tmp2_, make_tuple(number{}, number{}), {0, 0}); + }(); - // Allocate shared memory for block-level reduction - __shared__ char smem[BlockWelford::template GetSmemSize()]; - - // Block-level Welford reduction - ComputeDataType block_mean = thread_mean; - ComputeDataType block_var = thread_m2; - index_t block_count = thread_count; - - BlockWelford::template Run( - block_mean, block_var, block_count, smem); - - // Load scale (gamma) and bias (beta) for this channel - // Following old CK pattern: gamma/beta are ALWAYS provided (no nullptr checks) - // All threads load (efficient, no branching) - const ComputeDataType* p_gamma = static_cast(kargs.p_gamma); - const ComputeDataType* p_beta = static_cast(kargs.p_beta); - - const ComputeDataType gamma = p_gamma[c]; - const ComputeDataType beta = p_beta[c]; - - // Compute inverse standard deviation - ComputeDataType inv_std = type_convert(1) / - ck_tile::sqrt(block_var + epsilon); - - // Normalize and write output with scale and bias - // Formula: y = gamma * (x - mean) / std + beta - // = gamma * (x - mean) * inv_std + beta - for(index_t idx = thread_id; idx < per_channel_size; idx += kBlockSize) - { - const index_t n = idx / spatial_size; - const index_t hw = idx % spatial_size; + const auto gamma_window = [&]() { + const GammaDataType* p_gamma = static_cast(kargs.p_gamma); + const auto tmp_ = make_naive_tensor_view_packed( + p_gamma + c, + make_tuple(1), + number<1>{}); - const index_t offset = (n * C * H * W) + (c * H * W) + hw; - ComputeDataType val = type_convert(p_x[offset]); - - // Apply batch normalization with scale and bias - ComputeDataType normalized = gamma * ((val - block_mean) * inv_std) + beta; - - p_y[offset] = type_convert(normalized); - } + const auto tmp2_ = pad_tensor_view(tmp_, make_tuple(number{}), sequence{}); + return make_tile_window(tmp2_, make_tuple(number{}), {0}); + }(); - // Save mean and inverse std for backward pass (compile-time check) - if constexpr(Problem::Traits::kSaveMeanInvStd) - { - if(thread_id == 0) - { - using MeanVarDataType = typename Problem::MeanVarDataType; - MeanVarDataType* p_save_mean = static_cast(kargs.p_save_mean); - MeanVarDataType* p_save_inv_std = static_cast(kargs.p_save_inv_std); - - p_save_mean[c] = type_convert(block_mean); - p_save_inv_std[c] = type_convert(inv_std); - } - } + const auto beta_window = [&]() { + const BetaDataType* p_beta = static_cast(kargs.p_beta); + const auto tmp_ = make_naive_tensor_view_packed( + p_beta + c, + make_tuple(1), + number<1>{}); + + const auto tmp2_ = pad_tensor_view(tmp_, make_tuple(number{}), sequence{}); + return make_tile_window(tmp2_, make_tuple(number{}), {0}); + }(); - // Update running mean and variance (compile-time check) - if constexpr(Problem::Traits::kUpdateMovingAverage) - { - if(thread_id == 0) - { - using MeanVarDataType = typename Problem::MeanVarDataType; - MeanVarDataType* p_running_mean = static_cast(kargs.p_running_mean); - MeanVarDataType* p_running_var = static_cast(kargs.p_running_var); - - const ComputeDataType momentum = static_cast(kargs.momentum); - const ComputeDataType one_minus_momentum = type_convert(1) - momentum; - - // Exponential moving average: new = (1-momentum)*old + momentum*current - ComputeDataType old_mean = type_convert(p_running_mean[c]); - ComputeDataType old_var = type_convert(p_running_var[c]); - - p_running_mean[c] = type_convert(one_minus_momentum * old_mean + momentum * block_mean); - p_running_var[c] = type_convert(one_minus_momentum * old_var + momentum * block_var); - } - } + auto y_window = [&]() { + YDataType* p_y = static_cast(kargs.p_y); + const auto tmp_ = make_naive_tensor_view( + p_y + c * spatial_size, + make_tuple(N, spatial_size), + make_tuple(C * spatial_size, 1), + number<1>{}, + number<1>{}); + + const auto tmp2_ = pad_tensor_view( + tmp_, make_tuple(number{}, number{}), sequence{}); + + return make_tile_window(tmp2_, make_tuple(number{}, number{}), {0, 0}); + }(); + + // Allocate shared memory (use kernel's constexpr GetSmemSize) + __shared__ char smem[GetSmemSize()]; + + // Cast pointers for optional features + MeanVarDataType* p_running_mean = static_cast(kargs.p_running_mean); + MeanVarDataType* p_running_var = static_cast(kargs.p_running_var); + MeanVarDataType* p_save_mean = static_cast(kargs.p_save_mean); + MeanVarDataType* p_save_inv_std = static_cast(kargs.p_save_inv_std); + + // Call pipeline with properly distributed tile windows + Pipeline{}(x_window, + gamma_window, + beta_window, + y_window, + p_running_mean, + p_running_var, + p_save_mean, + p_save_inv_std, + static_cast(kargs.epsilon), + static_cast(kargs.momentum), + per_channel_size, + c, + smem); } // Validate arguments diff --git a/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp new file mode 100644 index 0000000000..acf5885229 --- /dev/null +++ b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/batchnorm/block/block_welford.hpp" +#include "ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp" + +namespace ck_tile { + +// BatchnormFwdPipeline: Computation logic for batch normalization +// Takes tile windows and performs Welford reduction + normalization +template +struct BatchnormFwdPipeline +{ + using Problem = remove_cvref_t; + using Policy = remove_cvref_t; + + using XDataType = typename Problem::XDataType; + using GammaDataType = typename Problem::GammaDataType; + using BetaDataType = typename Problem::BetaDataType; + using ComputeDataType = typename Problem::ComputeDataType; + using YDataType = typename Problem::YDataType; + using MeanVarDataType = typename Problem::MeanVarDataType; + + static constexpr index_t kBlockSize = Problem::BlockShape::BlockSize; + + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + return Policy::template GetSmemSize(); + } + + template + CK_TILE_DEVICE void operator()(const XWindow& x_window_, + const GammaWindow& gamma_window_, + const BetaWindow& beta_window_, + YWindow& y_window_, + MeanVarDataType* p_running_mean, + MeanVarDataType* p_running_var, + MeanVarDataType* p_save_mean, + MeanVarDataType* p_save_inv_std, + ComputeDataType epsilon, + ComputeDataType momentum, + [[maybe_unused]] index_t per_channel_size, + index_t channel_idx, + void* smem) const + { + const index_t thread_id = get_thread_id(); + + // Apply tile distributions (like layernorm2d does) + // Note: x_window and y_window are NOT const (need to move them) + auto x_window = + make_tile_window(x_window_, Policy::template MakeXBlockTileDistribution()); + const auto gamma_window = make_tile_window( + gamma_window_, Policy::template MakeGammaBetaBlockTileDistribution()); + const auto beta_window = make_tile_window( + beta_window_, Policy::template MakeGammaBetaBlockTileDistribution()); + auto y_window = + make_tile_window(y_window_, Policy::template MakeXBlockTileDistribution()); + + // Load gamma/beta once (constant per channel) + [[maybe_unused]]const auto gamma = load_tile(gamma_window); + [[maybe_unused]]const auto beta = load_tile(beta_window); + + // Calculate how many tiles needed (like layernorm2d two-pass) + constexpr index_t Block_N = Problem::BlockShape::Block_N; + index_t num_tile_iteration = integer_divide_ceil(per_channel_size, Block_N); + + // ========================================== + // PHASE 1: WELFORD REDUCTION OVER ALL TILES + // ========================================== + ComputeDataType thread_mean = type_convert(0); + ComputeDataType thread_m2 = type_convert(0); + index_t thread_count = 0; + + // Iterate over tiles for Welford accumulation + for(index_t tile_idx = 0; tile_idx < num_tile_iteration; ++tile_idx) + { + auto x = load_tile(x_window); + + sweep_tile(x, [&](auto idx) { + ComputeDataType val = type_convert(x[idx]); + thread_count++; + ComputeDataType delta = val - thread_mean; + thread_mean += delta / type_convert(thread_count); + ComputeDataType delta2 = val - thread_mean; + thread_m2 += delta * delta2; + }); + + // Move to next tile + if(tile_idx < num_tile_iteration - 1) + { + move_tile_window(x_window, {0, Block_N}); + } + } + + // Move x_window back to start + move_tile_window(x_window, {0, -static_cast(Block_N * (num_tile_iteration - 1))}); + + // Block-level reduction + ComputeDataType block_mean = thread_mean; + ComputeDataType block_var = thread_m2; + index_t block_count = thread_count; + + BlockWelford::template Run( + block_mean, block_var, block_count, smem); + + // ========================================== + // PHASE 2: COMPUTE INVERSE STD + // ========================================== + ComputeDataType inv_std = type_convert(1) / + ck_tile::sqrt(block_var + epsilon); + + // ========================================== + // PHASE 3: NORMALIZE AND STORE (ITERATE OVER TILES) + // ========================================== + for(index_t tile_idx = 0; tile_idx < num_tile_iteration; ++tile_idx) + { + auto x = load_tile(x_window); + auto y = make_static_distributed_tensor(x.get_tile_distribution()); + + sweep_tile(y, [&](auto idx) { + ComputeDataType x_val = type_convert(x[idx]); + + // y = (x - mean) / std (no gamma/beta for now) + ComputeDataType normalized = (x_val - block_mean) * inv_std; + y(idx) = type_convert(normalized); + }); + + store_tile(y_window, y); + + // Move to next tile + if(tile_idx < num_tile_iteration - 1) + { + move_tile_window(x_window, {0, Block_N}); + move_tile_window(y_window, {0, Block_N}); + } + } + + // ========================================== + // PHASE 6: SAVE STATISTICS (Optional) + // ========================================== + if constexpr(Problem::Traits::kSaveMeanInvStd) + { + if(thread_id == 0) + { + p_save_mean[channel_idx] = type_convert(block_mean); + p_save_inv_std[channel_idx] = type_convert(inv_std); + } + } + + // ========================================== + // PHASE 7: UPDATE RUNNING STATISTICS (Optional) + // ========================================== + if constexpr(Problem::Traits::kUpdateMovingAverage) + { + if(thread_id == 0) + { + ComputeDataType one_minus_momentum = type_convert(1) - momentum; + + ComputeDataType old_mean = type_convert(p_running_mean[channel_idx]); + ComputeDataType old_var = type_convert(p_running_var[channel_idx]); + + p_running_mean[channel_idx] = type_convert( + one_minus_momentum * old_mean + momentum * block_mean); + p_running_var[channel_idx] = type_convert( + one_minus_momentum * old_var + momentum * block_var); + } + } + } +}; + +} // 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 new file mode 100644 index 0000000000..a21fc32c35 --- /dev/null +++ b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_policy.hpp @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" + +namespace ck_tile { + +// Default policy for batchnorm forward pipeline +// Defines tile distributions and helper functions +struct BatchnormFwdPipelineDefaultPolicy +{ + // Tile distribution for input data (following layernorm2d pattern exactly) + template + CK_TILE_DEVICE static constexpr auto MakeXBlockTileDistribution() + { + using S = typename Problem::BlockShape; + + return make_static_tile_distribution( + tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 2>>, + sequence<1, 1, 2, 2>, + sequence<0, 3, 0, 3>>{}); + } + + // Tile distribution for gamma/beta parameters (following layernorm2d pattern) + template + CK_TILE_DEVICE static constexpr auto MakeGammaBetaBlockTileDistribution() + { + using S = typename Problem::BlockShape; + + return make_static_tile_distribution( + tile_distribution_encoding< + sequence, + tuple>, + tuple, sequence<0, 1>>, + tuple, sequence<1, 2>>, + sequence<1, 1>, + sequence<0, 3>>{}); + } + + // Calculate shared memory size + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + // For POC, use BlockWelford's smem requirement + using ComputeDataType = typename Problem::ComputeDataType; + constexpr index_t kBlockSize = Problem::BlockShape::BlockSize; + + return BlockWelford::template GetSmemSize(); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp index d46c6295ac..384806fb66 100644 --- a/include/ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp +++ b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_shape.hpp @@ -4,46 +4,14 @@ #pragma once #include "ck_tile/core.hpp" +#include "ck_tile/ops/common/generic_2d_block_shape.hpp" namespace ck_tile { -// BatchnormShape defines the block/warp/thread tile configuration -// Similar to reduce2d_shape but adapted for batchnorm -template -struct BatchnormShape -{ - using BlockWarps = remove_cvref_t; - using BlockTile = remove_cvref_t; - using WarpTile = remove_cvref_t; - using Vector = remove_cvref_t; - - static constexpr index_t kBlockWarps_M = BlockWarps::at(number<0>{}); - static constexpr index_t kBlockWarps_N = BlockWarps::at(number<1>{}); - - static constexpr index_t kBlockSize = kBlockWarps_M * kBlockWarps_N * get_warp_size(); - - static constexpr index_t Block_M = BlockTile::at(number<0>{}); - static constexpr index_t Block_N = BlockTile::at(number<1>{}); - - static constexpr index_t WarpTile_M = WarpTile::at(number<0>{}); - static constexpr index_t WarpTile_N = WarpTile::at(number<1>{}); - - static constexpr index_t VectorSize_M = Vector::at(number<0>{}); - static constexpr index_t VectorSize_N = Vector::at(number<1>{}); - - // Thread tile sizes - static constexpr index_t ThreadTile_M = WarpTile_M / kBlockWarps_M; - static constexpr index_t ThreadTile_N = WarpTile_N / kBlockWarps_N; - - // For batchnorm: - // - M dimension represents the batch*channel (merged N*C) - // - N dimension represents the spatial (merged H*W) - // We reduce over N (spatial) dimension per M (batch*channel) - - static constexpr auto BlockSize = number{}; -}; +// BatchnormShape using Generic2dBlockShape for proper tile distribution support +template + typename ThreadPerBlock_, // Threads along sequence + typename Vector_> // Vector size along sequence +using BatchnormShape = Generic2dBlockShape; } // namespace ck_tile