diff --git a/example/ck_tile/42_batchnorm/batchnorm.cpp b/example/ck_tile/42_batchnorm/batchnorm.cpp deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp index b81c1f923b..a150129f02 100644 --- a/example/ck_tile/42_batchnorm/batchnorm_simple.cpp +++ b/example/ck_tile/42_batchnorm/batchnorm_simple.cpp @@ -4,6 +4,7 @@ #include "ck_tile/host.hpp" #include "ck_tile/ops/batchnorm.hpp" #include +#include // Simple POC for batchnorm forward pass // Tests basic functionality with a small tensor @@ -26,8 +27,10 @@ auto create_args(int argc, char* argv[]) } // CPU reference implementation -template +template 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::index_t N, ck_tile::index_t C, @@ -73,10 +76,24 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor& x, } ComputeDataType variance = var_sum / static_cast(per_channel_size); - // Normalize all values in this channel + // Load gamma and beta for this channel + ComputeDataType gamma_val = static_cast(1.0); + ComputeDataType beta_val = static_cast(0.0); + + if(gamma != nullptr) + { + gamma_val = ck_tile::type_convert(gamma->mData[c]); + } + if(beta != nullptr) + { + beta_val = ck_tile::type_convert(beta->mData[c]); + } + + // Compute inverse standard deviation ComputeDataType inv_std = static_cast(1.0) / ck_tile::sqrt(variance + epsilon); + // Normalize all values in this channel with scale and bias for(ck_tile::index_t n = 0; n < N; ++n) { for(ck_tile::index_t h = 0; h < H; ++h) @@ -85,7 +102,7 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor& x, { ck_tile::index_t idx = n * C * H * W + c * H * W + h * W + w; ComputeDataType val = ck_tile::type_convert(x.mData[idx]); - ComputeDataType normalized = (val - mean) * inv_std; + ComputeDataType normalized = gamma_val * ((val - mean) * inv_std) + beta_val; y.mData[idx] = ck_tile::type_convert(normalized); } } @@ -115,17 +132,25 @@ bool run(const ck_tile::ArgParser& arg_parser) // Allocate host tensors ck_tile::index_t total_size = N * C * H * W; ck_tile::HostTensor x_host({N, C, H, W}); + ck_tile::HostTensor gamma_host({C}); + ck_tile::HostTensor beta_host({C}); ck_tile::HostTensor y_host_ref({N, C, H, W}); ck_tile::HostTensor y_host_dev({N, C, H, W}); // 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 // Allocate device memory ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes()); + 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()); x_buf.ToDevice(x_host.data()); + gamma_buf.ToDevice(gamma_host.data()); + beta_buf.ToDevice(beta_host.data()); // Define kernel configuration using BlockWarps = ck_tile::sequence<4, 1>; @@ -137,33 +162,44 @@ bool run(const ck_tile::ArgParser& arg_parser) using Problem = ck_tile::BatchnormProblem; using Kernel = ck_tile::BatchnormFwd; - const ck_tile::index_t kBlockSize = Kernel::BlockSize(); - const ck_tile::index_t kGridSize = C; // One block per channel + // Prepare host arguments + 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 + nullptr, // p_running_var + nullptr, // p_save_mean + nullptr, // p_save_inv_std + epsilon, // epsilon + 0.1f, // momentum (not used yet) + N, C, H, W, // dimensions + false, // update_moving_average + false // save_mean_inv_std + }; - std::cout << "Kernel config: BlockSize=" << kBlockSize << ", GridSize=" << kGridSize - << " (one block per channel, reducing over N×H×W=" << N*H*W << " elements)" << std::endl; - - if(!Kernel::IsSupportedArgument(N, C, H, W)) + // Validate arguments + if(!Kernel::IsSupportedArgument(hargs)) { std::cout << "Arguments not supported!" << std::endl; return false; } + // Get grid and block size + const auto grid_size = Kernel::GridSize(hargs); + const auto block_size = Kernel::BlockSize(); + + std::cout << "Kernel config: BlockSize=" << block_size << ", GridSize=" << grid_size.x + << " (one block per channel, reducing over N×H×W=" << N*H*W << " elements)" << std::endl; + + // Make kernel arguments + auto kargs = Kernel::MakeKernelArgs(hargs); + // Launch kernel float ave_time = ck_tile::launch_kernel( ck_tile::stream_config{nullptr, true, 0, warmup, repeat}, - ck_tile::make_kernel<1>( - Kernel{}, - kGridSize, - kBlockSize, - 0, - static_cast(x_buf.GetDeviceBuffer()), - static_cast(y_buf.GetDeviceBuffer()), - N, - C, - H, - W, - static_cast(epsilon))); + ck_tile::make_kernel<1>(Kernel{}, grid_size, block_size, 0, kargs)); std::size_t num_bytes = sizeof(XDataType) * total_size + sizeof(YDataType) * total_size; float gb_per_sec = num_bytes / 1.E6 / ave_time; @@ -174,13 +210,33 @@ bool run(const ck_tile::ArgParser& arg_parser) if(do_validation) { - // Compute reference - reference_batchnorm_fwd( - x_host, y_host_ref, N, C, H, W, epsilon); + // Compute reference with gamma and beta + reference_batchnorm_fwd( + x_host, &gamma_host, &beta_host, y_host_ref, N, C, H, W, epsilon); // Get device result y_buf.FromDevice(y_host_dev.mData.data()); + // Print sample outputs for each channel (2 samples per channel) + std::cout << "\n=== Sample Outputs (first 2 values per channel) ===" << std::endl; + for(ck_tile::index_t c = 0; c < C; ++c) + { + std::cout << "Channel " << c << ":" << std::endl; + + // Print 2 sample values from first sample (n=0) + for(ck_tile::index_t sample = 0; sample < 2 && sample < H * W; ++sample) + { + ck_tile::index_t idx = 0 * C * H * W + c * H * W + sample; + float ref_val = ck_tile::type_convert(y_host_ref.mData[idx]); + float dev_val = ck_tile::type_convert(y_host_dev.mData[idx]); + std::cout << " Sample[" << sample << "]: " + << "Ref=" << std::fixed << std::setprecision(6) << ref_val + << ", Kernel=" << dev_val + << ", Diff=" << std::abs(ref_val - dev_val) << std::endl; + } + } + std::cout << std::endl; + // Check error pass = ck_tile::check_err(y_host_dev, y_host_ref, "Error: Incorrect results!", 1e-2, 1e-2); 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 e206eeaa19..0634c93b3b 100644 --- a/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp +++ b/include/ck_tile/ops/batchnorm/kernel/batchnorm_fwd_kernel.hpp @@ -10,12 +10,30 @@ namespace ck_tile { +// Host-side arguments for batchnorm forward pass +struct BatchnormFwdHostArgs +{ + const void* p_x; // [N, C, H, W] input tensor (required) + const void* p_gamma; // [C] scale parameter (required, use all 1.0 if not needed) + const void* p_beta; // [C] bias parameter (required, use all 0.0 if not needed) + + void* p_y; // [N, C, H, W] output tensor (required) + + void* p_running_mean; // [C] running mean (nullptr if not used) + void* p_running_var; // [C] running variance (nullptr if not used) + void* p_save_mean; // [C] save mean for backward (nullptr if not used) + void* p_save_inv_std; // [C] save inv_std for backward (nullptr if not used) + + float epsilon; + float momentum; + + index_t N, C, H, W; + + bool update_moving_average; // If true, p_running_mean/var must be valid + bool save_mean_inv_std; // If true, p_save_mean/inv_std must be valid +}; + // BatchnormFwd: Forward pass batch normalization kernel -// Performs: -// 1. Welford reduction to compute mean and variance across spatial dimensions (H*W) -// 2. Normalization: y = (x - mean) / sqrt(variance + epsilon) -// -// For now, simplified version without scale/bias template struct BatchnormFwd { @@ -27,19 +45,74 @@ struct BatchnormFwd static constexpr index_t kBlockSize = BlockShape::kBlockSize; - CK_TILE_HOST static constexpr index_t BlockSize() { return kBlockSize; } - - CK_TILE_DEVICE void operator()(const XDataType* p_x, - YDataType* p_y, - index_t N, - index_t C, - index_t H, - index_t W, - ComputeDataType epsilon) const + // Kernel arguments + struct BatchnormFwdKargs { - // For batchnorm: input shape is [N, C, H, W] - // We reduce over N*H*W (batch and spatial) for EACH channel - // Each block handles ONE channel + 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; + + float epsilon; + float momentum; + + index_t N, C, H, W; + + bool update_moving_average; + bool save_mean_inv_std; + }; + + using Kargs = BatchnormFwdKargs; // Alias for convenience + using Hargs = BatchnormFwdHostArgs; + + // Convert host args to kernel args + CK_TILE_HOST static constexpr Kargs MakeKernelArgs(const Hargs& hargs) + { + return Kargs{hargs.p_x, + hargs.p_gamma, + hargs.p_beta, + hargs.p_y, + hargs.p_running_mean, + hargs.p_running_var, + hargs.p_save_mean, + hargs.p_save_inv_std, + hargs.epsilon, + hargs.momentum, + hargs.N, + hargs.C, + hargs.H, + hargs.W, + hargs.update_moving_average, + hargs.save_mean_inv_std}; + } + + // Grid size calculation + CK_TILE_HOST static constexpr auto GridSize(const Hargs& hargs) + { + return dim3(hargs.C); // One block per channel + } + + // Block size + CK_TILE_HOST static constexpr auto BlockSize() + { + return kBlockSize; + } + + 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); + + 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 @@ -89,11 +162,22 @@ struct BatchnormFwd BlockWelford::template Run( block_mean, block_var, block_count, smem); - // Now all threads have the same mean and variance for this channel - // Normalize and write output for ALL samples + // 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; @@ -101,23 +185,47 @@ struct BatchnormFwd const index_t offset = n * C * H * W + c * H * W + hw; ComputeDataType val = type_convert(p_x[offset]); - ComputeDataType normalized = (val - block_mean) * inv_std; + + // Apply batch normalization with scale and bias + ComputeDataType normalized = gamma * ((val - block_mean) * inv_std) + beta; + p_y[offset] = type_convert(normalized); } } // Validate arguments - CK_TILE_HOST static bool IsSupportedArgument(index_t N, - index_t C, - index_t H, - index_t W) + CK_TILE_HOST static bool IsSupportedArgument(const Hargs& hargs) { - // For POC, accept all sizes - // Later we can add alignment requirements - if(N <= 0 || C <= 0 || H <= 0 || W <= 0) + // 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 flags + if(hargs.update_moving_average) + { + if(hargs.p_running_mean == nullptr || hargs.p_running_var == nullptr) + { + return false; + } + } + + if(hargs.save_mean_inv_std) + { + if(hargs.p_save_mean == nullptr || hargs.p_save_inv_std == nullptr) + { + return false; + } + } + return true; } };