From a1bd7af3c5bc13c32b31bf6d3deb632ea34427e4 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 3 Sep 2026 11:25:29 -0700 Subject: [PATCH 1/6] Update batch_norm kernels. The batch_norm portion of a0890b364 ("Update kernels for conv_2d, batch_norm, pool_2d, concat, reshape, split, batch_matmul, transpose and upsample."), plus the create_4d_accessor_w_with_contents fix its GPU test needs. Co-Authored-By: Claude Opus 5.5 --- .../include/kernels/batch_norm_kernels.h | 71 ++-- .../include/kernels/batch_norm_kernels_cpu.h | 36 +-- .../include/kernels/batch_norm_kernels_gpu.h | 73 +++-- .../batch_norm_per_device_state.dtg.toml | 29 -- .../kernels/create_accessor_with_contents.h | 11 +- .../src/cuda/ops/batch_norm_kernels.cu | 302 +++++++++--------- lib/kernels/src/kernels/batch_norm_kernels.cc | 150 ++++----- .../src/kernels/batch_norm_kernels_cpu.cc | 32 +- .../src/kernels/batch_norm_kernels_gpu.cc | 225 +++++++++++++ .../test/src/test_batch_norm_kernel.cc | 101 ------ .../src/task-spec/ops/impl/batch_norm.cc | 106 +++--- 11 files changed, 631 insertions(+), 505 deletions(-) create mode 100644 lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc delete mode 100644 lib/kernels/test/src/test_batch_norm_kernel.cc diff --git a/lib/kernels/include/kernels/batch_norm_kernels.h b/lib/kernels/include/kernels/batch_norm_kernels.h index 47cb3d85a8..ee56bf7173 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels.h +++ b/lib/kernels/include/kernels/batch_norm_kernels.h @@ -1,47 +1,52 @@ #ifndef _FLEXFLOW_KERNELS_BATCH_NORM_KERNELS_H #define _FLEXFLOW_KERNELS_BATCH_NORM_KERNELS_H +#include "kernels/accessor.h" #include "kernels/allocation.h" #include "kernels/batch_norm_per_device_state.dtg.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/device_stream_t.dtg.h" -#include "kernels/ff_handle.h" +#include "op-attrs/ops/batch_norm_attrs.dtg.h" +#include "op-attrs/tensor_shape.dtg.h" +#include "pcg/device_type.dtg.h" -namespace FlexFlow::Kernels::BatchNorm { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - Allocator &allocator, - float *runningMean, - int output_n, - int output_c, - int output_h, - int output_w, - bool relu); - -void forward_kernel(device_stream_t const &stream, - BatchNormPerDeviceState const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr); - -void backward_kernel(device_stream_t const &stream, - BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements); - -void cleanup_kernel( + batch_norm_init_kernel(DeviceType device_type, + Allocator &allocator, + BatchNormAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); + +void batch_norm_forward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output); + +void batch_norm_backward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad); + +void batch_norm_cleanup_kernel( DeviceType device_type, Allocator &allocator, - std::optional const &per_device_state); + std::optional &per_device_state); + +} // namespace FlexFlow -} // namespace FlexFlow::Kernels::BatchNorm #endif diff --git a/lib/kernels/include/kernels/batch_norm_kernels_cpu.h b/lib/kernels/include/kernels/batch_norm_kernels_cpu.h index 8c564d6450..ecda6883a0 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels_cpu.h +++ b/lib/kernels/include/kernels/batch_norm_kernels_cpu.h @@ -1,28 +1,26 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_NORM_KERNELS_CPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_NORM_KERNELS_CPU_H -#include "kernels/allocation.h" -#include "kernels/batch_norm_per_device_state.dtg.h" -#include "kernels/device_stream_t.dtg.h" +#include "kernels/accessor.h" +#include "op-attrs/ops/batch_norm_attrs.dtg.h" -namespace FlexFlow::Kernels::BatchNorm { +namespace FlexFlow { -void cpu_forward_kernel(BatchNormPerDeviceState const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr); +void batch_norm_cpu_forward_kernel(BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output); -void cpu_backward_kernel(BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements); +void batch_norm_cpu_backward_kernel(BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad); -} // namespace FlexFlow::Kernels::BatchNorm +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/batch_norm_kernels_gpu.h b/lib/kernels/include/kernels/batch_norm_kernels_gpu.h index 41f9808bff..277c36a9ed 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels_gpu.h +++ b/lib/kernels/include/kernels/batch_norm_kernels_gpu.h @@ -1,43 +1,46 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_NORM_KERNELS_GPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_NORM_KERNELS_GPU_H +#include "kernels/accessor.h" #include "kernels/allocation.h" #include "kernels/batch_norm_per_device_state.dtg.h" #include "kernels/device.h" -#include "kernels/ff_handle.h" - -namespace FlexFlow::Kernels::BatchNorm { - -BatchNormPerDeviceState gpu_init_kernel(PerDeviceFFHandle const &handle, - Allocator &allocator, - float *runningMean, - int output_n, - int output_c, - int output_h, - int output_w, - bool relu); - -void gpu_forward_kernel(ffStream_t stream, - BatchNormPerDeviceState const &per_device_statem, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr); - -void gpu_backward_kernel(ffStream_t stream, - BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements); - -void gpu_cleanup_kernel(Allocator &allocator, - BatchNormPerDeviceState &per_device_state); - -} // namespace FlexFlow::Kernels::BatchNorm +#include "op-attrs/ops/batch_norm_attrs.dtg.h" + +namespace FlexFlow { + +BatchNormPerDeviceState + batch_norm_gpu_init_kernel(Allocator &allocator, + BatchNormAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); + +void batch_norm_gpu_forward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + BatchNormPerDeviceState const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output); + +void batch_norm_gpu_backward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + BatchNormPerDeviceState const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad); + +void batch_norm_gpu_cleanup_kernel(Allocator &allocator, + BatchNormPerDeviceState &per_device_state); + +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml b/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml index bdf9e1ed51..f1548e4891 100644 --- a/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml +++ b/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml @@ -5,13 +5,8 @@ features = [] includes = [ "kernels/device.h", - "kernels/ff_handle.h", ] -[[fields]] -name = "handle" -type = "::FlexFlow::PerDeviceFFHandle" - [[fields]] name = "inputTensor" type = "ffTensorDescriptor_t" @@ -24,10 +19,6 @@ type = "ffTensorDescriptor_t" name = "biasTensor" type = "ffTensorDescriptor_t" -[[fields]] -name = "actiDesc" -type = "ffActivationDescriptor_t" - [[fields]] name = "mode" type = "ffBatchNormMode_t" @@ -47,23 +38,3 @@ type = "float *" [[fields]] name = "saveVar" type = "float *" - -[[fields]] -name = "output_n" -type = "int" - -[[fields]] -name = "output_c" -type = "int" - -[[fields]] -name = "output_h" -type = "int" - -[[fields]] -name = "output_w" -type = "int" - -[[fields]] -name = "relu" -type = "bool" diff --git a/lib/kernels/include/kernels/create_accessor_with_contents.h b/lib/kernels/include/kernels/create_accessor_with_contents.h index dfe4428bdc..4e10ffec0c 100644 --- a/lib/kernels/include/kernels/create_accessor_with_contents.h +++ b/lib/kernels/include/kernels/create_accessor_with_contents.h @@ -155,7 +155,8 @@ GenericTensorAccessorW create_4d_accessor_w_with_contents( type_to_data_type_enum_v, }; - GenericTensorAccessorW accessor = allocator.allocate_tensor(shape); + Allocator cpu_allocator = create_local_cpu_memory_allocator(); + GenericTensorAccessorW cpu_accessor = cpu_allocator.allocate_tensor(shape); for (nonnegative_int dim0_idx : nonnegative_range(dim0_size.nonnegative_int_from_positive_int())) { @@ -165,7 +166,7 @@ GenericTensorAccessorW create_4d_accessor_w_with_contents( nonnegative_range(dim2_size.nonnegative_int_from_positive_int())) { for (nonnegative_int dim3_idx : nonnegative_range(dim3_size.nonnegative_int_from_positive_int())) { - accessor.at>(TensorDimsCoord{ + cpu_accessor.at>(TensorDimsCoord{ FFOrdered{dim0_idx, dim1_idx, dim2_idx, dim3_idx}}) = contents.at(dim0_idx.unwrap_nonnegative()) .at(dim1_idx.unwrap_nonnegative()) @@ -176,7 +177,11 @@ GenericTensorAccessorW create_4d_accessor_w_with_contents( } } - return accessor; + GenericTensorAccessorW result = allocator.allocate_tensor(shape); + copy_accessor_data_to_l_from_r( + result, read_only_accessor_from_write_accessor(cpu_accessor)); + + return result; } template diff --git a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu index 26234fd6e0..3e4514194c 100644 --- a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu +++ b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu @@ -14,170 +14,186 @@ */ #include "internal/device.h" -#include "kernels/allocation.h" -#include "kernels/batch_norm_kernels.h" -#include "kernels/ff_handle.h" -#include "utils/integer_conversions.h" +#include "kernels/batch_norm_kernels_gpu.h" +#include "op-attrs/tensor_dims.h" +#include "utils/containers/require_same.h" +#include namespace FlexFlow { -namespace Kernels { -namespace BatchNorm { - -void gpu_forward_kernel(cudaStream_t stream, - BatchNormPerDeviceState const &m, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr) { - checkCUDNN(cudnnSetStream(m.handle.dnn, stream)); - - float alpha = 1.0f, beta = 0.0f; - checkCUDNN(cudnnBatchNormalizationForwardTraining(m.handle.dnn, - m.mode, - &alpha, - &beta, - m.inputTensor, - input_ptr, - m.outputTensor, - output_ptr, - m.biasTensor, - scale_ptr, - bias_ptr, - 1.0, - m.runningMean, - m.runningVar, - CUDNN_BN_MIN_EPSILON, - m.saveMean, - m.saveVar)); -} -void gpu_backward_kernel(cudaStream_t stream, - BatchNormPerDeviceState const &m, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements) { - checkCUDNN(cudnnSetStream(m.handle.dnn, stream)); - - float alpha = 1.0f; - if (m.relu) { - reluBackward<<>>( - output_grad_ptr, output_ptr, numElements); - } - checkCUDNN(cudnnBatchNormalizationBackward(m.handle.dnn, - m.mode, - &alpha, - &alpha, - &alpha, - &alpha, - m.inputTensor, - input_ptr, - m.outputTensor, - output_grad_ptr, - m.inputTensor, - input_grad_ptr, - m.biasTensor, - scale_ptr, - scale_grad_ptr, - bias_grad_ptr, - CUDNN_BN_MIN_EPSILON, - m.saveMean, - m.saveVar)); +static positive_int get_num_channels(TensorShape const &shape) { + return dim_at_idx(shape.dims, ff_dim_t{1_n}); } -BatchNormPerDeviceState gpu_init_kernel(PerDeviceFFHandle const &handle, - Allocator &allocator, - float *runningMean, - int output_n, - int output_c, - int output_h, - int output_w, - bool relu) { +BatchNormPerDeviceState + batch_norm_gpu_init_kernel(Allocator &allocator, + BatchNormAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { + // Applying an activation as part of BatchNorm is not currently implemented + // (the forward kernel used to silently ignore it while the backward kernel + // applied it). If you need it, please create an issue. + ASSERT(!attrs.relu, + "BatchNorm does not currently support a fused relu activation"); + + ASSERT(attrs.affine, + "BatchNorm currently only supports attrs.affine = true. " + "If you need this feature, please create an issue."); + + TensorShape shape = require_same(input_shape, output_shape); + + ASSERT(get_num_dims(shape.dims) == num_tensor_dims_t{4_n}, + "BatchNorm currently only supports 4-dimensional (i.e., NCHW) " + "tensors. If you need this feature, please create an issue.", + shape); + + ASSERT(attrs.eps >= CUDNN_BN_MIN_EPSILON, + "cuDNN requires BatchNorm eps to be at least CUDNN_BN_MIN_EPSILON", + attrs.eps, + CUDNN_BN_MIN_EPSILON); + + int num_channels = get_num_channels(shape).int_from_positive_int(); + ffTensorDescriptor_t inputTensor; ffTensorDescriptor_t outputTensor; ffTensorDescriptor_t biasTensor; - ffActivationDescriptor_t actiDesc; - ffBatchNormMode_t mode; + checkCUDNN(cudnnCreateTensorDescriptor(&inputTensor)); - checkCUDNN(cudnnCreateTensorDescriptor(&biasTensor)); checkCUDNN(cudnnCreateTensorDescriptor(&outputTensor)); - mode = CUDNN_BATCHNORM_SPATIAL; + checkCUDNN(cudnnCreateTensorDescriptor(&biasTensor)); + + ffBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL; #if CUDNN_VERSION >= 7000 mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT; #endif - checkCUDNN(cudnnSetTensor4dDescriptor(inputTensor, - CUDNN_TENSOR_NCHW, - CUDNN_DATA_FLOAT, - output_n, - output_c, - output_h, - output_w)); - checkCUDNN(cudnnSetTensor4dDescriptor(outputTensor, + + checkCUDNN(cudnnSetTensorDescriptorFromTensorShape(inputTensor, input_shape)); + checkCUDNN( + cudnnSetTensorDescriptorFromTensorShape(outputTensor, output_shape)); + checkCUDNN(cudnnSetTensor4dDescriptor(biasTensor, CUDNN_TENSOR_NCHW, - CUDNN_DATA_FLOAT, - output_n, - output_c, - output_h, - output_w)); - checkCUDNN(cudnnSetTensor4dDescriptor( - biasTensor, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, output_c, 1, 1)); - // allocate memory for runningMean, runningVar, saveMean, saveVar - size_t totalSize = sizeof(float) * output_c * 4; - runningMean = (float *)allocator.allocate(totalSize); - float *runningVar = (float *)runningMean + output_c; - float *saveMean = (float *)runningVar + output_c; - float *saveVar = (float *)saveMean + output_c; - cudaStream_t stream; - checkCUDA(get_legion_stream(&stream)); - - assign_kernel<<>>( - runningMean, size_t_from_int(output_c), 0.0f); - assign_kernel<<>>( - runningVar, size_t_from_int(output_c), 0.0f); - - if (relu) { - checkCUDNN(cudnnCreateActivationDescriptor(&actiDesc)); - checkCUDNN(cudnnSetActivationDescriptor( - actiDesc, CUDNN_ACTIVATION_RELU, CUDNN_PROPAGATE_NAN, 0.0)); - } - - BatchNormPerDeviceState per_device_state = BatchNormPerDeviceState{ - handle, - inputTensor, - outputTensor, - biasTensor, - actiDesc, - mode, - runningMean, - runningVar, - saveMean, - saveVar, - output_n, - output_c, - output_h, - output_w, - relu, + ff_to_cudnn_datatype(shape.data_type), + /*n=*/1, + /*c=*/num_channels, + /*h=*/1, + /*w=*/1)); + + // Allocate memory for runningMean, runningVar, saveMean and saveVar as a + // single contiguous block (deallocated by batch_norm_gpu_cleanup_kernel). + float *runningMean = static_cast( + allocator.allocate(sizeof(float) * num_channels * 4)); + float *runningVar = runningMean + num_channels; + float *saveMean = runningVar + num_channels; + float *saveVar = saveMean + num_channels; + + // Match the PyTorch initialization of running_mean = 0 and running_var = 1. + std::vector initial_running_stats(num_channels * 2); + std::fill(initial_running_stats.begin(), + initial_running_stats.begin() + num_channels, + 0.0f); + std::fill(initial_running_stats.begin() + num_channels, + initial_running_stats.end(), + 1.0f); + checkCUDA(cudaMemcpy(runningMean, + initial_running_stats.data(), + sizeof(float) * num_channels * 2, + cudaMemcpyHostToDevice)); + + return BatchNormPerDeviceState{ + /*inputTensor=*/inputTensor, + /*outputTensor=*/outputTensor, + /*biasTensor=*/biasTensor, + /*mode=*/mode, + /*runningMean=*/runningMean, + /*runningVar=*/runningVar, + /*saveMean=*/saveMean, + /*saveVar=*/saveVar, }; +} + +void batch_norm_gpu_forward_kernel( + cudaStream_t stream, + PerDeviceFFHandle const &handle, + BatchNormPerDeviceState const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output) { + checkCUDNN(cudnnSetStream(handle.dnn, stream)); - checkCUDA(cudaStreamDestroy(stream)); - return per_device_state; + // NOTE: attrs.momentum = std::nullopt means "use a cumulative moving + // average", which cuDNN cannot express, so we fall back to fully replacing + // the running statistics on each call. The running statistics are currently + // never read back, so this only matters once inference mode is supported. + double exponential_average_factor = attrs.momentum.value_or(1.0); + + float alpha = 1.0f, beta_coeff = 0.0f; + checkCUDNN( + cudnnBatchNormalizationForwardTraining(handle.dnn, + per_device_state.mode, + &alpha, + &beta_coeff, + per_device_state.inputTensor, + input.ptr, + per_device_state.outputTensor, + output.ptr, + per_device_state.biasTensor, + gamma.ptr, + beta.ptr, + exponential_average_factor, + per_device_state.runningMean, + per_device_state.runningVar, + attrs.eps, + per_device_state.saveMean, + per_device_state.saveVar)); } -void gpu_cleanup_kernel(Allocator &allocator, - BatchNormPerDeviceState &per_device_state) { +void batch_norm_gpu_backward_kernel( + cudaStream_t stream, + PerDeviceFFHandle const &handle, + BatchNormPerDeviceState const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad) { + checkCUDNN(cudnnSetStream(handle.dnn, stream)); + + // NOTE: the beta coefficients are 1.0 so that the gradients are accumulated + // into rather than overwritten + float alpha_data = 1.0f, beta_data = 1.0f; + float alpha_param = 1.0f, beta_param = 1.0f; + checkCUDNN(cudnnBatchNormalizationBackward(handle.dnn, + per_device_state.mode, + &alpha_data, + &beta_data, + &alpha_param, + &beta_param, + per_device_state.inputTensor, + input.ptr, + per_device_state.outputTensor, + output_grad.ptr, + per_device_state.inputTensor, + input_grad.ptr, + per_device_state.biasTensor, + gamma.ptr, + gamma_grad.ptr, + beta_grad.ptr, + attrs.eps, + per_device_state.saveMean, + per_device_state.saveVar)); +} + +void batch_norm_gpu_cleanup_kernel(Allocator &allocator, + BatchNormPerDeviceState &per_device_state) { allocator.deallocate(per_device_state.runningMean); checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.inputTensor)); - checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.biasTensor)); checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.outputTensor)); - if (per_device_state.relu) { - checkCUDNN(cudnnDestroyActivationDescriptor(per_device_state.actiDesc)); - } + checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.biasTensor)); } -} // namespace BatchNorm -} // namespace Kernels } // namespace FlexFlow diff --git a/lib/kernels/src/kernels/batch_norm_kernels.cc b/lib/kernels/src/kernels/batch_norm_kernels.cc index e23f6a89e2..546aa1ce5b 100644 --- a/lib/kernels/src/kernels/batch_norm_kernels.cc +++ b/lib/kernels/src/kernels/batch_norm_kernels.cc @@ -1,107 +1,111 @@ #include "kernels/batch_norm_kernels.h" #include "kernels/batch_norm_kernels_cpu.h" #include "kernels/batch_norm_kernels_gpu.h" +#include "utils/optional.h" -namespace FlexFlow::Kernels::BatchNorm { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - Allocator &allocator, - float *runningMean, - int output_n, - int output_c, - int output_h, - int output_w, - bool relu) { + batch_norm_init_kernel(DeviceType device_type, + Allocator &allocator, + BatchNormAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { if (device_type == DeviceType::GPU) { - return gpu_init_kernel( - /*handle=*/handle.require_for_gpu(), + return batch_norm_gpu_init_kernel( /*allocator=*/allocator, - /*runningMean=*/runningMean, - /*output_n=*/output_n, - /*output_c=*/output_c, - /*output_h=*/output_h, - /*output_w=*/output_w, - /*relu=*/relu); + /*attrs=*/attrs, + /*input_shape=*/input_shape, + /*output_shape=*/output_shape); } else { ASSERT(device_type == DeviceType::CPU); - ASSERT(handle.is_for_cpu()); return std::nullopt; } } -void forward_kernel(device_stream_t const &stream, - BatchNormPerDeviceState const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr) { +void batch_norm_forward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + batch_norm_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state, - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr, - /*scale_ptr=*/scale_ptr, - /*bias_ptr=*/bias_ptr); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*input=*/input, + /*gamma=*/gamma, + /*beta=*/beta, + /*output=*/output); } else { ASSERT(stream.is_cpu()); - cpu_forward_kernel( - /*per_device_state=*/per_device_state, - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr, - /*scale_ptr=*/scale_ptr, - /*bias_ptr=*/bias_ptr); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + batch_norm_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*gamma=*/gamma, + /*beta=*/beta, + /*output=*/output); } } -void backward_kernel(device_stream_t const &stream, - BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements) { +void batch_norm_backward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + batch_norm_gpu_backward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state, - /*output_ptr=*/output_ptr, - /*output_grad_ptr=*/output_grad_ptr, - /*input_ptr=*/input_ptr, - /*input_grad_ptr=*/input_grad_ptr, - /*scale_ptr=*/scale_ptr, - /*scale_grad_ptr=*/scale_grad_ptr, - /*bias_grad_ptr=*/bias_grad_ptr, - /*numElements=*/numElements); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad, + /*gamma=*/gamma, + /*gamma_grad=*/gamma_grad, + /*beta_grad=*/beta_grad); } else { ASSERT(stream.is_cpu()); - cpu_backward_kernel( - /*per_device_state=*/per_device_state, - /*output_ptr=*/output_ptr, - /*output_grad_ptr=*/output_grad_ptr, - /*input_ptr=*/input_ptr, - /*input_grad_ptr=*/input_grad_ptr, - /*scale_ptr=*/scale_ptr, - /*scale_grad_ptr=*/scale_grad_ptr, - /*bias_grad_ptr=*/bias_grad_ptr, - /*numElements=*/numElements); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + batch_norm_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad, + /*gamma=*/gamma, + /*gamma_grad=*/gamma_grad, + /*beta_grad=*/beta_grad); } } -void cleanup_kernel(DeviceType device_type, - Allocator &allocator, - std::optional &per_device_state) { +void batch_norm_cleanup_kernel( + DeviceType device_type, + Allocator &allocator, + std::optional &per_device_state) { if (device_type == DeviceType::GPU) { - gpu_cleanup_kernel(allocator, per_device_state.value()); + batch_norm_gpu_cleanup_kernel(allocator, per_device_state.value()); } else { ASSERT(device_type == DeviceType::CPU); - ASSERT(per_device_state == std::nullopt); + ASSERT(!per_device_state.has_value()); } } -} // namespace FlexFlow::Kernels::BatchNorm +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc b/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc index be440ac4ea..5e16a11fe7 100644 --- a/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc +++ b/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc @@ -1,25 +1,25 @@ #include "kernels/batch_norm_kernels_cpu.h" +#include "utils/exception.h" -namespace FlexFlow::Kernels::BatchNorm { +namespace FlexFlow { -void cpu_forward_kernel(BatchNormPerDeviceState const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr) { +void batch_norm_cpu_forward_kernel(BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output) { NOT_IMPLEMENTED(); } -void cpu_backward_kernel(BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements) { +void batch_norm_cpu_backward_kernel(BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad) { NOT_IMPLEMENTED(); } -} // namespace FlexFlow::Kernels::BatchNorm +} // namespace FlexFlow diff --git a/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc new file mode 100644 index 0000000000..ad92f0ea56 --- /dev/null +++ b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc @@ -0,0 +1,225 @@ +#include "kernels/batch_norm_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +static BatchNormAttrs make_attrs() { + return BatchNormAttrs{ + /*relu=*/false, + /*affine=*/true, + /*eps=*/1e-5, + /*momentum=*/0.1, + }; +} + +// NCHW +static GenericTensorAccessorR make_input(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + {{1, 2}, {3, 4}}, + {{-1, -2}, {0.5, 0.25}}, + }, + { + {{5, 6}, {7, 8}}, + {{2, 0.125}, {-3, 1}}, + }, + }, + allocator); +} + +static GenericTensorAccessorR make_gamma(Allocator &allocator) { + return create_1d_accessor_r_with_contents({2, 0.5}, allocator); +} + +static GenericTensorAccessorR make_beta(Allocator &allocator) { + return create_1d_accessor_r_with_contents({-1, 3}, allocator); +} + +static GenericTensorAccessorR make_output(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + {{-4.055047512054443, -3.1821770668029785}, + {-2.3093061447143555, -1.4364354610443115}}, + {{2.760241985321045, 2.433763027191162}, + {3.249960422515869, 3.1683406829833984}}, + }, + { + {{-0.5635647177696228, 0.3093060255050659}, + {1.1821768283843994, 2.0550475120544434}}, + {{3.7396788597106934, 3.127530813217163}, + {2.1072840690612793, 3.4131999015808105}}, + }, + }, + allocator); +} + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("batch_norm_gpu_forward_kernel") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + BatchNormAttrs attrs = make_attrs(); + + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR gamma = make_gamma(allocator); + GenericTensorAccessorR beta = make_beta(allocator); + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(input.shape, allocator); + + BatchNormPerDeviceState per_device_state = + batch_norm_gpu_init_kernel(allocator, attrs, input.shape, output.shape); + + batch_norm_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*gamma=*/gamma, + /*beta=*/beta, + /*output=*/output); + + GenericTensorAccessorR correct = make_output(allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + + batch_norm_gpu_cleanup_kernel(allocator, per_device_state); + } + + TEST_CASE("batch_norm_gpu_backward_kernel") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + BatchNormAttrs attrs = make_attrs(); + + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR gamma = make_gamma(allocator); + GenericTensorAccessorR beta = make_beta(allocator); + + GenericTensorAccessorW forward_output = + create_random_filled_accessor_w(input.shape, allocator); + + BatchNormPerDeviceState per_device_state = batch_norm_gpu_init_kernel( + allocator, attrs, input.shape, forward_output.shape); + + // cudnnBatchNormalizationBackward reads the batch statistics saved by the + // forward pass, so the forward kernel has to run first + batch_norm_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*gamma=*/gamma, + /*beta=*/beta, + /*output=*/forward_output); + + GenericTensorAccessorR output = make_output(allocator); + + GenericTensorAccessorR output_grad = + create_4d_accessor_r_with_contents( + { + { + {{-0.875, -0.75}, {-0.625, -0.5}}, + {{-0.375, -0.25}, {-0.125, 0}}, + }, + { + {{0.125, 0.25}, {0.375, 0.5}}, + {{0.625, 0.75}, {0.875, 1}}, + }, + }, + allocator); + + // The gradients are accumulated into, so they need to start from known + // values + GenericTensorAccessorW input_grad = + create_4d_accessor_w_with_contents( + { + { + {{-1.75, -1.5}, {-1.25, -1}}, + {{-0.75, -0.5}, {-0.25, 0}}, + }, + { + {{0.25, 0.5}, {0.75, 1}}, + {{1.25, 1.5}, {1.75, 2}}, + }, + }, + allocator); + + GenericTensorAccessorW gamma_grad = + create_1d_accessor_w_with_contents({7, -3}, allocator); + + GenericTensorAccessorW beta_grad = + create_1d_accessor_w_with_contents({0.5, 11}, allocator); + + batch_norm_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad, + /*gamma=*/gamma, + /*gamma_grad=*/gamma_grad, + /*beta_grad=*/beta_grad); + + GenericTensorAccessorR correct_input_grad = + create_4d_accessor_r_with_contents( + { + { + {{-1.6772620677947998, -1.510392189025879}, + {-1.3435224294662476, -1.1766525506973267}}, + {{-0.9591808915138245, -0.6475732326507568}, + {-0.4087578058242798, -0.11274851113557816}}, + }, + { + {{0.42665261030197144, 0.5935224294662476}, + {0.7603921890258789, 0.927262008190155}}, + {{1.3049046993255615, 1.634710431098938}, + {1.9905133247375488, 2.198132038116455}}, + }, + }, + allocator); + + GenericTensorAccessorR correct_gamma_grad = + create_1d_accessor_r_with_contents( + {11.037027359008789, -2.2195112705230713}, allocator); + + GenericTensorAccessorR correct_beta_grad = + create_1d_accessor_r_with_contents({-1.0, 13.5}, allocator); + + CHECK_MESSAGE( + // cuDNN's batch norm data gradient is not bit-identical to PyTorch's + accessors_within_epsilon(input_grad, correct_input_grad, 1e-6), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + + CHECK_MESSAGE( + accessors_are_equal(gamma_grad, correct_gamma_grad), + check_kv("gamma_grad", format_accessor_w_contents(gamma_grad))); + + CHECK_MESSAGE(accessors_are_equal(beta_grad, correct_beta_grad), + check_kv("beta_grad", format_accessor_w_contents(beta_grad))); + + batch_norm_gpu_cleanup_kernel(allocator, per_device_state); + } +} diff --git a/lib/kernels/test/src/test_batch_norm_kernel.cc b/lib/kernels/test/src/test_batch_norm_kernel.cc deleted file mode 100644 index 1be45d8bbb..0000000000 --- a/lib/kernels/test/src/test_batch_norm_kernel.cc +++ /dev/null @@ -1,101 +0,0 @@ -#include "internal/test_utils.h" -#include "kernels/batch_norm_kernels_gpu.h" -#include "op-attrs/datatype_value.h" -#include - -using namespace ::FlexFlow; - -TEST_SUITE(FF_CUDA_TEST_SUITE) { - TEST_CASE("Test BatchNorm Kernel") { - positive_int output_n = 1_p; - positive_int output_c = 10_p; - positive_int output_h = 10_p; - positive_int output_w = 10_p; - - ManagedFFStream managed_stream{}; - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - - Allocator allocator = create_local_cuda_memory_allocator(); - - BatchNormPerDeviceState state = Kernels::BatchNorm::gpu_init_kernel( - /*handle=*/managed_handle.raw_handle(), - /*allocator=*/allocator, - /*runningMean=*/nullptr, - /*output_n=*/output_n.int_from_positive_int(), - /*output_c=*/output_c.int_from_positive_int(), - /*output_h=*/output_h.int_from_positive_int(), - /*output_w=*/output_w.int_from_positive_int(), - /*relu=*/true); - - TensorShape input_shape = TensorShape{ - TensorDims{FFOrdered{output_n, output_c, output_h, output_w}}, - DataType::FLOAT, - }; - TensorShape output_shape = TensorShape{ - TensorDims{FFOrdered{output_n, output_c, output_h, output_w}}, - DataType::FLOAT, - }; - TensorShape scale_shape = TensorShape{ - TensorDims{FFOrdered{output_n, output_c, output_h, output_w}}, - DataType::FLOAT, - }; - TensorShape bias_shape = TensorShape{ - TensorDims{FFOrdered{output_n, output_c, output_h, output_w}}, - DataType::FLOAT, - }; - - GenericTensorAccessorW input_accessor = - create_random_filled_accessor_w(input_shape, allocator); - GenericTensorAccessorW output_accessor = - create_random_filled_accessor_w(output_shape, allocator); - GenericTensorAccessorW scale_accessor = create_filled_accessor_w( - scale_shape, allocator, make_float_data_type_value(1)); - - SUBCASE("gpu_forward_kernel") { - GenericTensorAccessorW bias_accessor = create_filled_accessor_w( - bias_shape, allocator, make_float_data_type_value(0)); - - Kernels::BatchNorm::gpu_forward_kernel( - /*stream=*/managed_stream.raw_stream(), - /*per_device_state=*/state, - /*input_ptr=*/input_accessor.get_float_ptr(), - /*output_ptr=*/output_accessor.get_float_ptr(), - /*scale_ptr=*/scale_accessor.get_float_ptr(), - /*bias_ptr=*/bias_accessor.get_float_ptr()); - - CHECK(contains_non_zero(output_accessor)); - } - - SUBCASE("gpu_backward_kernel") { - GenericTensorAccessorW output_grad_accessor = - create_random_filled_accessor_w(output_shape, allocator); - GenericTensorAccessorW input_grad_accessor = - create_random_filled_accessor_w(input_shape, allocator); - GenericTensorAccessorW scale_grad_accessor = - create_random_filled_accessor_w(scale_shape, allocator); - GenericTensorAccessorW bias_grad_accessor = - create_random_filled_accessor_w(bias_shape, allocator); - - Kernels::BatchNorm::gpu_backward_kernel( - /*stream=*/managed_stream.raw_stream(), - /*per_device_state=*/state, - /*output_ptr=*/output_accessor.get_float_ptr(), - /*output_grad_ptr=*/output_grad_accessor.get_float_ptr(), - /*input_ptr=*/input_accessor.get_float_ptr(), - /*input_grad_ptr=*/input_grad_accessor.get_float_ptr(), - /*scale_ptr=*/scale_accessor.get_float_ptr(), - /*scale_grad_ptr=*/scale_grad_accessor.get_float_ptr(), - /*bias_grad_ptr=*/bias_grad_accessor.get_float_ptr(), - /*numElements=*/ - get_num_elements(input_accessor.shape.dims).int_from_positive_int()); - - CHECK(contains_non_zero(input_grad_accessor)); - CHECK(contains_non_zero(scale_grad_accessor)); - CHECK(contains_non_zero(bias_grad_accessor)); - } - - Kernels::BatchNorm::gpu_cleanup_kernel(allocator, state); - } -} diff --git a/lib/task-spec/src/task-spec/ops/impl/batch_norm.cc b/lib/task-spec/src/task-spec/ops/impl/batch_norm.cc index 5fb89f3eb9..ea029262eb 100644 --- a/lib/task-spec/src/task-spec/ops/impl/batch_norm.cc +++ b/lib/task-spec/src/task-spec/ops/impl/batch_norm.cc @@ -19,35 +19,18 @@ namespace FlexFlow { -using namespace FlexFlow::Kernels::BatchNorm; - static DeviceSpecificPerDeviceOpState init_task_impl(TaskArgumentAccessor const &acc) { Allocator allocator = acc.get_allocator(); - device_handle_t handle = acc.get_ff_handle(); - std::optional profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - - auto output = acc.get_tensor(TensorSlotName::OUTPUT); BatchNormAttrs attrs = acc.get_op_attrs().require_batch_norm(); - positive_int output_w = dim_at_idx(output.shape.dims, legion_dim_t{0_n}); - positive_int output_h = dim_at_idx(output.shape.dims, legion_dim_t{1_n}); - positive_int output_c = dim_at_idx(output.shape.dims, legion_dim_t{2_n}); - positive_int output_n = dim_at_idx(output.shape.dims, legion_dim_t{3_n}); - - float *runningMean; + TensorShape input_shape = acc.get_tensor_shape(TensorSlotName::INPUT); + TensorShape output_shape = acc.get_tensor_shape(TensorSlotName::OUTPUT); - std::optional per_device_state = init_kernel( - /*device_type=*/kernel_device_type, - /*handle=*/handle, - /*allocator=*/allocator, - /*runningMean=*/runningMean, - /*output_n=*/output_n.int_from_positive_int(), - /*output_c=*/output_c.int_from_positive_int(), - /*output_h=*/output_h.int_from_positive_int(), - /*output_w=*/output_w.int_from_positive_int(), - /*relu=*/attrs.relu); + std::optional per_device_state = + batch_norm_init_kernel( + kernel_device_type, allocator, attrs, input_shape, output_shape); return DeviceSpecificPerDeviceOpState{ acc.make_device_specific(per_device_state), @@ -56,56 +39,73 @@ static DeviceSpecificPerDeviceOpState static std::optional forward_task_impl(TaskArgumentAccessor const &acc) { - auto per_device_state = - acc.get_per_device_op_state().require_batch_norm().value(); std::optional profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); + device_handle_t handle = acc.get_ff_handle(); + std::optional per_device_state = + acc.get_per_device_op_state().require_batch_norm(); + BatchNormAttrs attrs = acc.get_op_attrs().require_batch_norm(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - auto scale = acc.get_tensor(TensorSlotName::SCALE); - auto bias = acc.get_tensor(TensorSlotName::BIAS); + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorR gamma = + acc.get_tensor(TensorSlotName::GAMMA); + GenericTensorAccessorR beta = + acc.get_tensor(TensorSlotName::BETA); + GenericTensorAccessorW output = + acc.get_tensor(TensorSlotName::OUTPUT); - return profile(forward_kernel, + return profile(batch_norm_forward_kernel, profiling, kernel_device_type, "[BatchNorm] forward_time = {:.2lf}ms\n", + handle, per_device_state, - input.get_float_ptr(), - output.get_float_ptr(), - scale.get_float_ptr(), - bias.get_float_ptr()); + attrs, + input, + gamma, + beta, + output); } static std::optional backward_task_impl(TaskArgumentAccessor const &acc) { - BatchNormPerDeviceState per_device_state = - acc.get_per_device_op_state().require_batch_norm().value(); std::optional profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); + device_handle_t handle = acc.get_ff_handle(); + std::optional per_device_state = + acc.get_per_device_op_state().require_batch_norm(); + BatchNormAttrs attrs = acc.get_op_attrs().require_batch_norm(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto input_grad = acc.get_tensor_grad(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - auto output_grad = - acc.get_tensor_grad(TensorSlotName::OUTPUT); - auto scale = acc.get_tensor(TensorSlotName::SCALE); - auto scale_grad = acc.get_tensor_grad(TensorSlotName::SCALE); - auto bias_grad = acc.get_tensor_grad(TensorSlotName::BIAS); - - return profile(backward_kernel, + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorW input_grad = + acc.get_tensor_grad(TensorSlotName::INPUT); + GenericTensorAccessorR output = + acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR output_grad = + acc.get_tensor_grad(TensorSlotName::OUTPUT); + GenericTensorAccessorR gamma = + acc.get_tensor(TensorSlotName::GAMMA); + GenericTensorAccessorW gamma_grad = + acc.get_tensor_grad(TensorSlotName::GAMMA); + GenericTensorAccessorW beta_grad = + acc.get_tensor_grad(TensorSlotName::BETA); + + return profile(batch_norm_backward_kernel, profiling, kernel_device_type, "[BatchNorm] backward_time = {:.2lf}ms\n", + handle, per_device_state, - output.get_float_ptr(), - output_grad.get_float_ptr(), - input.get_float_ptr(), - input_grad.get_float_ptr(), - scale.get_float_ptr(), - scale_grad.get_float_ptr(), - bias_grad.get_float_ptr(), - get_num_elements(output.shape.dims).int_from_positive_int()); + attrs, + output, + output_grad, + input, + input_grad, + gamma, + gamma_grad, + beta_grad); } TaskImplFunction get_batch_norm_init_task_impl() { From b06f3ee55863f5878a4248260347323e5569ccbe Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 3 Sep 2026 11:25:29 -0700 Subject: [PATCH 2/6] Implement SiLU in the element_unary GPU kernels. The element_unary portion of a0890b364 ("Update kernels for conv_2d, batch_norm, pool_2d, concat, reshape, split, batch_matmul, transpose and upsample."). Co-Authored-By: Claude Opus 5.5 --- .../src/cuda/ops/element_unary_kernels.cu | 30 ++- .../src/kernels/element_unary_kernels_gpu.cc | 196 ++++++++++++++---- 2 files changed, 176 insertions(+), 50 deletions(-) diff --git a/lib/kernels/src/cuda/ops/element_unary_kernels.cu b/lib/kernels/src/cuda/ops/element_unary_kernels.cu index 99ce5ba888..ecbb846797 100644 --- a/lib/kernels/src/cuda/ops/element_unary_kernels.cu +++ b/lib/kernels/src/cuda/ops/element_unary_kernels.cu @@ -40,12 +40,25 @@ static bool use_scalar(OperatorType op_type) { case OperatorType::SCALAR_SUB: case OperatorType::SCALAR_TRUE_DIV: case OperatorType::POW: + case OperatorType::SILU: return true; default: return false; } } +// SILU's scalar (beta) is optional and defaults to 1, matching the CPU +// kernel. The other scalar operators require a scalar. +static float require_scalar(OperatorType op_type, + std::optional const &scalar) { + if (op_type == OperatorType::SILU) { + return scalar.value_or(1.0f); + } + + ASSERT(scalar.has_value(), op_type); + return scalar.value(); +} + ElementUnaryPerDeviceState element_unary_gpu_init_kernel(ElementUnaryAttrs const &attrs, TensorShape const &input_shape, @@ -117,6 +130,10 @@ __global__ void elewise_scalar_unary_forward_kernel( out[i] = (T)(powf(in[i], scalar)); break; } + case OperatorType::SILU: { + out[i] = (T)(in[i] / (1.0f + expf(-1.0f * scalar * in[i]))); + break; + } default: assert(false); } @@ -191,6 +208,13 @@ __global__ void elewise_scalar_unary_backward_kernel(coord_t volume, (T)(output_grad[i] * scalar * powf(input[i], scalar - 1)); break; } + case OperatorType::SILU: { + float e_to_bx = expf(scalar * input[i]); + input_grad[i] += (T)(output_grad[i] * + (e_to_bx * (scalar * input[i] + e_to_bx + 1.0f)) / + ((e_to_bx + 1.0f) * (e_to_bx + 1.0f))); + break; + } default: assert(false); } @@ -266,11 +290,10 @@ struct ForwardKernel { m.outputTensor, output.get())); } else if (use_scalar(op_type)) { - assert(scalar.has_value()); elewise_scalar_unary_forward_kernel> <<>>( num_elements, - static_cast>(scalar.value()), + static_cast>(require_scalar(op_type, scalar)), op_type, input.get(), output.get()); @@ -312,11 +335,10 @@ struct BackwardKernel { m.inputTensor, input_grad.get())); } else if (use_scalar(op_type)) { - assert(scalar.has_value()); elewise_scalar_unary_backward_kernel> <<>>( num_elements, - static_cast>(scalar.value()), + static_cast>(require_scalar(op_type, scalar)), op_type, output.get(), output_grad.get(), diff --git a/lib/kernels/test/src/kernels/element_unary_kernels_gpu.cc b/lib/kernels/test/src/kernels/element_unary_kernels_gpu.cc index f7e21b7e1f..4c4ae5ca3d 100644 --- a/lib/kernels/test/src/kernels/element_unary_kernels_gpu.cc +++ b/lib/kernels/test/src/kernels/element_unary_kernels_gpu.cc @@ -1,40 +1,68 @@ +#include "kernels/element_unary_kernels_gpu.h" #include "internal/test_utils.h" #include "kernels/create_accessor_with_contents.h" -#include "kernels/element_unary_kernels_cpu.h" #include "kernels/format_accessor_contents.h" -#include "kernels/local_cpu_allocator.h" #include "op-attrs/ops/element_unary.h" #include "test/utils/doctest/check_kv.h" #include using namespace ::FlexFlow; +static ElementUnaryAttrs make_silu_attrs() { + return ElementUnaryAttrs{ + /*op_type=*/OperatorType::SILU, + /*scalar=*/std::nullopt, + }; +} + +static GenericTensorAccessorR make_input(Allocator &allocator) { + return create_2d_accessor_r_with_contents( + { + {3, -3, 6}, + {0, 1, 5}, + {1, 2, -2}, + {-8, 0.5, -3}, + }, + allocator); +} + +static GenericTensorAccessorR make_output_grad(Allocator &allocator) { + return create_2d_accessor_r_with_contents( + { + {1, 2, -1}, + {6, 4, 2}, + {0.5, 0.1, -2}, + {0, 0.5, 0}, + }, + allocator); +} + TEST_SUITE(FF_CUDA_TEST_SUITE) { TEST_CASE("element_unary_gpu_forward_kernel") { - SUBCASE("relu") { - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - ManagedFFStream managed_stream{}; + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; - Allocator allocator = create_local_cuda_memory_allocator(); + Allocator allocator = create_local_cuda_memory_allocator(); - ElementUnaryAttrs attrs = make_relu_attrs(); + GenericTensorAccessorR input = make_input(allocator); - GenericTensorAccessorR input = create_2d_accessor_r_with_contents( - { - {3, -3, 6}, - {0, 1, 5}, - {1, 2, -2}, - {-8, 0.5, -3}, - }, - allocator); + SUBCASE("relu") { + ElementUnaryAttrs attrs = make_relu_attrs(); - // Intentionally randomize this tensor so we can be confident we never read it + // Intentionally randomize this tensor so we can be confident we never + // read it GenericTensorAccessorW result = create_random_filled_accessor_w(input.shape, allocator); - element_unary_cpu_forward_kernel( + ElementUnaryPerDeviceState per_device_state = + element_unary_gpu_init_kernel(attrs, input.shape, result.shape); + + element_unary_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, /*attrs=*/attrs, /*input=*/input, /*output=*/result); @@ -52,52 +80,128 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { CHECK_MESSAGE(accessors_are_equal(result, correct), check_kv("result", format_accessor_w_contents(result))); } + + SUBCASE("silu") { + ElementUnaryAttrs attrs = make_silu_attrs(); + + // Intentionally randomize this tensor so we can be confident we never + // read it + GenericTensorAccessorW result = + create_random_filled_accessor_w(input.shape, allocator); + + ElementUnaryPerDeviceState per_device_state = + element_unary_gpu_init_kernel(attrs, input.shape, result.shape); + + element_unary_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*output=*/result); + + GenericTensorAccessorR correct = create_2d_accessor_r_with_contents< + float>( + { + {2.857722520828247, -0.14227761328220367, 5.985164642333984}, + {0.0, 0.7310585975646973, 4.966535568237305}, + {0.7310585975646973, 1.7615940570831299, -0.23840583860874176}, + {-0.00268280110321939, 0.3112296760082245, -0.14227761328220367}, + }, + allocator); + + // CUDA's expf is not bit-identical to PyTorch's CPU implementation + CHECK_MESSAGE(accessors_within_epsilon(result, correct, 1e-6), + check_kv("result", format_accessor_w_contents(result))); + } } TEST_CASE("element_unary_gpu_backward_kernel") { - SUBCASE("relu") { - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - ManagedFFStream managed_stream{}; + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; - Allocator allocator = create_local_cuda_memory_allocator(); + Allocator allocator = create_local_cuda_memory_allocator(); - ElementUnaryAttrs attrs = make_relu_attrs(); + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR output_grad = make_output_grad(allocator); - GenericTensorAccessorR input = create_2d_accessor_r_with_contents( + // input_grad is accumulated into, so it needs to start from a known value + auto make_input_grad = [&]() { + return create_2d_accessor_w_with_contents( { - {3, -3, 6}, - {0, 1, 5}, - {1, 2, -2}, - {-8, 0.5, -3}, + {10, 11, 12}, + {13, 14, 15}, + {16, 17, 18}, + {19, 20, 21}, }, allocator); + }; - // Intentionally randomize this tensor so we can be confident we never read it - GenericTensorAccessorW input_grad = create_random_filled_accessor_w( - get_tensor_shape_for_accessor_r(input), allocator); + SUBCASE("relu") { + ElementUnaryAttrs attrs = make_relu_attrs(); GenericTensorAccessorR output = create_2d_accessor_r_with_contents( { {3, 0, 6}, {0, 1, 5}, - {1, 2, -2}, + {1, 2, 0}, {0, 0.5, 0}, }, allocator); - GenericTensorAccessorR output_grad = + GenericTensorAccessorW input_grad = make_input_grad(); + + ElementUnaryPerDeviceState per_device_state = + element_unary_gpu_init_kernel(attrs, input.shape, output.shape); + + element_unary_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct_input_grad = create_2d_accessor_r_with_contents( { - {1, 2, -1}, - {6, 4, 2}, - {0.5, 0.1, -2}, - {0, 0.5, 0}, + {11, 11, 11}, + {13, 18, 17}, + {16.5, 17.1, 18}, + {19, 20.5, 21}, }, allocator); - element_unary_cpu_backward_kernel( + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct_input_grad), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + + SUBCASE("silu") { + ElementUnaryAttrs attrs = make_silu_attrs(); + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {2.857722520828247, -0.14227761328220367, 5.985164642333984}, + {0.0, 0.7310585975646973, 4.966535568237305}, + {0.7310585975646973, 1.7615940570831299, -0.23840583860874176}, + {-0.00268280110321939, 0.3112296760082245, -0.14227761328220367}, + }, + allocator); + + GenericTensorAccessorW input_grad = make_input_grad(); + + ElementUnaryPerDeviceState per_device_state = + element_unary_gpu_init_kernel(attrs, input.shape, output.shape); + + element_unary_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, /*attrs=*/attrs, /*output=*/output, /*output_grad=*/output_grad, @@ -107,15 +211,15 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { GenericTensorAccessorR correct_input_grad = create_2d_accessor_r_with_contents( { - {1.0f, 0.0f, -1.0f}, - {0.0f, 4.0f, 2.0f}, - {0.5f, 0.1f, 0.0f}, - {0.0f, 0.5f, 0.0f}, + {11.088104248046875, 10.82379150390625, 10.98767375946045}, + {16.0, 17.710681915283203, 17.0530948638916}, + {16.463834762573242, 17.109079360961914, 18.181568145751953}, + {19.0, 20.369979858398438, 21.0}, }, allocator); CHECK_MESSAGE( - accessors_are_equal(input_grad, correct_input_grad), + accessors_within_epsilon(input_grad, correct_input_grad, 1e-5), check_kv("input_grad", format_accessor_w_contents(input_grad))); } } From 8bdbab4fc32172647553921dac5a08a9d4a3e112 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 4 Sep 2026 11:08:28 -0700 Subject: [PATCH 3/6] Add BatchNormMode. --- .../src/cuda/ops/batch_norm_kernels.cu | 18 ++++++++++++++---- .../src/kernels/batch_norm_kernels_gpu.cc | 1 + lib/models/src/models/yolov10/yolov10.cc | 3 ++- .../op-attrs/ops/batch_norm_attrs.dtg.toml | 5 +++++ .../op-attrs/ops/batch_norm_mode.dtg.toml | 19 +++++++++++++++++++ .../op-attrs/computation_graph_op_attrs.cc | 1 + .../test/src/op-attrs/ops/batch_norm.cc | 4 ++++ .../include/pcg/computation_graph_builder.h | 1 + .../parallel_computation_graph_builder.h | 1 + lib/pcg/src/pcg/computation_graph_builder.cc | 2 ++ .../parallel_computation_graph_builder.cc | 2 ++ 11 files changed, 52 insertions(+), 5 deletions(-) create mode 100644 lib/op-attrs/include/op-attrs/ops/batch_norm_mode.dtg.toml diff --git a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu index 3e4514194c..2d14bc8ba3 100644 --- a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu +++ b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu @@ -62,10 +62,20 @@ BatchNormPerDeviceState checkCUDNN(cudnnCreateTensorDescriptor(&outputTensor)); checkCUDNN(cudnnCreateTensorDescriptor(&biasTensor)); - ffBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL; -#if CUDNN_VERSION >= 7000 - mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT; -#endif + ffBatchNormMode_t mode; + switch (attrs.mode) { + case BatchNormMode::PER_ACTIVATION: + mode = CUDNN_BATCHNORM_PER_ACTIVATION; + break; + case BatchNormMode::SPATIAL: + mode = CUDNN_BATCHNORM_SPATIAL; + break; + case BatchNormMode::SPATIAL_PERSISTENT: + mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT; + break; + default: + PANIC("Unknown BatchNormMode", attrs.mode); + } checkCUDNN(cudnnSetTensorDescriptorFromTensorShape(inputTensor, input_shape)); checkCUDNN( diff --git a/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc index ad92f0ea56..5d697ba9f8 100644 --- a/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc +++ b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc @@ -13,6 +13,7 @@ static BatchNormAttrs make_attrs() { /*affine=*/true, /*eps=*/1e-5, /*momentum=*/0.1, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }; } diff --git a/lib/models/src/models/yolov10/yolov10.cc b/lib/models/src/models/yolov10/yolov10.cc index 8d6b7dad34..99eccbd022 100644 --- a/lib/models/src/models/yolov10/yolov10.cc +++ b/lib/models/src/models/yolov10/yolov10.cc @@ -519,7 +519,8 @@ tensor_guid_t /*affine=*/true, /*activation=*/std::nullopt, /*eps=*/1e-3, - /*momentum=*/0.03); + /*momentum=*/0.03, + /*mode=*/BatchNormMode::SPATIAL); /** * https://github.com/ultralytics/ultralytics/blob/f8ad132a15b5f6818c2ce0647b40dc57e993bf0c/ultralytics/nn/modules/conv.py#L67 diff --git a/lib/op-attrs/include/op-attrs/ops/batch_norm_attrs.dtg.toml b/lib/op-attrs/include/op-attrs/ops/batch_norm_attrs.dtg.toml index a22fee4a64..0b89e6afec 100644 --- a/lib/op-attrs/include/op-attrs/ops/batch_norm_attrs.dtg.toml +++ b/lib/op-attrs/include/op-attrs/ops/batch_norm_attrs.dtg.toml @@ -13,6 +13,7 @@ features = [ includes = [ "", + "op-attrs/ops/batch_norm_mode.dtg.h", ] src_includes = [ @@ -36,3 +37,7 @@ type = "float" [[fields]] name = "momentum" type = "std::optional" + +[[fields]] +name = "mode" +type = "::FlexFlow::BatchNormMode" diff --git a/lib/op-attrs/include/op-attrs/ops/batch_norm_mode.dtg.toml b/lib/op-attrs/include/op-attrs/ops/batch_norm_mode.dtg.toml new file mode 100644 index 0000000000..bf2a920dde --- /dev/null +++ b/lib/op-attrs/include/op-attrs/ops/batch_norm_mode.dtg.toml @@ -0,0 +1,19 @@ +namespace = "FlexFlow" +name = "BatchNormMode" +type = "enum" + +features = [ + "hash", + "json", + "rapidcheck", + "fmt", +] + +[[values]] +name = "PER_ACTIVATION" + +[[values]] +name = "SPATIAL" + +[[values]] +name = "SPATIAL_PERSISTENT" diff --git a/lib/op-attrs/test/src/op-attrs/computation_graph_op_attrs.cc b/lib/op-attrs/test/src/op-attrs/computation_graph_op_attrs.cc index 84f1861f0b..2ac6578118 100644 --- a/lib/op-attrs/test/src/op-attrs/computation_graph_op_attrs.cc +++ b/lib/op-attrs/test/src/op-attrs/computation_graph_op_attrs.cc @@ -10,6 +10,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*affine=*/true, /*eps=*/1e-5, /*momentum=*/0.1, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }}; nlohmann::json j = correct; auto result = j.get(); diff --git a/lib/op-attrs/test/src/op-attrs/ops/batch_norm.cc b/lib/op-attrs/test/src/op-attrs/ops/batch_norm.cc index fe224c6674..20b7662c77 100644 --- a/lib/op-attrs/test/src/op-attrs/ops/batch_norm.cc +++ b/lib/op-attrs/test/src/op-attrs/ops/batch_norm.cc @@ -15,6 +15,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*affine=*/affine, /*eps=*/1.0, /*momentum=*/0.1, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }; }; @@ -63,6 +64,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*affine=*/true, /*eps=*/1.0, /*momentum=*/0.1, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }; BatchNormAttrs attrs_affine_false = [&] { @@ -143,6 +145,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*affine=*/true, /*eps=*/1.0, /*momentum=*/0.1, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }; BatchNormAttrs attrs_affine_false = [&] { @@ -347,6 +350,7 @@ TEST_SUITE(FF_TEST_SUITE) { /*affine=*/true, /*eps=*/1.0, /*momentum=*/0.1, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }; ParallelTensorShape input = ParallelTensorShape{ diff --git a/lib/pcg/include/pcg/computation_graph_builder.h b/lib/pcg/include/pcg/computation_graph_builder.h index 12e39e37fd..0c215b8c98 100644 --- a/lib/pcg/include/pcg/computation_graph_builder.h +++ b/lib/pcg/include/pcg/computation_graph_builder.h @@ -177,6 +177,7 @@ struct ComputationGraphBuilder { std::optional const &activation = std::nullopt, float eps = 1e-5f, std::optional const &momentum = 0.1f, + BatchNormMode mode = BatchNormMode::SPATIAL_PERSISTENT, std::optional const &name = std::nullopt); tensor_guid_t diff --git a/lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph_builder.h b/lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph_builder.h index 28ef8aa215..f9a3185f10 100644 --- a/lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph_builder.h +++ b/lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph_builder.h @@ -85,6 +85,7 @@ struct ParallelComputationGraphBuilder { std::optional const &activation, float eps, std::optional const &momentum, + BatchNormMode mode = BatchNormMode::SPATIAL_PERSISTENT, std::optional const &name = std::nullopt); parallel_tensor_guid_t diff --git a/lib/pcg/src/pcg/computation_graph_builder.cc b/lib/pcg/src/pcg/computation_graph_builder.cc index eb94f1f901..2227311ee7 100644 --- a/lib/pcg/src/pcg/computation_graph_builder.cc +++ b/lib/pcg/src/pcg/computation_graph_builder.cc @@ -720,6 +720,7 @@ tensor_guid_t ComputationGraphBuilder::batch_norm( std::optional const &activation, float eps, std::optional const &momentum, + BatchNormMode mode, std::optional const &maybe_name) { if (activation.has_value() && activation.value() != Activation::RELU) { @@ -736,6 +737,7 @@ tensor_guid_t ComputationGraphBuilder::batch_norm( /*affine=*/affine, /*eps=*/eps, /*momentum=*/momentum, + /*mode=*/mode, }; std::string name = diff --git a/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc b/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc index 0de4bca3dd..b531486399 100644 --- a/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc +++ b/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc @@ -342,6 +342,7 @@ parallel_tensor_guid_t ParallelComputationGraphBuilder::batch_norm( std::optional const &activation, float eps, std::optional const &momentum, + BatchNormMode mode, std::optional const &maybe_name) { if (activation.has_value() && activation.value() != Activation::RELU) { @@ -358,6 +359,7 @@ parallel_tensor_guid_t ParallelComputationGraphBuilder::batch_norm( /*affine=*/affine, /*eps=*/eps, /*momentum=*/momentum, + /*mode=*/mode, }; std::string name = From e04b17905872c0e453058c5d0cac0c784356b0e6 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 4 Sep 2026 17:40:20 -0700 Subject: [PATCH 4/6] Fix silu kernel backward pass. --- lib/kernels/src/cuda/ops/element_unary_kernels.cu | 14 ++++++++++---- .../src/kernels/element_unary_kernels_cpu.cc | 10 +++++++--- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/lib/kernels/src/cuda/ops/element_unary_kernels.cu b/lib/kernels/src/cuda/ops/element_unary_kernels.cu index ecbb846797..b540b5be3a 100644 --- a/lib/kernels/src/cuda/ops/element_unary_kernels.cu +++ b/lib/kernels/src/cuda/ops/element_unary_kernels.cu @@ -209,10 +209,16 @@ __global__ void elewise_scalar_unary_backward_kernel(coord_t volume, break; } case OperatorType::SILU: { - float e_to_bx = expf(scalar * input[i]); - input_grad[i] += (T)(output_grad[i] * - (e_to_bx * (scalar * input[i] + e_to_bx + 1.0f)) / - ((e_to_bx + 1.0f) * (e_to_bx + 1.0f))); + // Algebraically e^bx (bx + e^bx + 1) / (e^bx + 1)^2, but written in + // terms of the sigmoid so that it stays finite. exp(bx) overflows to + // infinity once bx is much above 88, and the expanded form then + // evaluates infinity/infinity = NaN. A zero output gradient does not + // rescue that, because 0 * NaN is NaN: a layer receiving no gradient + // at all would still poison everything upstream of it. + float bx = scalar * (float)input[i]; + float sigmoid = 1.0f / (1.0f + expf(-bx)); + input_grad[i] += + (T)(output_grad[i] * sigmoid * (1.0f + bx * (1.0f - sigmoid))); break; } default: diff --git a/lib/kernels/src/kernels/element_unary_kernels_cpu.cc b/lib/kernels/src/kernels/element_unary_kernels_cpu.cc index 183cdfd4cf..c14b472c08 100644 --- a/lib/kernels/src/kernels/element_unary_kernels_cpu.cc +++ b/lib/kernels/src/kernels/element_unary_kernels_cpu.cc @@ -114,9 +114,13 @@ void element_unary_cpu_backward_kernel( case OperatorType::SILU: { float beta = attrs.scalar.value_or(1.0f); return [beta](float x, float fx) -> float { - float e_to_bx = expf(beta * x); - return (e_to_bx * (beta * x + e_to_bx + 1.0f)) / - ((e_to_bx + 1) * (e_to_bx + 1)); + // Algebraically e^bx (bx + e^bx + 1) / (e^bx + 1)^2, but written in + // terms of the sigmoid so that it stays finite. exp(bx) overflows to + // infinity once bx is much above 88, and the expanded form then + // evaluates infinity/infinity = NaN. + float bx = beta * x; + float sigmoid = 1.0f / (1.0f + expf(-bx)); + return sigmoid * (1.0f + bx * (1.0f - sigmoid)); }; } case OperatorType::SIN: From 0715d15112b71700079b1bd1e96f20ce86105ad6 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Sat, 5 Sep 2026 18:12:36 -0700 Subject: [PATCH 5/6] Compute a batch norm and the activation after it in one kernel. A batch norm followed by an elementwise activation writes its output only for the activation to read it straight back and write it again, and because the backward pass then needs that value, it is read a third and fourth time there. On YOLOv10x those values come to 3145 MiB per iteration. Replaces BatchNormAttrs::relu, which was never implemented -- the forward kernel silently ignored it while the backward kernel applied it, so the kernels asserted it off -- with an optional Activation that the kernels do apply. The builders already took an activation argument to set the old flag from, so they now pass it through instead, and no longer refuse anything but relu. The forward kernel normalises and activates in registers, so the value between the two is never written down. The backward kernel then reconstructs it from the input and the saved statistics, at the price of a multiply-add rather than a pass over memory, which is what lets the forward kernel get away with not writing it. Measured across every batch-norm site in YOLOv10x: forward 13.70 -> 9.38 ms, backward 22.79 -> 13.80 ms. One block per channel, 512 threads. 512 beat 256 (by 32% forward, 21% backward) and 1024 (8%, 5%): below it there are not enough threads in flight to hide the memory latency, above it the block-wide reductions cost more than the occupancy buys. Splitting each channel across several blocks, with the same memory traffic but full occupancy, was tried and came out 14% slower. Statistics come out of Welford's method rather than E[x^2] - E[x]^2, whose cancellation loses most of the significant digits when the mean is large next to the standard deviation, and can return a negative variance and hence a NaN inverse standard deviation. Reconstructing the activation's input needs the shift, so beta joins the backward kernel's signature. The new test checks the fused kernels against the two operators they replace rather than against a table of numbers, since what matters is that folding them together did not change the answer. Which activations a batch norm can apply, and in which modes, lives in op-attrs so that the pass deciding to write one into the attrs and the kernel asked to run it cannot disagree. Sliced onto master without the change to have backward kernels overwrite their gradients, so the fused backward kernel accumulates into its gradients like the cuDNN path does, and its test starts every gradient at zero. Co-Authored-By: Claude Opus 5 Co-Authored-By: Claude Opus 5.5 --- .../include/kernels/batch_norm_kernels.h | 1 + .../include/kernels/batch_norm_kernels_cpu.h | 1 + .../include/kernels/batch_norm_kernels_gpu.h | 1 + .../src/cuda/ops/batch_norm_kernels.cu | 408 +++++++++++++++++- lib/kernels/src/kernels/batch_norm_kernels.cc | 3 + .../src/kernels/batch_norm_kernels_cpu.cc | 1 + .../src/kernels/batch_norm_kernels_gpu.cc | 207 ++++++++- .../include/op-attrs/ops/batch_norm.h | 21 + .../op-attrs/ops/batch_norm_attrs.dtg.toml | 10 +- lib/op-attrs/src/op-attrs/ops/batch_norm.cc | 22 + .../op-attrs/computation_graph_op_attrs.cc | 2 +- .../test/src/op-attrs/ops/batch_norm.cc | 8 +- lib/pcg/src/pcg/computation_graph_builder.cc | 11 +- .../parallel_computation_graph_builder.cc | 11 +- .../src/task-spec/ops/impl/batch_norm.cc | 3 + 15 files changed, 676 insertions(+), 34 deletions(-) diff --git a/lib/kernels/include/kernels/batch_norm_kernels.h b/lib/kernels/include/kernels/batch_norm_kernels.h index ee56bf7173..709a1d92e1 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels.h +++ b/lib/kernels/include/kernels/batch_norm_kernels.h @@ -39,6 +39,7 @@ void batch_norm_backward_kernel( GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad, GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, GenericTensorAccessorW const &gamma_grad, GenericTensorAccessorW const &beta_grad); diff --git a/lib/kernels/include/kernels/batch_norm_kernels_cpu.h b/lib/kernels/include/kernels/batch_norm_kernels_cpu.h index ecda6883a0..731756d5ff 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels_cpu.h +++ b/lib/kernels/include/kernels/batch_norm_kernels_cpu.h @@ -18,6 +18,7 @@ void batch_norm_cpu_backward_kernel(BatchNormAttrs const &attrs, GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad, GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, GenericTensorAccessorW const &gamma_grad, GenericTensorAccessorW const &beta_grad); diff --git a/lib/kernels/include/kernels/batch_norm_kernels_gpu.h b/lib/kernels/include/kernels/batch_norm_kernels_gpu.h index 277c36a9ed..ec6e0e5ec2 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels_gpu.h +++ b/lib/kernels/include/kernels/batch_norm_kernels_gpu.h @@ -35,6 +35,7 @@ void batch_norm_gpu_backward_kernel( GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad, GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, GenericTensorAccessorW const &gamma_grad, GenericTensorAccessorW const &beta_grad); diff --git a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu index 2d14bc8ba3..7d0d7c8d9c 100644 --- a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu +++ b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu @@ -15,6 +15,7 @@ #include "internal/device.h" #include "kernels/batch_norm_kernels_gpu.h" +#include "op-attrs/ops/batch_norm.h" #include "op-attrs/tensor_dims.h" #include "utils/containers/require_same.h" #include @@ -25,16 +26,367 @@ static positive_int get_num_channels(TensorShape const &shape) { return dim_at_idx(shape.dims, ff_dim_t{1_n}); } +namespace { + +/** + * @brief An NCHW shape as the fused kernels index it: one block per channel, + * walking the samples of that channel in turn. + */ +struct BatchNormExtents { + int num_samples; + int num_channels; + int spatial_size; +}; + +BatchNormExtents get_extents(TensorShape const &shape) { + return BatchNormExtents{ + /*num_samples=*/dim_at_idx(shape.dims, ff_dim_t{0_n}) + .int_from_positive_int(), + /*num_channels=*/ + dim_at_idx(shape.dims, ff_dim_t{1_n}).int_from_positive_int(), + /*spatial_size=*/ + (dim_at_idx(shape.dims, ff_dim_t{2_n}) * + dim_at_idx(shape.dims, ff_dim_t{3_n})) + .int_from_positive_int(), + }; +} + +} // namespace + +namespace { + +/** + * @brief The number of threads in each of the fused kernels' blocks. + * + * @details Each block reduces one channel, so this is also how much parallelism + * a channel gets. Measured across every batch-norm site in YOLOv10x, 512 beat + * both 256 (+32% forward, +21% backward) and 1024 (+8%, +5%): below it the GPU + * runs out of threads to hide memory latency with, above it the block-wide + * reductions start to cost more than the occupancy buys. + */ +constexpr int BATCH_NORM_FUSED_THREADS = 512; + +__device__ __forceinline__ float apply_activation(Activation activation, + float y) { + switch (activation) { + case Activation::RELU: + return y > 0.0f ? y : 0.0f; + case Activation::SIGMOID: + return 1.0f / (1.0f + expf(-y)); + case Activation::TANH: + return tanhf(y); + case Activation::SILU: + return y / (1.0f + expf(-y)); + default: + return y; + } +} + +/** + * @brief The gradient with respect to an activation's input, given \p dz, the + * gradient with respect to its output, and \p y, its input. + */ +__device__ __forceinline__ float + activation_backward(Activation activation, float dz, float y) { + switch (activation) { + case Activation::RELU: + return y > 0.0f ? dz : 0.0f; + case Activation::SIGMOID: { + float s = 1.0f / (1.0f + expf(-y)); + return dz * s * (1.0f - s); + } + case Activation::TANH: { + float t = tanhf(y); + return dz * (1.0f - t * t); + } + case Activation::SILU: { + // Written in terms of the sigmoid rather than the algebraically equal + // e^y (y + e^y + 1) / (e^y + 1)^2, which evaluates to NaN once y is much + // above 88 and e^y overflows. See element_unary_kernels.cu, which this + // has to agree with. + float s = 1.0f / (1.0f + expf(-y)); + return dz * s * (1.0f + y * (1.0f - s)); + } + default: + return dz; + } +} + +__device__ __forceinline__ void block_reduce_sum2(float &a, float &b) { + __shared__ float shared_a[BATCH_NORM_FUSED_THREADS / 32]; + __shared__ float shared_b[BATCH_NORM_FUSED_THREADS / 32]; + int lane = threadIdx.x % 32; + int warp = threadIdx.x / 32; + + for (int offset = 16; offset > 0; offset >>= 1) { + a += __shfl_down_sync(0xffffffff, a, offset); + b += __shfl_down_sync(0xffffffff, b, offset); + } + if (lane == 0) { + shared_a[warp] = a; + shared_b[warp] = b; + } + __syncthreads(); + + if (warp == 0) { + bool in_range = threadIdx.x < BATCH_NORM_FUSED_THREADS / 32; + a = in_range ? shared_a[threadIdx.x] : 0.0f; + b = in_range ? shared_b[threadIdx.x] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) { + a += __shfl_down_sync(0xffffffff, a, offset); + b += __shfl_down_sync(0xffffffff, b, offset); + } + if (threadIdx.x == 0) { + shared_a[0] = a; + shared_b[0] = b; + } + } + __syncthreads(); + a = shared_a[0]; + b = shared_b[0]; +} + +/** + * @brief Fold one set of Welford statistics into another. + * + * @details Welford's method is used in place of the textbook + * E[x^2] - E[x]^2 because that form subtracts two numbers that are + * nearly equal whenever the mean is large next to the standard deviation, + * which loses most of the significant digits and can even come out negative, + * making the inverse standard deviation NaN. + */ +__device__ __forceinline__ void welford_join( + float &count, float &mean, float &m2, float n, float m, float s) { + float total = count + n; + if (total == 0.0f) { + return; + } + float delta = m - mean; + mean += delta * (n / total); + m2 += s + delta * delta * (count * n / total); + count = total; +} + +__device__ __forceinline__ void + block_reduce_welford(float &count, float &mean, float &m2) { + __shared__ float shared_count[BATCH_NORM_FUSED_THREADS / 32]; + __shared__ float shared_mean[BATCH_NORM_FUSED_THREADS / 32]; + __shared__ float shared_m2[BATCH_NORM_FUSED_THREADS / 32]; + int lane = threadIdx.x % 32; + int warp = threadIdx.x / 32; + + for (int offset = 16; offset > 0; offset >>= 1) { + welford_join(count, + mean, + m2, + __shfl_down_sync(0xffffffff, count, offset), + __shfl_down_sync(0xffffffff, mean, offset), + __shfl_down_sync(0xffffffff, m2, offset)); + } + if (lane == 0) { + shared_count[warp] = count; + shared_mean[warp] = mean; + shared_m2[warp] = m2; + } + __syncthreads(); + + if (warp == 0) { + bool in_range = threadIdx.x < BATCH_NORM_FUSED_THREADS / 32; + count = in_range ? shared_count[threadIdx.x] : 0.0f; + mean = in_range ? shared_mean[threadIdx.x] : 0.0f; + m2 = in_range ? shared_m2[threadIdx.x] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) { + welford_join(count, + mean, + m2, + __shfl_down_sync(0xffffffff, count, offset), + __shfl_down_sync(0xffffffff, mean, offset), + __shfl_down_sync(0xffffffff, m2, offset)); + } + if (threadIdx.x == 0) { + shared_count[0] = count; + shared_mean[0] = mean; + shared_m2[0] = m2; + } + } + __syncthreads(); + count = shared_count[0]; + mean = shared_mean[0]; + m2 = shared_m2[0]; +} + +/** + * @brief Normalize and activate in one pass over the input, one block per + * channel. + * + * @details Saves the two passes over the output that a separate activation + * operator would cost: nothing ever writes or reads the normalized value, which + * exists only in registers. \ref batch_norm_fused_backward_kernel reconstructs + * it from \p input rather than being handed it, so it is not needed later + * either. + */ +template +__global__ void + batch_norm_fused_forward_kernel(int num_samples, + int num_channels, + int spatial_size, + float eps, + float exponential_average_factor, + float const *input, + float const *gamma, + float const *beta, + float *output, + float *save_mean, + float *save_invstd, + float *running_mean, + float *running_var) { + int channel = blockIdx.x; + + float count = 0.0f, mean = 0.0f, m2 = 0.0f; + for (int sample = 0; sample < num_samples; sample++) { + size_t base = ((size_t)sample * num_channels + channel) * spatial_size; + for (int i = threadIdx.x; i < spatial_size; i += BATCH_NORM_FUSED_THREADS) { + welford_join(count, mean, m2, 1.0f, input[base + i], 0.0f); + } + } + block_reduce_welford(count, mean, m2); + + float variance = m2 / count; + float invstd = rsqrtf(variance + eps); + + if (threadIdx.x == 0) { + save_mean[channel] = mean; + save_invstd[channel] = invstd; + + // Matching cuDNN, which keeps the running statistics in the units of the + // data and uses the unbiased variance for them even though the biased one + // is what normalizes this batch. + float unbiased_variance = count > 1.0f ? m2 / (count - 1.0f) : variance; + running_mean[channel] = + (1.0f - exponential_average_factor) * running_mean[channel] + + exponential_average_factor * mean; + running_var[channel] = + (1.0f - exponential_average_factor) * running_var[channel] + + exponential_average_factor * unbiased_variance; + } + + float scale = gamma[channel]; + float shift = beta[channel]; + for (int sample = 0; sample < num_samples; sample++) { + size_t base = ((size_t)sample * num_channels + channel) * spatial_size; + for (int i = threadIdx.x; i < spatial_size; i += BATCH_NORM_FUSED_THREADS) { + float normalized = (input[base + i] - mean) * invstd; + output[base + i] = + apply_activation(ACTIVATION, scale * normalized + shift); + } + } +} + +/** + * @brief Differentiate the activation and the normalization together, one block + * per channel. + * + * @details The value between the two, which a separate activation operator + * would have read from memory, is instead reconstructed from \p input and the + * statistics the forward pass saved. That trades one multiply-add for a whole + * pass over the tensor, and is what lets the forward pass get away with never + * writing it down. + * + * Like the cuDNN path, accumulates into the gradients rather than overwriting + * them. + */ +template +__global__ void batch_norm_fused_backward_kernel(int num_samples, + int num_channels, + int spatial_size, + float const *output_grad, + float const *input, + float const *gamma, + float const *beta, + float const *save_mean, + float const *save_invstd, + float *input_grad, + float *gamma_grad, + float *beta_grad) { + int channel = blockIdx.x; + float mean = save_mean[channel]; + float invstd = save_invstd[channel]; + float scale = gamma[channel]; + float shift = beta[channel]; + + float sum_grad = 0.0f, sum_grad_normalized = 0.0f; + for (int sample = 0; sample < num_samples; sample++) { + size_t base = ((size_t)sample * num_channels + channel) * spatial_size; + for (int i = threadIdx.x; i < spatial_size; i += BATCH_NORM_FUSED_THREADS) { + float normalized = (input[base + i] - mean) * invstd; + float grad = activation_backward( + ACTIVATION, output_grad[base + i], scale * normalized + shift); + sum_grad += grad; + sum_grad_normalized += grad * normalized; + } + } + block_reduce_sum2(sum_grad, sum_grad_normalized); + + if (threadIdx.x == 0) { + gamma_grad[channel] += sum_grad_normalized; + beta_grad[channel] += sum_grad; + } + + float count = (float)num_samples * spatial_size; + for (int sample = 0; sample < num_samples; sample++) { + size_t base = ((size_t)sample * num_channels + channel) * spatial_size; + for (int i = threadIdx.x; i < spatial_size; i += BATCH_NORM_FUSED_THREADS) { + float normalized = (input[base + i] - mean) * invstd; + float grad = activation_backward( + ACTIVATION, output_grad[base + i], scale * normalized + shift); + input_grad[base + i] += + scale * invstd * + (grad - (sum_grad + normalized * sum_grad_normalized) / count); + } + } +} + +/** + * @brief Call \p f with the activation as a template argument. + * + * @details The activation is fixed for the life of an operator, so resolving it + * once here keeps the branch out of the innermost loop. + */ +template +void dispatch_on_activation(Activation activation, F const &f) { + switch (activation) { + case Activation::RELU: + f(std::integral_constant{}); + return; + case Activation::SIGMOID: + f(std::integral_constant{}); + return; + case Activation::TANH: + f(std::integral_constant{}); + return; + case Activation::SILU: + f(std::integral_constant{}); + return; + default: + PANIC("BatchNorm cannot compute this activation itself", activation); + } +} + +} // namespace + BatchNormPerDeviceState batch_norm_gpu_init_kernel(Allocator &allocator, BatchNormAttrs const &attrs, TensorShape const &input_shape, TensorShape const &output_shape) { - // Applying an activation as part of BatchNorm is not currently implemented - // (the forward kernel used to silently ignore it while the backward kernel - // applied it). If you need it, please create an issue. - ASSERT(!attrs.relu, - "BatchNorm does not currently support a fused relu activation"); + if (attrs.activation.has_value()) { + ASSERT(batch_norm_supports_fused_activation(attrs.activation.value()), + "BatchNorm cannot compute this activation itself", + attrs.activation.value()); + ASSERT(batch_norm_mode_supports_fused_activation(attrs.mode), + "BatchNorm cannot compute a fused activation in this mode", + attrs.mode); + } ASSERT(attrs.affine, "BatchNorm currently only supports attrs.affine = true. " @@ -138,6 +490,29 @@ void batch_norm_gpu_forward_kernel( // never read back, so this only matters once inference mode is supported. double exponential_average_factor = attrs.momentum.value_or(1.0); + if (attrs.activation.has_value()) { + BatchNormExtents extents = get_extents(input.shape); + + dispatch_on_activation(attrs.activation.value(), [&](auto activation) { + batch_norm_fused_forward_kernel + <<>>( + extents.num_samples, + extents.num_channels, + extents.spatial_size, + attrs.eps, + static_cast(exponential_average_factor), + input.get_float_ptr(), + gamma.get_float_ptr(), + beta.get_float_ptr(), + output.get_float_ptr(), + per_device_state.saveMean, + per_device_state.saveVar, + per_device_state.runningMean, + per_device_state.runningVar); + }); + return; + } + float alpha = 1.0f, beta_coeff = 0.0f; checkCUDNN( cudnnBatchNormalizationForwardTraining(handle.dnn, @@ -169,10 +544,33 @@ void batch_norm_gpu_backward_kernel( GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad, GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, GenericTensorAccessorW const &gamma_grad, GenericTensorAccessorW const &beta_grad) { checkCUDNN(cudnnSetStream(handle.dnn, stream)); + if (attrs.activation.has_value()) { + BatchNormExtents extents = get_extents(input.shape); + + dispatch_on_activation(attrs.activation.value(), [&](auto activation) { + batch_norm_fused_backward_kernel + <<>>( + extents.num_samples, + extents.num_channels, + extents.spatial_size, + output_grad.get_float_ptr(), + input.get_float_ptr(), + gamma.get_float_ptr(), + beta.get_float_ptr(), + per_device_state.saveMean, + per_device_state.saveVar, + input_grad.get_float_ptr(), + gamma_grad.get_float_ptr(), + beta_grad.get_float_ptr()); + }); + return; + } + // NOTE: the beta coefficients are 1.0 so that the gradients are accumulated // into rather than overwritten float alpha_data = 1.0f, beta_data = 1.0f; diff --git a/lib/kernels/src/kernels/batch_norm_kernels.cc b/lib/kernels/src/kernels/batch_norm_kernels.cc index 546aa1ce5b..1a40f3ae25 100644 --- a/lib/kernels/src/kernels/batch_norm_kernels.cc +++ b/lib/kernels/src/kernels/batch_norm_kernels.cc @@ -65,6 +65,7 @@ void batch_norm_backward_kernel( GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad, GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, GenericTensorAccessorW const &gamma_grad, GenericTensorAccessorW const &beta_grad) { if (stream.is_gpu()) { @@ -78,6 +79,7 @@ void batch_norm_backward_kernel( /*input=*/input, /*input_grad=*/input_grad, /*gamma=*/gamma, + /*beta=*/beta, /*gamma_grad=*/gamma_grad, /*beta_grad=*/beta_grad); } else { @@ -91,6 +93,7 @@ void batch_norm_backward_kernel( /*input=*/input, /*input_grad=*/input_grad, /*gamma=*/gamma, + /*beta=*/beta, /*gamma_grad=*/gamma_grad, /*beta_grad=*/beta_grad); } diff --git a/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc b/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc index 5e16a11fe7..d4071556a0 100644 --- a/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc +++ b/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc @@ -17,6 +17,7 @@ void batch_norm_cpu_backward_kernel(BatchNormAttrs const &attrs, GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad, GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, GenericTensorAccessorW const &gamma_grad, GenericTensorAccessorW const &beta_grad) { NOT_IMPLEMENTED(); diff --git a/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc index 5d697ba9f8..f863f19d94 100644 --- a/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc +++ b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc @@ -1,15 +1,17 @@ #include "kernels/batch_norm_kernels_gpu.h" #include "internal/test_utils.h" #include "kernels/create_accessor_with_contents.h" +#include "kernels/element_unary_kernels_gpu.h" #include "kernels/format_accessor_contents.h" #include "test/utils/doctest/check_kv.h" #include using namespace ::FlexFlow; -static BatchNormAttrs make_attrs() { +static BatchNormAttrs + make_attrs(std::optional activation = std::nullopt) { return BatchNormAttrs{ - /*relu=*/false, + /*activation=*/activation, /*affine=*/true, /*eps=*/1e-5, /*momentum=*/0.1, @@ -181,6 +183,7 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { /*input=*/input, /*input_grad=*/input_grad, /*gamma=*/gamma, + /*beta=*/beta, /*gamma_grad=*/gamma_grad, /*beta_grad=*/beta_grad); @@ -223,4 +226,204 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { batch_norm_gpu_cleanup_kernel(allocator, per_device_state); } + + // Applying an activation as part of batch norm has to compute what the two + // separate operators it replaces would have computed. Comparing against those + // operators rather than against a table of expected numbers is the point: + // what matters is that folding them together did not change the answer. + TEST_CASE("batch_norm_gpu fused activation matches the separate operators") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + Allocator allocator = create_local_cuda_memory_allocator(); + + ElementUnaryAttrs silu_attrs = ElementUnaryAttrs{ + /*op_type=*/OperatorType::SILU, + /*scalar=*/std::nullopt, + }; + + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR gamma = make_gamma(allocator); + GenericTensorAccessorR beta = make_beta(allocator); + + ElementUnaryPerDeviceState silu_state = + element_unary_gpu_init_kernel(silu_attrs, input.shape, input.shape); + + SUBCASE("forward") { + BatchNormAttrs separate_attrs = make_attrs(); + BatchNormAttrs fused_attrs = make_attrs(Activation::SILU); + + GenericTensorAccessorW normalized = + create_random_filled_accessor_w(input.shape, allocator); + GenericTensorAccessorW separate = + create_random_filled_accessor_w(input.shape, allocator); + GenericTensorAccessorW fused = + create_random_filled_accessor_w(input.shape, allocator); + + BatchNormPerDeviceState separate_state = + batch_norm_gpu_init_kernel(allocator, + separate_attrs, + input.shape, + input.shape); + batch_norm_gpu_forward_kernel(managed_stream.raw_stream(), + managed_handle.raw_handle(), + separate_state, + separate_attrs, + input, + gamma, + beta, + normalized); + element_unary_gpu_forward_kernel( + managed_stream.raw_stream(), + managed_handle.raw_handle(), + silu_state, + silu_attrs, + read_only_accessor_from_write_accessor(normalized), + separate); + + BatchNormPerDeviceState fused_state = + batch_norm_gpu_init_kernel(allocator, + fused_attrs, + input.shape, + input.shape); + batch_norm_gpu_forward_kernel(managed_stream.raw_stream(), + managed_handle.raw_handle(), + fused_state, + fused_attrs, + input, + gamma, + beta, + fused); + + CHECK_MESSAGE(accessors_within_epsilon(fused, separate, 1e-5), + check_kv("fused", format_accessor_w_contents(fused)), + check_kv("separate", format_accessor_w_contents(separate))); + + batch_norm_gpu_cleanup_kernel(allocator, separate_state); + batch_norm_gpu_cleanup_kernel(allocator, fused_state); + } + + SUBCASE("backward") { + BatchNormAttrs separate_attrs = make_attrs(); + BatchNormAttrs fused_attrs = make_attrs(Activation::SILU); + + GenericTensorAccessorR output_grad = + create_random_filled_accessor_r(input.shape, allocator); + + // The backward kernels accumulate into their gradients, so every + // gradient starts at zero for the two paths to be comparable. + + // The separate path: batch norm writes the value the activation reads, + // and the activation's backward turns the output gradient into the + // gradient of that value. + GenericTensorAccessorW normalized = + create_random_filled_accessor_w(input.shape, allocator); + GenericTensorAccessorW activated = + create_random_filled_accessor_w(input.shape, allocator); + GenericTensorAccessorW normalized_grad = + create_zero_filled_accessor_w(input.shape, allocator); + GenericTensorAccessorW separate_input_grad = + create_zero_filled_accessor_w(input.shape, allocator); + GenericTensorAccessorW separate_gamma_grad = + create_zero_filled_accessor_w(gamma.shape, allocator); + GenericTensorAccessorW separate_beta_grad = + create_zero_filled_accessor_w(gamma.shape, allocator); + + BatchNormPerDeviceState separate_state = + batch_norm_gpu_init_kernel(allocator, + separate_attrs, + input.shape, + input.shape); + batch_norm_gpu_forward_kernel(managed_stream.raw_stream(), + managed_handle.raw_handle(), + separate_state, + separate_attrs, + input, + gamma, + beta, + normalized); + element_unary_gpu_forward_kernel( + managed_stream.raw_stream(), + managed_handle.raw_handle(), + silu_state, + silu_attrs, + read_only_accessor_from_write_accessor(normalized), + activated); + element_unary_gpu_backward_kernel( + managed_stream.raw_stream(), + managed_handle.raw_handle(), + silu_state, + silu_attrs, + read_only_accessor_from_write_accessor(activated), + output_grad, + read_only_accessor_from_write_accessor(normalized), + normalized_grad); + batch_norm_gpu_backward_kernel( + managed_stream.raw_stream(), + managed_handle.raw_handle(), + separate_state, + separate_attrs, + read_only_accessor_from_write_accessor(normalized), + read_only_accessor_from_write_accessor(normalized_grad), + input, + separate_input_grad, + gamma, + beta, + separate_gamma_grad, + separate_beta_grad); + + // The fused path, which never writes the value between them. + GenericTensorAccessorW fused_output = + create_random_filled_accessor_w(input.shape, allocator); + GenericTensorAccessorW fused_input_grad = + create_zero_filled_accessor_w(input.shape, allocator); + GenericTensorAccessorW fused_gamma_grad = + create_zero_filled_accessor_w(gamma.shape, allocator); + GenericTensorAccessorW fused_beta_grad = + create_zero_filled_accessor_w(gamma.shape, allocator); + + BatchNormPerDeviceState fused_state = + batch_norm_gpu_init_kernel(allocator, + fused_attrs, + input.shape, + input.shape); + batch_norm_gpu_forward_kernel(managed_stream.raw_stream(), + managed_handle.raw_handle(), + fused_state, + fused_attrs, + input, + gamma, + beta, + fused_output); + batch_norm_gpu_backward_kernel( + managed_stream.raw_stream(), + managed_handle.raw_handle(), + fused_state, + fused_attrs, + read_only_accessor_from_write_accessor(fused_output), + output_grad, + input, + fused_input_grad, + gamma, + beta, + fused_gamma_grad, + fused_beta_grad); + + CHECK_MESSAGE( + accessors_within_epsilon(fused_input_grad, separate_input_grad, 1e-5), + check_kv("fused", format_accessor_w_contents(fused_input_grad)), + check_kv("separate", + format_accessor_w_contents(separate_input_grad))); + CHECK_MESSAGE( + accessors_within_epsilon(fused_gamma_grad, separate_gamma_grad, 1e-5), + check_kv("fused", format_accessor_w_contents(fused_gamma_grad))); + CHECK_MESSAGE( + accessors_within_epsilon(fused_beta_grad, separate_beta_grad, 1e-5), + check_kv("fused", format_accessor_w_contents(fused_beta_grad))); + + batch_norm_gpu_cleanup_kernel(allocator, separate_state); + batch_norm_gpu_cleanup_kernel(allocator, fused_state); + } + } } diff --git a/lib/op-attrs/include/op-attrs/ops/batch_norm.h b/lib/op-attrs/include/op-attrs/ops/batch_norm.h index 50dcd4bb5c..ae34f37544 100644 --- a/lib/op-attrs/include/op-attrs/ops/batch_norm.h +++ b/lib/op-attrs/include/op-attrs/ops/batch_norm.h @@ -1,6 +1,7 @@ #ifndef _FLEXFLOW_OP_ATTRS_INCLUDE_OP_ATTRS_OPS_BATCH_NORM_H #define _FLEXFLOW_OP_ATTRS_INCLUDE_OP_ATTRS_OPS_BATCH_NORM_H +#include "op-attrs/activation.dtg.h" #include "op-attrs/incoming_tensor_role.dtg.h" #include "op-attrs/initializer_attrs.dtg.h" #include "op-attrs/ops/batch_norm_attrs.dtg.h" @@ -12,6 +13,26 @@ namespace FlexFlow { +/** + * \brief Whether a batch-norm kernel can compute \p activation itself, as part + * of normalizing, rather than leaving it to a separate operator. + * + * \details Both sides of that arrangement ask here: whatever writes an + * activation into a \ref BatchNormAttrs, to decide whether the result would be + * runnable, and the kernels, to check that whoever built the attrs did ask. + */ +bool batch_norm_supports_fused_activation(Activation); + +/** + * \brief Whether a batch-norm kernel can compute a fused activation in \p mode. + * + * \details The spatial modes take statistics per channel over the samples and + * both spatial dimensions, which is what the fused kernels compute. + * PER_ACTIVATION takes them per position over the samples alone, and would need + * a kernel of its own. + */ +bool batch_norm_mode_supports_fused_activation(BatchNormMode); + std::map get_batch_norm_incoming_tensor_roles(BatchNormAttrs const &); diff --git a/lib/op-attrs/include/op-attrs/ops/batch_norm_attrs.dtg.toml b/lib/op-attrs/include/op-attrs/ops/batch_norm_attrs.dtg.toml index 0b89e6afec..212a11f1a5 100644 --- a/lib/op-attrs/include/op-attrs/ops/batch_norm_attrs.dtg.toml +++ b/lib/op-attrs/include/op-attrs/ops/batch_norm_attrs.dtg.toml @@ -13,6 +13,7 @@ features = [ includes = [ "", + "op-attrs/activation.dtg.h", "op-attrs/ops/batch_norm_mode.dtg.h", ] @@ -23,8 +24,13 @@ src_includes = [ ] [[fields]] -name = "relu" -type = "bool" +name = "activation" +type = "std::optional<::FlexFlow::Activation>" +docstring = ''' +\brief An elementwise activation applied to this operator's output, computed by the same kernel rather than by a separate operator. + +Either written by the model author, or introduced by folding an adjacent elementwise operator into this one, so that the value between them is never materialized. +''' [[fields]] name = "affine" diff --git a/lib/op-attrs/src/op-attrs/ops/batch_norm.cc b/lib/op-attrs/src/op-attrs/ops/batch_norm.cc index 95a12146ea..c6953eb64e 100644 --- a/lib/op-attrs/src/op-attrs/ops/batch_norm.cc +++ b/lib/op-attrs/src/op-attrs/ops/batch_norm.cc @@ -10,6 +10,28 @@ namespace FlexFlow { +bool batch_norm_supports_fused_activation(Activation activation) { + switch (activation) { + case Activation::RELU: + case Activation::SIGMOID: + case Activation::TANH: + case Activation::SILU: + return true; + default: + return false; + } +} + +bool batch_norm_mode_supports_fused_activation(BatchNormMode mode) { + switch (mode) { + case BatchNormMode::SPATIAL: + case BatchNormMode::SPATIAL_PERSISTENT: + return true; + default: + return false; + } +} + std::map get_batch_norm_incoming_tensor_roles(BatchNormAttrs const &attrs) { std::map result = { diff --git a/lib/op-attrs/test/src/op-attrs/computation_graph_op_attrs.cc b/lib/op-attrs/test/src/op-attrs/computation_graph_op_attrs.cc index 2ac6578118..5af11f93da 100644 --- a/lib/op-attrs/test/src/op-attrs/computation_graph_op_attrs.cc +++ b/lib/op-attrs/test/src/op-attrs/computation_graph_op_attrs.cc @@ -6,7 +6,7 @@ using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("ComputationGraphOpAttrs to/from json") { ComputationGraphOpAttrs correct = ComputationGraphOpAttrs{BatchNormAttrs{ - /*relu=*/false, + /*activation=*/std::nullopt, /*affine=*/true, /*eps=*/1e-5, /*momentum=*/0.1, diff --git a/lib/op-attrs/test/src/op-attrs/ops/batch_norm.cc b/lib/op-attrs/test/src/op-attrs/ops/batch_norm.cc index 20b7662c77..7ff658a0fd 100644 --- a/lib/op-attrs/test/src/op-attrs/ops/batch_norm.cc +++ b/lib/op-attrs/test/src/op-attrs/ops/batch_norm.cc @@ -11,7 +11,7 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("get_batch_norm_incoming_tensor_roles(BatchNormAttrs)") { auto make_attrs = [](bool affine) { return BatchNormAttrs{ - /*relu=*/false, + /*activation=*/std::nullopt, /*affine=*/affine, /*eps=*/1.0, /*momentum=*/0.1, @@ -60,7 +60,7 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("shape inference (BatchNorm)") { BatchNormAttrs attrs_affine_true = BatchNormAttrs{ - /*relu=*/false, + /*activation=*/std::nullopt, /*affine=*/true, /*eps=*/1.0, /*momentum=*/0.1, @@ -141,7 +141,7 @@ TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("parallel dim degree inference (BatchNormAttrs)") { BatchNormAttrs attrs_affine_true = BatchNormAttrs{ - /*relu=*/false, + /*activation=*/std::nullopt, /*affine=*/true, /*eps=*/1.0, /*momentum=*/0.1, @@ -346,7 +346,7 @@ TEST_SUITE(FF_TEST_SUITE) { // here we just do a basic check that they compose BatchNormAttrs attrs = BatchNormAttrs{ - /*relu=*/true, + /*activation=*/Activation::RELU, /*affine=*/true, /*eps=*/1.0, /*momentum=*/0.1, diff --git a/lib/pcg/src/pcg/computation_graph_builder.cc b/lib/pcg/src/pcg/computation_graph_builder.cc index 2227311ee7..0788c37b1b 100644 --- a/lib/pcg/src/pcg/computation_graph_builder.cc +++ b/lib/pcg/src/pcg/computation_graph_builder.cc @@ -723,17 +723,8 @@ tensor_guid_t ComputationGraphBuilder::batch_norm( BatchNormMode mode, std::optional const &maybe_name) { - if (activation.has_value() && activation.value() != Activation::RELU) { - PANIC(fmt::format( - "batch_norm currently only supports (1) no activation function, or (2) " - "relu activation function, but received {}. " - "If you need support for additional activation functions, please " - "create an issue.", - activation)); - } - BatchNormAttrs attrs = BatchNormAttrs{ - /*relu=*/activation.has_value(), + /*activation=*/activation, /*affine=*/affine, /*eps=*/eps, /*momentum=*/momentum, diff --git a/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc b/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc index b531486399..ffcdf19447 100644 --- a/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc +++ b/lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph_builder.cc @@ -345,17 +345,8 @@ parallel_tensor_guid_t ParallelComputationGraphBuilder::batch_norm( BatchNormMode mode, std::optional const &maybe_name) { - if (activation.has_value() && activation.value() != Activation::RELU) { - throw mk_runtime_error(fmt::format( - "batch_norm currently only supports (1) no activation function, or (2) " - "relu activation function, but received {}. " - "If you need support for additional activation functions, please " - "create an issue.", - activation)); - } - BatchNormAttrs attrs = BatchNormAttrs{ - /*relu=*/activation.has_value(), + /*activation=*/activation, /*affine=*/affine, /*eps=*/eps, /*momentum=*/momentum, diff --git a/lib/task-spec/src/task-spec/ops/impl/batch_norm.cc b/lib/task-spec/src/task-spec/ops/impl/batch_norm.cc index ea029262eb..ca44492dbb 100644 --- a/lib/task-spec/src/task-spec/ops/impl/batch_norm.cc +++ b/lib/task-spec/src/task-spec/ops/impl/batch_norm.cc @@ -87,6 +87,8 @@ static std::optional acc.get_tensor_grad(TensorSlotName::OUTPUT); GenericTensorAccessorR gamma = acc.get_tensor(TensorSlotName::GAMMA); + GenericTensorAccessorR beta = + acc.get_tensor(TensorSlotName::BETA); GenericTensorAccessorW gamma_grad = acc.get_tensor_grad(TensorSlotName::GAMMA); GenericTensorAccessorW beta_grad = @@ -104,6 +106,7 @@ static std::optional input, input_grad, gamma, + beta, gamma_grad, beta_grad); } From 4ee795c41b4b7b1bc0b333375fcd627a5f2f7837 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 17 Sep 2026 13:03:46 -0700 Subject: [PATCH 6/6] Give batch norm a grid the device can fill. The fused kernels put one block on each channel, which is all the parallelism the statistics reduction can have but more than the normalization that follows it needs. At YOLOv10x's 80-channel sites that left 80 blocks on 82 SMs: a single wave, sixteen warps per SM, and 58-60% of the bandwidth the same access pattern reaches elsewhere in the network, against 88-126% where the channel count is larger. Two changes, each chosen per call site from the extents: - below two blocks per SM, split the pass in two, so the normalization runs over a grid sized to the tensor rather than to the channel count. It is two thirds of the traffic and needs none of the reduction's structure. - walk four floats at a time where the spatial size leaves every thread work. Below that the vector walk divides the work by four and loses. Batch norm goes from 28.95 to 21.01 ms per iteration, and the iteration from 153.09 to 146.68 ms (4.2%), measured with evaluation/benchmark_flexflow.sh on the deploy build. The validation suite passes all four stages. Sliced onto master without the change to have backward kernels overwrite their gradients, so these kernels accumulate into their gradients too. That means the split backward pass can no longer hand its per-channel sums from the stats kernel to the apply kernel through gamma_grad and beta_grad, so they go through a gradSums scratch buffer in BatchNormPerDeviceState, allocated alongside the saved statistics. A new test runs the fused kernels at sizes that reach the vectorized and split paths, which the existing one is too small for, starting both the fused and separate paths from the same non-zero gradients. Co-Authored-By: Claude Opus 5.5 --- .../batch_norm_per_device_state.dtg.toml | 9 + .../src/cuda/ops/batch_norm_kernels.cu | 524 ++++++++++++++++-- .../src/kernels/batch_norm_kernels_gpu.cc | 195 ++++++- 3 files changed, 655 insertions(+), 73 deletions(-) diff --git a/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml b/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml index f1548e4891..662ce1e34b 100644 --- a/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml +++ b/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml @@ -38,3 +38,12 @@ type = "float *" [[fields]] name = "saveVar" type = "float *" + +[[fields]] +name = "gradSums" +type = "float *" +docstring = ''' +Scratch for the fused backward pass: the per-channel sums of the gradient and +of the gradient times the normalized input, handed from the kernel that +reduces them to the kernel that applies them. +''' diff --git a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu index 7d0d7c8d9c..9f49a476ea 100644 --- a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu +++ b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu @@ -225,7 +225,19 @@ __device__ __forceinline__ void * it from \p input rather than being handed it, so it is not needed later * either. */ -template +/** + * @brief How a channel's elements are walked. + * + * @details A thread loading four contiguous floats at a time keeps more loads + * in flight, which is what the kernels are short of, but it also divides the + * work by four -- so below \ref BATCH_NORM_MIN_VECTOR_SPATIAL there are no + * longer enough of them to go round and the scalar walk wins. + */ +enum class BatchNormWalk { SCALAR, VECTOR }; + +constexpr int BATCH_NORM_MIN_VECTOR_SPATIAL = 1024; + +template __global__ void batch_norm_fused_forward_kernel(int num_samples, int num_channels, @@ -245,8 +257,21 @@ __global__ void float count = 0.0f, mean = 0.0f, m2 = 0.0f; for (int sample = 0; sample < num_samples; sample++) { size_t base = ((size_t)sample * num_channels + channel) * spatial_size; - for (int i = threadIdx.x; i < spatial_size; i += BATCH_NORM_FUSED_THREADS) { - welford_join(count, mean, m2, 1.0f, input[base + i], 0.0f); + if constexpr (WALK == BatchNormWalk::VECTOR) { + float4 const *in4 = reinterpret_cast(input + base); + for (int i = threadIdx.x; i < spatial_size / 4; + i += BATCH_NORM_FUSED_THREADS) { + float4 v = in4[i]; + welford_join(count, mean, m2, 1.0f, v.x, 0.0f); + welford_join(count, mean, m2, 1.0f, v.y, 0.0f); + welford_join(count, mean, m2, 1.0f, v.z, 0.0f); + welford_join(count, mean, m2, 1.0f, v.w, 0.0f); + } + } else { + for (int i = threadIdx.x; i < spatial_size; + i += BATCH_NORM_FUSED_THREADS) { + welford_join(count, mean, m2, 1.0f, input[base + i], 0.0f); + } } } block_reduce_welford(count, mean, m2); @@ -274,10 +299,29 @@ __global__ void float shift = beta[channel]; for (int sample = 0; sample < num_samples; sample++) { size_t base = ((size_t)sample * num_channels + channel) * spatial_size; - for (int i = threadIdx.x; i < spatial_size; i += BATCH_NORM_FUSED_THREADS) { - float normalized = (input[base + i] - mean) * invstd; - output[base + i] = - apply_activation(ACTIVATION, scale * normalized + shift); + if constexpr (WALK == BatchNormWalk::VECTOR) { + float4 const *in4 = reinterpret_cast(input + base); + float4 *out4 = reinterpret_cast(output + base); + for (int i = threadIdx.x; i < spatial_size / 4; + i += BATCH_NORM_FUSED_THREADS) { + float4 v = in4[i]; + out4[i] = make_float4( + apply_activation(ACTIVATION, + scale * ((v.x - mean) * invstd) + shift), + apply_activation(ACTIVATION, + scale * ((v.y - mean) * invstd) + shift), + apply_activation(ACTIVATION, + scale * ((v.z - mean) * invstd) + shift), + apply_activation(ACTIVATION, + scale * ((v.w - mean) * invstd) + shift)); + } + } else { + for (int i = threadIdx.x; i < spatial_size; + i += BATCH_NORM_FUSED_THREADS) { + float normalized = (input[base + i] - mean) * invstd; + output[base + i] = + apply_activation(ACTIVATION, scale * normalized + shift); + } } } } @@ -295,7 +339,7 @@ __global__ void * Like the cuDNN path, accumulates into the gradients rather than overwriting * them. */ -template +template __global__ void batch_norm_fused_backward_kernel(int num_samples, int num_channels, int spatial_size, @@ -317,12 +361,40 @@ __global__ void batch_norm_fused_backward_kernel(int num_samples, float sum_grad = 0.0f, sum_grad_normalized = 0.0f; for (int sample = 0; sample < num_samples; sample++) { size_t base = ((size_t)sample * num_channels + channel) * spatial_size; - for (int i = threadIdx.x; i < spatial_size; i += BATCH_NORM_FUSED_THREADS) { - float normalized = (input[base + i] - mean) * invstd; - float grad = activation_backward( - ACTIVATION, output_grad[base + i], scale * normalized + shift); - sum_grad += grad; - sum_grad_normalized += grad * normalized; + if constexpr (WALK == BatchNormWalk::VECTOR) { + float4 const *in4 = reinterpret_cast(input + base); + float4 const *grad4 = + reinterpret_cast(output_grad + base); + for (int i = threadIdx.x; i < spatial_size / 4; + i += BATCH_NORM_FUSED_THREADS) { + float4 x = in4[i], d = grad4[i]; + float normalized, grad; + normalized = (x.x - mean) * invstd; + grad = activation_backward(ACTIVATION, d.x, scale * normalized + shift); + sum_grad += grad; + sum_grad_normalized += grad * normalized; + normalized = (x.y - mean) * invstd; + grad = activation_backward(ACTIVATION, d.y, scale * normalized + shift); + sum_grad += grad; + sum_grad_normalized += grad * normalized; + normalized = (x.z - mean) * invstd; + grad = activation_backward(ACTIVATION, d.z, scale * normalized + shift); + sum_grad += grad; + sum_grad_normalized += grad * normalized; + normalized = (x.w - mean) * invstd; + grad = activation_backward(ACTIVATION, d.w, scale * normalized + shift); + sum_grad += grad; + sum_grad_normalized += grad * normalized; + } + } else { + for (int i = threadIdx.x; i < spatial_size; + i += BATCH_NORM_FUSED_THREADS) { + float normalized = (input[base + i] - mean) * invstd; + float grad = activation_backward( + ACTIVATION, output_grad[base + i], scale * normalized + shift); + sum_grad += grad; + sum_grad_normalized += grad * normalized; + } } } block_reduce_sum2(sum_grad, sum_grad_normalized); @@ -335,15 +407,277 @@ __global__ void batch_norm_fused_backward_kernel(int num_samples, float count = (float)num_samples * spatial_size; for (int sample = 0; sample < num_samples; sample++) { size_t base = ((size_t)sample * num_channels + channel) * spatial_size; - for (int i = threadIdx.x; i < spatial_size; i += BATCH_NORM_FUSED_THREADS) { - float normalized = (input[base + i] - mean) * invstd; - float grad = activation_backward( - ACTIVATION, output_grad[base + i], scale * normalized + shift); - input_grad[base + i] += - scale * invstd * - (grad - (sum_grad + normalized * sum_grad_normalized) / count); + if constexpr (WALK == BatchNormWalk::VECTOR) { + float4 const *in4 = reinterpret_cast(input + base); + float4 const *grad4 = + reinterpret_cast(output_grad + base); + float4 *dst4 = reinterpret_cast(input_grad + base); + for (int i = threadIdx.x; i < spatial_size / 4; + i += BATCH_NORM_FUSED_THREADS) { + float4 x = in4[i], d = grad4[i], r = dst4[i]; + float normalized, grad; + normalized = (x.x - mean) * invstd; + grad = activation_backward(ACTIVATION, d.x, scale * normalized + shift); + r.x += scale * invstd * + (grad - (sum_grad + normalized * sum_grad_normalized) / count); + normalized = (x.y - mean) * invstd; + grad = activation_backward(ACTIVATION, d.y, scale * normalized + shift); + r.y += scale * invstd * + (grad - (sum_grad + normalized * sum_grad_normalized) / count); + normalized = (x.z - mean) * invstd; + grad = activation_backward(ACTIVATION, d.z, scale * normalized + shift); + r.z += scale * invstd * + (grad - (sum_grad + normalized * sum_grad_normalized) / count); + normalized = (x.w - mean) * invstd; + grad = activation_backward(ACTIVATION, d.w, scale * normalized + shift); + r.w += scale * invstd * + (grad - (sum_grad + normalized * sum_grad_normalized) / count); + dst4[i] = r; + } + } else { + for (int i = threadIdx.x; i < spatial_size; + i += BATCH_NORM_FUSED_THREADS) { + float normalized = (input[base + i] - mean) * invstd; + float grad = activation_backward( + ACTIVATION, output_grad[base + i], scale * normalized + shift); + input_grad[base + i] += + scale * invstd * + (grad - (sum_grad + normalized * sum_grad_normalized) / count); + } + } + } +} + +/** + * @brief The forward pass split in two, for channel counts that cannot fill the + * device. + * + * @details One block per channel is all the parallelism the reduction can have, + * but the normalization that follows it needs none of that structure. When + * there are fewer channels than SMs the fused kernel leaves most of the device + * idle for both halves; splitting lets the second half, which is two thirds of + * the traffic, run over a grid sized to the tensor instead of to the channel + * count. Measured on YOLOv10x's 80-channel sites this is worth 1.3-1.5x; above + * \ref BATCH_NORM_MIN_SPLIT_BLOCKS_PER_SM the fused kernel is already busy + * enough that the extra pass over the input costs more than the grid buys. + */ +constexpr int BATCH_NORM_MIN_SPLIT_BLOCKS_PER_SM = 2; +constexpr int BATCH_NORM_MIN_SPLIT_SPATIAL = 1600; +constexpr int BATCH_NORM_APPLY_BLOCKS = 4096; +constexpr int BATCH_NORM_APPLY_THREADS = 256; + +template +__global__ void batch_norm_stats_kernel(int num_samples, + int num_channels, + int spatial_size, + float eps, + float exponential_average_factor, + float const *input, + float *save_mean, + float *save_invstd, + float *running_mean, + float *running_var) { + int channel = blockIdx.x; + + float count = 0.0f, mean = 0.0f, m2 = 0.0f; + for (int sample = 0; sample < num_samples; sample++) { + size_t base = ((size_t)sample * num_channels + channel) * spatial_size; + if constexpr (WALK == BatchNormWalk::VECTOR) { + float4 const *in4 = reinterpret_cast(input + base); + for (int i = threadIdx.x; i < spatial_size / 4; + i += BATCH_NORM_FUSED_THREADS) { + float4 v = in4[i]; + welford_join(count, mean, m2, 1.0f, v.x, 0.0f); + welford_join(count, mean, m2, 1.0f, v.y, 0.0f); + welford_join(count, mean, m2, 1.0f, v.z, 0.0f); + welford_join(count, mean, m2, 1.0f, v.w, 0.0f); + } + } else { + for (int i = threadIdx.x; i < spatial_size; + i += BATCH_NORM_FUSED_THREADS) { + welford_join(count, mean, m2, 1.0f, input[base + i], 0.0f); + } + } + } + block_reduce_welford(count, mean, m2); + + if (threadIdx.x == 0) { + float variance = m2 / count; + save_mean[channel] = mean; + save_invstd[channel] = rsqrtf(variance + eps); + + float unbiased_variance = count > 1.0f ? m2 / (count - 1.0f) : variance; + running_mean[channel] = + (1.0f - exponential_average_factor) * running_mean[channel] + + exponential_average_factor * mean; + running_var[channel] = + (1.0f - exponential_average_factor) * running_var[channel] + + exponential_average_factor * unbiased_variance; + } +} + +template +__global__ void batch_norm_apply_kernel(size_t num_vectors, + int num_channels, + int spatial_size, + float const *input, + float const *gamma, + float const *beta, + float const *save_mean, + float const *save_invstd, + float *output) { + int vectors_per_channel = spatial_size / 4; + float4 const *in4 = reinterpret_cast(input); + float4 *out4 = reinterpret_cast(output); + + for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; + i < num_vectors; + i += (size_t)blockDim.x * gridDim.x) { + int channel = (int)((i / vectors_per_channel) % num_channels); + float mean = save_mean[channel], invstd = save_invstd[channel]; + float scale = gamma[channel], shift = beta[channel]; + float4 v = in4[i]; + out4[i] = make_float4( + apply_activation(ACTIVATION, scale * ((v.x - mean) * invstd) + shift), + apply_activation(ACTIVATION, scale * ((v.y - mean) * invstd) + shift), + apply_activation(ACTIVATION, scale * ((v.z - mean) * invstd) + shift), + apply_activation(ACTIVATION, scale * ((v.w - mean) * invstd) + shift)); + } +} + +template +__global__ void batch_norm_bwd_stats_kernel(int num_samples, + int num_channels, + int spatial_size, + float const *output_grad, + float const *input, + float const *gamma, + float const *beta, + float const *save_mean, + float const *save_invstd, + float *grad_sums, + float *gamma_grad, + float *beta_grad) { + int channel = blockIdx.x; + float mean = save_mean[channel]; + float invstd = save_invstd[channel]; + float scale = gamma[channel]; + float shift = beta[channel]; + + float sum_grad = 0.0f, sum_grad_normalized = 0.0f; + for (int sample = 0; sample < num_samples; sample++) { + size_t base = ((size_t)sample * num_channels + channel) * spatial_size; + if constexpr (WALK == BatchNormWalk::VECTOR) { + float4 const *in4 = reinterpret_cast(input + base); + float4 const *grad4 = + reinterpret_cast(output_grad + base); + for (int i = threadIdx.x; i < spatial_size / 4; + i += BATCH_NORM_FUSED_THREADS) { + float4 x = in4[i], d = grad4[i]; + float const xs[4] = {x.x, x.y, x.z, x.w}; + float const ds[4] = {d.x, d.y, d.z, d.w}; +#pragma unroll + for (int k = 0; k < 4; k++) { + float normalized = (xs[k] - mean) * invstd; + float grad = activation_backward( + ACTIVATION, ds[k], scale * normalized + shift); + sum_grad += grad; + sum_grad_normalized += grad * normalized; + } + } + } else { + for (int i = threadIdx.x; i < spatial_size; + i += BATCH_NORM_FUSED_THREADS) { + float normalized = (input[base + i] - mean) * invstd; + float grad = activation_backward( + ACTIVATION, output_grad[base + i], scale * normalized + shift); + sum_grad += grad; + sum_grad_normalized += grad * normalized; + } } } + block_reduce_sum2(sum_grad, sum_grad_normalized); + + if (threadIdx.x == 0) { + grad_sums[channel] = sum_grad; + grad_sums[num_channels + channel] = sum_grad_normalized; + gamma_grad[channel] += sum_grad_normalized; + beta_grad[channel] += sum_grad; + } +} + +/** + * @brief The second half of the backward pass, reading the two per-channel sums + * the stats kernel left in \p grad_sums. + * + * @details The sums cannot be read back out of the parameter gradients, since + * those are accumulated into rather than overwritten. + */ +template +__global__ void batch_norm_bwd_apply_kernel(size_t num_vectors, + int num_samples, + int num_channels, + int spatial_size, + float const *output_grad, + float const *input, + float const *gamma, + float const *beta, + float const *save_mean, + float const *save_invstd, + float const *grad_sums, + float *input_grad) { + int vectors_per_channel = spatial_size / 4; + float count = (float)num_samples * spatial_size; + float4 const *in4 = reinterpret_cast(input); + float4 const *grad4 = reinterpret_cast(output_grad); + float4 *dst4 = reinterpret_cast(input_grad); + + for (size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; + i < num_vectors; + i += (size_t)blockDim.x * gridDim.x) { + int channel = (int)((i / vectors_per_channel) % num_channels); + float mean = save_mean[channel], invstd = save_invstd[channel]; + float scale = gamma[channel], shift = beta[channel]; + float sum_grad = grad_sums[channel]; + float sum_grad_normalized = grad_sums[num_channels + channel]; + + float4 x = in4[i], d = grad4[i], r = dst4[i]; + float const xs[4] = {x.x, x.y, x.z, x.w}; + float const ds[4] = {d.x, d.y, d.z, d.w}; + float rs[4]; +#pragma unroll + for (int k = 0; k < 4; k++) { + float normalized = (xs[k] - mean) * invstd; + float grad = + activation_backward(ACTIVATION, ds[k], scale * normalized + shift); + rs[k] = scale * invstd * + (grad - (sum_grad + normalized * sum_grad_normalized) / count); + } + r = make_float4(r.x + rs[0], r.y + rs[1], r.z + rs[2], r.w + rs[3]); + dst4[i] = r; + } +} + +/** + * @brief Whether the split kernels above are the better shape for \p extents. + */ +bool batch_norm_should_split(BatchNormExtents const &extents) { + // Queried once: this runs on the launch path of every batch norm in the + // graph, twice per iteration. + static int const sm_count = []() { + int count; + checkCUDA(cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, /*device=*/0)); + return count; + }(); + return extents.num_channels < BATCH_NORM_MIN_SPLIT_BLOCKS_PER_SM * sm_count && + extents.spatial_size >= BATCH_NORM_MIN_SPLIT_SPATIAL && + extents.spatial_size % 4 == 0; +} + +bool batch_norm_should_vectorize(BatchNormExtents const &extents) { + return extents.spatial_size % 4 == 0 && + extents.spatial_size >= BATCH_NORM_MIN_VECTOR_SPATIAL; } /** @@ -440,13 +774,15 @@ BatchNormPerDeviceState /*h=*/1, /*w=*/1)); - // Allocate memory for runningMean, runningVar, saveMean and saveVar as a - // single contiguous block (deallocated by batch_norm_gpu_cleanup_kernel). + // Allocate memory for runningMean, runningVar, saveMean, saveVar and + // gradSums (two per channel) as a single contiguous block (deallocated by + // batch_norm_gpu_cleanup_kernel). float *runningMean = static_cast( - allocator.allocate(sizeof(float) * num_channels * 4)); + allocator.allocate(sizeof(float) * num_channels * 6)); float *runningVar = runningMean + num_channels; float *saveMean = runningVar + num_channels; float *saveVar = saveMean + num_channels; + float *gradSums = saveVar + num_channels; // Match the PyTorch initialization of running_mean = 0 and running_var = 1. std::vector initial_running_stats(num_channels * 2); @@ -470,6 +806,7 @@ BatchNormPerDeviceState /*runningVar=*/runningVar, /*saveMean=*/saveMean, /*saveVar=*/saveVar, + /*gradSums=*/gradSums, }; } @@ -494,21 +831,59 @@ void batch_norm_gpu_forward_kernel( BatchNormExtents extents = get_extents(input.shape); dispatch_on_activation(attrs.activation.value(), [&](auto activation) { - batch_norm_fused_forward_kernel - <<>>( - extents.num_samples, - extents.num_channels, - extents.spatial_size, - attrs.eps, - static_cast(exponential_average_factor), - input.get_float_ptr(), - gamma.get_float_ptr(), - beta.get_float_ptr(), - output.get_float_ptr(), - per_device_state.saveMean, - per_device_state.saveVar, - per_device_state.runningMean, - per_device_state.runningVar); + constexpr Activation ACT = decltype(activation)::value; + size_t num_vectors = (size_t)extents.num_samples * extents.num_channels * + extents.spatial_size / 4; + + if (batch_norm_should_split(extents)) { + batch_norm_stats_kernel + <<>>( + extents.num_samples, + extents.num_channels, + extents.spatial_size, + attrs.eps, + static_cast(exponential_average_factor), + input.get_float_ptr(), + per_device_state.saveMean, + per_device_state.saveVar, + per_device_state.runningMean, + per_device_state.runningVar); + batch_norm_apply_kernel + <<>>( + num_vectors, + extents.num_channels, + extents.spatial_size, + input.get_float_ptr(), + gamma.get_float_ptr(), + beta.get_float_ptr(), + per_device_state.saveMean, + per_device_state.saveVar, + output.get_float_ptr()); + return; + } + + auto launch = [&](auto walk) { + batch_norm_fused_forward_kernel + <<>>( + extents.num_samples, + extents.num_channels, + extents.spatial_size, + attrs.eps, + static_cast(exponential_average_factor), + input.get_float_ptr(), + gamma.get_float_ptr(), + beta.get_float_ptr(), + output.get_float_ptr(), + per_device_state.saveMean, + per_device_state.saveVar, + per_device_state.runningMean, + per_device_state.runningVar); + }; + if (batch_norm_should_vectorize(extents)) { + launch(std::integral_constant{}); + } else { + launch(std::integral_constant{}); + } }); return; } @@ -553,20 +928,63 @@ void batch_norm_gpu_backward_kernel( BatchNormExtents extents = get_extents(input.shape); dispatch_on_activation(attrs.activation.value(), [&](auto activation) { - batch_norm_fused_backward_kernel - <<>>( - extents.num_samples, - extents.num_channels, - extents.spatial_size, - output_grad.get_float_ptr(), - input.get_float_ptr(), - gamma.get_float_ptr(), - beta.get_float_ptr(), - per_device_state.saveMean, - per_device_state.saveVar, - input_grad.get_float_ptr(), - gamma_grad.get_float_ptr(), - beta_grad.get_float_ptr()); + constexpr Activation ACT = decltype(activation)::value; + size_t num_vectors = (size_t)extents.num_samples * extents.num_channels * + extents.spatial_size / 4; + + if (batch_norm_should_split(extents)) { + batch_norm_bwd_stats_kernel + <<>>( + extents.num_samples, + extents.num_channels, + extents.spatial_size, + output_grad.get_float_ptr(), + input.get_float_ptr(), + gamma.get_float_ptr(), + beta.get_float_ptr(), + per_device_state.saveMean, + per_device_state.saveVar, + per_device_state.gradSums, + gamma_grad.get_float_ptr(), + beta_grad.get_float_ptr()); + batch_norm_bwd_apply_kernel + <<>>( + num_vectors, + extents.num_samples, + extents.num_channels, + extents.spatial_size, + output_grad.get_float_ptr(), + input.get_float_ptr(), + gamma.get_float_ptr(), + beta.get_float_ptr(), + per_device_state.saveMean, + per_device_state.saveVar, + per_device_state.gradSums, + input_grad.get_float_ptr()); + return; + } + + auto launch = [&](auto walk) { + batch_norm_fused_backward_kernel + <<>>( + extents.num_samples, + extents.num_channels, + extents.spatial_size, + output_grad.get_float_ptr(), + input.get_float_ptr(), + gamma.get_float_ptr(), + beta.get_float_ptr(), + per_device_state.saveMean, + per_device_state.saveVar, + input_grad.get_float_ptr(), + gamma_grad.get_float_ptr(), + beta_grad.get_float_ptr()); + }; + if (batch_norm_should_vectorize(extents)) { + launch(std::integral_constant{}); + } else { + launch(std::integral_constant{}); + } }); return; } diff --git a/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc index f863f19d94..0a39d6f041 100644 --- a/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc +++ b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc @@ -261,11 +261,8 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { GenericTensorAccessorW fused = create_random_filled_accessor_w(input.shape, allocator); - BatchNormPerDeviceState separate_state = - batch_norm_gpu_init_kernel(allocator, - separate_attrs, - input.shape, - input.shape); + BatchNormPerDeviceState separate_state = batch_norm_gpu_init_kernel( + allocator, separate_attrs, input.shape, input.shape); batch_norm_gpu_forward_kernel(managed_stream.raw_stream(), managed_handle.raw_handle(), separate_state, @@ -282,11 +279,8 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { read_only_accessor_from_write_accessor(normalized), separate); - BatchNormPerDeviceState fused_state = - batch_norm_gpu_init_kernel(allocator, - fused_attrs, - input.shape, - input.shape); + BatchNormPerDeviceState fused_state = batch_norm_gpu_init_kernel( + allocator, fused_attrs, input.shape, input.shape); batch_norm_gpu_forward_kernel(managed_stream.raw_stream(), managed_handle.raw_handle(), fused_state, @@ -330,11 +324,8 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { GenericTensorAccessorW separate_beta_grad = create_zero_filled_accessor_w(gamma.shape, allocator); - BatchNormPerDeviceState separate_state = - batch_norm_gpu_init_kernel(allocator, - separate_attrs, - input.shape, - input.shape); + BatchNormPerDeviceState separate_state = batch_norm_gpu_init_kernel( + allocator, separate_attrs, input.shape, input.shape); batch_norm_gpu_forward_kernel(managed_stream.raw_stream(), managed_handle.raw_handle(), separate_state, @@ -383,11 +374,8 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { GenericTensorAccessorW fused_beta_grad = create_zero_filled_accessor_w(gamma.shape, allocator); - BatchNormPerDeviceState fused_state = - batch_norm_gpu_init_kernel(allocator, - fused_attrs, - input.shape, - input.shape); + BatchNormPerDeviceState fused_state = batch_norm_gpu_init_kernel( + allocator, fused_attrs, input.shape, input.shape); batch_norm_gpu_forward_kernel(managed_stream.raw_stream(), managed_handle.raw_handle(), fused_state, @@ -426,4 +414,171 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { batch_norm_gpu_cleanup_kernel(allocator, fused_state); } } + + // The case above is too small to reach the vectorized walk or the split + // kernels, which only run once the spatial size (and, for the split, a + // channel count too small to fill the device) makes them worth it. Here the + // gradients also start from the same non-zero values on both paths, since + // the fused kernels have to accumulate into them just as cuDNN does. + TEST_CASE("batch_norm_gpu fused activation matches the separate operators at " + "sizes that reach the vectorized and split kernels") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + Allocator allocator = create_local_cuda_memory_allocator(); + + ffStream_t stream = managed_stream.raw_stream(); + PerDeviceFFHandle handle = managed_handle.raw_handle(); + + ElementUnaryAttrs silu_attrs = ElementUnaryAttrs{ + /*op_type=*/OperatorType::SILU, + /*scalar=*/std::nullopt, + }; + BatchNormAttrs separate_attrs = make_attrs(); + BatchNormAttrs fused_attrs = make_attrs(Activation::SILU); + + auto check = [&](TensorShape const &input_shape) { + TensorShape channel_shape = TensorShape{ + TensorDims{FFOrdered{dim_at_idx(input_shape.dims, ff_dim_t{1_n})}}, + DataType::FLOAT, + }; + + GenericTensorAccessorR input = + create_random_filled_accessor_r(input_shape, allocator); + GenericTensorAccessorR gamma = + create_random_filled_accessor_r(channel_shape, allocator); + GenericTensorAccessorR beta = + create_random_filled_accessor_r(channel_shape, allocator); + GenericTensorAccessorR output_grad = + create_random_filled_accessor_r(input_shape, allocator); + + auto copy_of = [&](GenericTensorAccessorW const &accessor) { + GenericTensorAccessorW result = + allocator.allocate_tensor(accessor.shape); + copy_accessor_data_to_l_from_r( + result, read_only_accessor_from_write_accessor(accessor)); + return result; + }; + + ElementUnaryPerDeviceState silu_state = + element_unary_gpu_init_kernel(silu_attrs, input_shape, input_shape); + + // The separate path. + GenericTensorAccessorW normalized = + create_random_filled_accessor_w(input_shape, allocator); + GenericTensorAccessorW separate_output = + create_random_filled_accessor_w(input_shape, allocator); + GenericTensorAccessorW normalized_grad = + create_zero_filled_accessor_w(input_shape, allocator); + GenericTensorAccessorW separate_input_grad = + create_random_filled_accessor_w(input_shape, allocator); + GenericTensorAccessorW separate_gamma_grad = + create_random_filled_accessor_w(channel_shape, allocator); + GenericTensorAccessorW separate_beta_grad = + create_random_filled_accessor_w(channel_shape, allocator); + + // The fused path, starting from the same gradients. + GenericTensorAccessorW fused_output = + create_random_filled_accessor_w(input_shape, allocator); + GenericTensorAccessorW fused_input_grad = copy_of(separate_input_grad); + GenericTensorAccessorW fused_gamma_grad = copy_of(separate_gamma_grad); + GenericTensorAccessorW fused_beta_grad = copy_of(separate_beta_grad); + + BatchNormPerDeviceState separate_state = batch_norm_gpu_init_kernel( + allocator, separate_attrs, input_shape, input_shape); + batch_norm_gpu_forward_kernel(stream, + handle, + separate_state, + separate_attrs, + input, + gamma, + beta, + normalized); + element_unary_gpu_forward_kernel( + stream, + handle, + silu_state, + silu_attrs, + read_only_accessor_from_write_accessor(normalized), + separate_output); + element_unary_gpu_backward_kernel( + stream, + handle, + silu_state, + silu_attrs, + read_only_accessor_from_write_accessor(separate_output), + output_grad, + read_only_accessor_from_write_accessor(normalized), + normalized_grad); + batch_norm_gpu_backward_kernel( + stream, + handle, + separate_state, + separate_attrs, + read_only_accessor_from_write_accessor(normalized), + read_only_accessor_from_write_accessor(normalized_grad), + input, + separate_input_grad, + gamma, + beta, + separate_gamma_grad, + separate_beta_grad); + + BatchNormPerDeviceState fused_state = batch_norm_gpu_init_kernel( + allocator, fused_attrs, input_shape, input_shape); + batch_norm_gpu_forward_kernel(stream, + handle, + fused_state, + fused_attrs, + input, + gamma, + beta, + fused_output); + batch_norm_gpu_backward_kernel( + stream, + handle, + fused_state, + fused_attrs, + read_only_accessor_from_write_accessor(fused_output), + output_grad, + input, + fused_input_grad, + gamma, + beta, + fused_gamma_grad, + fused_beta_grad); + + CHECK(accessors_within_epsilon(fused_output, separate_output, 1e-4)); + CHECK(accessors_within_epsilon( + fused_input_grad, separate_input_grad, 1e-4)); + CHECK_MESSAGE( + accessors_within_epsilon(fused_gamma_grad, separate_gamma_grad, 1e-3), + check_kv("fused", format_accessor_w_contents(fused_gamma_grad)), + check_kv("separate", + format_accessor_w_contents(separate_gamma_grad))); + CHECK_MESSAGE( + accessors_within_epsilon(fused_beta_grad, separate_beta_grad, 1e-3), + check_kv("fused", format_accessor_w_contents(fused_beta_grad)), + check_kv("separate", format_accessor_w_contents(separate_beta_grad))); + + batch_norm_gpu_cleanup_kernel(allocator, separate_state); + batch_norm_gpu_cleanup_kernel(allocator, fused_state); + }; + + SUBCASE("vectorized") { + // Enough channels to fill any current device, so not split. + check(TensorShape{ + TensorDims{FFOrdered{2_p, 512_p, 32_p, 32_p}}, + DataType::FLOAT, + }); + } + + SUBCASE("split") { + check(TensorShape{ + TensorDims{FFOrdered{2_p, 4_p, 40_p, 40_p}}, + DataType::FLOAT, + }); + } + } }