diff --git a/lib/kernels/include/kernels/batch_norm_kernels.h b/lib/kernels/include/kernels/batch_norm_kernels.h index 47cb3d85a8..709a1d92e1 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels.h +++ b/lib/kernels/include/kernels/batch_norm_kernels.h @@ -1,47 +1,53 @@ #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, + GenericTensorAccessorR const &beta, + 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..731756d5ff 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels_cpu.h +++ b/lib/kernels/include/kernels/batch_norm_kernels_cpu.h @@ -1,28 +1,27 @@ #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, + GenericTensorAccessorR const &beta, + 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..ec6e0e5ec2 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels_gpu.h +++ b/lib/kernels/include/kernels/batch_norm_kernels_gpu.h @@ -1,43 +1,47 @@ #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, + GenericTensorAccessorR const &beta, + 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..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 @@ -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" @@ -49,21 +40,10 @@ 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" +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/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..9f49a476ea 100644 --- a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu +++ b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu @@ -14,170 +14,1012 @@ */ #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/ops/batch_norm.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)); -} - -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) { + +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. + */ +/** + * @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, + 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; + 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); + + 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; + 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); + } + } + } +} + +/** + * @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; + 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); + + 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; + 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; +} + +/** + * @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) { + 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. " + "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; -#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(cudnnCreateTensorDescriptor(&biasTensor)); + + 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( + 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, 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 * 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); + 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, + /*gradSums=*/gradSums, }; +} + +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); + + if (attrs.activation.has_value()) { + BatchNormExtents extents = get_extents(input.shape); + + dispatch_on_activation(attrs.activation.value(), [&](auto activation) { + 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; + } + + 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 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, + 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) { + 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; + } + + // 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 gpu_cleanup_kernel(Allocator &allocator, - BatchNormPerDeviceState &per_device_state) { +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/cuda/ops/element_unary_kernels.cu b/lib/kernels/src/cuda/ops/element_unary_kernels.cu index 99ce5ba888..b540b5be3a 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,19 @@ __global__ void elewise_scalar_unary_backward_kernel(coord_t volume, (T)(output_grad[i] * scalar * powf(input[i], scalar - 1)); break; } + case OperatorType::SILU: { + // 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: assert(false); } @@ -266,11 +296,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 +341,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/src/kernels/batch_norm_kernels.cc b/lib/kernels/src/kernels/batch_norm_kernels.cc index e23f6a89e2..1a40f3ae25 100644 --- a/lib/kernels/src/kernels/batch_norm_kernels.cc +++ b/lib/kernels/src/kernels/batch_norm_kernels.cc @@ -1,107 +1,114 @@ #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, + GenericTensorAccessorR const &beta, + 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, + /*beta=*/beta, + /*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, + /*beta=*/beta, + /*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..d4071556a0 100644 --- a/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc +++ b/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc @@ -1,25 +1,26 @@ #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, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad) { NOT_IMPLEMENTED(); } -} // namespace FlexFlow::Kernels::BatchNorm +} // namespace FlexFlow 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: 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..0a39d6f041 --- /dev/null +++ b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc @@ -0,0 +1,584 @@ +#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(std::optional activation = std::nullopt) { + return BatchNormAttrs{ + /*activation=*/activation, + /*affine=*/true, + /*eps=*/1e-5, + /*momentum=*/0.1, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, + }; +} + +// 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, + /*beta=*/beta, + /*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); + } + + // 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); + } + } + + // 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, + }); + } + } +} 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))); } } 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/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.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 a22fee4a64..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,8 @@ features = [ includes = [ "", + "op-attrs/activation.dtg.h", + "op-attrs/ops/batch_norm_mode.dtg.h", ] src_includes = [ @@ -22,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" @@ -36,3 +43,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/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 84f1861f0b..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,10 +6,11 @@ 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, + /*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..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,10 +11,11 @@ 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, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }; }; @@ -59,10 +60,11 @@ 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, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }; BatchNormAttrs attrs_affine_false = [&] { @@ -139,10 +141,11 @@ 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, + /*mode=*/BatchNormMode::SPATIAL_PERSISTENT, }; BatchNormAttrs attrs_affine_false = [&] { @@ -343,10 +346,11 @@ 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, + /*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..0788c37b1b 100644 --- a/lib/pcg/src/pcg/computation_graph_builder.cc +++ b/lib/pcg/src/pcg/computation_graph_builder.cc @@ -720,22 +720,15 @@ 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) { - 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, + /*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..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 @@ -342,22 +342,15 @@ 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) { - 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, + /*mode=*/mode, }; std::string name = 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..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 @@ -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,76 @@ 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); + GenericTensorAccessorR beta = + acc.get_tensor(TensorSlotName::BETA); + 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, + beta, + gamma_grad, + beta_grad); } TaskImplFunction get_batch_norm_init_task_impl() {