diff --git a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp index 51bcf01254..2731b43cc9 100644 --- a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp +++ b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp @@ -142,12 +142,9 @@ bool run(const ck_tile::ArgParser& arg_parser) // Fill input with random data ck_tile::FillUniformDistribution{-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(1.0); - beta_host.mData[c] = static_cast(0.0); - } + // Fill gamma and beta with random values (test scale/bias) + ck_tile::FillUniformDistribution{0.8f, 1.2f}(gamma_host); // Scale around 1 + ck_tile::FillUniformDistribution{-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; 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 f342f44be7..9c7a0f9f29 100644 --- a/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp +++ b/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp @@ -139,28 +139,6 @@ struct BatchnormFwd return make_tile_window(tmp2_, make_tuple(number{}, number{}), {0, 0}); }(); - 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 auto tmp2_ = pad_tensor_view(tmp_, make_tuple(number{}), sequence{}); - return make_tile_window(tmp2_, make_tuple(number{}), {0}); - }(); - - 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}); - }(); - auto y_window = [&]() { YDataType* p_y = static_cast(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(kargs.p_save_mean); MeanVarDataType* p_save_inv_std = static_cast(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(kargs.p_gamma); + const BetaDataType* p_beta = static_cast(kargs.p_beta); + Pipeline{}(x_window, - gamma_window, - beta_window, + p_gamma, + p_beta, y_window, p_running_mean, p_running_var, diff --git a/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp index 25bd00ed7f..b8736d3318 100644 --- a/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp +++ b/include/ck_tile/ops/batchnorm/pipeline/batchnorm_fwd_pipeline.hpp @@ -32,12 +32,10 @@ struct BatchnormFwdPipeline } template 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()); - // 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(p_gamma[channel_idx]); + const ComputeDataType beta_val = type_convert(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(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(normalized); });