add gamma, bias to the simple kernel

example/ck_tile/42_batchnorm/batchnorm_simple.cpp

@@ -4,6 +4,7 @@
 #include "ck_tile/host.hpp"
 #include "ck_tile/ops/batchnorm.hpp"
 #include <cstring>
+#include <iomanip>
 
 // 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 <typename XDataType, typename YDataType, typename ComputeDataType>
+template <typename XDataType, typename YDataType, typename ComputeDataType, typename GammaDataType, typename BetaDataType>
 void reference_batchnorm_fwd(const ck_tile::HostTensor<XDataType>& x,
+                             const ck_tile::HostTensor<GammaDataType>* gamma,
+                             const ck_tile::HostTensor<BetaDataType>* beta,
                              ck_tile::HostTensor<YDataType>& y,
                              ck_tile::index_t N,
                              ck_tile::index_t C,
@@ -73,10 +76,24 @@ void reference_batchnorm_fwd(const ck_tile::HostTensor<XDataType>& x,
         }
         ComputeDataType variance = var_sum / static_cast<ComputeDataType>(per_channel_size);
         
-        // Normalize all values in this channel
+        // Load gamma and beta for this channel
+        ComputeDataType gamma_val = static_cast<ComputeDataType>(1.0);
+        ComputeDataType beta_val = static_cast<ComputeDataType>(0.0);
+        
+        if(gamma != nullptr)
+        {
+            gamma_val = ck_tile::type_convert<ComputeDataType>(gamma->mData[c]);
+        }
+        if(beta != nullptr)
+        {
+            beta_val = ck_tile::type_convert<ComputeDataType>(beta->mData[c]);
+        }
+        
+        // Compute inverse standard deviation
         ComputeDataType inv_std = static_cast<ComputeDataType>(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<XDataType>& x,
                 {
                     ck_tile::index_t idx = n * C * H * W + c * H * W + h * W + w;
                     ComputeDataType val = ck_tile::type_convert<ComputeDataType>(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<YDataType>(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<XDataType> x_host({N, C, H, W});
+    ck_tile::HostTensor<ComputeDataType> gamma_host({C});
+    ck_tile::HostTensor<ComputeDataType> beta_host({C});
     ck_tile::HostTensor<YDataType> y_host_ref({N, C, H, W});
     ck_tile::HostTensor<YDataType> y_host_dev({N, C, H, W});
 
     // Fill input with random data
     ck_tile::FillUniformDistribution<XDataType>{-5.f, 5.f}(x_host);
+    ck_tile::FillUniformDistribution<ComputeDataType>{0.8f, 1.2f}(gamma_host);  // Scale around 1.0
+    ck_tile::FillUniformDistribution<ComputeDataType>{-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<XDataType, ComputeDataType, YDataType, Shape>;
     using Kernel = ck_tile::BatchnormFwd<Problem>;
 
-    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<const XDataType*>(x_buf.GetDeviceBuffer()),
-            static_cast<YDataType*>(y_buf.GetDeviceBuffer()),
-            N,
-            C,
-            H,
-            W,
-            static_cast<ComputeDataType>(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<XDataType, YDataType, ComputeDataType>(
-            x_host, y_host_ref, N, C, H, W, epsilon);
+        // Compute reference with gamma and beta
+        reference_batchnorm_fwd<XDataType, YDataType, ComputeDataType, ComputeDataType, ComputeDataType>(
+            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<float>(y_host_ref.mData[idx]);
+                float dev_val = ck_tile::type_convert<float>(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);
         

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 <typename Problem_>
 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<const XDataType*>(kargs.p_x);
+        YDataType* p_y = static_cast<YDataType*>(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<ComputeDataType>(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<ComputeDataType>::template Run<index_t, kBlockSize>(
             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<const ComputeDataType*>(kargs.p_gamma);
+        const ComputeDataType* p_beta = static_cast<const ComputeDataType*>(kargs.p_beta);
+        
+        const ComputeDataType gamma = p_gamma[c];
+        const ComputeDataType beta = p_beta[c];
+        
+        // Compute inverse standard deviation
         ComputeDataType inv_std = type_convert<ComputeDataType>(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<ComputeDataType>(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<YDataType>(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;
     }
 };