From cf328b4fea76536fec812ff6db74528c41c90bd5 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 3 Sep 2026 11:25:29 -0700 Subject: [PATCH 1/2] Update kernels for conv_2d, batch_norm, pool_2d, concat, reshape, split, batch_matmul, transpose and upsample. --- .../include/kernels/batch_matmul_kernels.h | 3 + .../kernels/batch_matmul_kernels_gpu.h | 3 + .../include/kernels/batch_norm_kernels.h | 71 ++- .../include/kernels/batch_norm_kernels_cpu.h | 36 +- .../include/kernels/batch_norm_kernels_gpu.h | 73 ++- .../batch_norm_per_device_state.dtg.toml | 29 - lib/kernels/include/kernels/concat_kernels.h | 24 +- .../include/kernels/concat_kernels_cpu.h | 22 +- .../include/kernels/concat_kernels_gpu.h | 25 +- lib/kernels/include/kernels/conv_2d_kernels.h | 76 ++- .../include/kernels/conv_2d_kernels_cpu.h | 36 +- .../include/kernels/conv_2d_kernels_gpu.h | 77 ++- .../kernels/conv_2d_per_device_state.dtg.toml | 9 - .../kernels/create_accessor_with_contents.h | 11 +- lib/kernels/include/kernels/pool_2d_kernels.h | 64 +- .../include/kernels/pool_2d_kernels_cpu.h | 20 +- .../include/kernels/pool_2d_kernels_gpu.h | 71 +-- .../kernels/pool_2d_per_device_state.dtg.toml | 13 - lib/kernels/include/kernels/reshape_kernels.h | 18 +- .../include/kernels/reshape_kernels_cpu.h | 14 +- .../include/kernels/reshape_kernels_gpu.h | 18 +- lib/kernels/include/kernels/split_kernels.h | 31 +- .../include/kernels/split_kernels_cpu.h | 29 +- .../include/kernels/split_kernels_gpu.h | 32 +- .../kernels/tensor_accessor_binary_ops.h | 7 +- .../include/kernels/transpose_kernels.h | 22 +- .../include/kernels/transpose_kernels_cpu.h | 18 +- .../include/kernels/transpose_kernels_gpu.h | 22 +- .../include/kernels/upsample_kernels_cpu.h | 2 + .../include/kernels/upsample_kernels_gpu.h | 1 - .../src/cuda/ops/batch_matmul_kernels.cu | 143 ++++- .../src/cuda/ops/batch_norm_kernels.cu | 302 ++++----- lib/kernels/src/cuda/ops/concat_kernels.cu | 105 +-- lib/kernels/src/cuda/ops/conv_2d_kernels.cu | 600 ++++++++---------- .../src/cuda/ops/element_unary_kernels.cu | 30 +- lib/kernels/src/cuda/ops/pool_2d_kernels.cu | 183 +++--- lib/kernels/src/cuda/ops/reshape_kernels.cu | 77 ++- lib/kernels/src/cuda/ops/split_kernels.cu | 100 ++- lib/kernels/src/cuda/ops/transpose_kernels.cu | 92 +-- lib/kernels/src/cuda/ops/upsample_kernels.cu | 161 ++++- .../src/kernels/batch_matmul_kernels.cc | 6 + lib/kernels/src/kernels/batch_norm_kernels.cc | 150 ++--- .../src/kernels/batch_norm_kernels_cpu.cc | 32 +- lib/kernels/src/kernels/concat_kernels.cc | 51 +- lib/kernels/src/kernels/concat_kernels_cpu.cc | 120 +++- lib/kernels/src/kernels/conv_2d_kernels.cc | 155 +++-- .../src/kernels/conv_2d_kernels_cpu.cc | 32 +- lib/kernels/src/kernels/pool_2d_kernels.cc | 129 ++-- .../src/kernels/pool_2d_kernels_cpu.cc | 17 +- lib/kernels/src/kernels/reshape_kernels.cc | 34 +- .../src/kernels/reshape_kernels_cpu.cc | 59 +- lib/kernels/src/kernels/split_kernels.cc | 79 +-- lib/kernels/src/kernels/split_kernels_cpu.cc | 133 +++- lib/kernels/src/kernels/transpose_kernels.cc | 42 +- .../src/kernels/transpose_kernels_cpu.cc | 85 ++- lib/kernels/src/kernels/upsample_kernels.cc | 2 + .../src/kernels/upsample_kernels_cpu.cc | 4 +- .../src/kernels/batch_matmul_kernels_cpu.cc | 2 +- .../src/kernels/batch_matmul_kernels_gpu.cc | 193 ++++++ .../src/kernels/batch_norm_kernels_gpu.cc | 225 +++++++ .../test/src/kernels/concat_kernels_cpu.cc | 279 ++++++++ .../test/src/kernels/concat_kernels_gpu.cc | 292 +++++++++ .../test/src/kernels/conv_2d_kernels_gpu.cc | 334 ++++++++++ .../src/kernels/element_unary_kernels_gpu.cc | 196 ++++-- .../test/src/kernels/pool_2d_kernels_gpu.cc | 284 +++++++++ .../test/src/kernels/reshape_kernels_cpu.cc | 101 +++ .../test/src/kernels/reshape_kernels_gpu.cc | 112 ++++ .../test/src/kernels/split_kernels_cpu.cc | 293 +++++++++ .../test/src/kernels/split_kernels_gpu.cc | 306 +++++++++ .../test/src/kernels/transpose_kernels_cpu.cc | 157 +++++ .../test/src/kernels/transpose_kernels_gpu.cc | 168 +++++ .../test/src/kernels/upsample_kernels_cpu.cc | 179 ++++++ .../test/src/kernels/upsample_kernels_gpu.cc | 313 +++++++++ .../test/src/test_batch_norm_kernel.cc | 101 --- lib/kernels/test/src/test_concat_kernel.cc | 113 ---- lib/kernels/test/src/test_pool_2d_kernels.cc | 91 --- lib/kernels/test/src/test_reshape_kernel.cc | 49 -- lib/kernels/test/src/test_split_kernel.cc | 72 --- lib/kernels/test/src/test_transpose_kernel.cc | 56 -- .../src/task-spec/ops/impl/batch_matmul.cc | 4 + .../src/task-spec/ops/impl/batch_norm.cc | 106 ++-- .../src/task-spec/ops/impl/concat.cc | 49 +- .../src/task-spec/ops/impl/conv_2d.cc | 115 ++-- .../src/task-spec/ops/impl/pool_2d.cc | 105 ++- .../src/task-spec/ops/impl/reshape.cc | 25 +- lib/task-spec/src/task-spec/ops/impl/split.cc | 101 ++- .../src/task-spec/ops/impl/transpose.cc | 28 +- 87 files changed, 5832 insertions(+), 2417 deletions(-) create mode 100644 lib/kernels/test/src/kernels/batch_matmul_kernels_gpu.cc create mode 100644 lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc create mode 100644 lib/kernels/test/src/kernels/concat_kernels_cpu.cc create mode 100644 lib/kernels/test/src/kernels/concat_kernels_gpu.cc create mode 100644 lib/kernels/test/src/kernels/conv_2d_kernels_gpu.cc create mode 100644 lib/kernels/test/src/kernels/pool_2d_kernels_gpu.cc create mode 100644 lib/kernels/test/src/kernels/reshape_kernels_cpu.cc create mode 100644 lib/kernels/test/src/kernels/reshape_kernels_gpu.cc create mode 100644 lib/kernels/test/src/kernels/split_kernels_cpu.cc create mode 100644 lib/kernels/test/src/kernels/split_kernels_gpu.cc create mode 100644 lib/kernels/test/src/kernels/transpose_kernels_cpu.cc create mode 100644 lib/kernels/test/src/kernels/transpose_kernels_gpu.cc create mode 100644 lib/kernels/test/src/kernels/upsample_kernels_cpu.cc create mode 100644 lib/kernels/test/src/kernels/upsample_kernels_gpu.cc delete mode 100644 lib/kernels/test/src/test_batch_norm_kernel.cc delete mode 100644 lib/kernels/test/src/test_concat_kernel.cc delete mode 100644 lib/kernels/test/src/test_pool_2d_kernels.cc delete mode 100644 lib/kernels/test/src/test_reshape_kernel.cc delete mode 100644 lib/kernels/test/src/test_split_kernel.cc delete mode 100644 lib/kernels/test/src/test_transpose_kernel.cc diff --git a/lib/kernels/include/kernels/batch_matmul_kernels.h b/lib/kernels/include/kernels/batch_matmul_kernels.h index b2902e3bdf..b6a6b8e970 100644 --- a/lib/kernels/include/kernels/batch_matmul_kernels.h +++ b/lib/kernels/include/kernels/batch_matmul_kernels.h @@ -2,17 +2,20 @@ #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_MATMUL_KERNELS_H #include "kernels/accessor.h" +#include "kernels/device_handle_t.dtg.h" #include "kernels/device_stream_t.dtg.h" #include "op-attrs/ops/batch_matmul_attrs.dtg.h" namespace FlexFlow { void batch_matmul_forward_kernel(device_stream_t const &stream, + device_handle_t const &handle, GenericTensorAccessorR const &input_lhs, GenericTensorAccessorR const &input_rhs, GenericTensorAccessorW const &output); void batch_matmul_backward_kernel(device_stream_t const &stream, + device_handle_t const &handle, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input_lhs, diff --git a/lib/kernels/include/kernels/batch_matmul_kernels_gpu.h b/lib/kernels/include/kernels/batch_matmul_kernels_gpu.h index c8bb9432da..83c5390de4 100644 --- a/lib/kernels/include/kernels/batch_matmul_kernels_gpu.h +++ b/lib/kernels/include/kernels/batch_matmul_kernels_gpu.h @@ -2,17 +2,20 @@ #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_MATMUL_KERNELS_GPU_H #include "kernels/accessor.h" +#include "kernels/device.h" #include "op-attrs/ops/batch_matmul_attrs.dtg.h" namespace FlexFlow { void batch_matmul_gpu_forward_kernel(ffStream_t stream, + PerDeviceFFHandle const &handle, GenericTensorAccessorR const &input_lhs, GenericTensorAccessorR const &input_rhs, GenericTensorAccessorW const &output); void batch_matmul_gpu_backward_kernel( ffStream_t stream, + PerDeviceFFHandle const &handle, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input_lhs, diff --git a/lib/kernels/include/kernels/batch_norm_kernels.h b/lib/kernels/include/kernels/batch_norm_kernels.h index 47cb3d85a8..ee56bf7173 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels.h +++ b/lib/kernels/include/kernels/batch_norm_kernels.h @@ -1,47 +1,52 @@ #ifndef _FLEXFLOW_KERNELS_BATCH_NORM_KERNELS_H #define _FLEXFLOW_KERNELS_BATCH_NORM_KERNELS_H +#include "kernels/accessor.h" #include "kernels/allocation.h" #include "kernels/batch_norm_per_device_state.dtg.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/device_stream_t.dtg.h" -#include "kernels/ff_handle.h" +#include "op-attrs/ops/batch_norm_attrs.dtg.h" +#include "op-attrs/tensor_shape.dtg.h" +#include "pcg/device_type.dtg.h" -namespace FlexFlow::Kernels::BatchNorm { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - Allocator &allocator, - float *runningMean, - int output_n, - int output_c, - int output_h, - int output_w, - bool relu); - -void forward_kernel(device_stream_t const &stream, - BatchNormPerDeviceState const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr); - -void backward_kernel(device_stream_t const &stream, - BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements); - -void cleanup_kernel( + batch_norm_init_kernel(DeviceType device_type, + Allocator &allocator, + BatchNormAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); + +void batch_norm_forward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output); + +void batch_norm_backward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad); + +void batch_norm_cleanup_kernel( DeviceType device_type, Allocator &allocator, - std::optional const &per_device_state); + std::optional &per_device_state); + +} // namespace FlexFlow -} // namespace FlexFlow::Kernels::BatchNorm #endif diff --git a/lib/kernels/include/kernels/batch_norm_kernels_cpu.h b/lib/kernels/include/kernels/batch_norm_kernels_cpu.h index 8c564d6450..ecda6883a0 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels_cpu.h +++ b/lib/kernels/include/kernels/batch_norm_kernels_cpu.h @@ -1,28 +1,26 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_NORM_KERNELS_CPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_NORM_KERNELS_CPU_H -#include "kernels/allocation.h" -#include "kernels/batch_norm_per_device_state.dtg.h" -#include "kernels/device_stream_t.dtg.h" +#include "kernels/accessor.h" +#include "op-attrs/ops/batch_norm_attrs.dtg.h" -namespace FlexFlow::Kernels::BatchNorm { +namespace FlexFlow { -void cpu_forward_kernel(BatchNormPerDeviceState const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr); +void batch_norm_cpu_forward_kernel(BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output); -void cpu_backward_kernel(BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements); +void batch_norm_cpu_backward_kernel(BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad); -} // namespace FlexFlow::Kernels::BatchNorm +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/batch_norm_kernels_gpu.h b/lib/kernels/include/kernels/batch_norm_kernels_gpu.h index 41f9808bff..277c36a9ed 100644 --- a/lib/kernels/include/kernels/batch_norm_kernels_gpu.h +++ b/lib/kernels/include/kernels/batch_norm_kernels_gpu.h @@ -1,43 +1,46 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_NORM_KERNELS_GPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_NORM_KERNELS_GPU_H +#include "kernels/accessor.h" #include "kernels/allocation.h" #include "kernels/batch_norm_per_device_state.dtg.h" #include "kernels/device.h" -#include "kernels/ff_handle.h" - -namespace FlexFlow::Kernels::BatchNorm { - -BatchNormPerDeviceState gpu_init_kernel(PerDeviceFFHandle const &handle, - Allocator &allocator, - float *runningMean, - int output_n, - int output_c, - int output_h, - int output_w, - bool relu); - -void gpu_forward_kernel(ffStream_t stream, - BatchNormPerDeviceState const &per_device_statem, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr); - -void gpu_backward_kernel(ffStream_t stream, - BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements); - -void gpu_cleanup_kernel(Allocator &allocator, - BatchNormPerDeviceState &per_device_state); - -} // namespace FlexFlow::Kernels::BatchNorm +#include "op-attrs/ops/batch_norm_attrs.dtg.h" + +namespace FlexFlow { + +BatchNormPerDeviceState + batch_norm_gpu_init_kernel(Allocator &allocator, + BatchNormAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); + +void batch_norm_gpu_forward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + BatchNormPerDeviceState const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output); + +void batch_norm_gpu_backward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + BatchNormPerDeviceState const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad); + +void batch_norm_gpu_cleanup_kernel(Allocator &allocator, + BatchNormPerDeviceState &per_device_state); + +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml b/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml index bdf9e1ed51..f1548e4891 100644 --- a/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml +++ b/lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml @@ -5,13 +5,8 @@ features = [] includes = [ "kernels/device.h", - "kernels/ff_handle.h", ] -[[fields]] -name = "handle" -type = "::FlexFlow::PerDeviceFFHandle" - [[fields]] name = "inputTensor" type = "ffTensorDescriptor_t" @@ -24,10 +19,6 @@ type = "ffTensorDescriptor_t" name = "biasTensor" type = "ffTensorDescriptor_t" -[[fields]] -name = "actiDesc" -type = "ffActivationDescriptor_t" - [[fields]] name = "mode" type = "ffBatchNormMode_t" @@ -47,23 +38,3 @@ type = "float *" [[fields]] name = "saveVar" type = "float *" - -[[fields]] -name = "output_n" -type = "int" - -[[fields]] -name = "output_c" -type = "int" - -[[fields]] -name = "output_h" -type = "int" - -[[fields]] -name = "output_w" -type = "int" - -[[fields]] -name = "relu" -type = "bool" diff --git a/lib/kernels/include/kernels/concat_kernels.h b/lib/kernels/include/kernels/concat_kernels.h index 793bf52505..415ca25fbb 100644 --- a/lib/kernels/include/kernels/concat_kernels.h +++ b/lib/kernels/include/kernels/concat_kernels.h @@ -3,19 +3,23 @@ #include "kernels/accessor.h" #include "kernels/device_stream_t.dtg.h" +#include "op-attrs/ops/concat_attrs.dtg.h" -namespace FlexFlow::Kernels::Concat { +namespace FlexFlow { -void forward_kernel(device_stream_t const &stream, - GenericTensorAccessorW const &output, - std::vector const &inputs, - ff_dim_t axis); +void concat_forward_kernel(device_stream_t const &stream, + ConcatAttrs const &attrs, + std::vector const &inputs, + GenericTensorAccessorW const &output); -void backward_kernel(device_stream_t const &stream, - GenericTensorAccessorR const &output_grad, - std::vector const &input_grads, - ff_dim_t axis); +void concat_backward_kernel( + device_stream_t const &stream, + ConcatAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + std::vector const &inputs, + std::vector const &input_grads); -} // namespace FlexFlow::Kernels::Concat +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/concat_kernels_cpu.h b/lib/kernels/include/kernels/concat_kernels_cpu.h index 4a7f9fd3c8..596e15f07a 100644 --- a/lib/kernels/include/kernels/concat_kernels_cpu.h +++ b/lib/kernels/include/kernels/concat_kernels_cpu.h @@ -2,18 +2,22 @@ #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_CONCAT_KERNELS_CPU_H #include "kernels/accessor.h" -#include "kernels/device.h" +#include "op-attrs/ops/concat_attrs.dtg.h" -namespace FlexFlow::Kernels::Concat { +namespace FlexFlow { -void cpu_forward_kernel(GenericTensorAccessorW const &output, - std::vector const &inputs, - ff_dim_t axis); +void concat_cpu_forward_kernel( + ConcatAttrs const &attrs, + std::vector const &inputs, + GenericTensorAccessorW const &output); -void cpu_backward_kernel(GenericTensorAccessorR const &output_grad, - std::vector const &input_grads, - ff_dim_t axis); +void concat_cpu_backward_kernel( + ConcatAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + std::vector const &inputs, + std::vector const &input_grads); -} // namespace FlexFlow::Kernels::Concat +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/concat_kernels_gpu.h b/lib/kernels/include/kernels/concat_kernels_gpu.h index 3aaf3fbe2c..1bbdee5f87 100644 --- a/lib/kernels/include/kernels/concat_kernels_gpu.h +++ b/lib/kernels/include/kernels/concat_kernels_gpu.h @@ -3,19 +3,24 @@ #include "kernels/accessor.h" #include "kernels/device.h" +#include "op-attrs/ops/concat_attrs.dtg.h" -namespace FlexFlow::Kernels::Concat { +namespace FlexFlow { -void gpu_forward_kernel(ffStream_t stream, - GenericTensorAccessorW const &output, - std::vector const &inputs, - ff_dim_t axis); +void concat_gpu_forward_kernel( + ffStream_t stream, + ConcatAttrs const &attrs, + std::vector const &inputs, + GenericTensorAccessorW const &output); -void gpu_backward_kernel(ffStream_t stream, - GenericTensorAccessorR const &output_grad, - std::vector const &input_grads, - ff_dim_t axis); +void concat_gpu_backward_kernel( + ffStream_t stream, + ConcatAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + std::vector const &inputs, + std::vector const &input_grads); -} // namespace FlexFlow::Kernels::Concat +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/conv_2d_kernels.h b/lib/kernels/include/kernels/conv_2d_kernels.h index eb7cd7327a..1d429aadfe 100644 --- a/lib/kernels/include/kernels/conv_2d_kernels.h +++ b/lib/kernels/include/kernels/conv_2d_kernels.h @@ -5,50 +5,46 @@ #include "kernels/conv_2d_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/activation.dtg.h" +#include "op-attrs/ops/conv_2d_attrs.dtg.h" +#include "op-attrs/tensor_shape.dtg.h" +#include "pcg/device_type.dtg.h" -namespace FlexFlow::Kernels::Conv2D { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - std::optional activation, - int kernel_h, - int kernel_w, - int groups, - int padding_h, - int padding_w, - int stride_h, - int stride_w, - GenericTensorAccessorW const &input, - GenericTensorAccessorW const &output, - float const *filter_ptr, - float *filter_grad_ptr); - -void forward_kernel(device_stream_t const &stream, - std::optional const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *filter_ptr, - float const *bias_ptr, - std::optional activation); - -void backward_kernel( + conv_2d_init_kernel(DeviceType device_type, + device_handle_t const &handle, + Conv2DAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); + +void conv_2d_forward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &filter, + std::optional const &bias, + GenericTensorAccessorW const &output); + +void conv_2d_backward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *filter_ptr, - float *filter_grad_ptr, - float *bias_grad_ptr, - std::optional activation); - -void cleanup_kernel(DeviceType device_type, - std::optional &per_device_state); - -} // namespace FlexFlow::Kernels::Conv2D + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &filter, + GenericTensorAccessorW const &filter_grad, + std::optional const &bias_grad); + +void conv_2d_cleanup_kernel( + DeviceType device_type, + std::optional &per_device_state); + +} // namespace FlexFlow #endif // _FLEXFLOW_OPS_KERNELS_CONV_2D_KERNELS_H diff --git a/lib/kernels/include/kernels/conv_2d_kernels_cpu.h b/lib/kernels/include/kernels/conv_2d_kernels_cpu.h index 3a783a395f..b00ac4e2b9 100644 --- a/lib/kernels/include/kernels/conv_2d_kernels_cpu.h +++ b/lib/kernels/include/kernels/conv_2d_kernels_cpu.h @@ -1,26 +1,28 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_CONV_2D_KERNELS_CPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_CONV_2D_KERNELS_CPU_H -#include "op-attrs/activation.dtg.h" -#include +#include "kernels/accessor.h" +#include "op-attrs/ops/conv_2d_attrs.dtg.h" -namespace FlexFlow::Kernels::Conv2D { +namespace FlexFlow { -void cpu_forward_kernel(float const *input_ptr, - float *output_ptr, - float const *filter_ptr, - float const *bias_ptr, - std::optional const &activation); +void conv_2d_cpu_forward_kernel( + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &filter, + std::optional const &bias, + GenericTensorAccessorW const &output); -void cpu_backward_kernel(float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *filter_ptr, - float *filter_grad_ptr, - float *bias_grad_ptr, - std::optional const &activation); +void conv_2d_cpu_backward_kernel( + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &filter, + GenericTensorAccessorW const &filter_grad, + std::optional const &bias_grad); -} // namespace FlexFlow::Kernels::Conv2D +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/conv_2d_kernels_gpu.h b/lib/kernels/include/kernels/conv_2d_kernels_gpu.h index 9084838e9d..c06dcd3af0 100644 --- a/lib/kernels/include/kernels/conv_2d_kernels_gpu.h +++ b/lib/kernels/include/kernels/conv_2d_kernels_gpu.h @@ -1,44 +1,43 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_CONV_2D_KERNELS_GPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_CONV_2D_KERNELS_GPU_H -namespace FlexFlow::Kernels::Conv2D { - -Conv2DPerDeviceState - gpu_init_kernel(PerDeviceFFHandle const &handle, - std::optional const &activation, - int kernel_h, - int kernel_w, - int groups, - int padding_h, - int padding_w, - int stride_h, - int stride_w, - GenericTensorAccessorW const &input, - GenericTensorAccessorW const &output, - float const *filter_ptr, - float *filter_grad_ptr); - -void gpu_forward_kernel(ffStream_t stream, - Conv2DPerDeviceState const &m, - float const *input_ptr, - float *output_ptr, - float const *filter_ptr, - float const *bias_ptr, - std::optional activation); - -void gpu_backward_kernel(ffStream_t stream, - Conv2DPerDeviceState const &m, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *filter_ptr, - float *filter_grad_ptr, - float *bias_grad_ptr, - std::optional activation); - -void gpu_cleanup_kernel(Conv2DPerDeviceState &per_device_state); - -} // namespace FlexFlow::Kernels::Conv2D +#include "kernels/accessor.h" +#include "kernels/conv_2d_per_device_state.dtg.h" +#include "kernels/device.h" +#include "op-attrs/ops/conv_2d_attrs.dtg.h" + +namespace FlexFlow { + +Conv2DPerDeviceState conv_2d_gpu_init_kernel(PerDeviceFFHandle const &handle, + Conv2DAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); + +void conv_2d_gpu_forward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + Conv2DPerDeviceState const &per_device_state, + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &filter, + std::optional const &bias, + GenericTensorAccessorW const &output); + +void conv_2d_gpu_backward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + Conv2DPerDeviceState const &per_device_state, + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &filter, + GenericTensorAccessorW const &filter_grad, + std::optional const &bias_grad); + +void conv_2d_gpu_cleanup_kernel(Conv2DPerDeviceState &per_device_state); + +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/conv_2d_per_device_state.dtg.toml b/lib/kernels/include/kernels/conv_2d_per_device_state.dtg.toml index cdbb2fea38..c4001f016d 100644 --- a/lib/kernels/include/kernels/conv_2d_per_device_state.dtg.toml +++ b/lib/kernels/include/kernels/conv_2d_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" @@ -28,10 +23,6 @@ type = "ffTensorDescriptor_t" name = "filterDesc" type = "ffFilterDescriptor_t" -[[fields]] -name = "actiDesc" -type = "ffActivationDescriptor_t" - [[fields]] name = "convDesc" type = "ffConvolutionDescriptor_t" 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/include/kernels/pool_2d_kernels.h b/lib/kernels/include/kernels/pool_2d_kernels.h index c18ff92289..12261d9427 100644 --- a/lib/kernels/include/kernels/pool_2d_kernels.h +++ b/lib/kernels/include/kernels/pool_2d_kernels.h @@ -1,52 +1,44 @@ #ifndef _FLEXFLOW_OPS_KERNELS_POOL_2D_KERNELS_H #define _FLEXFLOW_OPS_KERNELS_POOL_2D_KERNELS_H +#include "kernels/accessor.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/device_stream_t.dtg.h" -#include "kernels/ff_handle.h" #include "kernels/pool_2d_per_device_state.dtg.h" -#include "op-attrs/activation.dtg.h" -#include "op-attrs/ops/pool_2d.h" +#include "op-attrs/ops/pool_2d_attrs.dtg.h" +#include "op-attrs/tensor_shape.dtg.h" #include "pcg/device_type.dtg.h" -namespace FlexFlow::Kernels::Pool2D { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - std::optional activation, - int input_w, - int input_h, - int input_c, - int input_n, - int output_w, - int output_h, - int output_c, - int output_n, - int pad_h, - int pad_w, - int kernel_h, - int kernel_w, - int stride_h, - int stride_w, - PoolOp pool_type); - -void forward_kernel(device_stream_t const &stream, - std::optional const &per_device_state, - void const *input_ptr, - void *output_ptr); - -void backward_kernel( + pool_2d_init_kernel(DeviceType device_type, + Pool2DAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); + +void pool_2d_forward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + Pool2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); + +void pool_2d_backward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - void const *output_ptr, - void const *output_grad_ptr, - void const *input_ptr, - void *input_grad_ptr); + Pool2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -void cleanup_kernel(DeviceType device_type, - std::optional &per_device_state); +void pool_2d_cleanup_kernel( + DeviceType device_type, + std::optional &per_device_state); -} // namespace FlexFlow::Kernels::Pool2D +} // namespace FlexFlow #endif // _FLEXFLOW_OPS_KERNELS_POOL_2D_KERNELS_H diff --git a/lib/kernels/include/kernels/pool_2d_kernels_cpu.h b/lib/kernels/include/kernels/pool_2d_kernels_cpu.h index aa13e913b6..22e12050c4 100644 --- a/lib/kernels/include/kernels/pool_2d_kernels_cpu.h +++ b/lib/kernels/include/kernels/pool_2d_kernels_cpu.h @@ -1,15 +1,21 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_POOL_2D_KERNELS_CPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_POOL_2D_KERNELS_CPU_H -namespace FlexFlow::Kernels::Pool2D { +#include "kernels/accessor.h" +#include "op-attrs/ops/pool_2d_attrs.dtg.h" -void cpu_forward_kernel(void const *input_ptr, void *output_ptr); +namespace FlexFlow { -void cpu_backward_kernel(void const *output_ptr, - void const *output_grad_ptr, - void const *input_ptr, - void *input_grad_ptr); +void pool_2d_cpu_forward_kernel(Pool2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); -} // namespace FlexFlow::Kernels::Pool2D +void pool_2d_cpu_backward_kernel(Pool2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); + +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/pool_2d_kernels_gpu.h b/lib/kernels/include/kernels/pool_2d_kernels_gpu.h index 8a1499e97e..7d4c51aaf5 100644 --- a/lib/kernels/include/kernels/pool_2d_kernels_gpu.h +++ b/lib/kernels/include/kernels/pool_2d_kernels_gpu.h @@ -1,46 +1,35 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_POOL_2D_KERNELS_GPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_POOL_2D_KERNELS_GPU_H +#include "kernels/accessor.h" #include "kernels/device.h" -#include "kernels/ff_handle.h" #include "kernels/pool_2d_per_device_state.dtg.h" -#include "op-attrs/activation.dtg.h" -#include "op-attrs/ops/pool_2d.h" - -namespace FlexFlow::Kernels::Pool2D { - -Pool2DPerDeviceState gpu_init_kernel(PerDeviceFFHandle handle, - std::optional activation, - int input_w, - int input_h, - int input_c, - int input_n, - int output_w, - int output_h, - int output_c, - int output_n, - int pad_h, - int pad_w, - int kernel_h, - int kernel_w, - int stride_h, - int stride_w, - PoolOp pool_type); - -void gpu_forward_kernel(ffStream_t stream, - Pool2DPerDeviceState const &per_device_state, - void const *input_ptr, - void *output_ptr); - -void gpu_backward_kernel(ffStream_t stream, - Pool2DPerDeviceState const &per_device_state, - void const *output_ptr, - void const *output_grad_ptr, - void const *input_ptr, - void *input_grad_ptr); - -void gpu_cleanup_kernel(Pool2DPerDeviceState &per_device_state); - -} // namespace FlexFlow::Kernels::Pool2D - -#endif +#include "op-attrs/ops/pool_2d_attrs.dtg.h" + +namespace FlexFlow { + +Pool2DPerDeviceState pool_2d_gpu_init_kernel(Pool2DAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); + +void pool_2d_gpu_forward_kernel(ffStream_t stream, + PerDeviceFFHandle const &handle, + Pool2DPerDeviceState const &per_device_state, + Pool2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); + +void pool_2d_gpu_backward_kernel(ffStream_t stream, + PerDeviceFFHandle const &handle, + Pool2DPerDeviceState const &per_device_state, + Pool2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); + +void pool_2d_gpu_cleanup_kernel(Pool2DPerDeviceState &per_device_state); + +} // namespace FlexFlow + +#endif // _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_POOL_2D_KERNELS_GPU_H diff --git a/lib/kernels/include/kernels/pool_2d_per_device_state.dtg.toml b/lib/kernels/include/kernels/pool_2d_per_device_state.dtg.toml index afc8bad21a..57119a758b 100644 --- a/lib/kernels/include/kernels/pool_2d_per_device_state.dtg.toml +++ b/lib/kernels/include/kernels/pool_2d_per_device_state.dtg.toml @@ -4,14 +4,9 @@ type = "struct" features = [] includes = [ - "kernels/ff_handle.h", "kernels/device.h", ] -[[fields]] -name = "handle" -type = "::FlexFlow::PerDeviceFFHandle" - [[fields]] name = "inputTensor" type = "ffTensorDescriptor_t" @@ -20,14 +15,6 @@ type = "ffTensorDescriptor_t" name = "outputTensor" type = "ffTensorDescriptor_t" -[[fields]] -name = "actiDesc" -type = "ffActivationDescriptor_t" - [[fields]] name = "poolDesc" type = "ffPoolingDescriptor_t" - -[[fields]] -name = "relu" -type = "bool" diff --git a/lib/kernels/include/kernels/reshape_kernels.h b/lib/kernels/include/kernels/reshape_kernels.h index 310b349473..bab512a273 100644 --- a/lib/kernels/include/kernels/reshape_kernels.h +++ b/lib/kernels/include/kernels/reshape_kernels.h @@ -4,16 +4,18 @@ #include "kernels/accessor.h" #include "kernels/device_stream_t.dtg.h" -namespace FlexFlow::Kernels::Reshape { +namespace FlexFlow { -void forward_kernel(device_stream_t const &stream, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output); +void reshape_forward_kernel(device_stream_t const &stream, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); -void backward_kernel(device_stream_t const &stream, - GenericTensorAccessorR const &output, - GenericTensorAccessorW const &input); +void reshape_backward_kernel(device_stream_t const &stream, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -} // namespace FlexFlow::Kernels::Reshape +} // namespace FlexFlow #endif // _FLEXFLOW_OPS_KERNELS_RESHAPE_KERNELS_H diff --git a/lib/kernels/include/kernels/reshape_kernels_cpu.h b/lib/kernels/include/kernels/reshape_kernels_cpu.h index a81ea639f6..a842fe1362 100644 --- a/lib/kernels/include/kernels/reshape_kernels_cpu.h +++ b/lib/kernels/include/kernels/reshape_kernels_cpu.h @@ -3,14 +3,16 @@ #include "kernels/accessor.h" -namespace FlexFlow::Kernels::Reshape { +namespace FlexFlow { -void cpu_forward_kernel(GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output); +void reshape_cpu_forward_kernel(GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); -void cpu_backward_kernel(GenericTensorAccessorR const &output, - GenericTensorAccessorW const &input); +void reshape_cpu_backward_kernel(GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -} // namespace FlexFlow::Kernels::Reshape +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/reshape_kernels_gpu.h b/lib/kernels/include/kernels/reshape_kernels_gpu.h index 1454ce56ee..64e038dfa6 100644 --- a/lib/kernels/include/kernels/reshape_kernels_gpu.h +++ b/lib/kernels/include/kernels/reshape_kernels_gpu.h @@ -4,16 +4,18 @@ #include "kernels/accessor.h" #include "kernels/device.h" -namespace FlexFlow::Kernels::Reshape { +namespace FlexFlow { -void gpu_forward_kernel(ffStream_t stream, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output); +void reshape_gpu_forward_kernel(ffStream_t stream, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); -void gpu_backward_kernel(ffStream_t stream, - GenericTensorAccessorR const &output, - GenericTensorAccessorW const &input); +void reshape_gpu_backward_kernel(ffStream_t stream, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -} // namespace FlexFlow::Kernels::Reshape +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/split_kernels.h b/lib/kernels/include/kernels/split_kernels.h index 6c3d576f29..ca46cecbb8 100644 --- a/lib/kernels/include/kernels/split_kernels.h +++ b/lib/kernels/include/kernels/split_kernels.h @@ -1,26 +1,25 @@ #ifndef _FLEXFLOW_OPS_KERNELS_SPLIT_KERNELS_H #define _FLEXFLOW_OPS_KERNELS_SPLIT_KERNELS_H +#include "kernels/accessor.h" #include "kernels/device_stream_t.dtg.h" +#include "op-attrs/ops/split_attrs.dtg.h" -namespace FlexFlow::Kernels::Split { +namespace FlexFlow { -void forward_kernel(device_stream_t const &stream, - float **out_ptrs, - float const *in_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs); +void split_forward_kernel(device_stream_t const &stream, + SplitAttrs const &attrs, + GenericTensorAccessorR const &input, + std::vector const &outputs); -void backward_kernel(device_stream_t const &stream, - float *in_grad_ptr, - float const **out_grad_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs); +void split_backward_kernel( + device_stream_t const &stream, + SplitAttrs const &attrs, + std::vector const &outputs, + std::vector const &output_grads, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -} // namespace FlexFlow::Kernels::Split +} // namespace FlexFlow #endif // _FLEXFLOW_OPS_KERNELS_SPLIT_KERNELS_H diff --git a/lib/kernels/include/kernels/split_kernels_cpu.h b/lib/kernels/include/kernels/split_kernels_cpu.h index 7f50804dff..2f5165c1e1 100644 --- a/lib/kernels/include/kernels/split_kernels_cpu.h +++ b/lib/kernels/include/kernels/split_kernels_cpu.h @@ -1,22 +1,23 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_SPLIT_KERNELS_CPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_SPLIT_KERNELS_CPU_H -namespace FlexFlow::Kernels::Split { +#include "kernels/accessor.h" +#include "op-attrs/ops/split_attrs.dtg.h" -void cpu_forward_kernel(float **out_ptrs, - float const *in_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs); +namespace FlexFlow { -void cpu_backward_kernel(float *in_grad_ptr, - float const **out_grad_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs); +void split_cpu_forward_kernel( + SplitAttrs const &attrs, + GenericTensorAccessorR const &input, + std::vector const &outputs); -} // namespace FlexFlow::Kernels::Split +void split_cpu_backward_kernel( + SplitAttrs const &attrs, + std::vector const &outputs, + std::vector const &output_grads, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); + +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/split_kernels_gpu.h b/lib/kernels/include/kernels/split_kernels_gpu.h index e6bfc5454c..def1468ace 100644 --- a/lib/kernels/include/kernels/split_kernels_gpu.h +++ b/lib/kernels/include/kernels/split_kernels_gpu.h @@ -1,26 +1,26 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_SPLIT_KERNELS_GPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_SPLIT_KERNELS_GPU_H +#include "kernels/accessor.h" #include "kernels/device.h" +#include "op-attrs/ops/split_attrs.dtg.h" -namespace FlexFlow::Kernels::Split { +namespace FlexFlow { -void gpu_forward_kernel(ffStream_t stream, - float **out_ptrs, - float const *in_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs); +void split_gpu_forward_kernel( + ffStream_t stream, + SplitAttrs const &attrs, + GenericTensorAccessorR const &input, + std::vector const &outputs); -void gpu_backward_kernel(ffStream_t stream, - float *in_grad_ptr, - float const **out_grad_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs); +void split_gpu_backward_kernel( + ffStream_t stream, + SplitAttrs const &attrs, + std::vector const &outputs, + std::vector const &output_grads, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -} // namespace FlexFlow::Kernels::Split +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/tensor_accessor_binary_ops.h b/lib/kernels/include/kernels/tensor_accessor_binary_ops.h index f5cf2233ba..72039ddc91 100644 --- a/lib/kernels/include/kernels/tensor_accessor_binary_ops.h +++ b/lib/kernels/include/kernels/tensor_accessor_binary_ops.h @@ -46,9 +46,10 @@ void tensor_accessor_matmul_to(GenericTensorAccessorR const &lhs, GenericTensorAccessorR const &rhs, GenericTensorAccessorW const &output); -void tensor_accessor_batch_matmul(GenericTensorAccessorR const &lhs, - GenericTensorAccessorR const &rhs, - GenericTensorAccessorW const &output); +GenericTensorAccessorW + tensor_accessor_batch_matmul(GenericTensorAccessorR const &lhs, + GenericTensorAccessorR const &rhs, + Allocator &output_allocator); void tensor_accessor_batch_matmul_to(GenericTensorAccessorR const &lhs, GenericTensorAccessorR const &rhs, diff --git a/lib/kernels/include/kernels/transpose_kernels.h b/lib/kernels/include/kernels/transpose_kernels.h index 96b0a9c4aa..c0b2e6da0b 100644 --- a/lib/kernels/include/kernels/transpose_kernels.h +++ b/lib/kernels/include/kernels/transpose_kernels.h @@ -5,18 +5,20 @@ #include "kernels/device_stream_t.dtg.h" #include "op-attrs/ops/transpose_attrs.dtg.h" -namespace FlexFlow::Kernels::Transpose { +namespace FlexFlow { -void forward_kernel(device_stream_t const &stream, - TransposeAttrs const &attrs, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output); +void transpose_forward_kernel(device_stream_t const &stream, + TransposeAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); -void backward_kernel(device_stream_t const &stream, - TransposeAttrs const &attrs, - GenericTensorAccessorR const &out_grad, - GenericTensorAccessorW const &in_grad); +void transpose_backward_kernel(device_stream_t const &stream, + TransposeAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -} // namespace FlexFlow::Kernels::Transpose +} // namespace FlexFlow #endif // _FLEXFLOW_OPS_KERNELS_TRANSPOSE_KERNELS_H diff --git a/lib/kernels/include/kernels/transpose_kernels_cpu.h b/lib/kernels/include/kernels/transpose_kernels_cpu.h index dd8963d5e4..f4d53abd83 100644 --- a/lib/kernels/include/kernels/transpose_kernels_cpu.h +++ b/lib/kernels/include/kernels/transpose_kernels_cpu.h @@ -4,16 +4,18 @@ #include "kernels/accessor.h" #include "op-attrs/ops/transpose_attrs.dtg.h" -namespace FlexFlow::Kernels::Transpose { +namespace FlexFlow { -void cpu_forward_kernel(TransposeAttrs const &attrs, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output); +void transpose_cpu_forward_kernel(TransposeAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); -void cpu_backward_kernel(TransposeAttrs const &attrs, - GenericTensorAccessorR const &out_grad, - GenericTensorAccessorW const &in_grad); +void transpose_cpu_backward_kernel(TransposeAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -} // namespace FlexFlow::Kernels::Transpose +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/transpose_kernels_gpu.h b/lib/kernels/include/kernels/transpose_kernels_gpu.h index 67f6e48665..8c359db0c6 100644 --- a/lib/kernels/include/kernels/transpose_kernels_gpu.h +++ b/lib/kernels/include/kernels/transpose_kernels_gpu.h @@ -5,18 +5,20 @@ #include "kernels/device.h" #include "op-attrs/ops/transpose_attrs.dtg.h" -namespace FlexFlow::Kernels::Transpose { +namespace FlexFlow { -void gpu_forward_kernel(ffStream_t stream, - TransposeAttrs const &attrs, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output); +void transpose_gpu_forward_kernel(ffStream_t stream, + TransposeAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output); -void gpu_backward_kernel(ffStream_t stream, - TransposeAttrs const &attrs, - GenericTensorAccessorR const &out_grad, - GenericTensorAccessorW const &in_grad); +void transpose_gpu_backward_kernel(ffStream_t stream, + TransposeAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad); -} // namespace FlexFlow::Kernels::Transpose +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/upsample_kernels_cpu.h b/lib/kernels/include/kernels/upsample_kernels_cpu.h index 4d51691c8c..4eace8489a 100644 --- a/lib/kernels/include/kernels/upsample_kernels_cpu.h +++ b/lib/kernels/include/kernels/upsample_kernels_cpu.h @@ -11,7 +11,9 @@ void upsample_cpu_forward_kernel(UpsampleAttrs const &attrs, GenericTensorAccessorW const &output); void upsample_cpu_backward_kernel(UpsampleAttrs const &attrs, + GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad); } // namespace FlexFlow diff --git a/lib/kernels/include/kernels/upsample_kernels_gpu.h b/lib/kernels/include/kernels/upsample_kernels_gpu.h index ac6d9c30f6..c06a8ce57c 100644 --- a/lib/kernels/include/kernels/upsample_kernels_gpu.h +++ b/lib/kernels/include/kernels/upsample_kernels_gpu.h @@ -3,7 +3,6 @@ #include "kernels/accessor.h" #include "kernels/device.h" -#include "kernels/ff_handle.h" #include "op-attrs/ops/upsample_attrs.dtg.h" namespace FlexFlow { diff --git a/lib/kernels/src/cuda/ops/batch_matmul_kernels.cu b/lib/kernels/src/cuda/ops/batch_matmul_kernels.cu index f8524c604d..e8e4363dd8 100644 --- a/lib/kernels/src/cuda/ops/batch_matmul_kernels.cu +++ b/lib/kernels/src/cuda/ops/batch_matmul_kernels.cu @@ -1,24 +1,157 @@ #include "internal/device.h" -#include "kernels/upsample_kernels_gpu.h" +#include "kernels/batch_matmul_kernels_gpu.h" +#include "op-attrs/tensor_dims.h" +#include "utils/containers/require_same.h" namespace FlexFlow { -void batch_matmul_gpu_forward_kernel(ffStream_t stream, +// The dimensions of a batched matmul `output[b] = input_lhs[b] * input_rhs[b]`, +// where input_lhs is (batch_size, lhs_rows, inner), input_rhs is (batch_size, +// inner, rhs_cols) and output is (batch_size, lhs_rows, rhs_cols). +struct BatchMatmulDims { + int batch_size; + int lhs_rows; + int inner; + int rhs_cols; +}; + +static BatchMatmulDims get_batch_matmul_dims(TensorShape const &lhs_shape, + TensorShape const &rhs_shape, + TensorShape const &output_shape) { + ASSERT(get_num_dims(lhs_shape.dims) == num_tensor_dims_t{3_n}, + "BatchMatmul expects 3-dimensional tensors", + lhs_shape); + + require_same(lhs_shape.data_type, rhs_shape.data_type); + ASSERT(require_same(lhs_shape.data_type, output_shape.data_type) == + DataType::FLOAT, + "BatchMatmul currently only supports data_type = FLOAT. " + "If you need this feature, please create an issue.", + lhs_shape.data_type); + + positive_int batch_size = + require_same(dim_at_idx(lhs_shape.dims, ff_dim_t{0_n}), + dim_at_idx(rhs_shape.dims, ff_dim_t{0_n}), + dim_at_idx(output_shape.dims, ff_dim_t{0_n})); + + positive_int lhs_rows = + require_same(dim_at_idx(lhs_shape.dims, ff_dim_t{1_n}), + dim_at_idx(output_shape.dims, ff_dim_t{1_n})); + + positive_int inner = require_same(dim_at_idx(lhs_shape.dims, ff_dim_t{2_n}), + dim_at_idx(rhs_shape.dims, ff_dim_t{1_n})); + + positive_int rhs_cols = + require_same(dim_at_idx(rhs_shape.dims, ff_dim_t{2_n}), + dim_at_idx(output_shape.dims, ff_dim_t{2_n})); + + return BatchMatmulDims{ + /*batch_size=*/batch_size.int_from_positive_int(), + /*lhs_rows=*/lhs_rows.int_from_positive_int(), + /*inner=*/inner.int_from_positive_int(), + /*rhs_cols=*/rhs_cols.int_from_positive_int(), + }; +} + +void batch_matmul_gpu_forward_kernel(cudaStream_t stream, + PerDeviceFFHandle const &handle, GenericTensorAccessorR const &input_lhs, GenericTensorAccessorR const &input_rhs, GenericTensorAccessorW const &output) { - NOT_IMPLEMENTED(); + BatchMatmulDims dims = + get_batch_matmul_dims(input_lhs.shape, input_rhs.shape, output.shape); + + checkCUBLAS(cublasSetStream(handle.blas, stream)); + + // Our tensors are row-major but cuBLAS is column-major, so we compute + // `output^T = input_rhs^T * input_lhs^T` instead. A row-major (r, c) matrix + // is bit-for-bit a column-major (c, r) matrix, so this needs no transposes + // and no data movement: it is just a matter of swapping the operands. + float alpha = 1.0f, beta = 0.0f; + checkCUBLAS(cublasSgemmStridedBatched(handle.blas, + CUBLAS_OP_N, + CUBLAS_OP_N, + /*m=*/dims.rhs_cols, + /*n=*/dims.lhs_rows, + /*k=*/dims.inner, + &alpha, + /*A=*/input_rhs.get_float_ptr(), + /*lda=*/dims.rhs_cols, + /*strideA=*/dims.inner * dims.rhs_cols, + /*B=*/input_lhs.get_float_ptr(), + /*ldb=*/dims.inner, + /*strideB=*/dims.lhs_rows * dims.inner, + &beta, + /*C=*/output.get_float_ptr(), + /*ldc=*/dims.rhs_cols, + /*strideC=*/ + dims.lhs_rows * dims.rhs_cols, + /*batchCount=*/dims.batch_size)); } void batch_matmul_gpu_backward_kernel( - ffStream_t stream, + cudaStream_t stream, + PerDeviceFFHandle const &handle, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input_lhs, GenericTensorAccessorW const &input_lhs_grad, GenericTensorAccessorR const &input_rhs, GenericTensorAccessorW const &input_rhs_grad) { - NOT_IMPLEMENTED(); + require_same(input_lhs.shape, input_lhs_grad.shape); + require_same(input_rhs.shape, input_rhs_grad.shape); + require_same(output.shape, output_grad.shape); + + BatchMatmulDims dims = get_batch_matmul_dims( + input_lhs.shape, input_rhs.shape, output_grad.shape); + + checkCUBLAS(cublasSetStream(handle.blas, stream)); + + // NOTE: beta is 0 so that the gradients are overwritten rather than + // accumulated into, matching batch_matmul_cpu_backward_kernel. + float alpha = 1.0f, beta = 0.0f; + + // input_lhs_grad = output_grad * input_rhs^T + checkCUBLAS(cublasSgemmStridedBatched(handle.blas, + CUBLAS_OP_T, + CUBLAS_OP_N, + /*m=*/dims.inner, + /*n=*/dims.lhs_rows, + /*k=*/dims.rhs_cols, + &alpha, + /*A=*/input_rhs.get_float_ptr(), + /*lda=*/dims.rhs_cols, + /*strideA=*/dims.inner * dims.rhs_cols, + /*B=*/output_grad.get_float_ptr(), + /*ldb=*/dims.rhs_cols, + /*strideB=*/ + dims.lhs_rows * dims.rhs_cols, + &beta, + /*C=*/input_lhs_grad.get_float_ptr(), + /*ldc=*/dims.inner, + /*strideC=*/dims.lhs_rows * dims.inner, + /*batchCount=*/dims.batch_size)); + + // input_rhs_grad = input_lhs^T * output_grad + checkCUBLAS(cublasSgemmStridedBatched(handle.blas, + CUBLAS_OP_N, + CUBLAS_OP_T, + /*m=*/dims.rhs_cols, + /*n=*/dims.inner, + /*k=*/dims.lhs_rows, + &alpha, + /*A=*/output_grad.get_float_ptr(), + /*lda=*/dims.rhs_cols, + /*strideA=*/ + dims.lhs_rows * dims.rhs_cols, + /*B=*/input_lhs.get_float_ptr(), + /*ldb=*/dims.inner, + /*strideB=*/dims.lhs_rows * dims.inner, + &beta, + /*C=*/input_rhs_grad.get_float_ptr(), + /*ldc=*/dims.rhs_cols, + /*strideC=*/dims.inner * dims.rhs_cols, + /*batchCount=*/dims.batch_size)); } } // namespace FlexFlow diff --git a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu index 26234fd6e0..3e4514194c 100644 --- a/lib/kernels/src/cuda/ops/batch_norm_kernels.cu +++ b/lib/kernels/src/cuda/ops/batch_norm_kernels.cu @@ -14,170 +14,186 @@ */ #include "internal/device.h" -#include "kernels/allocation.h" -#include "kernels/batch_norm_kernels.h" -#include "kernels/ff_handle.h" -#include "utils/integer_conversions.h" +#include "kernels/batch_norm_kernels_gpu.h" +#include "op-attrs/tensor_dims.h" +#include "utils/containers/require_same.h" +#include namespace FlexFlow { -namespace Kernels { -namespace BatchNorm { - -void gpu_forward_kernel(cudaStream_t stream, - BatchNormPerDeviceState const &m, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr) { - checkCUDNN(cudnnSetStream(m.handle.dnn, stream)); - - float alpha = 1.0f, beta = 0.0f; - checkCUDNN(cudnnBatchNormalizationForwardTraining(m.handle.dnn, - m.mode, - &alpha, - &beta, - m.inputTensor, - input_ptr, - m.outputTensor, - output_ptr, - m.biasTensor, - scale_ptr, - bias_ptr, - 1.0, - m.runningMean, - m.runningVar, - CUDNN_BN_MIN_EPSILON, - m.saveMean, - m.saveVar)); -} -void gpu_backward_kernel(cudaStream_t stream, - BatchNormPerDeviceState const &m, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements) { - checkCUDNN(cudnnSetStream(m.handle.dnn, stream)); - - float alpha = 1.0f; - if (m.relu) { - reluBackward<<>>( - output_grad_ptr, output_ptr, numElements); - } - checkCUDNN(cudnnBatchNormalizationBackward(m.handle.dnn, - m.mode, - &alpha, - &alpha, - &alpha, - &alpha, - m.inputTensor, - input_ptr, - m.outputTensor, - output_grad_ptr, - m.inputTensor, - input_grad_ptr, - m.biasTensor, - scale_ptr, - scale_grad_ptr, - bias_grad_ptr, - CUDNN_BN_MIN_EPSILON, - m.saveMean, - m.saveVar)); +static positive_int get_num_channels(TensorShape const &shape) { + return dim_at_idx(shape.dims, ff_dim_t{1_n}); } -BatchNormPerDeviceState gpu_init_kernel(PerDeviceFFHandle const &handle, - Allocator &allocator, - float *runningMean, - int output_n, - int output_c, - int output_h, - int output_w, - bool relu) { +BatchNormPerDeviceState + batch_norm_gpu_init_kernel(Allocator &allocator, + BatchNormAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { + // Applying an activation as part of BatchNorm is not currently implemented + // (the forward kernel used to silently ignore it while the backward kernel + // applied it). If you need it, please create an issue. + ASSERT(!attrs.relu, + "BatchNorm does not currently support a fused relu activation"); + + ASSERT(attrs.affine, + "BatchNorm currently only supports attrs.affine = true. " + "If you need this feature, please create an issue."); + + TensorShape shape = require_same(input_shape, output_shape); + + ASSERT(get_num_dims(shape.dims) == num_tensor_dims_t{4_n}, + "BatchNorm currently only supports 4-dimensional (i.e., NCHW) " + "tensors. If you need this feature, please create an issue.", + shape); + + ASSERT(attrs.eps >= CUDNN_BN_MIN_EPSILON, + "cuDNN requires BatchNorm eps to be at least CUDNN_BN_MIN_EPSILON", + attrs.eps, + CUDNN_BN_MIN_EPSILON); + + int num_channels = get_num_channels(shape).int_from_positive_int(); + ffTensorDescriptor_t inputTensor; ffTensorDescriptor_t outputTensor; ffTensorDescriptor_t biasTensor; - ffActivationDescriptor_t actiDesc; - ffBatchNormMode_t mode; + checkCUDNN(cudnnCreateTensorDescriptor(&inputTensor)); - checkCUDNN(cudnnCreateTensorDescriptor(&biasTensor)); checkCUDNN(cudnnCreateTensorDescriptor(&outputTensor)); - mode = CUDNN_BATCHNORM_SPATIAL; + checkCUDNN(cudnnCreateTensorDescriptor(&biasTensor)); + + ffBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL; #if CUDNN_VERSION >= 7000 mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT; #endif - checkCUDNN(cudnnSetTensor4dDescriptor(inputTensor, - CUDNN_TENSOR_NCHW, - CUDNN_DATA_FLOAT, - output_n, - output_c, - output_h, - output_w)); - checkCUDNN(cudnnSetTensor4dDescriptor(outputTensor, + + checkCUDNN(cudnnSetTensorDescriptorFromTensorShape(inputTensor, input_shape)); + checkCUDNN( + cudnnSetTensorDescriptorFromTensorShape(outputTensor, output_shape)); + checkCUDNN(cudnnSetTensor4dDescriptor(biasTensor, CUDNN_TENSOR_NCHW, - CUDNN_DATA_FLOAT, - output_n, - output_c, - output_h, - output_w)); - checkCUDNN(cudnnSetTensor4dDescriptor( - biasTensor, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, output_c, 1, 1)); - // allocate memory for runningMean, runningVar, saveMean, saveVar - size_t totalSize = sizeof(float) * output_c * 4; - runningMean = (float *)allocator.allocate(totalSize); - float *runningVar = (float *)runningMean + output_c; - float *saveMean = (float *)runningVar + output_c; - float *saveVar = (float *)saveMean + output_c; - cudaStream_t stream; - checkCUDA(get_legion_stream(&stream)); - - assign_kernel<<>>( - runningMean, size_t_from_int(output_c), 0.0f); - assign_kernel<<>>( - runningVar, size_t_from_int(output_c), 0.0f); - - if (relu) { - checkCUDNN(cudnnCreateActivationDescriptor(&actiDesc)); - checkCUDNN(cudnnSetActivationDescriptor( - actiDesc, CUDNN_ACTIVATION_RELU, CUDNN_PROPAGATE_NAN, 0.0)); - } - - BatchNormPerDeviceState per_device_state = BatchNormPerDeviceState{ - handle, - inputTensor, - outputTensor, - biasTensor, - actiDesc, - mode, - runningMean, - runningVar, - saveMean, - saveVar, - output_n, - output_c, - output_h, - output_w, - relu, + ff_to_cudnn_datatype(shape.data_type), + /*n=*/1, + /*c=*/num_channels, + /*h=*/1, + /*w=*/1)); + + // Allocate memory for runningMean, runningVar, saveMean and saveVar as a + // single contiguous block (deallocated by batch_norm_gpu_cleanup_kernel). + float *runningMean = static_cast( + allocator.allocate(sizeof(float) * num_channels * 4)); + float *runningVar = runningMean + num_channels; + float *saveMean = runningVar + num_channels; + float *saveVar = saveMean + num_channels; + + // Match the PyTorch initialization of running_mean = 0 and running_var = 1. + std::vector initial_running_stats(num_channels * 2); + std::fill(initial_running_stats.begin(), + initial_running_stats.begin() + num_channels, + 0.0f); + std::fill(initial_running_stats.begin() + num_channels, + initial_running_stats.end(), + 1.0f); + checkCUDA(cudaMemcpy(runningMean, + initial_running_stats.data(), + sizeof(float) * num_channels * 2, + cudaMemcpyHostToDevice)); + + return BatchNormPerDeviceState{ + /*inputTensor=*/inputTensor, + /*outputTensor=*/outputTensor, + /*biasTensor=*/biasTensor, + /*mode=*/mode, + /*runningMean=*/runningMean, + /*runningVar=*/runningVar, + /*saveMean=*/saveMean, + /*saveVar=*/saveVar, }; +} + +void batch_norm_gpu_forward_kernel( + cudaStream_t stream, + PerDeviceFFHandle const &handle, + BatchNormPerDeviceState const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output) { + checkCUDNN(cudnnSetStream(handle.dnn, stream)); - checkCUDA(cudaStreamDestroy(stream)); - return per_device_state; + // NOTE: attrs.momentum = std::nullopt means "use a cumulative moving + // average", which cuDNN cannot express, so we fall back to fully replacing + // the running statistics on each call. The running statistics are currently + // never read back, so this only matters once inference mode is supported. + double exponential_average_factor = attrs.momentum.value_or(1.0); + + float alpha = 1.0f, beta_coeff = 0.0f; + checkCUDNN( + cudnnBatchNormalizationForwardTraining(handle.dnn, + per_device_state.mode, + &alpha, + &beta_coeff, + per_device_state.inputTensor, + input.ptr, + per_device_state.outputTensor, + output.ptr, + per_device_state.biasTensor, + gamma.ptr, + beta.ptr, + exponential_average_factor, + per_device_state.runningMean, + per_device_state.runningVar, + attrs.eps, + per_device_state.saveMean, + per_device_state.saveVar)); } -void gpu_cleanup_kernel(Allocator &allocator, - BatchNormPerDeviceState &per_device_state) { +void batch_norm_gpu_backward_kernel( + cudaStream_t stream, + PerDeviceFFHandle const &handle, + BatchNormPerDeviceState const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad) { + checkCUDNN(cudnnSetStream(handle.dnn, stream)); + + // NOTE: the beta coefficients are 1.0 so that the gradients are accumulated + // into rather than overwritten + float alpha_data = 1.0f, beta_data = 1.0f; + float alpha_param = 1.0f, beta_param = 1.0f; + checkCUDNN(cudnnBatchNormalizationBackward(handle.dnn, + per_device_state.mode, + &alpha_data, + &beta_data, + &alpha_param, + &beta_param, + per_device_state.inputTensor, + input.ptr, + per_device_state.outputTensor, + output_grad.ptr, + per_device_state.inputTensor, + input_grad.ptr, + per_device_state.biasTensor, + gamma.ptr, + gamma_grad.ptr, + beta_grad.ptr, + attrs.eps, + per_device_state.saveMean, + per_device_state.saveVar)); +} + +void batch_norm_gpu_cleanup_kernel(Allocator &allocator, + BatchNormPerDeviceState &per_device_state) { allocator.deallocate(per_device_state.runningMean); checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.inputTensor)); - checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.biasTensor)); checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.outputTensor)); - if (per_device_state.relu) { - checkCUDNN(cudnnDestroyActivationDescriptor(per_device_state.actiDesc)); - } + checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.biasTensor)); } -} // namespace BatchNorm -} // namespace Kernels } // namespace FlexFlow diff --git a/lib/kernels/src/cuda/ops/concat_kernels.cu b/lib/kernels/src/cuda/ops/concat_kernels.cu index 667a7c0b74..5415558a29 100644 --- a/lib/kernels/src/cuda/ops/concat_kernels.cu +++ b/lib/kernels/src/cuda/ops/concat_kernels.cu @@ -15,83 +15,84 @@ #include "internal/device.h" #include "kernels/concat_kernels_gpu.h" -#include - -namespace FlexFlow::Kernels::Concat { - -static void calc_blk_size(size_t &num_blocks, - size_t &blk_size, - TensorShape const &shape, - ff_dim_t axis) { - blk_size = get_num_elements(slice_tensor_dims(shape.dims, axis, std::nullopt)) - .int_from_positive_int(); - num_blocks = - get_num_elements(slice_tensor_dims(shape.dims, ff_dim_t{0_n}, axis)) - .int_from_positive_int(); +#include "op-attrs/tensor_dims.h" + +namespace FlexFlow { + +// The number of contiguous elements spanned by `axis` and everything after it. +static int get_blk_size(TensorShape const &shape, ff_dim_t axis) { + return get_num_elements(slice_tensor_dims(shape.dims, axis, std::nullopt)) + .int_from_positive_int(); } -void gpu_forward_kernel(cudaStream_t stream, - GenericTensorAccessorW const &output, - std::vector const &inputs, - ff_dim_t axis) { - assert(inputs.size() <= MAX_NUM_INPUTS); - size_t num_blocks = 1, output_blk_size = 1; - calc_blk_size(num_blocks, output_blk_size, output.shape, axis); - off_t offset = 0; +// The number of such blocks, i.e. the product of the dims before `axis`. +static int get_num_blks(TensorShape const &shape, ff_dim_t axis) { + return get_num_elements(slice_tensor_dims(shape.dims, ff_dim_t{0_n}, axis)) + .int_from_positive_int(); +} - for (GenericTensorAccessorR const &input : inputs) { - size_t input_num_blocks = 1, input_blk_size = 1; - calc_blk_size(input_num_blocks, input_blk_size, input.shape, axis); - assert(input_num_blocks == num_blocks || output_blk_size == input_blk_size); +void concat_gpu_forward_kernel( + cudaStream_t stream, + ConcatAttrs const &attrs, + std::vector const &inputs, + GenericTensorAccessorW const &output) { + ASSERT(inputs.size() == attrs.num_inputs.int_from_int_ge_two()); + ASSERT(inputs.size() <= MAX_NUM_INPUTS); - int blocks_to_copy = - (output_blk_size == input_blk_size) ? input_num_blocks : num_blocks; + int num_blks = get_num_blks(output.shape, attrs.axis); + int output_blk_size = get_blk_size(output.shape, attrs.axis); - copy_with_stride<<>>(output.get_float_ptr() + offset, input.get_float_ptr(), - blocks_to_copy, + num_blks, output_blk_size, input_blk_size); - offset += (output_blk_size == input_blk_size) - ? input_blk_size * input_num_blocks - : input_blk_size; + offset += input_blk_size; } + + ASSERT(offset == output_blk_size); } -void gpu_backward_kernel(cudaStream_t stream, - GenericTensorAccessorR const &output_grad, - std::vector const &input_grads, - ff_dim_t axis) { - assert(input_grads.size() <= MAX_NUM_INPUTS); - size_t num_blocks = 1, output_blk_size = 1; - calc_blk_size(num_blocks, output_blk_size, output_grad.shape, axis); - off_t offset = 0; +void concat_gpu_backward_kernel( + cudaStream_t stream, + ConcatAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + std::vector const &inputs, + std::vector const &input_grads) { + ASSERT(input_grads.size() == attrs.num_inputs.int_from_int_ge_two()); + ASSERT(input_grads.size() <= MAX_NUM_INPUTS); - for (auto &input_grad : input_grads) { - size_t input_num_blocks = 1, input_blk_size = 1; - calc_blk_size(input_num_blocks, input_blk_size, input_grad.shape, axis); - assert(input_num_blocks == num_blocks || output_blk_size == input_blk_size); + int num_blks = get_num_blks(output_grad.shape, attrs.axis); + int output_blk_size = get_blk_size(output_grad.shape, attrs.axis); - int blocks_to_add = - (output_blk_size == input_blk_size) ? input_num_blocks : num_blocks; + int offset = 0; + for (GenericTensorAccessorW const &input_grad : input_grads) { + ASSERT(get_num_blks(input_grad.shape, attrs.axis) == num_blks); + int input_blk_size = get_blk_size(input_grad.shape, attrs.axis); - add_with_stride<<>>(input_grad.get_float_ptr(), output_grad.get_float_ptr() + offset, - blocks_to_add, + num_blks, input_blk_size, output_blk_size); - offset += (output_blk_size == input_blk_size) - ? input_blk_size * input_num_blocks - : input_blk_size; + offset += input_blk_size; } + + ASSERT(offset == output_blk_size); } -} // namespace FlexFlow::Kernels::Concat +} // namespace FlexFlow diff --git a/lib/kernels/src/cuda/ops/conv_2d_kernels.cu b/lib/kernels/src/cuda/ops/conv_2d_kernels.cu index 92046b30ae..3489575fce 100644 --- a/lib/kernels/src/cuda/ops/conv_2d_kernels.cu +++ b/lib/kernels/src/cuda/ops/conv_2d_kernels.cu @@ -1,195 +1,109 @@ #include "internal/device.h" -#include "kernels/conv_2d_kernels.h" +#include "kernels/conv_2d_kernels_gpu.h" +#include "op-attrs/ops/conv_2d.h" +#include "op-attrs/tensor_dims.h" +#include namespace FlexFlow { -namespace Kernels { -namespace Conv2D { - -cudnnConvolutionBwdDataAlgo_t selectConvolutionBackwardDataAlgorithm( - cudnnHandle_t handle, - const cudnnFilterDescriptor_t wDesc, - void const *w, - const cudnnTensorDescriptor_t dyDesc, - void const *dy, - const cudnnConvolutionDescriptor_t convDesc, - void *workSpace, - size_t workSpaceSize, - const cudnnTensorDescriptor_t dxDesc, - void *dx, - float *time) { - int const reqAlgCnt = 8; - int cnt = 0; - cudnnConvolutionBwdDataAlgoPerf_t perfResults[reqAlgCnt]; - checkCUDNN(cudnnFindConvolutionBackwardDataAlgorithmEx(handle, - wDesc, - w, - dyDesc, - dy, - convDesc, - dxDesc, - dx, - reqAlgCnt, - &cnt, - perfResults, - workSpace, - workSpaceSize)); - assert(cnt > 0); - checkCUDNN(perfResults[0].status); - if (time != nullptr) { - *time = perfResults[0].time; - } - return perfResults[0].algo; -} -cudnnConvolutionFwdAlgo_t selectConvolutionForwardAlgorithm( - cudnnHandle_t handle, - const cudnnTensorDescriptor_t xDesc, - void const *x, - const cudnnFilterDescriptor_t wDesc, - void const *w, - const cudnnConvolutionDescriptor_t convDesc, - void *workSpace, - size_t workSpaceSize, - const cudnnTensorDescriptor_t yDesc, - void *y, - float *time) { - int const reqAlgCnt = 8; - int cnt = 0; - cudnnConvolutionFwdAlgoPerf_t perfResults[reqAlgCnt]; - checkCUDNN(cudnnFindConvolutionForwardAlgorithmEx(handle, - xDesc, - x, - wDesc, - w, - convDesc, - yDesc, - y, - reqAlgCnt, - &cnt, - perfResults, - workSpace, - workSpaceSize)); - assert(cnt > 0); - checkCUDNN(perfResults[0].status); - if (time != nullptr) { - *time = perfResults[0].time; +// Picks the fastest algorithm reported by cuDNN's heuristics that both +// succeeds and fits in the available workspace. +template +static int select_algorithm_idx(std::vector const &perf_results, + int num_results, + size_t workspace_size) { + for (int i = 0; i < num_results; i++) { + if (perf_results.at(i).status == CUDNN_STATUS_SUCCESS && + perf_results.at(i).memory <= workspace_size) { + return i; + } } - return perfResults[0].algo; -} -cudnnConvolutionBwdFilterAlgo_t selectConvolutionBackwardFilterAlgorithm( - cudnnHandle_t handle, - const cudnnTensorDescriptor_t xDesc, - void const *x, - const cudnnTensorDescriptor_t dyDesc, - void const *dy, - const cudnnConvolutionDescriptor_t convDesc, - void *workSpace, - size_t workSpaceSize, - const cudnnFilterDescriptor_t dwDesc, - void *dw, - float *time) { - int const reqAlgCnt = 8; - int cnt = 0; - cudnnConvolutionBwdFilterAlgoPerf_t perfResults[reqAlgCnt]; - checkCUDNN(cudnnFindConvolutionBackwardFilterAlgorithmEx(handle, - xDesc, - x, - dyDesc, - dy, - convDesc, - dwDesc, - dw, - reqAlgCnt, - &cnt, - perfResults, - workSpace, - workSpaceSize)); - assert(cnt > 0); - checkCUDNN(perfResults[0].status); - if (time != nullptr) { - *time = perfResults[0].time; - } - return perfResults[0].algo; + PANIC("No cuDNN convolution algorithm fits within the available workspace", + num_results, + workspace_size); } -Conv2DPerDeviceState - gpu_init_kernel(PerDeviceFFHandle const &handle, - std::optional const &activation, - int kernel_h, - int kernel_w, - int groups, - int pad_h, - int pad_w, - int stride_h, - int stride_w, - GenericTensorAccessorW const &input, - GenericTensorAccessorW const &output, - float const *filter_ptr, - float *filter_grad_ptr) { +Conv2DPerDeviceState conv_2d_gpu_init_kernel(PerDeviceFFHandle const &handle, + Conv2DAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { + // Applying an activation as part of Conv2D is not currently implemented. The + // previous implementation hardcoded relu (regardless of which activation was + // requested) and computed the backward pass by destructively modifying the + // output gradient. If you need it, please create an issue. + ASSERT(!attrs.activation.has_value(), + "Conv2D does not currently support fused activations", + attrs.activation); + + ASSERT(get_num_dims(input_shape.dims) == num_tensor_dims_t{4_n}, + "Conv2D expects 4-dimensional (i.e., NCHW) input tensors", + input_shape); + ASSERT(conv2d_get_output_shape(attrs, input_shape) == output_shape, + "Conv2D output shape does not match the shape implied by its " + "attributes and input shape", + attrs, + input_shape, + output_shape); + + TensorShape filter_shape = conv2d_get_kernel_shape(attrs, input_shape); + + positive_int input_c = dim_at_idx(input_shape.dims, ff_dim_t{1_n}); + positive_int output_c = dim_at_idx(output_shape.dims, ff_dim_t{1_n}); + + ASSERT(input_c % attrs.groups == 0, + "Conv2D requires the number of input channels to be divisible by the " + "number of groups", + input_c, + attrs.groups); ffTensorDescriptor_t inputTensor; ffTensorDescriptor_t biasTensor; ffTensorDescriptor_t outputTensor; ffFilterDescriptor_t filterDesc; - ffActivationDescriptor_t actiDesc; ffConvolutionDescriptor_t convDesc; - ffConvolutionFwdAlgo_t fwdAlgo; - ffConvolutionBwdFilterAlgo_t bwdFilterAlgo; - ffConvolutionBwdDataAlgo_t bwdDataAlgo; - - int input_w = - dim_at_idx(input.shape.dims, legion_dim_t{0_n}).int_from_positive_int(); - int input_h = - dim_at_idx(input.shape.dims, legion_dim_t{1_n}).int_from_positive_int(); - int input_c = - dim_at_idx(input.shape.dims, legion_dim_t{2_n}).int_from_positive_int(); - int input_n = - dim_at_idx(input.shape.dims, legion_dim_t{3_n}).int_from_positive_int(); - - int output_w = - dim_at_idx(output.shape.dims, legion_dim_t{0_n}).int_from_positive_int(); - int output_h = - dim_at_idx(output.shape.dims, legion_dim_t{1_n}).int_from_positive_int(); - int output_c = - dim_at_idx(output.shape.dims, legion_dim_t{2_n}).int_from_positive_int(); - int output_n = - dim_at_idx(output.shape.dims, legion_dim_t{3_n}).int_from_positive_int(); checkCUDNN(cudnnCreateTensorDescriptor(&inputTensor)); checkCUDNN(cudnnCreateTensorDescriptor(&biasTensor)); checkCUDNN(cudnnCreateTensorDescriptor(&outputTensor)); checkCUDNN(cudnnCreateFilterDescriptor(&filterDesc)); checkCUDNN(cudnnCreateConvolutionDescriptor(&convDesc)); - checkCUDNN(cudnnCreateActivationDescriptor(&actiDesc)); - - checkCUDNN(cudnnSetTensorDescriptorFromTensorShape(inputTensor, input.shape)); - - checkCUDNN(cudnnSetTensor4dDescriptor( - biasTensor, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, output_c, 1, 1)); - - // Require that input_c is divisible by conv->groups - assert(input_c % groups == 0); - checkCUDNN(cudnnSetFilter4dDescriptor(filterDesc, - CUDNN_DATA_FLOAT, - CUDNN_TENSOR_NCHW, - output_c, - input_c / groups, - kernel_h, - kernel_w)); - - checkCUDNN(cudnnSetConvolution2dDescriptor(convDesc, - pad_h, - pad_w, - stride_h, - stride_w, - 1 /*upscale_x*/, - 1 /*upscale_y*/, - CUDNN_CROSS_CORRELATION, - CUDNN_DATA_FLOAT)); - if (groups != 1) { - checkCUDNN(cudnnSetConvolutionGroupCount(convDesc, groups)); - } + + checkCUDNN(cudnnSetTensorDescriptorFromTensorShape(inputTensor, input_shape)); + checkCUDNN( + cudnnSetTensorDescriptorFromTensorShape(outputTensor, output_shape)); + + checkCUDNN( + cudnnSetTensor4dDescriptor(biasTensor, + CUDNN_TENSOR_NCHW, + ff_to_cudnn_datatype(output_shape.data_type), + /*n=*/1, + /*c=*/output_c.int_from_positive_int(), + /*h=*/1, + /*w=*/1)); + + checkCUDNN(cudnnSetFilter4dDescriptor( + filterDesc, + ff_to_cudnn_datatype(filter_shape.data_type), + CUDNN_TENSOR_NCHW, + dim_at_idx(filter_shape.dims, ff_dim_t{0_n}).int_from_positive_int(), + dim_at_idx(filter_shape.dims, ff_dim_t{1_n}).int_from_positive_int(), + dim_at_idx(filter_shape.dims, ff_dim_t{2_n}).int_from_positive_int(), + dim_at_idx(filter_shape.dims, ff_dim_t{3_n}).int_from_positive_int())); + + checkCUDNN( + cudnnSetConvolution2dDescriptor(convDesc, + attrs.padding_h.unwrap_nonnegative(), + attrs.padding_w.unwrap_nonnegative(), + attrs.stride_h.int_from_positive_int(), + attrs.stride_w.int_from_positive_int(), + /*dilation_h=*/1, + /*dilation_w=*/1, + CUDNN_CROSS_CORRELATION, + CUDNN_DATA_FLOAT)); + + checkCUDNN(cudnnSetConvolutionGroupCount( + convDesc, attrs.groups.int_from_positive_int())); // enable tensor core when possible if (handle.allowTensorOpMathConversion) { @@ -199,202 +113,202 @@ Conv2DPerDeviceState checkCUDNN(cudnnSetConvolutionMathType(convDesc, CUDNN_TENSOR_OP_MATH)); } - int n, c, h, w; - checkCUDNN(cudnnGetConvolution2dForwardOutputDim( - convDesc, inputTensor, filterDesc, &n, &c, &h, &w)); - assert(n == output_n); - assert(c == output_c); - assert(h == output_h); - assert(w == output_w); - - checkCUDNN(cudnnSetTensor4dDescriptor( - outputTensor, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, n, c, h, w)); - - // select forward algorithm - fwdAlgo = selectConvolutionForwardAlgorithm( - handle.dnn, - inputTensor, - static_cast(input.get_float_ptr()), - filterDesc, - filter_ptr, - convDesc, - handle.workSpace, - handle.workSpaceSize, - outputTensor, - output.get_float_ptr(), - nullptr); - - // select backward filter algorithm - bwdFilterAlgo = selectConvolutionBackwardFilterAlgorithm( - handle.dnn, - inputTensor, - static_cast(input.get_float_ptr()), - outputTensor, - output.get_float_ptr(), - convDesc, - handle.workSpace, - handle.workSpaceSize, - filterDesc, - filter_grad_ptr, - nullptr); + { + int n, c, h, w; + checkCUDNN(cudnnGetConvolution2dForwardOutputDim( + convDesc, inputTensor, filterDesc, &n, &c, &h, &w)); + ASSERT(dim_at_idx(output_shape.dims, ff_dim_t{0_n}) == positive_int{n}); + ASSERT(dim_at_idx(output_shape.dims, ff_dim_t{1_n}) == positive_int{c}); + ASSERT(dim_at_idx(output_shape.dims, ff_dim_t{2_n}) == positive_int{h}); + ASSERT(dim_at_idx(output_shape.dims, ff_dim_t{3_n}) == positive_int{w}); + } - // select backward data algorithm - bwdDataAlgo = selectConvolutionBackwardDataAlgorithm( - handle.dnn, - filterDesc, - filter_ptr, - outputTensor, - output.get_float_ptr(), - convDesc, - handle.workSpace, - handle.workSpaceSize, - inputTensor, - static_cast(input.get_float_ptr()), - nullptr); - if (activation.has_value()) { - checkCUDNN(cudnnSetActivationDescriptor( - actiDesc, CUDNN_ACTIVATION_RELU, CUDNN_PROPAGATE_NAN, 0.0)); + // NOTE: we use cuDNN's heuristics (rather than the cudnnFind*Ex family) to + // select algorithms, as the latter requires (and destructively writes to) + // real input, output and gradient buffers. + ffConvolutionFwdAlgo_t fwdAlgo; + { + int max_num_results; + checkCUDNN(cudnnGetConvolutionForwardAlgorithmMaxCount(handle.dnn, + &max_num_results)); + std::vector perf_results(max_num_results); + int num_results = 0; + checkCUDNN(cudnnGetConvolutionForwardAlgorithm_v7(handle.dnn, + inputTensor, + filterDesc, + convDesc, + outputTensor, + max_num_results, + &num_results, + perf_results.data())); + fwdAlgo = perf_results + .at(select_algorithm_idx( + perf_results, num_results, handle.workSpaceSize)) + .algo; } - Conv2DPerDeviceState per_device_state = Conv2DPerDeviceState{ - handle, - inputTensor, - biasTensor, - outputTensor, - filterDesc, - actiDesc, - convDesc, - fwdAlgo, - bwdFilterAlgo, - bwdDataAlgo, + ffConvolutionBwdFilterAlgo_t bwdFilterAlgo; + { + int max_num_results; + checkCUDNN(cudnnGetConvolutionBackwardFilterAlgorithmMaxCount( + handle.dnn, &max_num_results)); + std::vector perf_results( + max_num_results); + int num_results = 0; + checkCUDNN( + cudnnGetConvolutionBackwardFilterAlgorithm_v7(handle.dnn, + inputTensor, + outputTensor, + convDesc, + filterDesc, + max_num_results, + &num_results, + perf_results.data())); + bwdFilterAlgo = perf_results + .at(select_algorithm_idx( + perf_results, num_results, handle.workSpaceSize)) + .algo; + } + + ffConvolutionBwdDataAlgo_t bwdDataAlgo; + { + int max_num_results; + checkCUDNN(cudnnGetConvolutionBackwardDataAlgorithmMaxCount( + handle.dnn, &max_num_results)); + std::vector perf_results( + max_num_results); + int num_results = 0; + checkCUDNN( + cudnnGetConvolutionBackwardDataAlgorithm_v7(handle.dnn, + filterDesc, + outputTensor, + convDesc, + inputTensor, + max_num_results, + &num_results, + perf_results.data())); + bwdDataAlgo = perf_results + .at(select_algorithm_idx( + perf_results, num_results, handle.workSpaceSize)) + .algo; + } + + return Conv2DPerDeviceState{ + /*inputTensor=*/inputTensor, + /*biasTensor=*/biasTensor, + /*outputTensor=*/outputTensor, + /*filterDesc=*/filterDesc, + /*convDesc=*/convDesc, + /*fwdAlgo=*/fwdAlgo, + /*bwdFilterAlgo=*/bwdFilterAlgo, + /*bwdDataAlgo=*/bwdDataAlgo, }; - return per_device_state; } -void gpu_forward_kernel(ffStream_t stream, - Conv2DPerDeviceState const &m, - float const *input_ptr, - float *output_ptr, - float const *filter_ptr, - float const *bias_ptr, - std::optional activation) { - checkCUDNN(cudnnSetStream(m.handle.dnn, stream)); +void conv_2d_gpu_forward_kernel( + cudaStream_t stream, + PerDeviceFFHandle const &handle, + Conv2DPerDeviceState const &per_device_state, + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &filter, + std::optional const &bias, + GenericTensorAccessorW const &output) { + ASSERT(bias.has_value() == attrs.use_bias); + + checkCUDNN(cudnnSetStream(handle.dnn, stream)); float alpha = 1.0f, beta = 0.0f; - checkCUDNN(cudnnConvolutionForward(m.handle.dnn, + checkCUDNN(cudnnConvolutionForward(handle.dnn, &alpha, - m.inputTensor, - input_ptr, - m.filterDesc, - filter_ptr, - m.convDesc, - m.fwdAlgo, - m.handle.workSpace, - m.handle.workSpaceSize, + per_device_state.inputTensor, + input.ptr, + per_device_state.filterDesc, + filter.ptr, + per_device_state.convDesc, + per_device_state.fwdAlgo, + handle.workSpace, + handle.workSpaceSize, &beta, - m.outputTensor, - output_ptr)); + per_device_state.outputTensor, + output.ptr)); - if (bias_ptr != NULL) { - checkCUDNN(cudnnAddTensor(m.handle.dnn, + if (bias.has_value()) { + checkCUDNN(cudnnAddTensor(handle.dnn, &alpha, - m.biasTensor, - bias_ptr, + per_device_state.biasTensor, + bias.value().ptr, &alpha, - m.outputTensor, - output_ptr)); - } - if (activation.has_value()) { - checkCUDNN(cudnnActivationForward(m.handle.dnn, - m.actiDesc, - &alpha, - m.outputTensor, - output_ptr, - &beta, - m.outputTensor, - output_ptr)); + per_device_state.outputTensor, + output.ptr)); } } -void gpu_backward_kernel(ffStream_t stream, - Conv2DPerDeviceState const &m, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *filter_ptr, - float *filter_grad_ptr, - float *bias_grad_ptr, - std::optional activation) { - checkCUDNN(cudnnSetStream(m.handle.dnn, stream)); +void conv_2d_gpu_backward_kernel( + cudaStream_t stream, + PerDeviceFFHandle const &handle, + Conv2DPerDeviceState const &per_device_state, + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &filter, + GenericTensorAccessorW const &filter_grad, + std::optional const &bias_grad) { + ASSERT(bias_grad.has_value() == attrs.use_bias); + + checkCUDNN(cudnnSetStream(handle.dnn, stream)); + // NOTE: alpha is used for beta as well so that the gradients are accumulated + // into rather than overwritten float alpha = 1.0f; - // float beta = 0.0f; - if (activation.has_value()) { - cudnnDataType_t dataType; - int n, c, h, w, nStride, cStride, hStride, wStride; - checkCUDNN(cudnnGetTensor4dDescriptor(m.outputTensor, - &dataType, - &n, - &c, - &h, - &w, - &nStride, - &cStride, - &hStride, - &wStride)); - reluBackward<<>>( - output_grad_ptr, output_ptr, n * c * h * w); - } - // Compute filter gradiant - // NOTE: we use alpha for kernel_grad to accumulate gradients - checkCUDNN(cudnnConvolutionBackwardFilter(m.handle.dnn, - &alpha, - m.inputTensor, - input_ptr, - m.outputTensor, - output_grad_ptr, - m.convDesc, - m.bwdFilterAlgo, - m.handle.workSpace, - m.handle.workSpaceSize, - &alpha, - m.filterDesc, - filter_grad_ptr)); - // Compute bias gradiant - // NOTE: we use alpha for bias_grad to accumulate gradients - if (bias_grad_ptr != NULL) { - checkCUDNN(cudnnConvolutionBackwardBias(m.handle.dnn, + + checkCUDNN(cudnnConvolutionBackwardFilter(handle.dnn, &alpha, - m.outputTensor, - output_grad_ptr, + per_device_state.inputTensor, + input.ptr, + per_device_state.outputTensor, + output_grad.ptr, + per_device_state.convDesc, + per_device_state.bwdFilterAlgo, + handle.workSpace, + handle.workSpaceSize, &alpha, - m.biasTensor, - bias_grad_ptr)); - } - // Compute data gradiant - // NOTE: we use alpha for input_grad to accumulate gradients - if (input_grad_ptr != NULL) { - checkCUDNN(cudnnConvolutionBackwardData(m.handle.dnn, + per_device_state.filterDesc, + filter_grad.ptr)); + + if (bias_grad.has_value()) { + checkCUDNN(cudnnConvolutionBackwardBias(handle.dnn, &alpha, - m.filterDesc, - filter_ptr, - m.outputTensor, - output_grad_ptr, - m.convDesc, - m.bwdDataAlgo, - m.handle.workSpace, - m.handle.workSpaceSize, + per_device_state.outputTensor, + output_grad.ptr, &alpha, - m.inputTensor, - input_grad_ptr)); + per_device_state.biasTensor, + bias_grad.value().ptr)); } + + checkCUDNN(cudnnConvolutionBackwardData(handle.dnn, + &alpha, + per_device_state.filterDesc, + filter.ptr, + per_device_state.outputTensor, + output_grad.ptr, + per_device_state.convDesc, + per_device_state.bwdDataAlgo, + handle.workSpace, + handle.workSpaceSize, + &alpha, + per_device_state.inputTensor, + input_grad.ptr)); } -void gpu_cleanup_kernel(Conv2DPerDeviceState &per_device_state) { - NOT_IMPLEMENTED(); +void conv_2d_gpu_cleanup_kernel(Conv2DPerDeviceState &per_device_state) { + checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.inputTensor)); + checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.biasTensor)); + checkCUDNN(cudnnDestroyTensorDescriptor(per_device_state.outputTensor)); + checkCUDNN(cudnnDestroyFilterDescriptor(per_device_state.filterDesc)); + checkCUDNN(cudnnDestroyConvolutionDescriptor(per_device_state.convDesc)); } -} // namespace Conv2D -} // 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..ecbb846797 100644 --- a/lib/kernels/src/cuda/ops/element_unary_kernels.cu +++ b/lib/kernels/src/cuda/ops/element_unary_kernels.cu @@ -40,12 +40,25 @@ static bool use_scalar(OperatorType op_type) { case OperatorType::SCALAR_SUB: case OperatorType::SCALAR_TRUE_DIV: case OperatorType::POW: + case OperatorType::SILU: return true; default: return false; } } +// SILU's scalar (beta) is optional and defaults to 1, matching the CPU +// kernel. The other scalar operators require a scalar. +static float require_scalar(OperatorType op_type, + std::optional const &scalar) { + if (op_type == OperatorType::SILU) { + return scalar.value_or(1.0f); + } + + ASSERT(scalar.has_value(), op_type); + return scalar.value(); +} + ElementUnaryPerDeviceState element_unary_gpu_init_kernel(ElementUnaryAttrs const &attrs, TensorShape const &input_shape, @@ -117,6 +130,10 @@ __global__ void elewise_scalar_unary_forward_kernel( out[i] = (T)(powf(in[i], scalar)); break; } + case OperatorType::SILU: { + out[i] = (T)(in[i] / (1.0f + expf(-1.0f * scalar * in[i]))); + break; + } default: assert(false); } @@ -191,6 +208,13 @@ __global__ void elewise_scalar_unary_backward_kernel(coord_t volume, (T)(output_grad[i] * scalar * powf(input[i], scalar - 1)); break; } + case OperatorType::SILU: { + float e_to_bx = expf(scalar * input[i]); + input_grad[i] += (T)(output_grad[i] * + (e_to_bx * (scalar * input[i] + e_to_bx + 1.0f)) / + ((e_to_bx + 1.0f) * (e_to_bx + 1.0f))); + break; + } default: assert(false); } @@ -266,11 +290,10 @@ struct ForwardKernel { m.outputTensor, output.get())); } else if (use_scalar(op_type)) { - assert(scalar.has_value()); elewise_scalar_unary_forward_kernel> <<>>( num_elements, - static_cast>(scalar.value()), + static_cast>(require_scalar(op_type, scalar)), op_type, input.get(), output.get()); @@ -312,11 +335,10 @@ struct BackwardKernel { m.inputTensor, input_grad.get())); } else if (use_scalar(op_type)) { - assert(scalar.has_value()); elewise_scalar_unary_backward_kernel> <<>>( num_elements, - static_cast>(scalar.value()), + static_cast>(require_scalar(op_type, scalar)), op_type, output.get(), output_grad.get(), diff --git a/lib/kernels/src/cuda/ops/pool_2d_kernels.cu b/lib/kernels/src/cuda/ops/pool_2d_kernels.cu index 4e06f2da02..b1763d34f6 100644 --- a/lib/kernels/src/cuda/ops/pool_2d_kernels.cu +++ b/lib/kernels/src/cuda/ops/pool_2d_kernels.cu @@ -15,136 +15,125 @@ #include "internal/device.h" #include "kernels/pool_2d_kernels_gpu.h" +#include "op-attrs/tensor_dims.h" #include "utils/exception.h" namespace FlexFlow { -namespace Kernels { -namespace Pool2D { - -Pool2DPerDeviceState gpu_init_kernel(PerDeviceFFHandle handle, - std::optional activation, - int input_w, - int input_h, - int input_c, - int input_n, - int output_w, - int output_h, - int output_c, - int output_n, - int pad_h, - int pad_w, - int kernel_h, - int kernel_w, - int stride_h, - int stride_w, - PoolOp pool_type) { +static cudnnPoolingMode_t cudnn_pooling_mode_from_pool_op(PoolOp pool_type) { + switch (pool_type) { + case PoolOp::MAX: + return CUDNN_POOLING_MAX; + case PoolOp::AVG: + return CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING; + default: + NOT_IMPLEMENTED(); + } +} + +Pool2DPerDeviceState pool_2d_gpu_init_kernel(Pool2DAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { + // Applying an activation as part of Pool2D is not currently implemented (it + // used to be silently ignored). If you need it, please create an issue. + ASSERT(!attrs.activation.has_value(), + "Pool2D does not currently support fused activations", + attrs.activation); + + ASSERT(get_num_dims(input_shape.dims) == num_tensor_dims_t{4_n}, + "Pool2D expects 4-dimensional (i.e., NCHW) input tensors", + input_shape); + ASSERT(get_num_dims(output_shape.dims) == num_tensor_dims_t{4_n}, + "Pool2D expects 4-dimensional (i.e., NCHW) output tensors", + output_shape); + ffTensorDescriptor_t inputTensor; ffTensorDescriptor_t outputTensor; - ffActivationDescriptor_t actiDesc; ffPoolingDescriptor_t poolDesc; checkCUDNN(cudnnCreateTensorDescriptor(&inputTensor)); checkCUDNN(cudnnCreateTensorDescriptor(&outputTensor)); - checkCUDNN(cudnnCreateActivationDescriptor(&actiDesc)); checkCUDNN(cudnnCreatePoolingDescriptor(&poolDesc)); - checkCUDNN(cudnnSetTensor4dDescriptor(inputTensor, - CUDNN_TENSOR_NCHW, - CUDNN_DATA_FLOAT, - input_n, - input_c, - input_h, - input_w)); - cudnnPoolingMode_t mode; - if (pool_type == PoolOp::MAX) { - mode = CUDNN_POOLING_MAX; - } else { - assert(pool_type == PoolOp::AVG); - mode = CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING; - } + checkCUDNN(cudnnSetTensorDescriptorFromTensorShape(inputTensor, input_shape)); - checkCUDNN(cudnnSetPooling2dDescriptor(poolDesc, - mode, - CUDNN_PROPAGATE_NAN, - kernel_h, - kernel_w, - pad_h, - pad_w, - stride_h, - stride_w)); + checkCUDNN(cudnnSetPooling2dDescriptor( + poolDesc, + cudnn_pooling_mode_from_pool_op(attrs.pool_type), + CUDNN_PROPAGATE_NAN, + attrs.kernel_h.int_from_positive_int(), + attrs.kernel_w.int_from_positive_int(), + attrs.padding_h.unwrap_nonnegative(), + attrs.padding_w.unwrap_nonnegative(), + attrs.stride_h.int_from_positive_int(), + attrs.stride_w.int_from_positive_int())); int n, c, h, w; checkCUDNN( cudnnGetPooling2dForwardOutputDim(poolDesc, inputTensor, &n, &c, &h, &w)); - assert(n == output_n); - assert(c == output_c); - assert(h == output_h); - assert(w == output_w); - - checkCUDNN(cudnnSetTensor4dDescriptor( - outputTensor, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, n, c, h, w)); - bool relu = false; - if (activation == Activation::RELU) { - relu = true; - } - Pool2DPerDeviceState state = Pool2DPerDeviceState{ - /*handle=*/handle, + + ASSERT(dim_at_idx(output_shape.dims, ff_dim_t{0_n}) == positive_int{n}); + ASSERT(dim_at_idx(output_shape.dims, ff_dim_t{1_n}) == positive_int{c}); + ASSERT(dim_at_idx(output_shape.dims, ff_dim_t{2_n}) == positive_int{h}); + ASSERT(dim_at_idx(output_shape.dims, ff_dim_t{3_n}) == positive_int{w}); + + checkCUDNN( + cudnnSetTensorDescriptorFromTensorShape(outputTensor, output_shape)); + + return Pool2DPerDeviceState{ /*inputTensor=*/inputTensor, /*outputTensor=*/outputTensor, - /*actiDesc=*/actiDesc, /*poolDesc=*/poolDesc, - /*relu=*/relu, }; - return state; } -void gpu_forward_kernel(cudaStream_t stream, - Pool2DPerDeviceState const &m, - void const *input_ptr, - void *output_ptr) { - - checkCUDNN(cudnnSetStream(m.handle.dnn, stream)); +void pool_2d_gpu_forward_kernel(cudaStream_t stream, + PerDeviceFFHandle const &handle, + Pool2DPerDeviceState const &per_device_state, + Pool2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { + checkCUDNN(cudnnSetStream(handle.dnn, stream)); float alpha = 1.0f, beta = 0.0f; - checkCUDNN(cudnnPoolingForward(m.handle.dnn, - m.poolDesc, + checkCUDNN(cudnnPoolingForward(handle.dnn, + per_device_state.poolDesc, &alpha, - m.inputTensor, - input_ptr, + per_device_state.inputTensor, + input.ptr, &beta, - m.outputTensor, - output_ptr)); + per_device_state.outputTensor, + output.ptr)); } -void gpu_backward_kernel(cudaStream_t stream, - Pool2DPerDeviceState const &m, - void const *output_ptr, - void const *output_grad_ptr, - void const *input_ptr, - void *input_grad_ptr) { - - checkCUDNN(cudnnSetStream(m.handle.dnn, stream)); - - float alpha = 1.0f; - checkCUDNN(cudnnPoolingBackward(m.handle.dnn, - m.poolDesc, - &alpha, - m.outputTensor, - output_ptr, - m.outputTensor, - output_grad_ptr, - m.inputTensor, - input_ptr, +void pool_2d_gpu_backward_kernel(cudaStream_t stream, + PerDeviceFFHandle const &handle, + Pool2DPerDeviceState const &per_device_state, + Pool2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { + checkCUDNN(cudnnSetStream(handle.dnn, stream)); + + // NOTE: beta is 1.0 so that input_grad is accumulated into + float alpha = 1.0f, beta = 1.0f; + checkCUDNN(cudnnPoolingBackward(handle.dnn, + per_device_state.poolDesc, &alpha, - m.inputTensor, - input_grad_ptr)); + per_device_state.outputTensor, + output.ptr, + per_device_state.outputTensor, + output_grad.ptr, + per_device_state.inputTensor, + input.ptr, + &beta, + per_device_state.inputTensor, + input_grad.ptr)); } -void gpu_cleanup_kernel(Pool2DPerDeviceState &per_device_state) { +void pool_2d_gpu_cleanup_kernel(Pool2DPerDeviceState &per_device_state) { NOT_IMPLEMENTED(); } -} // namespace Pool2D -} // namespace Kernels } // namespace FlexFlow diff --git a/lib/kernels/src/cuda/ops/reshape_kernels.cu b/lib/kernels/src/cuda/ops/reshape_kernels.cu index 1414706ef0..2efd482b86 100644 --- a/lib/kernels/src/cuda/ops/reshape_kernels.cu +++ b/lib/kernels/src/cuda/ops/reshape_kernels.cu @@ -16,53 +16,60 @@ #include "internal/device.h" #include "kernels/datatype_dispatch.h" #include "kernels/reshape_kernels_gpu.h" +#include "op-attrs/tensor_shape.h" +#include "utils/containers/require_same.h" namespace FlexFlow { -namespace Kernels { -namespace Reshape { +void reshape_gpu_forward_kernel(cudaStream_t stream, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { + require_same(input.shape.data_type, output.shape.data_type); + require_same(get_num_elements(input.shape.dims), + get_num_elements(output.shape.dims)); -template -__global__ void apply_add_with_scale2(DT *data_ptr, - DTGrad const *grad_ptr, - size_t size, - DT scale) { - CUDA_KERNEL_LOOP(i, size) { - data_ptr[i] += grad_ptr[i] * scale; + size_t num_bytes = + get_size_in_bytes(output.shape).unwrap_num_bytes().unwrap_nonnegative(); + + checkCUDA(cudaMemcpyAsync( + output.ptr, input.ptr, num_bytes, cudaMemcpyDeviceToDevice, stream)); +} + +template +__global__ void reshape_accumulate_kernel(T *input_grad, + T const *output_grad, + size_t num_elements) { + CUDA_KERNEL_LOOP(i, num_elements) { + input_grad[i] += output_grad[i]; } } -template -struct BackwardKernel { +template +struct ReshapeGPUBackwardKernel { void operator()(cudaStream_t stream, - GenericTensorAccessorR const &output, - GenericTensorAccessorW const &input) { - float alpha = 1.0f; - apply_add_with_scale2, real_type_t> - <<>>(input.get(), - output.get(), - get_num_elements(input.shape.dims).int_from_positive_int(), - static_cast>(alpha)); + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorW const &input_grad) const { + int num_elements = + get_num_elements(input_grad.shape.dims).int_from_positive_int(); + + reshape_accumulate_kernel> + <<>>( + input_grad.get
(), output_grad.get
(), num_elements); } }; -void gpu_forward_kernel(cudaStream_t stream, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output) { - copy_accessor_data_to_l_from_r(output, input); -} +void reshape_gpu_backward_kernel(cudaStream_t stream, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { + DataType data_type = + require_same(output_grad.shape.data_type, input_grad.shape.data_type); + require_same(get_num_elements(output_grad.shape.dims), + get_num_elements(input_grad.shape.dims)); -void gpu_backward_kernel(cudaStream_t stream, - GenericTensorAccessorR const &output, - GenericTensorAccessorW const &input) { - DataTypeDispatch2{}( - input.shape.data_type, output.shape.data_type, stream, output, input); + DataTypeDispatch1{}( + data_type, stream, output_grad, input_grad); } -} // namespace Reshape -} // namespace Kernels } // namespace FlexFlow diff --git a/lib/kernels/src/cuda/ops/split_kernels.cu b/lib/kernels/src/cuda/ops/split_kernels.cu index ca953bd7b1..96587ba1c3 100644 --- a/lib/kernels/src/cuda/ops/split_kernels.cu +++ b/lib/kernels/src/cuda/ops/split_kernels.cu @@ -15,49 +15,91 @@ #include "internal/device.h" #include "kernels/split_kernels_gpu.h" +#include "op-attrs/tensor_dims.h" namespace FlexFlow { -namespace Kernels { -namespace Split { +// The number of contiguous elements spanned by `axis` and everything after it. +static int get_blk_size(TensorShape const &shape, ff_dim_t axis) { + return get_num_elements(slice_tensor_dims(shape.dims, axis, std::nullopt)) + .int_from_positive_int(); +} + +// The number of such blocks, i.e. the product of the dims before `axis`. +static int get_num_blks(TensorShape const &shape, ff_dim_t axis) { + return get_num_elements(slice_tensor_dims(shape.dims, ff_dim_t{0_n}, axis)) + .int_from_positive_int(); +} + +void split_gpu_forward_kernel( + cudaStream_t stream, + SplitAttrs const &attrs, + GenericTensorAccessorR const &input, + std::vector const &outputs) { + ASSERT(outputs.size() == attrs.splits.size()); + ASSERT(outputs.size() <= MAX_NUM_OUTPUTS); -void gpu_forward_kernel(cudaStream_t stream, - float **out_ptrs, - float const *in_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs) { + int num_blks = get_num_blks(input.shape, attrs.axis); + int input_blk_size = get_blk_size(input.shape, attrs.axis); - for (int i = 0; i < numOutputs; i++) { - copy_with_stride<<>>( - out_ptrs[i], in_ptr, num_blks, out_blk_sizes[i], in_blk_size); - in_ptr += out_blk_sizes[i]; + stream>>>(output.get_float_ptr(), + input.get_float_ptr() + offset, + num_blks, + output_blk_size, + input_blk_size); + + offset += output_blk_size; } + + ASSERT(offset == input_blk_size); } -void gpu_backward_kernel(cudaStream_t stream, - float *in_grad_ptr, - float const **out_grad_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs) { +void split_gpu_backward_kernel( + cudaStream_t stream, + SplitAttrs const &attrs, + std::vector const &outputs, + std::vector const &output_grads, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { + ASSERT(output_grads.size() == attrs.splits.size()); + ASSERT(output_grads.size() <= MAX_NUM_OUTPUTS); + + int num_blks = get_num_blks(input_grad.shape, attrs.axis); + int input_blk_size = get_blk_size(input_grad.shape, attrs.axis); + + int offset = 0; + for (int i = 0; i < output_grads.size(); i++) { + GenericTensorAccessorR const &output_grad = output_grads.at(i); - for (int i = 0; i < numOutputs; i++) { - add_with_stride<<>>( - in_grad_ptr, out_grad_ptr[i], num_blks, in_blk_size, out_blk_sizes[i]); - in_grad_ptr += out_blk_sizes[i]; + stream>>>(input_grad.get_float_ptr() + offset, + output_grad.get_float_ptr(), + num_blks, + input_blk_size, + output_blk_size); + + offset += output_blk_size; } - // checkCUDA(cudaDeviceSynchronize()); + + ASSERT(offset == input_blk_size); } -} // namespace Split -} // namespace Kernels } // namespace FlexFlow diff --git a/lib/kernels/src/cuda/ops/transpose_kernels.cu b/lib/kernels/src/cuda/ops/transpose_kernels.cu index 5becb2190d..2c20636a5c 100644 --- a/lib/kernels/src/cuda/ops/transpose_kernels.cu +++ b/lib/kernels/src/cuda/ops/transpose_kernels.cu @@ -17,6 +17,7 @@ #include "kernels/accessor.h" #include "kernels/legion_ordered/legion_ordered_transform.h" #include "kernels/transpose_kernels_gpu.h" +#include "op-attrs/tensor_dim_permutation.h" #include "utils/exception.h" #include "utils/nonnegative_int/num_elements.h" @@ -28,9 +29,8 @@ struct TransposeStrides { perm[MAX_TENSOR_DIM]; }; -namespace Kernels { -namespace Transpose { - +// Computes out_ptr[o] = out_ptr[o] * beta + in_ptr[i], where the coordinates +// of o and i are related by `info.perm` (see `make_strides`). __global__ void transpose_simple_kernel(std::size_t volume, float const *in_ptr, float *out_ptr, @@ -44,17 +44,21 @@ __global__ void transpose_simple_kernel(std::size_t volume, t -= ratio * info.out_strides[i]; i_idx += ratio * info.in_strides[info.perm[i]]; } - out_ptr[o_idx] += out_ptr[o_idx] * beta + in_ptr[i_idx]; + out_ptr[o_idx] = out_ptr[o_idx] * beta + in_ptr[i_idx]; } } -static TransposeStrides make_strides(TransposeAttrs const &m, +// Builds the stride/permutation info needed to compute +// `output[coord] = input[permuted_coord]`, where +// `permuted_coord[permutation.at_l(d)] = coord[d]`. +static TransposeStrides make_strides(TensorDimPermutation const &permutation, TensorDims const &input_dims, TensorDims const &output_dims) { - ASSERT(get_num_dims(input_dims) == m.permutation.num_tensor_dims()); + ASSERT(get_num_dims(input_dims) == permutation.num_tensor_dims()); + ASSERT(get_num_dims(output_dims) == permutation.num_tensor_dims()); TransposeStrides info; - num_tensor_dims_t num_dims = m.permutation.num_tensor_dims(); + num_tensor_dims_t num_dims = permutation.num_tensor_dims(); info.num_dim = num_dims.int_from_num_tensor_dims(); for (int i = 0; i < info.num_dim; i++) { @@ -65,15 +69,16 @@ static TransposeStrides make_strides(TransposeAttrs const &m, info.in_strides[i] = 1; info.out_strides[i] = 1; } else { + legion_dim_t prev_legion_dim = legion_dim_t{nonnegative_int{i - 1}}; int in_dim_size = - dim_at_idx(input_dims, legion_dim).int_from_positive_int(); + dim_at_idx(input_dims, prev_legion_dim).int_from_positive_int(); int out_dim_size = - dim_at_idx(output_dims, legion_dim).int_from_positive_int(); + dim_at_idx(output_dims, prev_legion_dim).int_from_positive_int(); info.in_strides[i] = info.in_strides[i - 1] * in_dim_size; info.out_strides[i] = info.out_strides[i - 1] * out_dim_size; } - ff_dim_t ff_permuted_dim = m.permutation.at_l(ff_dim); + ff_dim_t ff_permuted_dim = permutation.at_l(ff_dim); legion_dim_t legion_permuted_dim = legion_dim_from_ff_dim(ff_permuted_dim, num_dims); info.perm[i] = legion_permuted_dim.value.unwrap_nonnegative(); @@ -82,38 +87,45 @@ static TransposeStrides make_strides(TransposeAttrs const &m, return info; } -void gpu_forward_kernel(cudaStream_t stream, - TransposeAttrs const &m, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output) { - transpose_simple_kernel<<< - GET_BLOCKS(get_num_elements(output.shape.dims).int_from_positive_int()), - CUDA_NUM_THREADS, - 0, - stream>>>(get_num_elements(output.shape.dims).int_from_positive_int(), - input.get_float_ptr(), - output.get_float_ptr(), - make_strides(m, input.shape.dims, output.shape.dims), - /*beta=*/0.0f); +void transpose_gpu_forward_kernel(cudaStream_t stream, + TransposeAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { + ASSERT(permute_tensor_dims(attrs.permutation, input.shape.dims) == + output.shape.dims); + + int volume = get_num_elements(output.shape.dims).int_from_positive_int(); + + transpose_simple_kernel<<>>( + volume, + input.get_float_ptr(), + output.get_float_ptr(), + make_strides(attrs.permutation, input.shape.dims, output.shape.dims), + /*beta=*/0.0f); } -void gpu_backward_kernel(cudaStream_t stream, - TransposeAttrs const &m, - GenericTensorAccessorR const &out_grad, - GenericTensorAccessorW const &in_grad) { - ASSERT(get_num_dims(in_grad.shape.dims) == m.permutation.num_tensor_dims()); - - transpose_simple_kernel<<< - GET_BLOCKS(get_num_elements(in_grad.shape.dims).int_from_positive_int()), - CUDA_NUM_THREADS, - 0, - stream>>>(get_num_elements(in_grad.shape.dims).int_from_positive_int(), - out_grad.get_float_ptr(), - in_grad.get_float_ptr(), - make_strides(m, out_grad.shape.dims, in_grad.shape.dims), - /*beta=*/1.0f); +void transpose_gpu_backward_kernel(cudaStream_t stream, + TransposeAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { + ASSERT(permute_tensor_dims(attrs.permutation, input_grad.shape.dims) == + output_grad.shape.dims); + + // The gradient of a permutation is the inverse permutation. + TensorDimPermutation inverse_permutation = + invert_tensor_dim_permutation(attrs.permutation); + + int volume = get_num_elements(input_grad.shape.dims).int_from_positive_int(); + + transpose_simple_kernel<<>>( + volume, + output_grad.get_float_ptr(), + input_grad.get_float_ptr(), + make_strides( + inverse_permutation, output_grad.shape.dims, input_grad.shape.dims), + /*beta=*/1.0f); } -} // namespace Transpose -} // namespace Kernels } // namespace FlexFlow diff --git a/lib/kernels/src/cuda/ops/upsample_kernels.cu b/lib/kernels/src/cuda/ops/upsample_kernels.cu index ce91429704..11cb5999e4 100644 --- a/lib/kernels/src/cuda/ops/upsample_kernels.cu +++ b/lib/kernels/src/cuda/ops/upsample_kernels.cu @@ -1,22 +1,175 @@ #include "internal/device.h" +#include "kernels/datatype_dispatch.h" #include "kernels/upsample_kernels_gpu.h" +#include "op-attrs/tensor_dims.h" +#include "utils/containers/require_same.h" namespace FlexFlow { -void upsample_gpu_forward_kernel(ffStream_t stream, +// The shape of an NCHW tensor, flattened into the quantities the kernels below +// need in order to turn a linear index into coordinates. +struct UpsampleNCHWDims { + int num_samples; + int num_channels; + int height; + int width; +}; + +static UpsampleNCHWDims get_nchw_dims(TensorShape const &shape) { + return UpsampleNCHWDims{ + /*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(), + /*height=*/dim_at_idx(shape.dims, ff_dim_t{2_n}).int_from_positive_int(), + /*width=*/dim_at_idx(shape.dims, ff_dim_t{3_n}).int_from_positive_int(), + }; +} + +static void check_shapes(UpsampleAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { + ASSERT(get_num_dims(input_shape.dims) == num_tensor_dims_t{4_n}, + "Currently Upsample only supports 4-dimensional input tensors (i.e., " + "NCHW tensors). " + "If you need other support for other tensor shapes, please create an " + "issue.", + input_shape); + + ASSERT(attrs.mode == UpsampleMode::NEAREST, + "Currently Upsample is only supports mode {}. " + "If you need other support for other modes, please create an issue.", + UpsampleMode::NEAREST); + + UpsampleNCHWDims input_dims = get_nchw_dims(input_shape); + UpsampleNCHWDims output_dims = get_nchw_dims(output_shape); + int scale_factor = attrs.scale_factor.int_from_int_ge_two(); + + ASSERT(output_dims.num_samples == input_dims.num_samples); + ASSERT(output_dims.num_channels == input_dims.num_channels); + ASSERT(output_dims.height == input_dims.height * scale_factor); + ASSERT(output_dims.width == input_dims.width * scale_factor); +} + +// Each thread computes one output element by reading the input element the +// output coordinate maps back to. +template +__global__ void upsample_nearest_forward_kernel(size_t num_output_elements, + T const *input, + T *output, + UpsampleNCHWDims input_dims, + int scale_factor) { + CUDA_KERNEL_LOOP(output_idx, num_output_elements) { + int output_width = input_dims.width * scale_factor; + int output_height = input_dims.height * scale_factor; + + int w = output_idx % output_width; + int h = (output_idx / output_width) % output_height; + int nc = output_idx / (output_width * output_height); + + int input_idx = + (nc * input_dims.height + h / scale_factor) * input_dims.width + + w / scale_factor; + + output[output_idx] = input[input_idx]; + } +} + +// Each thread accumulates one input element by summing over the block of +// output elements that read from it. Iterating over inputs (rather than +// outputs) keeps the accumulation deterministic and avoids atomics. +template +__global__ void upsample_nearest_backward_kernel(size_t num_input_elements, + T const *output_grad, + T *input_grad, + UpsampleNCHWDims input_dims, + int scale_factor) { + CUDA_KERNEL_LOOP(input_idx, num_input_elements) { + int output_width = input_dims.width * scale_factor; + int output_height = input_dims.height * scale_factor; + + int w = input_idx % input_dims.width; + int h = (input_idx / input_dims.width) % input_dims.height; + int nc = input_idx / (input_dims.width * input_dims.height); + + T sum = 0; + for (int dh = 0; dh < scale_factor; dh++) { + for (int dw = 0; dw < scale_factor; dw++) { + int output_idx = + (nc * output_height + h * scale_factor + dh) * output_width + + w * scale_factor + dw; + sum += output_grad[output_idx]; + } + } + + input_grad[input_idx] += sum; + } +} + +template +struct UpsampleGPUForwardKernel { + void operator()(cudaStream_t stream, + UpsampleAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) const { + int num_output_elements = + get_num_elements(output.shape.dims).int_from_positive_int(); + + upsample_nearest_forward_kernel> + <<>>( + num_output_elements, + input.get
(), + output.get
(), + get_nchw_dims(input.shape), + attrs.scale_factor.int_from_int_ge_two()); + } +}; + +void upsample_gpu_forward_kernel(cudaStream_t stream, UpsampleAttrs const &attrs, GenericTensorAccessorR const &input, GenericTensorAccessorW const &output) { - NOT_IMPLEMENTED(); + check_shapes(attrs, input.shape, output.shape); + + DataType data_type = + require_same(input.shape.data_type, output.shape.data_type); + + DataTypeDispatch1{}( + data_type, stream, attrs, input, output); } -void upsample_gpu_backward_kernel(ffStream_t stream, +template +struct UpsampleGPUBackwardKernel { + void operator()(cudaStream_t stream, + UpsampleAttrs const &attrs, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorW const &input_grad) const { + int num_input_elements = + get_num_elements(input_grad.shape.dims).int_from_positive_int(); + + upsample_nearest_backward_kernel> + <<>>( + num_input_elements, + output_grad.get
(), + input_grad.get
(), + get_nchw_dims(input_grad.shape), + attrs.scale_factor.int_from_int_ge_two()); + } +}; + +void upsample_gpu_backward_kernel(cudaStream_t stream, UpsampleAttrs const &attrs, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad) { - NOT_IMPLEMENTED(); + check_shapes(attrs, input_grad.shape, output_grad.shape); + + DataType data_type = + require_same(output_grad.shape.data_type, input_grad.shape.data_type); + + DataTypeDispatch1{}( + data_type, stream, attrs, output_grad, input_grad); } } // namespace FlexFlow diff --git a/lib/kernels/src/kernels/batch_matmul_kernels.cc b/lib/kernels/src/kernels/batch_matmul_kernels.cc index 724bff7136..f30bfa3f34 100644 --- a/lib/kernels/src/kernels/batch_matmul_kernels.cc +++ b/lib/kernels/src/kernels/batch_matmul_kernels.cc @@ -5,17 +5,20 @@ namespace FlexFlow { void batch_matmul_forward_kernel(device_stream_t const &stream, + device_handle_t const &handle, GenericTensorAccessorR const &input_lhs, GenericTensorAccessorR const &input_rhs, GenericTensorAccessorW const &output) { if (stream.is_gpu()) { batch_matmul_gpu_forward_kernel( /*stream=*/stream.require_gpu(), + /*handle=*/handle.require_for_gpu(), /*input_lhs=*/input_lhs, /*input_rhs=*/input_rhs, /*output=*/output); } else { ASSERT(stream.is_cpu()); + ASSERT(handle.is_for_cpu()); batch_matmul_cpu_forward_kernel( /*input_lhs=*/input_lhs, @@ -26,6 +29,7 @@ void batch_matmul_forward_kernel(device_stream_t const &stream, void batch_matmul_backward_kernel( device_stream_t const &stream, + device_handle_t const &handle, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input_lhs, @@ -35,6 +39,7 @@ void batch_matmul_backward_kernel( if (stream.is_gpu()) { batch_matmul_gpu_backward_kernel( /*stream=*/stream.require_gpu(), + /*handle=*/handle.require_for_gpu(), /*output=*/output, /*output_grad=*/output_grad, /*input_lhs=*/input_lhs, @@ -43,6 +48,7 @@ void batch_matmul_backward_kernel( /*input_rhs_grad=*/input_rhs_grad); } else { ASSERT(stream.is_cpu()); + ASSERT(handle.is_for_cpu()); batch_matmul_cpu_backward_kernel( /*output=*/output, diff --git a/lib/kernels/src/kernels/batch_norm_kernels.cc b/lib/kernels/src/kernels/batch_norm_kernels.cc index e23f6a89e2..546aa1ce5b 100644 --- a/lib/kernels/src/kernels/batch_norm_kernels.cc +++ b/lib/kernels/src/kernels/batch_norm_kernels.cc @@ -1,107 +1,111 @@ #include "kernels/batch_norm_kernels.h" #include "kernels/batch_norm_kernels_cpu.h" #include "kernels/batch_norm_kernels_gpu.h" +#include "utils/optional.h" -namespace FlexFlow::Kernels::BatchNorm { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - Allocator &allocator, - float *runningMean, - int output_n, - int output_c, - int output_h, - int output_w, - bool relu) { + batch_norm_init_kernel(DeviceType device_type, + Allocator &allocator, + BatchNormAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { if (device_type == DeviceType::GPU) { - return gpu_init_kernel( - /*handle=*/handle.require_for_gpu(), + return batch_norm_gpu_init_kernel( /*allocator=*/allocator, - /*runningMean=*/runningMean, - /*output_n=*/output_n, - /*output_c=*/output_c, - /*output_h=*/output_h, - /*output_w=*/output_w, - /*relu=*/relu); + /*attrs=*/attrs, + /*input_shape=*/input_shape, + /*output_shape=*/output_shape); } else { ASSERT(device_type == DeviceType::CPU); - ASSERT(handle.is_for_cpu()); return std::nullopt; } } -void forward_kernel(device_stream_t const &stream, - BatchNormPerDeviceState const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr) { +void batch_norm_forward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + batch_norm_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state, - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr, - /*scale_ptr=*/scale_ptr, - /*bias_ptr=*/bias_ptr); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*input=*/input, + /*gamma=*/gamma, + /*beta=*/beta, + /*output=*/output); } else { ASSERT(stream.is_cpu()); - cpu_forward_kernel( - /*per_device_state=*/per_device_state, - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr, - /*scale_ptr=*/scale_ptr, - /*bias_ptr=*/bias_ptr); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + batch_norm_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*gamma=*/gamma, + /*beta=*/beta, + /*output=*/output); } } -void backward_kernel(device_stream_t const &stream, - BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements) { +void batch_norm_backward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + batch_norm_gpu_backward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state, - /*output_ptr=*/output_ptr, - /*output_grad_ptr=*/output_grad_ptr, - /*input_ptr=*/input_ptr, - /*input_grad_ptr=*/input_grad_ptr, - /*scale_ptr=*/scale_ptr, - /*scale_grad_ptr=*/scale_grad_ptr, - /*bias_grad_ptr=*/bias_grad_ptr, - /*numElements=*/numElements); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad, + /*gamma=*/gamma, + /*gamma_grad=*/gamma_grad, + /*beta_grad=*/beta_grad); } else { ASSERT(stream.is_cpu()); - cpu_backward_kernel( - /*per_device_state=*/per_device_state, - /*output_ptr=*/output_ptr, - /*output_grad_ptr=*/output_grad_ptr, - /*input_ptr=*/input_ptr, - /*input_grad_ptr=*/input_grad_ptr, - /*scale_ptr=*/scale_ptr, - /*scale_grad_ptr=*/scale_grad_ptr, - /*bias_grad_ptr=*/bias_grad_ptr, - /*numElements=*/numElements); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + batch_norm_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad, + /*gamma=*/gamma, + /*gamma_grad=*/gamma_grad, + /*beta_grad=*/beta_grad); } } -void cleanup_kernel(DeviceType device_type, - Allocator &allocator, - std::optional &per_device_state) { +void batch_norm_cleanup_kernel( + DeviceType device_type, + Allocator &allocator, + std::optional &per_device_state) { if (device_type == DeviceType::GPU) { - gpu_cleanup_kernel(allocator, per_device_state.value()); + batch_norm_gpu_cleanup_kernel(allocator, per_device_state.value()); } else { ASSERT(device_type == DeviceType::CPU); - ASSERT(per_device_state == std::nullopt); + ASSERT(!per_device_state.has_value()); } } -} // namespace FlexFlow::Kernels::BatchNorm +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc b/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc index be440ac4ea..5e16a11fe7 100644 --- a/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc +++ b/lib/kernels/src/kernels/batch_norm_kernels_cpu.cc @@ -1,25 +1,25 @@ #include "kernels/batch_norm_kernels_cpu.h" +#include "utils/exception.h" -namespace FlexFlow::Kernels::BatchNorm { +namespace FlexFlow { -void cpu_forward_kernel(BatchNormPerDeviceState const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *scale_ptr, - float const *bias_ptr) { +void batch_norm_cpu_forward_kernel(BatchNormAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorR const &beta, + GenericTensorAccessorW const &output) { NOT_IMPLEMENTED(); } -void cpu_backward_kernel(BatchNormPerDeviceState const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *scale_ptr, - float *scale_grad_ptr, - float *bias_grad_ptr, - size_t numElements) { +void batch_norm_cpu_backward_kernel(BatchNormAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &gamma, + GenericTensorAccessorW const &gamma_grad, + GenericTensorAccessorW const &beta_grad) { NOT_IMPLEMENTED(); } -} // namespace FlexFlow::Kernels::BatchNorm +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/concat_kernels.cc b/lib/kernels/src/kernels/concat_kernels.cc index 8d0c3112aa..c64ebece74 100644 --- a/lib/kernels/src/kernels/concat_kernels.cc +++ b/lib/kernels/src/kernels/concat_kernels.cc @@ -2,44 +2,51 @@ #include "kernels/concat_kernels_cpu.h" #include "kernels/concat_kernels_gpu.h" -namespace FlexFlow::Kernels::Concat { +namespace FlexFlow { -void forward_kernel(device_stream_t const &stream, - GenericTensorAccessorW const &output, - std::vector const &inputs, - ff_dim_t axis) { +void concat_forward_kernel(device_stream_t const &stream, + ConcatAttrs const &attrs, + std::vector const &inputs, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + concat_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*output=*/output, + /*attrs=*/attrs, /*inputs=*/inputs, - /*axis=*/axis); + /*output=*/output); } else { ASSERT(stream.is_cpu()); - cpu_forward_kernel( - /*output=*/output, + concat_cpu_forward_kernel( + /*attrs=*/attrs, /*inputs=*/inputs, - /*axis=*/axis); + /*output=*/output); } } -void backward_kernel(device_stream_t const &stream, - GenericTensorAccessorR const &output_grad, - std::vector const &input_grads, - ff_dim_t axis) { +void concat_backward_kernel( + device_stream_t const &stream, + ConcatAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + std::vector const &inputs, + std::vector const &input_grads) { if (stream.is_gpu()) { - gpu_backward_kernel( + concat_gpu_backward_kernel( /*stream=*/stream.require_gpu(), + /*attrs=*/attrs, + /*output=*/output, /*output_grad=*/output_grad, - /*input_grads=*/input_grads, - /*axis=*/axis); + /*inputs=*/inputs, + /*input_grads=*/input_grads); } else { ASSERT(stream.is_cpu()); - cpu_backward_kernel( + concat_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, /*output_grad=*/output_grad, - /*input_grads=*/input_grads, - /*axis=*/axis); + /*inputs=*/inputs, + /*input_grads=*/input_grads); } } -} // namespace FlexFlow::Kernels::Concat +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/concat_kernels_cpu.cc b/lib/kernels/src/kernels/concat_kernels_cpu.cc index 03bbff67bb..4a1817c239 100644 --- a/lib/kernels/src/kernels/concat_kernels_cpu.cc +++ b/lib/kernels/src/kernels/concat_kernels_cpu.cc @@ -1,17 +1,117 @@ #include "kernels/concat_kernels_cpu.h" +#include "kernels/datatype_dispatch.h" +#include "op-attrs/tensor_dims.h" +#include "op-attrs/tensor_dims_coord.h" +#include "utils/containers/require_same.h" -namespace FlexFlow::Kernels::Concat { +namespace FlexFlow { -void cpu_forward_kernel(GenericTensorAccessorW const &output, - std::vector const &inputs, - ff_dim_t axis) { - NOT_IMPLEMENTED(); +// The coordinate in the concatenated tensor that `input_coord` (a coordinate +// in the input starting at `offset` along `axis`) corresponds to. +static TensorDimsCoord output_coord_from_input_coord( + TensorDimsCoord const &input_coord, ff_dim_t axis, nonnegative_int offset) { + TensorDimsCoord output_coord = input_coord; + tensor_dims_coord_at_idx(output_coord, axis) += offset; + return output_coord; } -void cpu_backward_kernel(GenericTensorAccessorR const &output_grad, - std::vector const &input_grads, - ff_dim_t axis) { - NOT_IMPLEMENTED(); +// Checks that the inputs agree with the output on every dim but `axis`, and +// that their sizes along `axis` add up to the output's. +static void check_shapes(ConcatAttrs const &attrs, + std::vector const &input_shapes, + TensorShape const &output_shape) { + ASSERT(input_shapes.size() == attrs.num_inputs.int_from_int_ge_two()); + + nonnegative_int total = 0_n; + for (TensorShape const &input_shape : input_shapes) { + require_same(input_shape.data_type, output_shape.data_type); + ASSERT( + tensor_dims_drop_dims(input_shape.dims, + [&](ff_dim_t d) { return d == attrs.axis; }) == + tensor_dims_drop_dims(output_shape.dims, + [&](ff_dim_t d) { return d == attrs.axis; }), + input_shape, + output_shape, + attrs.axis); + total += dim_at_idx(input_shape.dims, attrs.axis) + .nonnegative_int_from_positive_int(); + } + + ASSERT(total == dim_at_idx(output_shape.dims, attrs.axis) + .nonnegative_int_from_positive_int()); +} + +template +struct ConcatCPUForwardKernel { + void operator()(ConcatAttrs const &attrs, + std::vector const &inputs, + GenericTensorAccessorW const &output) const { + nonnegative_int offset = 0_n; + + for (GenericTensorAccessorR const &input : inputs) { + for (TensorDimsCoord const &input_coord : + get_tensor_dims_coord_set(input.shape.dims)) { + output.at
(output_coord_from_input_coord( + input_coord, attrs.axis, offset)) = input.at
(input_coord); + } + + offset += dim_at_idx(input.shape.dims, attrs.axis) + .nonnegative_int_from_positive_int(); + } + } +}; + +void concat_cpu_forward_kernel( + ConcatAttrs const &attrs, + std::vector const &inputs, + GenericTensorAccessorW const &output) { + check_shapes(attrs, + transform(inputs, + [](GenericTensorAccessorR const &input) { + return input.shape; + }), + output.shape); + + DataTypeDispatch1{}( + output.shape.data_type, attrs, inputs, output); +} + +template +struct ConcatCPUBackwardKernel { + void + operator()(ConcatAttrs const &attrs, + GenericTensorAccessorR const &output_grad, + std::vector const &input_grads) const { + nonnegative_int offset = 0_n; + + for (GenericTensorAccessorW const &input_grad : input_grads) { + for (TensorDimsCoord const &input_coord : + get_tensor_dims_coord_set(input_grad.shape.dims)) { + input_grad.at
(input_coord) += output_grad.at
( + output_coord_from_input_coord(input_coord, attrs.axis, offset)); + } + + offset += dim_at_idx(input_grad.shape.dims, attrs.axis) + .nonnegative_int_from_positive_int(); + } + } +}; + +void concat_cpu_backward_kernel( + ConcatAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + std::vector const &inputs, + std::vector const &input_grads) { + check_shapes(attrs, + transform(input_grads, + [](GenericTensorAccessorW const &input_grad) { + return input_grad.shape; + }), + output_grad.shape); + + DataTypeDispatch1{}( + output_grad.shape.data_type, attrs, output_grad, input_grads); } -} // namespace FlexFlow::Kernels::Concat +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/conv_2d_kernels.cc b/lib/kernels/src/kernels/conv_2d_kernels.cc index 3008e7d1c0..6a612aab39 100644 --- a/lib/kernels/src/kernels/conv_2d_kernels.cc +++ b/lib/kernels/src/kernels/conv_2d_kernels.cc @@ -1,39 +1,22 @@ #include "kernels/conv_2d_kernels.h" #include "kernels/conv_2d_kernels_cpu.h" #include "kernels/conv_2d_kernels_gpu.h" +#include "utils/optional.h" -namespace FlexFlow::Kernels::Conv2D { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - std::optional activation, - int kernel_h, - int kernel_w, - int groups, - int padding_h, - int padding_w, - int stride_h, - int stride_w, - GenericTensorAccessorW const &input, - GenericTensorAccessorW const &output, - float const *filter_ptr, - float *filter_grad_ptr) { + conv_2d_init_kernel(DeviceType device_type, + device_handle_t const &handle, + Conv2DAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { if (device_type == DeviceType::GPU) { - return gpu_init_kernel( + return conv_2d_gpu_init_kernel( /*handle=*/handle.require_for_gpu(), - /*activation=*/activation, - /*kernel_h=*/kernel_h, - /*kernel_w=*/kernel_w, - /*groups=*/groups, - /*padding_h=*/padding_h, - /*padding_w=*/padding_w, - /*stride_h=*/stride_h, - /*stride_w=*/stride_w, - /*input=*/input, - /*output=*/output, - /*filter_ptr=*/filter_ptr, - /*filter_grad_ptr=*/filter_grad_ptr); + /*attrs=*/attrs, + /*input_shape=*/input_shape, + /*output_shape=*/output_shape); } else { ASSERT(device_type == DeviceType::CPU); ASSERT(handle.is_for_cpu()); @@ -41,78 +24,88 @@ std::optional } } -void forward_kernel(device_stream_t const &stream, - std::optional const &per_device_state, - float const *input_ptr, - float *output_ptr, - float const *filter_ptr, - float const *bias_ptr, - std::optional activation) { +void conv_2d_forward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &filter, + std::optional const &bias, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + conv_2d_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr, - /*filter_ptr=*/filter_ptr, - /*bias_ptr=*/bias_ptr, - /*activation=*/activation); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*input=*/input, + /*filter=*/filter, + /*bias=*/bias, + /*output=*/output); } else { ASSERT(stream.is_cpu()); - cpu_forward_kernel( - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr, - /*filter_ptr=*/filter_ptr, - /*bias_ptr=*/bias_ptr, - /*activation=*/activation); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + conv_2d_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*filter=*/filter, + /*bias=*/bias, + /*output=*/output); } } -void backward_kernel( +void conv_2d_backward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *filter_ptr, - float *filter_grad_ptr, - float *bias_grad_ptr, - std::optional activation) { + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &filter, + GenericTensorAccessorW const &filter_grad, + std::optional const &bias_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + conv_2d_gpu_backward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*output_ptr=*/output_ptr, - /*output_grad_ptr=*/output_grad_ptr, - /*input_ptr=*/input_ptr, - /*input_grad_ptr=*/input_grad_ptr, - /*filter_ptr=*/filter_ptr, - /*filter_grad_ptr=*/filter_grad_ptr, - /*bias_grad_ptr=*/bias_grad_ptr, - /*activation=*/activation); + /*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, + /*filter=*/filter, + /*filter_grad=*/filter_grad, + /*bias_grad=*/bias_grad); } else { ASSERT(stream.is_cpu()); - cpu_backward_kernel( - /*output_ptr=*/output_ptr, - /*output_grad_ptr=*/output_grad_ptr, - /*input_ptr=*/input_ptr, - /*input_grad_ptr=*/input_grad_ptr, - /*filter_ptr=*/filter_ptr, - /*filter_grad_ptr=*/filter_grad_ptr, - /*bias_grad_ptr=*/bias_grad_ptr, - /*activation=*/activation); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + conv_2d_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad, + /*filter=*/filter, + /*filter_grad=*/filter_grad, + /*bias_grad=*/bias_grad); } } -void cleanup_kernel(DeviceType device_type, - std::optional &per_device_state) { +void conv_2d_cleanup_kernel( + DeviceType device_type, + std::optional &per_device_state) { if (device_type == DeviceType::GPU) { - gpu_cleanup_kernel(per_device_state.value()); + conv_2d_gpu_cleanup_kernel(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::Conv2D +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/conv_2d_kernels_cpu.cc b/lib/kernels/src/kernels/conv_2d_kernels_cpu.cc index c595ecb586..7267159d59 100644 --- a/lib/kernels/src/kernels/conv_2d_kernels_cpu.cc +++ b/lib/kernels/src/kernels/conv_2d_kernels_cpu.cc @@ -1,25 +1,27 @@ #include "kernels/conv_2d_kernels_cpu.h" #include "utils/exception.h" -namespace FlexFlow::Kernels::Conv2D { +namespace FlexFlow { -void cpu_forward_kernel(float const *input_ptr, - float *output_ptr, - float const *filter_ptr, - float const *bias_ptr, - std::optional const &activation) { +void conv_2d_cpu_forward_kernel( + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorR const &filter, + std::optional const &bias, + GenericTensorAccessorW const &output) { NOT_IMPLEMENTED(); } -void cpu_backward_kernel(float const *output_ptr, - float *output_grad_ptr, - float const *input_ptr, - float *input_grad_ptr, - float const *filter_ptr, - float *filter_grad_ptr, - float *bias_grad_ptr, - std::optional const &activation) { +void conv_2d_cpu_backward_kernel( + Conv2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad, + GenericTensorAccessorR const &filter, + GenericTensorAccessorW const &filter_grad, + std::optional const &bias_grad) { NOT_IMPLEMENTED(); } -} // namespace FlexFlow::Kernels::Conv2D +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/pool_2d_kernels.cc b/lib/kernels/src/kernels/pool_2d_kernels.cc index f8f5571716..8dfe1e7ae1 100644 --- a/lib/kernels/src/kernels/pool_2d_kernels.cc +++ b/lib/kernels/src/kernels/pool_2d_kernels.cc @@ -1,106 +1,93 @@ #include "kernels/pool_2d_kernels.h" #include "kernels/pool_2d_kernels_cpu.h" #include "kernels/pool_2d_kernels_gpu.h" -#include +#include "utils/optional.h" -namespace FlexFlow::Kernels::Pool2D { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - std::optional activation, - int input_w, - int input_h, - int input_c, - int input_n, - int output_w, - int output_h, - int output_c, - int output_n, - int pad_h, - int pad_w, - int kernel_h, - int kernel_w, - int stride_h, - int stride_w, - PoolOp pool_type) { + pool_2d_init_kernel(DeviceType device_type, + Pool2DAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { if (device_type == DeviceType::GPU) { - return gpu_init_kernel( - /*handle=*/handle.require_for_gpu(), - /*activation=*/activation, - /*input_w=*/input_w, - /*input_h=*/input_h, - /*input_c=*/input_c, - /*input_n=*/input_n, - /*output_w=*/output_w, - /*output_h=*/output_h, - /*output_c=*/output_c, - /*output_n=*/output_n, - /*pad_h=*/pad_h, - /*pad_w=*/pad_w, - /*kernel_h=*/kernel_h, - /*kernel_w=*/kernel_w, - /*stride_h=*/stride_h, - /*stride_w=*/stride_w, - /*pool_type=*/pool_type); + return pool_2d_gpu_init_kernel( + /*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, - std::optional const &per_device_state, - void const *input_ptr, - void *output_ptr) { +void pool_2d_forward_kernel( + device_stream_t const &stream, + device_handle_t const &handle, + std::optional const &per_device_state, + Pool2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + pool_2d_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); } else { ASSERT(stream.is_cpu()); - cpu_forward_kernel( - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + pool_2d_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); } } -void backward_kernel( +void pool_2d_backward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - void const *output_ptr, - void const *output_grad_ptr, - void const *input_ptr, - void *input_grad_ptr) { + Pool2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + pool_2d_gpu_backward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*output_ptr=*/output_ptr, - /*output_grad_ptr=*/output_grad_ptr, - /*input_ptr=*/input_ptr, - /*input_grad_ptr=*/input_grad_ptr); + /*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); } else { ASSERT(stream.is_cpu()); - cpu_backward_kernel( - /*output_ptr=*/output_ptr, - /*output_grad_ptr=*/output_grad_ptr, - /*input_ptr=*/input_ptr, - /*input_grad_ptr=*/input_grad_ptr); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + pool_2d_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); } } -void cleanup_kernel(DeviceType device_type, - std::optional &per_device_state) { +void pool_2d_cleanup_kernel( + DeviceType device_type, + std::optional &per_device_state) { if (device_type == DeviceType::GPU) { - gpu_cleanup_kernel(per_device_state.value()); + pool_2d_gpu_cleanup_kernel(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::Pool2D +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/pool_2d_kernels_cpu.cc b/lib/kernels/src/kernels/pool_2d_kernels_cpu.cc index f2d2141e96..c688ef8129 100644 --- a/lib/kernels/src/kernels/pool_2d_kernels_cpu.cc +++ b/lib/kernels/src/kernels/pool_2d_kernels_cpu.cc @@ -1,17 +1,20 @@ #include "kernels/pool_2d_kernels_cpu.h" #include "utils/exception.h" -namespace FlexFlow::Kernels::Pool2D { +namespace FlexFlow { -void cpu_forward_kernel(void const *input_ptr, void *output_ptr) { +void pool_2d_cpu_forward_kernel(Pool2DAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { NOT_IMPLEMENTED(); } -void cpu_backward_kernel(void const *output_ptr, - void const *output_grad_ptr, - void const *input_ptr, - void *input_grad_ptr) { +void pool_2d_cpu_backward_kernel(Pool2DAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { NOT_IMPLEMENTED(); } -} // namespace FlexFlow::Kernels::Pool2D +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/reshape_kernels.cc b/lib/kernels/src/kernels/reshape_kernels.cc index 2ac90352bb..25d79ff5a6 100644 --- a/lib/kernels/src/kernels/reshape_kernels.cc +++ b/lib/kernels/src/kernels/reshape_kernels.cc @@ -2,38 +2,44 @@ #include "kernels/reshape_kernels_cpu.h" #include "kernels/reshape_kernels_gpu.h" -namespace FlexFlow::Kernels::Reshape { +namespace FlexFlow { -void forward_kernel(device_stream_t const &stream, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output) { +void reshape_forward_kernel(device_stream_t const &stream, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + reshape_gpu_forward_kernel( /*stream=*/stream.require_gpu(), /*input=*/input, /*output=*/output); } else { ASSERT(stream.is_cpu()); - cpu_forward_kernel( + reshape_cpu_forward_kernel( /*input=*/input, /*output=*/output); } } -void backward_kernel(device_stream_t const &stream, - GenericTensorAccessorR const &output, - GenericTensorAccessorW const &input) { +void reshape_backward_kernel(device_stream_t const &stream, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + reshape_gpu_backward_kernel( /*stream=*/stream.require_gpu(), /*output=*/output, - /*input=*/input); + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); } else { ASSERT(stream.is_cpu()); - cpu_backward_kernel( + reshape_cpu_backward_kernel( /*output=*/output, - /*input=*/input); + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); } } -} // namespace FlexFlow::Kernels::Reshape +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/reshape_kernels_cpu.cc b/lib/kernels/src/kernels/reshape_kernels_cpu.cc index b48272cdde..5ac5407244 100644 --- a/lib/kernels/src/kernels/reshape_kernels_cpu.cc +++ b/lib/kernels/src/kernels/reshape_kernels_cpu.cc @@ -1,15 +1,58 @@ #include "kernels/reshape_kernels_cpu.h" +#include "kernels/datatype_dispatch.h" +#include "op-attrs/tensor_dims.h" +#include "utils/containers/require_same.h" -namespace FlexFlow::Kernels::Reshape { +namespace FlexFlow { -void cpu_forward_kernel(GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output) { - NOT_IMPLEMENTED(); +// Reshape does not move any data: the input and output have different shapes +// but the same contiguous layout, so viewing one with the other's shape lines +// their elements up. +static GenericTensorAccessorW view_with_shape(GenericTensorAccessorW const &acc, + TensorShape const &shape) { + return GenericTensorAccessorW{ + /*shape=*/shape, + /*ptr=*/acc.ptr, + /*device_type=*/acc.device_type, + }; } -void cpu_backward_kernel(GenericTensorAccessorR const &output, - GenericTensorAccessorW const &input) { - NOT_IMPLEMENTED(); +static void check_shapes(TensorShape const &input_shape, + TensorShape const &output_shape) { + require_same(input_shape.data_type, output_shape.data_type); + require_same(get_num_elements(input_shape.dims), + get_num_elements(output_shape.dims)); } -} // namespace FlexFlow::Kernels::Reshape +void reshape_cpu_forward_kernel(GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { + check_shapes(input.shape, output.shape); + + copy_accessor_data_to_l_from_r(output, input); +} + +template +struct ReshapeCPUBackwardKernel { + void operator()(GenericTensorAccessorR const &output_grad, + GenericTensorAccessorW const &input_grad) const { + GenericTensorAccessorW input_grad_view = + view_with_shape(input_grad, output_grad.shape); + + for (TensorDimsCoord const &coord : + get_tensor_dims_coord_set(output_grad.shape.dims)) { + input_grad_view.at
(coord) += output_grad.at
(coord); + } + } +}; + +void reshape_cpu_backward_kernel(GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { + check_shapes(input_grad.shape, output_grad.shape); + + DataTypeDispatch1{}( + output_grad.shape.data_type, output_grad, input_grad); +} + +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/split_kernels.cc b/lib/kernels/src/kernels/split_kernels.cc index f38ae2e8af..4443c31764 100644 --- a/lib/kernels/src/kernels/split_kernels.cc +++ b/lib/kernels/src/kernels/split_kernels.cc @@ -1,63 +1,52 @@ #include "kernels/split_kernels.h" #include "kernels/split_kernels_cpu.h" #include "kernels/split_kernels_gpu.h" -#include -namespace FlexFlow::Kernels::Split { +namespace FlexFlow { -void forward_kernel(device_stream_t const &stream, - float **out_ptrs, - float const *in_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs) { +void split_forward_kernel(device_stream_t const &stream, + SplitAttrs const &attrs, + GenericTensorAccessorR const &input, + std::vector const &outputs) { if (stream.is_gpu()) { - gpu_forward_kernel( + split_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*out_ptrs=*/out_ptrs, - /*in_ptr=*/in_ptr, - /*out_blk_sizes=*/out_blk_sizes, - /*in_blk_size=*/in_blk_size, - /*num_blks=*/num_blks, - /*numOutputs=*/numOutputs); + /*attrs=*/attrs, + /*input=*/input, + /*outputs=*/outputs); } else { - cpu_forward_kernel( - /*out_ptrs=*/out_ptrs, - /*in_ptr=*/in_ptr, - /*out_blk_sizes=*/out_blk_sizes, - /*in_blk_size=*/in_blk_size, - /*num_blks=*/num_blks, - /*numOutputs=*/numOutputs); + ASSERT(stream.is_cpu()); + split_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*outputs=*/outputs); } } -void backward_kernel(device_stream_t const &stream, - float *in_grad_ptr, - float const **out_grad_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs) { +void split_backward_kernel( + device_stream_t const &stream, + SplitAttrs const &attrs, + std::vector const &outputs, + std::vector const &output_grads, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + split_gpu_backward_kernel( /*stream=*/stream.require_gpu(), - /*in_grad_ptr=*/in_grad_ptr, - /*out_grad_ptr=*/out_grad_ptr, - /*out_blk_sizes=*/out_blk_sizes, - /*in_blk_size=*/in_blk_size, - /*num_blks=*/num_blks, - /*numOutputs=*/numOutputs); + /*attrs=*/attrs, + /*outputs=*/outputs, + /*output_grads=*/output_grads, + /*input=*/input, + /*input_grad=*/input_grad); } else { ASSERT(stream.is_cpu()); - cpu_backward_kernel( - /*in_grad_ptr=*/in_grad_ptr, - /*out_grad_ptr=*/out_grad_ptr, - /*out_blk_sizes=*/out_blk_sizes, - /*in_blk_size=*/in_blk_size, - /*num_blks=*/num_blks, - /*numOutputs=*/numOutputs); + split_cpu_backward_kernel( + /*attrs=*/attrs, + /*outputs=*/outputs, + /*output_grads=*/output_grads, + /*input=*/input, + /*input_grad=*/input_grad); } } -} // namespace FlexFlow::Kernels::Split +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/split_kernels_cpu.cc b/lib/kernels/src/kernels/split_kernels_cpu.cc index 1639848ef4..cd4161fe2a 100644 --- a/lib/kernels/src/kernels/split_kernels_cpu.cc +++ b/lib/kernels/src/kernels/split_kernels_cpu.cc @@ -1,24 +1,123 @@ #include "kernels/split_kernels_cpu.h" -#include "utils/exception.h" +#include "kernels/datatype_dispatch.h" +#include "op-attrs/tensor_dims.h" +#include "op-attrs/tensor_dims_coord.h" +#include "utils/containers/require_same.h" -namespace FlexFlow::Kernels::Split { +namespace FlexFlow { -void cpu_forward_kernel(float **out_ptrs, - float const *in_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs) { - NOT_IMPLEMENTED(); +// The coordinate in the tensor being split that `output_coord` (a coordinate +// in the output starting at `offset` along `axis`) corresponds to. +static TensorDimsCoord + input_coord_from_output_coord(TensorDimsCoord const &output_coord, + ff_dim_t axis, + nonnegative_int offset) { + TensorDimsCoord input_coord = output_coord; + tensor_dims_coord_at_idx(input_coord, axis) += offset; + return input_coord; } -void cpu_backward_kernel(float *in_grad_ptr, - float const **out_grad_ptr, - int const *out_blk_sizes, - int in_blk_size, - int num_blks, - int numOutputs) { - NOT_IMPLEMENTED(); +// Checks that the outputs agree with the input on every dim but `axis`, and +// that their sizes along `axis` are the requested splits and add up to the +// input's. +static void check_shapes(SplitAttrs const &attrs, + TensorShape const &input_shape, + std::vector const &output_shapes) { + ASSERT(output_shapes.size() == attrs.splits.size()); + + nonnegative_int total = 0_n; + for (int i = 0; i < output_shapes.size(); i++) { + TensorShape const &output_shape = output_shapes.at(i); + + require_same(input_shape.data_type, output_shape.data_type); + ASSERT( + tensor_dims_drop_dims(output_shape.dims, + [&](ff_dim_t d) { return d == attrs.axis; }) == + tensor_dims_drop_dims(input_shape.dims, + [&](ff_dim_t d) { return d == attrs.axis; }), + input_shape, + output_shape, + attrs.axis); + ASSERT(dim_at_idx(output_shape.dims, attrs.axis) == attrs.splits.at(i)); + + total += dim_at_idx(output_shape.dims, attrs.axis) + .nonnegative_int_from_positive_int(); + } + + ASSERT(total == dim_at_idx(input_shape.dims, attrs.axis) + .nonnegative_int_from_positive_int()); +} + +template +struct SplitCPUForwardKernel { + void operator()(SplitAttrs const &attrs, + GenericTensorAccessorR const &input, + std::vector const &outputs) const { + nonnegative_int offset = 0_n; + + for (GenericTensorAccessorW const &output : outputs) { + for (TensorDimsCoord const &output_coord : + get_tensor_dims_coord_set(output.shape.dims)) { + output.at
(output_coord) = input.at
( + input_coord_from_output_coord(output_coord, attrs.axis, offset)); + } + + offset += dim_at_idx(output.shape.dims, attrs.axis) + .nonnegative_int_from_positive_int(); + } + } +}; + +void split_cpu_forward_kernel( + SplitAttrs const &attrs, + GenericTensorAccessorR const &input, + std::vector const &outputs) { + check_shapes(attrs, + input.shape, + transform(outputs, [](GenericTensorAccessorW const &output) { + return output.shape; + })); + + DataTypeDispatch1{}( + input.shape.data_type, attrs, input, outputs); +} + +template +struct SplitCPUBackwardKernel { + void operator()(SplitAttrs const &attrs, + std::vector const &output_grads, + GenericTensorAccessorW const &input_grad) const { + nonnegative_int offset = 0_n; + + for (GenericTensorAccessorR const &output_grad : output_grads) { + for (TensorDimsCoord const &output_coord : + get_tensor_dims_coord_set(output_grad.shape.dims)) { + input_grad.at
( + input_coord_from_output_coord(output_coord, attrs.axis, offset)) += + output_grad.at
(output_coord); + } + + offset += dim_at_idx(output_grad.shape.dims, attrs.axis) + .nonnegative_int_from_positive_int(); + } + } +}; + +void split_cpu_backward_kernel( + SplitAttrs const &attrs, + std::vector const &outputs, + std::vector const &output_grads, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { + check_shapes( + attrs, + input_grad.shape, + transform(output_grads, [](GenericTensorAccessorR const &output_grad) { + return output_grad.shape; + })); + + DataTypeDispatch1{}( + input_grad.shape.data_type, attrs, output_grads, input_grad); } -} // namespace FlexFlow::Kernels::Split +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/transpose_kernels.cc b/lib/kernels/src/kernels/transpose_kernels.cc index bb3775a073..6e29ef42a6 100644 --- a/lib/kernels/src/kernels/transpose_kernels.cc +++ b/lib/kernels/src/kernels/transpose_kernels.cc @@ -2,44 +2,50 @@ #include "kernels/transpose_kernels_cpu.h" #include "kernels/transpose_kernels_gpu.h" -namespace FlexFlow::Kernels::Transpose { +namespace FlexFlow { -void forward_kernel(device_stream_t const &stream, - TransposeAttrs const &attrs, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output) { +void transpose_forward_kernel(device_stream_t const &stream, + TransposeAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + transpose_gpu_forward_kernel( /*stream=*/stream.require_gpu(), /*attrs=*/attrs, /*input=*/input, /*output=*/output); } else { ASSERT(stream.is_cpu()); - cpu_forward_kernel( + transpose_cpu_forward_kernel( /*attrs=*/attrs, /*input=*/input, /*output=*/output); } } -void backward_kernel(device_stream_t const &stream, - TransposeAttrs const &attrs, - GenericTensorAccessorR const &out_grad, - GenericTensorAccessorW const &in_grad) { +void transpose_backward_kernel(device_stream_t const &stream, + TransposeAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + transpose_gpu_backward_kernel( /*stream=*/stream.require_gpu(), /*attrs=*/attrs, - /*out_grad=*/out_grad, - /*in_grad=*/in_grad); + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); } else { ASSERT(stream.is_cpu()); - cpu_forward_kernel( + transpose_cpu_backward_kernel( /*attrs=*/attrs, - /*out_grad=*/out_grad, - /*in_grad=*/in_grad); + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); } } -} // namespace FlexFlow::Kernels::Transpose +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/transpose_kernels_cpu.cc b/lib/kernels/src/kernels/transpose_kernels_cpu.cc index 7950e71eb4..c48e48394d 100644 --- a/lib/kernels/src/kernels/transpose_kernels_cpu.cc +++ b/lib/kernels/src/kernels/transpose_kernels_cpu.cc @@ -1,17 +1,82 @@ #include "kernels/transpose_kernels_cpu.h" +#include "kernels/datatype_dispatch.h" +#include "op-attrs/ff_ordered/ff_ordered_from_map.h" +#include "op-attrs/ff_ordered/map_from_ff_ordered.h" +#include "op-attrs/tensor_dim_permutation.h" +#include "op-attrs/tensor_dims.h" +#include "utils/containers/map_keys.h" +#include "utils/containers/require_same.h" -namespace FlexFlow::Kernels::Transpose { +namespace FlexFlow { -void cpu_forward_kernel(TransposeAttrs const &attrs, - GenericTensorAccessorR const &input, - GenericTensorAccessorW const &output) { - NOT_IMPLEMENTED(); +// The coordinate in the transposed tensor that `input_coord` corresponds to, +// i.e. `output_coord[d] = input_coord[permutation.at_l(d)]`. This is the same +// relabelling `permute_tensor_dims` applies to the dims themselves. +static TensorDimsCoord + permute_tensor_dims_coord(TensorDimPermutation const &permutation, + TensorDimsCoord const &input_coord) { + return TensorDimsCoord{ + ff_ordered_from_map( + map_keys(map_from_ff_ordered(input_coord.ff_ordered), + [&](ff_dim_t d) { return permutation.at_r(d); })), + }; } -void cpu_backward_kernel(TransposeAttrs const &attrs, - GenericTensorAccessorR const &out_grad, - GenericTensorAccessorW const &in_grad) { - NOT_IMPLEMENTED(); +static void check_shapes(TransposeAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { + require_same(input_shape.data_type, output_shape.data_type); + ASSERT(permute_tensor_dims(attrs.permutation, input_shape.dims) == + output_shape.dims, + attrs, + input_shape, + output_shape); } -} // namespace FlexFlow::Kernels::Transpose +template +struct TransposeCPUForwardKernel { + void operator()(TransposeAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) const { + for (TensorDimsCoord const &input_coord : + get_tensor_dims_coord_set(input.shape.dims)) { + output.at
(permute_tensor_dims_coord(attrs.permutation, input_coord)) = + input.at
(input_coord); + } + } +}; + +void transpose_cpu_forward_kernel(TransposeAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { + check_shapes(attrs, input.shape, output.shape); + + DataTypeDispatch1{}( + input.shape.data_type, attrs, input, output); +} + +template +struct TransposeCPUBackwardKernel { + void operator()(TransposeAttrs const &attrs, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorW const &input_grad) const { + for (TensorDimsCoord const &input_coord : + get_tensor_dims_coord_set(input_grad.shape.dims)) { + input_grad.at
(input_coord) += output_grad.at
( + permute_tensor_dims_coord(attrs.permutation, input_coord)); + } + } +}; + +void transpose_cpu_backward_kernel(TransposeAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { + check_shapes(attrs, input_grad.shape, output_grad.shape); + + DataTypeDispatch1{}( + input_grad.shape.data_type, attrs, output_grad, input_grad); +} + +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/upsample_kernels.cc b/lib/kernels/src/kernels/upsample_kernels.cc index ad01dd6911..cacc4c0a45 100644 --- a/lib/kernels/src/kernels/upsample_kernels.cc +++ b/lib/kernels/src/kernels/upsample_kernels.cc @@ -43,7 +43,9 @@ void upsample_backward_kernel(device_stream_t const &stream, upsample_cpu_backward_kernel( /*attrs=*/attrs, + /*output=*/output, /*output_grad=*/output_grad, + /*input=*/input, /*input_grad=*/input_grad); } } diff --git a/lib/kernels/src/kernels/upsample_kernels_cpu.cc b/lib/kernels/src/kernels/upsample_kernels_cpu.cc index 83def0cf86..c41a2bf61d 100644 --- a/lib/kernels/src/kernels/upsample_kernels_cpu.cc +++ b/lib/kernels/src/kernels/upsample_kernels_cpu.cc @@ -70,13 +70,15 @@ struct UpsampleCPUBackwardKernel { get_tensor_dims_coord_set(output_grad.shape.dims)) { TensorDimsCoord input_coord = input_coord_from_output_coord(output_coord); - input_grad.at
(output_coord) += output_grad.at
(input_coord); + input_grad.at
(input_coord) += output_grad.at
(output_coord); } } }; void upsample_cpu_backward_kernel(UpsampleAttrs const &attrs, + GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad) { ASSERT(get_num_dims(input_grad.shape.dims) == num_tensor_dims_t{4_n}, "Currently Upsample only supports 4-dimensional input tensors (i.e., " diff --git a/lib/kernels/test/src/kernels/batch_matmul_kernels_cpu.cc b/lib/kernels/test/src/kernels/batch_matmul_kernels_cpu.cc index 06f316b14a..992dd3469c 100644 --- a/lib/kernels/test/src/kernels/batch_matmul_kernels_cpu.cc +++ b/lib/kernels/test/src/kernels/batch_matmul_kernels_cpu.cc @@ -116,7 +116,7 @@ TEST_SUITE(FF_TEST_SUITE) { { {-29.8, 2.5}, {47.6, 7.5}, - {-2.6, 7.7475}, + {-2.6, 7.475}, {-3.0, 19.5}, }, }, diff --git a/lib/kernels/test/src/kernels/batch_matmul_kernels_gpu.cc b/lib/kernels/test/src/kernels/batch_matmul_kernels_gpu.cc new file mode 100644 index 0000000000..051faca789 --- /dev/null +++ b/lib/kernels/test/src/kernels/batch_matmul_kernels_gpu.cc @@ -0,0 +1,193 @@ +#include "kernels/batch_matmul_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +static GenericTensorAccessorR make_input_lhs(Allocator &allocator) { + return create_3d_accessor_r_with_contents( + { + { + {3, 3, 6}, + {2, 1, 5}, + {1, 2, -2}, + {8, 0.5, -3}, + }, + { + {5, 1, -2}, + {-8, 0, -1}, + {0.25, -0.3, -2}, + {0, -1, -5}, + }, + }, + allocator); +} + +static GenericTensorAccessorR make_input_rhs(Allocator &allocator) { + return create_3d_accessor_r_with_contents( + { + { + {1.0, 0.5}, + {2.0, 4.0}, + {1.5, -1.0}, + }, + { + {-6.0, -0.5}, + {1.0, -2.0}, + {0.4, -3.5}, + }, + }, + allocator); +} + +static GenericTensorAccessorR make_output(Allocator &allocator) { + return create_3d_accessor_r_with_contents( + { + { + {18.0, 7.5}, + {11.5, 0.0}, + {2.0, 10.5}, + {4.5, 9.0}, + }, + { + {-29.799999237060547, 2.5}, + {47.599998474121094, 7.5}, + {-2.5999999046325684, 7.474999904632568}, + {-3.0, 19.5}, + }, + }, + allocator); +} + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("batch_matmul_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(); + + GenericTensorAccessorR input_lhs = make_input_lhs(allocator); + GenericTensorAccessorR input_rhs = make_input_rhs(allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{2_p, 4_p, 2_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW result = + create_random_filled_accessor_w(output_shape, allocator); + + batch_matmul_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*input_lhs=*/input_lhs, + /*input_rhs=*/input_rhs, + /*output=*/result); + + GenericTensorAccessorR correct = make_output(allocator); + + CHECK_MESSAGE(accessors_are_equal(result, correct), + check_kv("result", format_accessor_w_contents(result))); + } + + TEST_CASE("batch_matmul_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(); + + GenericTensorAccessorR input_lhs = make_input_lhs(allocator); + GenericTensorAccessorR input_rhs = make_input_rhs(allocator); + GenericTensorAccessorR output = make_output(allocator); + + GenericTensorAccessorR output_grad = + create_3d_accessor_r_with_contents( + { + { + {1.0, 2.0}, + {0.1, -1.0}, + {0.5, 1.5}, + {4.5, -3.0}, + }, + { + {0.0, -2.5}, + {1.0, -0.5}, + {-2.0, -0.5}, + {-3.0, 0.25}, + }, + }, + allocator); + + // The gradients are overwritten rather than accumulated into (matching + // batch_matmul_cpu_backward_kernel), so randomize them to be confident + // their previous contents are never read + GenericTensorAccessorW input_lhs_grad = create_random_filled_accessor_w( + get_tensor_shape_for_accessor_r(input_lhs), allocator); + + GenericTensorAccessorW input_rhs_grad = create_random_filled_accessor_w( + get_tensor_shape_for_accessor_r(input_rhs), allocator); + + batch_matmul_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*output=*/output, + /*output_grad=*/output_grad, + /*input_lhs=*/input_lhs, + /*input_lhs_grad=*/input_lhs_grad, + /*input_rhs=*/input_rhs, + /*input_rhs_grad=*/input_rhs_grad); + + GenericTensorAccessorR correct_input_lhs_grad = + create_3d_accessor_r_with_contents( + { + { + {2.0, 10.0, -0.5}, + {-0.4000000059604645, + -3.799999952316284, + 1.149999976158142}, + {1.25, 7.0, -0.75}, + {3.0, -3.0, 9.75}, + }, + { + {1.25, 5.0, 8.75}, + {-5.75, 2.0, 2.1500000953674316}, + {12.25, -1.0, 0.949999988079071}, + {17.875, -3.5, -2.075000047683716}, + }, + }, + allocator); + + GenericTensorAccessorR correct_input_rhs_grad = + create_3d_accessor_r_with_contents( + { + { + {39.70000076293945, -18.5}, + {6.349999904632568, 6.5}, + {-8.0, 13.0}, + }, + { + {-8.5, -8.625}, + {3.5999999046325684, -2.5999999046325684}, + {18.0, 5.25}, + }, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_lhs_grad, correct_input_lhs_grad), + check_kv("input_lhs_grad", format_accessor_w_contents(input_lhs_grad))); + + CHECK_MESSAGE( + accessors_are_equal(input_rhs_grad, correct_input_rhs_grad), + check_kv("input_rhs_grad", format_accessor_w_contents(input_rhs_grad))); + } +} diff --git a/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc new file mode 100644 index 0000000000..ad92f0ea56 --- /dev/null +++ b/lib/kernels/test/src/kernels/batch_norm_kernels_gpu.cc @@ -0,0 +1,225 @@ +#include "kernels/batch_norm_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +static BatchNormAttrs make_attrs() { + return BatchNormAttrs{ + /*relu=*/false, + /*affine=*/true, + /*eps=*/1e-5, + /*momentum=*/0.1, + }; +} + +// NCHW +static GenericTensorAccessorR make_input(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + {{1, 2}, {3, 4}}, + {{-1, -2}, {0.5, 0.25}}, + }, + { + {{5, 6}, {7, 8}}, + {{2, 0.125}, {-3, 1}}, + }, + }, + allocator); +} + +static GenericTensorAccessorR make_gamma(Allocator &allocator) { + return create_1d_accessor_r_with_contents({2, 0.5}, allocator); +} + +static GenericTensorAccessorR make_beta(Allocator &allocator) { + return create_1d_accessor_r_with_contents({-1, 3}, allocator); +} + +static GenericTensorAccessorR make_output(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + {{-4.055047512054443, -3.1821770668029785}, + {-2.3093061447143555, -1.4364354610443115}}, + {{2.760241985321045, 2.433763027191162}, + {3.249960422515869, 3.1683406829833984}}, + }, + { + {{-0.5635647177696228, 0.3093060255050659}, + {1.1821768283843994, 2.0550475120544434}}, + {{3.7396788597106934, 3.127530813217163}, + {2.1072840690612793, 3.4131999015808105}}, + }, + }, + allocator); +} + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("batch_norm_gpu_forward_kernel") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + BatchNormAttrs attrs = make_attrs(); + + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR gamma = make_gamma(allocator); + GenericTensorAccessorR beta = make_beta(allocator); + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(input.shape, allocator); + + BatchNormPerDeviceState per_device_state = + batch_norm_gpu_init_kernel(allocator, attrs, input.shape, output.shape); + + batch_norm_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*gamma=*/gamma, + /*beta=*/beta, + /*output=*/output); + + GenericTensorAccessorR correct = make_output(allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + + batch_norm_gpu_cleanup_kernel(allocator, per_device_state); + } + + TEST_CASE("batch_norm_gpu_backward_kernel") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + BatchNormAttrs attrs = make_attrs(); + + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR gamma = make_gamma(allocator); + GenericTensorAccessorR beta = make_beta(allocator); + + GenericTensorAccessorW forward_output = + create_random_filled_accessor_w(input.shape, allocator); + + BatchNormPerDeviceState per_device_state = batch_norm_gpu_init_kernel( + allocator, attrs, input.shape, forward_output.shape); + + // cudnnBatchNormalizationBackward reads the batch statistics saved by the + // forward pass, so the forward kernel has to run first + batch_norm_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*gamma=*/gamma, + /*beta=*/beta, + /*output=*/forward_output); + + GenericTensorAccessorR output = make_output(allocator); + + GenericTensorAccessorR output_grad = + create_4d_accessor_r_with_contents( + { + { + {{-0.875, -0.75}, {-0.625, -0.5}}, + {{-0.375, -0.25}, {-0.125, 0}}, + }, + { + {{0.125, 0.25}, {0.375, 0.5}}, + {{0.625, 0.75}, {0.875, 1}}, + }, + }, + allocator); + + // The gradients are accumulated into, so they need to start from known + // values + GenericTensorAccessorW input_grad = + create_4d_accessor_w_with_contents( + { + { + {{-1.75, -1.5}, {-1.25, -1}}, + {{-0.75, -0.5}, {-0.25, 0}}, + }, + { + {{0.25, 0.5}, {0.75, 1}}, + {{1.25, 1.5}, {1.75, 2}}, + }, + }, + allocator); + + GenericTensorAccessorW gamma_grad = + create_1d_accessor_w_with_contents({7, -3}, allocator); + + GenericTensorAccessorW beta_grad = + create_1d_accessor_w_with_contents({0.5, 11}, allocator); + + batch_norm_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad, + /*gamma=*/gamma, + /*gamma_grad=*/gamma_grad, + /*beta_grad=*/beta_grad); + + GenericTensorAccessorR correct_input_grad = + create_4d_accessor_r_with_contents( + { + { + {{-1.6772620677947998, -1.510392189025879}, + {-1.3435224294662476, -1.1766525506973267}}, + {{-0.9591808915138245, -0.6475732326507568}, + {-0.4087578058242798, -0.11274851113557816}}, + }, + { + {{0.42665261030197144, 0.5935224294662476}, + {0.7603921890258789, 0.927262008190155}}, + {{1.3049046993255615, 1.634710431098938}, + {1.9905133247375488, 2.198132038116455}}, + }, + }, + allocator); + + GenericTensorAccessorR correct_gamma_grad = + create_1d_accessor_r_with_contents( + {11.037027359008789, -2.2195112705230713}, allocator); + + GenericTensorAccessorR correct_beta_grad = + create_1d_accessor_r_with_contents({-1.0, 13.5}, allocator); + + CHECK_MESSAGE( + // cuDNN's batch norm data gradient is not bit-identical to PyTorch's + accessors_within_epsilon(input_grad, correct_input_grad, 1e-6), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + + CHECK_MESSAGE( + accessors_are_equal(gamma_grad, correct_gamma_grad), + check_kv("gamma_grad", format_accessor_w_contents(gamma_grad))); + + CHECK_MESSAGE(accessors_are_equal(beta_grad, correct_beta_grad), + check_kv("beta_grad", format_accessor_w_contents(beta_grad))); + + batch_norm_gpu_cleanup_kernel(allocator, per_device_state); + } +} diff --git a/lib/kernels/test/src/kernels/concat_kernels_cpu.cc b/lib/kernels/test/src/kernels/concat_kernels_cpu.cc new file mode 100644 index 0000000000..08b3bfb207 --- /dev/null +++ b/lib/kernels/test/src/kernels/concat_kernels_cpu.cc @@ -0,0 +1,279 @@ +#include "kernels/concat_kernels_cpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "kernels/local_cpu_allocator.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("concat_cpu_forward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + std::vector inputs = { + create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {-1, -2, -3}, + {-4, -5, -6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {0.5, 0.25, 0.125}, + {10, 20, 30}, + }, + allocator), + }; + + SUBCASE("axis = 0") { + ConcatAttrs attrs = ConcatAttrs{ + /*axis=*/ff_dim_t{0_n}, + /*num_inputs=*/int_ge_two{3}, + }; + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{6_p, 3_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never + // read it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + concat_cpu_forward_kernel( + /*attrs=*/attrs, + /*inputs=*/inputs, + /*output=*/output); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {-1, -2, -3}, + {-4, -5, -6}, + {0.5, 0.25, 0.125}, + {10, 20, 30}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + SUBCASE("axis = 1") { + ConcatAttrs attrs = ConcatAttrs{ + /*axis=*/ff_dim_t{1_n}, + /*num_inputs=*/int_ge_two{3}, + }; + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{2_p, 9_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never + // read it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + concat_cpu_forward_kernel( + /*attrs=*/attrs, + /*inputs=*/inputs, + /*output=*/output); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {1, 2, 3, -1, -2, -3, 0.5, 0.25, 0.125}, + {4, 5, 6, -4, -5, -6, 10, 20, 30}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + } + + TEST_CASE("concat_cpu_backward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + std::vector inputs = { + create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {-1, -2, -3}, + {-4, -5, -6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {0.5, 0.25, 0.125}, + {10, 20, 30}, + }, + allocator), + }; + + // input_grads are accumulated into, so they need to start from known values + auto make_input_grads = [&]() { + return std::vector{ + create_2d_accessor_w_with_contents( + { + {100, 101, 102}, + {103, 104, 105}, + }, + allocator), + create_2d_accessor_w_with_contents( + { + {200, 201, 202}, + {203, 204, 205}, + }, + allocator), + create_2d_accessor_w_with_contents( + { + {300, 301, 302}, + {303, 304, 305}, + }, + allocator), + }; + }; + + SUBCASE("axis = 0") { + ConcatAttrs attrs = ConcatAttrs{ + /*axis=*/ff_dim_t{0_n}, + /*num_inputs=*/int_ge_two{3}, + }; + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {-1, -2, -3}, + {-4, -5, -6}, + {0.5, 0.25, 0.125}, + {10, 20, 30}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_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, 2.25, 2.5}, + }, + allocator); + + std::vector input_grads = make_input_grads(); + + concat_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*inputs=*/inputs, + /*input_grads=*/input_grads); + + std::vector correct = { + create_2d_accessor_r_with_contents( + { + {98.25, 99.5, 100.75}, + {102, 103.25, 104.5}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {199.75, 201, 202.25}, + {203.5, 204.75, 206}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {301.25, 302.5, 303.75}, + {305, 306.25, 307.5}, + }, + allocator), + }; + + for (int i = 0; i < input_grads.size(); i++) { + CHECK_MESSAGE(accessors_are_equal(input_grads.at(i), correct.at(i)), + check_kv("input_grad", + format_accessor_w_contents(input_grads.at(i)))); + } + } + + SUBCASE("axis = 1") { + ConcatAttrs attrs = ConcatAttrs{ + /*axis=*/ff_dim_t{1_n}, + /*num_inputs=*/int_ge_two{3}, + }; + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {1, 2, 3, -1, -2, -3, 0.5, 0.25, 0.125}, + {4, 5, 6, -4, -5, -6, 10, 20, 30}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_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, 2.25, 2.5}, + }, + allocator); + + std::vector input_grads = make_input_grads(); + + concat_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*inputs=*/inputs, + /*input_grads=*/input_grads); + + std::vector correct = { + create_2d_accessor_r_with_contents( + { + {98.25, 99.5, 100.75}, + {103.5, 104.75, 106}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {199, 200.25, 201.5}, + {204.25, 205.5, 206.75}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {299.75, 301, 302.25}, + {305, 306.25, 307.5}, + }, + allocator), + }; + + for (int i = 0; i < input_grads.size(); i++) { + CHECK_MESSAGE(accessors_are_equal(input_grads.at(i), correct.at(i)), + check_kv("input_grad", + format_accessor_w_contents(input_grads.at(i)))); + } + } + } +} diff --git a/lib/kernels/test/src/kernels/concat_kernels_gpu.cc b/lib/kernels/test/src/kernels/concat_kernels_gpu.cc new file mode 100644 index 0000000000..4728c22f7b --- /dev/null +++ b/lib/kernels/test/src/kernels/concat_kernels_gpu.cc @@ -0,0 +1,292 @@ +#include "kernels/concat_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("concat_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(); + + std::vector inputs = { + create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {-1, -2, -3}, + {-4, -5, -6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {0.5, 0.25, 0.125}, + {10, 20, 30}, + }, + allocator), + }; + + SUBCASE("axis = 0") { + ConcatAttrs attrs = ConcatAttrs{ + /*axis=*/ff_dim_t{0_n}, + /*num_inputs=*/int_ge_two{3}, + }; + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{6_p, 3_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never + // read it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + concat_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*inputs=*/inputs, + /*output=*/output); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {-1, -2, -3}, + {-4, -5, -6}, + {0.5, 0.25, 0.125}, + {10, 20, 30}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + SUBCASE("axis = 1") { + ConcatAttrs attrs = ConcatAttrs{ + /*axis=*/ff_dim_t{1_n}, + /*num_inputs=*/int_ge_two{3}, + }; + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{2_p, 9_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never + // read it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + concat_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*inputs=*/inputs, + /*output=*/output); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {1, 2, 3, -1, -2, -3, 0.5, 0.25, 0.125}, + {4, 5, 6, -4, -5, -6, 10, 20, 30}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + } + + TEST_CASE("concat_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(); + + std::vector inputs = { + create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {-1, -2, -3}, + {-4, -5, -6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {0.5, 0.25, 0.125}, + {10, 20, 30}, + }, + allocator), + }; + + // input_grads are accumulated into, so they need to start from known values + auto make_input_grads = [&]() { + return std::vector{ + create_2d_accessor_w_with_contents( + { + {100, 101, 102}, + {103, 104, 105}, + }, + allocator), + create_2d_accessor_w_with_contents( + { + {200, 201, 202}, + {203, 204, 205}, + }, + allocator), + create_2d_accessor_w_with_contents( + { + {300, 301, 302}, + {303, 304, 305}, + }, + allocator), + }; + }; + + SUBCASE("axis = 0") { + ConcatAttrs attrs = ConcatAttrs{ + /*axis=*/ff_dim_t{0_n}, + /*num_inputs=*/int_ge_two{3}, + }; + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {-1, -2, -3}, + {-4, -5, -6}, + {0.5, 0.25, 0.125}, + {10, 20, 30}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_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, 2.25, 2.5}, + }, + allocator); + + std::vector input_grads = make_input_grads(); + + concat_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*inputs=*/inputs, + /*input_grads=*/input_grads); + + std::vector correct = { + create_2d_accessor_r_with_contents( + { + {98.25, 99.5, 100.75}, + {102, 103.25, 104.5}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {199.75, 201, 202.25}, + {203.5, 204.75, 206}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {301.25, 302.5, 303.75}, + {305, 306.25, 307.5}, + }, + allocator), + }; + + for (int i = 0; i < input_grads.size(); i++) { + CHECK_MESSAGE(accessors_are_equal(input_grads.at(i), correct.at(i)), + check_kv("input_grad", + format_accessor_w_contents(input_grads.at(i)))); + } + } + + SUBCASE("axis = 1") { + ConcatAttrs attrs = ConcatAttrs{ + /*axis=*/ff_dim_t{1_n}, + /*num_inputs=*/int_ge_two{3}, + }; + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {1, 2, 3, -1, -2, -3, 0.5, 0.25, 0.125}, + {4, 5, 6, -4, -5, -6, 10, 20, 30}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_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, 2.25, 2.5}, + }, + allocator); + + std::vector input_grads = make_input_grads(); + + concat_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*inputs=*/inputs, + /*input_grads=*/input_grads); + + std::vector correct = { + create_2d_accessor_r_with_contents( + { + {98.25, 99.5, 100.75}, + {103.5, 104.75, 106}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {199, 200.25, 201.5}, + {204.25, 205.5, 206.75}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {299.75, 301, 302.25}, + {305, 306.25, 307.5}, + }, + allocator), + }; + + for (int i = 0; i < input_grads.size(); i++) { + CHECK_MESSAGE(accessors_are_equal(input_grads.at(i), correct.at(i)), + check_kv("input_grad", + format_accessor_w_contents(input_grads.at(i)))); + } + } + } +} diff --git a/lib/kernels/test/src/kernels/conv_2d_kernels_gpu.cc b/lib/kernels/test/src/kernels/conv_2d_kernels_gpu.cc new file mode 100644 index 0000000000..8fd34689c2 --- /dev/null +++ b/lib/kernels/test/src/kernels/conv_2d_kernels_gpu.cc @@ -0,0 +1,334 @@ +#include "kernels/conv_2d_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +static Conv2DAttrs make_attrs() { + return Conv2DAttrs{ + /*out_channels=*/3_p, + /*kernel_h=*/3_p, + /*kernel_w=*/3_p, + /*stride_h=*/1_p, + /*stride_w=*/1_p, + /*padding_h=*/1_n, + /*padding_w=*/1_n, + /*groups=*/1_p, + /*activation=*/std::nullopt, + /*use_bias=*/true, + }; +} + +// NCHW +static GenericTensorAccessorR make_input(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + { + {1, 2, 3, 4}, + {0.5, 0.25, -1, -2}, + {-3, 0.125, 2, 1}, + {4, -0.5, 0.75, 0}, + }, + { + {-1, 0.5, 2, -2}, + {3, 1, 0.25, 0.5}, + {-0.25, -1.5, 1.25, 2.5}, + {0.75, 3, -4, 1}, + }, + }, + }, + allocator); +} + +static GenericTensorAccessorR make_filter(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + {{1, 0, -1}, {0.5, 1, 0.5}, {-1, 0, 1}}, + {{0.25, -0.5, 0.75}, {1, -1, 0.5}, {0, 0.25, -0.25}}, + }, + { + {{-1, 2, 0.5}, {0, 1, -1}, {0.25, 0.5, 0.75}}, + {{1, 1, 1}, {-0.5, -0.5, -0.5}, {0.25, 0, -0.25}}, + }, + { + {{0.5, 0.5, 0.5}, {0.25, 0.25, 0.25}, {-1, -1, -1}}, + {{2, -1, 0}, {0.5, 0.75, -0.25}, {1, 0.5, 0.25}}, + }, + }, + allocator); +} + +static GenericTensorAccessorR make_bias(Allocator &allocator) { + return create_1d_accessor_r_with_contents({0.5, -1, 2}, allocator); +} + +static GenericTensorAccessorR make_output_grad(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + { + {-0.9375, -0.875, -0.8125, -0.75}, + {-0.6875, -0.625, -0.5625, -0.5}, + {-0.4375, -0.375, -0.3125, -0.25}, + {-0.1875, -0.125, -0.0625, 0}, + }, + { + {0.0625, 0.125, 0.1875, 0.25}, + {0.3125, 0.375, 0.4375, 0.5}, + {0.5625, 0.625, 0.6875, 0.75}, + {0.8125, 0.875, 0.9375, 1}, + }, + { + {1.0625, 1.125, 1.1875, 1.25}, + {1.3125, 1.375, 1.4375, 1.5}, + {1.5625, 1.625, 1.6875, 1.75}, + {1.8125, 1.875, 1.9375, 2}, + }, + }, + }, + allocator); +} + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("conv_2d_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(); + + Conv2DAttrs attrs = make_attrs(); + + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR filter = make_filter(allocator); + GenericTensorAccessorR bias = make_bias(allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{1_p, 3_p, 4_p, 4_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + Conv2DPerDeviceState per_device_state = conv_2d_gpu_init_kernel( + managed_handle.raw_handle(), attrs, input.shape, output.shape); + + conv_2d_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*filter=*/filter, + /*bias=*/bias, + /*output=*/output); + + GenericTensorAccessorR correct = create_4d_accessor_r_with_contents( + { + { + { + {4.5, 2.6875, 1.6875, 11.125}, + {-2.0625, 5.9375, -4.1875, 0.875}, + {-5, 2.4375, 3.5625, -0.4375}, + {3.875, -5.25, 8.5, -3.0625}, + }, + { + {-1.5625, -2.5625, -4.0625, 1.8125}, + {-1.28125, 4.5625, 6.40625, 2.9375}, + {2.75, 3.625, -1.875, -4.9375}, + {-6.0625, 1.625, 6.375, 4.25}, + }, + { + {2.875, 6.6875, 10.5, 6.75}, + {9.0625, 4.8125, 2, 10.75}, + {-3.53125, 2.09375, 3.71875, -0.5}, + {1.5, 7.25, -2.375, 2.4375}, + }, + }, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + + conv_2d_gpu_cleanup_kernel(per_device_state); + } + + TEST_CASE("conv_2d_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(); + + Conv2DAttrs attrs = make_attrs(); + + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR filter = make_filter(allocator); + GenericTensorAccessorR output_grad = make_output_grad(allocator); + + GenericTensorAccessorR output = create_4d_accessor_r_with_contents( + { + { + { + {4.5, 2.6875, 1.6875, 11.125}, + {-2.0625, 5.9375, -4.1875, 0.875}, + {-5, 2.4375, 3.5625, -0.4375}, + {3.875, -5.25, 8.5, -3.0625}, + }, + { + {-1.5625, -2.5625, -4.0625, 1.8125}, + {-1.28125, 4.5625, 6.40625, 2.9375}, + {2.75, 3.625, -1.875, -4.9375}, + {-6.0625, 1.625, 6.375, 4.25}, + }, + { + {2.875, 6.6875, 10.5, 6.75}, + {9.0625, 4.8125, 2, 10.75}, + {-3.53125, 2.09375, 3.71875, -0.5}, + {1.5, 7.25, -2.375, 2.4375}, + }, + }, + }, + allocator); + + // The gradients are accumulated into, so they need to start from known + // values + GenericTensorAccessorW input_grad = + create_4d_accessor_w_with_contents( + { + { + { + {-2.75, -2.5, -2.25, -2}, + {-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}, + {2.25, 2.5, 2.75, 3}, + {3.25, 3.5, 3.75, 4}, + {4.25, 4.5, 4.75, 5}, + }, + }, + }, + allocator); + + GenericTensorAccessorW filter_grad = create_4d_accessor_w_with_contents< + float>( + { + { + {{-1.875, -1.75, -1.625}, + {-1.5, -1.375, -1.25}, + {-1.125, -1, -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}, {1.125, 1.25, 1.375}}, + {{1.5, 1.625, 1.75}, {1.875, 2, 2.125}, {2.25, 2.375, 2.5}}, + }, + { + {{2.625, 2.75, 2.875}, {3, 3.125, 3.25}, {3.375, 3.5, 3.625}}, + {{3.75, 3.875, 4}, {4.125, 4.25, 4.375}, {4.5, 4.625, 4.75}}, + }, + }, + allocator); + + GenericTensorAccessorW bias_grad = + create_1d_accessor_w_with_contents({10, -5, 0.25}, allocator); + + Conv2DPerDeviceState per_device_state = conv_2d_gpu_init_kernel( + managed_handle.raw_handle(), attrs, input.shape, output.shape); + + conv_2d_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, + /*filter=*/filter, + /*filter_grad=*/filter_grad, + /*bias_grad=*/bias_grad); + + GenericTensorAccessorR correct_input_grad = + create_4d_accessor_r_with_contents( + { + { + { + {-2.546875, -0.6875, -0.078125, 0.765625}, + {-1.296875, -1.59375, -1.078125, -0.09375}, + {0.640625, 0.46875, 0.984375, 2.09375}, + {-0.640625, -2.375, -2.046875, -0.84375}, + }, + { + {4.890625, 4.359375, 4.890625, 2.03125}, + {8.09375, 8.546875, 9.1875, 3.984375}, + {10.34375, 11.109375, 11.75, 5.171875}, + {8.21875, 8.015625, 8.375, 6.15625}, + }, + }, + }, + allocator); + + GenericTensorAccessorR correct_filter_grad = + create_4d_accessor_r_with_contents( + { + { + {{-4.7734375, -6.453125, -6.5703125}, + {-6.0859375, -8.78125, -8.5390625}, + {-1.7421875, -0.8125, -0.5703125}}, + {{-2.78125, -2.109375, -1.6875}, + {-4, -3.5625, -2.5625}, + {-3.46875, -5.4375, -2.59375}}, + }, + { + {{2.3515625, 3.671875, 5.0546875}, + {5.2890625, 5.59375, 3.3359375}, + {3.6328125, 3.5625, 2.3046875}}, + {{4.71875, 6.390625, 5.0625}, + {3.25, 5.6875, 4.1875}, + {2.28125, 4.3125, 3.65625}}, + }, + { + {{9.4765625, 13.796875, 16.6796875}, + {16.6640625, 19.96875, 15.2109375}, + {9.0078125, 7.9375, 5.1796875}}, + {{12.21875, 14.890625, 11.8125}, + {10.5, 14.9375, 10.9375}, + {8.03125, 14.0625, 9.90625}}, + }, + }, + allocator); + + GenericTensorAccessorR correct_bias_grad = + create_1d_accessor_r_with_contents({2.5, 3.5, 24.75}, allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct_input_grad), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + + CHECK_MESSAGE( + accessors_are_equal(filter_grad, correct_filter_grad), + check_kv("filter_grad", format_accessor_w_contents(filter_grad))); + + CHECK_MESSAGE(accessors_are_equal(bias_grad, correct_bias_grad), + check_kv("bias_grad", format_accessor_w_contents(bias_grad))); + + conv_2d_gpu_cleanup_kernel(per_device_state); + } +} 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/kernels/pool_2d_kernels_gpu.cc b/lib/kernels/test/src/kernels/pool_2d_kernels_gpu.cc new file mode 100644 index 0000000000..47cc2ff95d --- /dev/null +++ b/lib/kernels/test/src/kernels/pool_2d_kernels_gpu.cc @@ -0,0 +1,284 @@ +#include "kernels/pool_2d_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +// NCHW +static GenericTensorAccessorR make_input(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + { + {1, 2, 3, 4}, + {5, 6, 7, 8}, + {-1, -2, -3, -4}, + {0.5, 0.25, 0.125, 0.0625}, + }, + { + {9, 1, 0.5, -7}, + {2, 3, 0.25, 0.125}, + {-8, 4, 10, 20}, + {6, -5, 30, 40}, + }, + }, + }, + allocator); +} + +// input_grad is accumulated into, so it needs to start from a known value +static GenericTensorAccessorW make_input_grad(Allocator &allocator) { + return create_4d_accessor_w_with_contents( + { + { + { + {-4.5, -4, -3.5, -3}, + {-2.5, -2, -1.5, -1}, + {-0.5, 0, 0.5, 1}, + {1.5, 2, 2.5, 3}, + }, + { + {3.5, 4, 4.5, 5}, + {5.5, 6, 6.5, 7}, + {7.5, 8, 8.5, 9}, + {9.5, 10, 10.5, 11}, + }, + }, + }, + allocator); +} + +static GenericTensorAccessorR make_output_grad(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + { + {-0.75, -0.5}, + {-0.25, 0}, + }, + { + {0.25, 0.5}, + {0.75, 1}, + }, + }, + }, + allocator); +} + +static Pool2DAttrs make_attrs(PoolOp pool_type) { + return Pool2DAttrs{ + /*kernel_h=*/2_p, + /*kernel_w=*/2_p, + /*stride_h=*/2_p, + /*stride_w=*/2_p, + /*padding_h=*/0_n, + /*padding_w=*/0_n, + /*pool_type=*/pool_type, + /*activation=*/std::nullopt, + }; +} + +static TensorShape output_shape() { + return TensorShape{ + TensorDims{FFOrdered{1_p, 2_p, 2_p, 2_p}}, + DataType::FLOAT, + }; +} + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("pool_2d_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(); + + GenericTensorAccessorR input = make_input(allocator); + + SUBCASE("PoolOp::MAX") { + Pool2DAttrs attrs = make_attrs(PoolOp::MAX); + + // Intentionally randomize this tensor so we can be confident we never + // read it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape(), allocator); + + Pool2DPerDeviceState per_device_state = + pool_2d_gpu_init_kernel(attrs, input.shape, output.shape); + + pool_2d_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = + create_4d_accessor_r_with_contents( + { + { + {{6, 8}, {0.5, 0.125}}, + {{9, 0.5}, {6, 40}}, + }, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + SUBCASE("PoolOp::AVG") { + Pool2DAttrs attrs = make_attrs(PoolOp::AVG); + + // Intentionally randomize this tensor so we can be confident we never + // read it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape(), allocator); + + Pool2DPerDeviceState per_device_state = + pool_2d_gpu_init_kernel(attrs, input.shape, output.shape); + + pool_2d_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = + create_4d_accessor_r_with_contents( + { + { + {{3.5, 5.5}, {-0.5625, -1.703125}}, + {{3.75, -1.53125}, {-0.75, 25}}, + }, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + } + + TEST_CASE("pool_2d_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(); + + GenericTensorAccessorR input = make_input(allocator); + GenericTensorAccessorR output_grad = make_output_grad(allocator); + + SUBCASE("PoolOp::MAX") { + Pool2DAttrs attrs = make_attrs(PoolOp::MAX); + + GenericTensorAccessorR output = create_4d_accessor_r_with_contents( + { + { + {{6, 8}, {0.5, 0.125}}, + {{9, 0.5}, {6, 40}}, + }, + }, + allocator); + + GenericTensorAccessorW input_grad = make_input_grad(allocator); + + Pool2DPerDeviceState per_device_state = + pool_2d_gpu_init_kernel(attrs, input.shape, output.shape); + + pool_2d_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 = + create_4d_accessor_r_with_contents( + { + { + { + {-4.5, -4, -3.5, -3}, + {-2.5, -2.75, -1.5, -1.5}, + {-0.5, 0, 0.5, 1}, + {1.25, 2, 2.5, 3}, + }, + { + {3.75, 4, 5, 5}, + {5.5, 6, 6.5, 7}, + {7.5, 8, 8.5, 9}, + {10.25, 10, 10.5, 12}, + }, + }, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + + SUBCASE("PoolOp::AVG") { + Pool2DAttrs attrs = make_attrs(PoolOp::AVG); + + GenericTensorAccessorR output = create_4d_accessor_r_with_contents( + { + { + {{3.5, 5.5}, {-0.5625, -1.703125}}, + {{3.75, -1.53125}, {-0.75, 25}}, + }, + }, + allocator); + + GenericTensorAccessorW input_grad = make_input_grad(allocator); + + Pool2DPerDeviceState per_device_state = + pool_2d_gpu_init_kernel(attrs, input.shape, output.shape); + + pool_2d_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 = + create_4d_accessor_r_with_contents( + { + { + { + {-4.6875, -4.1875, -3.625, -3.125}, + {-2.6875, -2.1875, -1.625, -1.125}, + {-0.5625, -0.0625, 0.5, 1}, + {1.4375, 1.9375, 2.5, 3}, + }, + { + {3.5625, 4.0625, 4.625, 5.125}, + {5.5625, 6.0625, 6.625, 7.125}, + {7.6875, 8.1875, 8.75, 9.25}, + {9.6875, 10.1875, 10.75, 11.25}, + }, + }, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + } +} diff --git a/lib/kernels/test/src/kernels/reshape_kernels_cpu.cc b/lib/kernels/test/src/kernels/reshape_kernels_cpu.cc new file mode 100644 index 0000000000..9989764a8c --- /dev/null +++ b/lib/kernels/test/src/kernels/reshape_kernels_cpu.cc @@ -0,0 +1,101 @@ +#include "kernels/reshape_kernels_cpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "kernels/local_cpu_allocator.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("reshape_cpu_forward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + GenericTensorAccessorR input = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + }, + allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{3_p, 2_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + reshape_cpu_forward_kernel( + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = create_2d_accessor_r_with_contents( + { + {1, 2}, + {3, 4}, + {5, 6}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + TEST_CASE("reshape_cpu_backward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + GenericTensorAccessorR input = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + }, + allocator); + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {1, 2}, + {3, 4}, + {5, 6}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_with_contents( + { + {1, 2}, + {-3, 4}, + {0.5, 0.25}, + }, + allocator); + + // input_grad is accumulated into, so it needs to start from a known value + GenericTensorAccessorW input_grad = + create_2d_accessor_w_with_contents( + { + {10, 20, 30}, + {40, 50, 60}, + }, + allocator); + + reshape_cpu_backward_kernel( + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = create_2d_accessor_r_with_contents( + { + {11, 22, 27}, + {44, 50.5, 60.25}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } +} diff --git a/lib/kernels/test/src/kernels/reshape_kernels_gpu.cc b/lib/kernels/test/src/kernels/reshape_kernels_gpu.cc new file mode 100644 index 0000000000..7e7ae65ecf --- /dev/null +++ b/lib/kernels/test/src/kernels/reshape_kernels_gpu.cc @@ -0,0 +1,112 @@ +#include "kernels/reshape_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("reshape_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(); + + GenericTensorAccessorR input = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + }, + allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{3_p, 2_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + reshape_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = create_2d_accessor_r_with_contents( + { + {1, 2}, + {3, 4}, + {5, 6}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + TEST_CASE("reshape_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(); + + GenericTensorAccessorR input = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + }, + allocator); + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {1, 2}, + {3, 4}, + {5, 6}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_with_contents( + { + {1, 2}, + {-3, 4}, + {0.5, 0.25}, + }, + allocator); + + // input_grad is accumulated into, so it needs to start from a known value + GenericTensorAccessorW input_grad = + create_2d_accessor_w_with_contents( + { + {10, 20, 30}, + {40, 50, 60}, + }, + allocator); + + reshape_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = create_2d_accessor_r_with_contents( + { + {11, 22, 27}, + {44, 50.5, 60.25}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } +} diff --git a/lib/kernels/test/src/kernels/split_kernels_cpu.cc b/lib/kernels/test/src/kernels/split_kernels_cpu.cc new file mode 100644 index 0000000000..8827367acf --- /dev/null +++ b/lib/kernels/test/src/kernels/split_kernels_cpu.cc @@ -0,0 +1,293 @@ +#include "kernels/split_kernels_cpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "kernels/local_cpu_allocator.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +static GenericTensorAccessorR make_input(Allocator &allocator) { + return create_2d_accessor_r_with_contents( + { + {1, 2, 3, 4, 5, 6}, + {-1, -2, -3, -4, -5, -6}, + {0.5, 0.25, 0.125, 10, 20, 30}, + {7, 8, 9, 0.75, 0.5, 0.25}, + }, + allocator); +} + +// input_grad is accumulated into, so it needs to start from a known value +static GenericTensorAccessorW make_input_grad(Allocator &allocator) { + return create_2d_accessor_w_with_contents( + { + {-3, -1, 1, 3, 5, 7}, + {9, 11, 13, 15, 17, 19}, + {21, 23, 25, 27, 29, 31}, + {33, 35, 37, 39, 41, 43}, + }, + allocator); +} + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("split_cpu_forward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + GenericTensorAccessorR input = make_input(allocator); + + SUBCASE("axis = 0") { + SplitAttrs attrs = SplitAttrs{ + /*splits=*/{1_p, 3_p}, + /*axis=*/ff_dim_t{0_n}, + }; + + // Intentionally randomize these tensors so we can be confident we never + // read them + std::vector outputs = { + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{1_p, 6_p}}, DataType::FLOAT}, + allocator), + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{3_p, 6_p}}, DataType::FLOAT}, + allocator), + }; + + split_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*outputs=*/outputs); + + std::vector correct = { + create_2d_accessor_r_with_contents( + { + {1, 2, 3, 4, 5, 6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {-1, -2, -3, -4, -5, -6}, + {0.5, 0.25, 0.125, 10, 20, 30}, + {7, 8, 9, 0.75, 0.5, 0.25}, + }, + allocator), + }; + + for (int i = 0; i < outputs.size(); i++) { + CHECK_MESSAGE( + accessors_are_equal(outputs.at(i), correct.at(i)), + check_kv("output", format_accessor_w_contents(outputs.at(i)))); + } + } + + SUBCASE("axis = 1") { + SplitAttrs attrs = SplitAttrs{ + /*splits=*/{2_p, 1_p, 3_p}, + /*axis=*/ff_dim_t{1_n}, + }; + + // Intentionally randomize these tensors so we can be confident we never + // read them + std::vector outputs = { + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{4_p, 2_p}}, DataType::FLOAT}, + allocator), + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{4_p, 1_p}}, DataType::FLOAT}, + allocator), + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{4_p, 3_p}}, DataType::FLOAT}, + allocator), + }; + + split_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*outputs=*/outputs); + + std::vector correct = { + create_2d_accessor_r_with_contents( + { + {1, 2}, + {-1, -2}, + {0.5, 0.25}, + {7, 8}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {3}, + {-3}, + {0.125}, + {9}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {4, 5, 6}, + {-4, -5, -6}, + {10, 20, 30}, + {0.75, 0.5, 0.25}, + }, + allocator), + }; + + for (int i = 0; i < outputs.size(); i++) { + CHECK_MESSAGE( + accessors_are_equal(outputs.at(i), correct.at(i)), + check_kv("output", format_accessor_w_contents(outputs.at(i)))); + } + } + } + + TEST_CASE("split_cpu_backward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + GenericTensorAccessorR input = make_input(allocator); + + SUBCASE("axis = 0") { + SplitAttrs attrs = SplitAttrs{ + /*splits=*/{1_p, 3_p}, + /*axis=*/ff_dim_t{0_n}, + }; + + std::vector outputs = { + create_2d_accessor_r_with_contents( + { + {1, 2, 3, 4, 5, 6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {-1, -2, -3, -4, -5, -6}, + {0.5, 0.25, 0.125, 10, 20, 30}, + {7, 8, 9, 0.75, 0.5, 0.25}, + }, + allocator), + }; + + std::vector output_grads = { + create_2d_accessor_r_with_contents( + { + {-0.75, -0.5, -0.25, 0, 0.25, 0.5}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {9.25, 9.5, 9.75, 10, 10.25, 10.5}, + {10.75, 11, 11.25, 11.5, 11.75, 12}, + {12.25, 12.5, 12.75, 13, 13.25, 13.5}, + }, + allocator), + }; + + GenericTensorAccessorW input_grad = make_input_grad(allocator); + + split_cpu_backward_kernel( + /*attrs=*/attrs, + /*outputs=*/outputs, + /*output_grads=*/output_grads, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {-3.75, -1.5, 0.75, 3, 5.25, 7.5}, + {18.25, 20.5, 22.75, 25, 27.25, 29.5}, + {31.75, 34, 36.25, 38.5, 40.75, 43}, + {45.25, 47.5, 49.75, 52, 54.25, 56.5}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + + SUBCASE("axis = 1") { + SplitAttrs attrs = SplitAttrs{ + /*splits=*/{2_p, 1_p, 3_p}, + /*axis=*/ff_dim_t{1_n}, + }; + + std::vector outputs = { + create_2d_accessor_r_with_contents( + { + {1, 2}, + {-1, -2}, + {0.5, 0.25}, + {7, 8}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {3}, + {-3}, + {0.125}, + {9}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {4, 5, 6}, + {-4, -5, -6}, + {10, 20, 30}, + {0.75, 0.5, 0.25}, + }, + allocator), + }; + + std::vector output_grads = { + create_2d_accessor_r_with_contents( + { + {-0.75, -0.5}, + {-0.25, 0}, + {0.25, 0.5}, + {0.75, 1}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {9.25}, + {9.5}, + {9.75}, + {10}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {19.25, 19.5, 19.75}, + {20, 20.25, 20.5}, + {20.75, 21, 21.25}, + {21.5, 21.75, 22}, + }, + allocator), + }; + + GenericTensorAccessorW input_grad = make_input_grad(allocator); + + split_cpu_backward_kernel( + /*attrs=*/attrs, + /*outputs=*/outputs, + /*output_grads=*/output_grads, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {-3.75, -1.5, 10.25, 22.25, 24.5, 26.75}, + {8.75, 11, 22.5, 35, 37.25, 39.5}, + {21.25, 23.5, 34.75, 47.75, 50, 52.25}, + {33.75, 36, 47, 60.5, 62.75, 65}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + } +} diff --git a/lib/kernels/test/src/kernels/split_kernels_gpu.cc b/lib/kernels/test/src/kernels/split_kernels_gpu.cc new file mode 100644 index 0000000000..f1918f28f2 --- /dev/null +++ b/lib/kernels/test/src/kernels/split_kernels_gpu.cc @@ -0,0 +1,306 @@ +#include "kernels/split_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +static GenericTensorAccessorR make_input(Allocator &allocator) { + return create_2d_accessor_r_with_contents( + { + {1, 2, 3, 4, 5, 6}, + {-1, -2, -3, -4, -5, -6}, + {0.5, 0.25, 0.125, 10, 20, 30}, + {7, 8, 9, 0.75, 0.5, 0.25}, + }, + allocator); +} + +// input_grad is accumulated into, so it needs to start from a known value +static GenericTensorAccessorW make_input_grad(Allocator &allocator) { + return create_2d_accessor_w_with_contents( + { + {-3, -1, 1, 3, 5, 7}, + {9, 11, 13, 15, 17, 19}, + {21, 23, 25, 27, 29, 31}, + {33, 35, 37, 39, 41, 43}, + }, + allocator); +} + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("split_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(); + + GenericTensorAccessorR input = make_input(allocator); + + SUBCASE("axis = 0") { + SplitAttrs attrs = SplitAttrs{ + /*splits=*/{1_p, 3_p}, + /*axis=*/ff_dim_t{0_n}, + }; + + // Intentionally randomize these tensors so we can be confident we never + // read them + std::vector outputs = { + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{1_p, 6_p}}, DataType::FLOAT}, + allocator), + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{3_p, 6_p}}, DataType::FLOAT}, + allocator), + }; + + split_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*input=*/input, + /*outputs=*/outputs); + + std::vector correct = { + create_2d_accessor_r_with_contents( + { + {1, 2, 3, 4, 5, 6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {-1, -2, -3, -4, -5, -6}, + {0.5, 0.25, 0.125, 10, 20, 30}, + {7, 8, 9, 0.75, 0.5, 0.25}, + }, + allocator), + }; + + for (int i = 0; i < outputs.size(); i++) { + CHECK_MESSAGE( + accessors_are_equal(outputs.at(i), correct.at(i)), + check_kv("output", format_accessor_w_contents(outputs.at(i)))); + } + } + + SUBCASE("axis = 1") { + SplitAttrs attrs = SplitAttrs{ + /*splits=*/{2_p, 1_p, 3_p}, + /*axis=*/ff_dim_t{1_n}, + }; + + // Intentionally randomize these tensors so we can be confident we never + // read them + std::vector outputs = { + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{4_p, 2_p}}, DataType::FLOAT}, + allocator), + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{4_p, 1_p}}, DataType::FLOAT}, + allocator), + create_random_filled_accessor_w( + TensorShape{TensorDims{FFOrdered{4_p, 3_p}}, DataType::FLOAT}, + allocator), + }; + + split_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*input=*/input, + /*outputs=*/outputs); + + std::vector correct = { + create_2d_accessor_r_with_contents( + { + {1, 2}, + {-1, -2}, + {0.5, 0.25}, + {7, 8}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {3}, + {-3}, + {0.125}, + {9}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {4, 5, 6}, + {-4, -5, -6}, + {10, 20, 30}, + {0.75, 0.5, 0.25}, + }, + allocator), + }; + + for (int i = 0; i < outputs.size(); i++) { + CHECK_MESSAGE( + accessors_are_equal(outputs.at(i), correct.at(i)), + check_kv("output", format_accessor_w_contents(outputs.at(i)))); + } + } + } + + TEST_CASE("split_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(); + + GenericTensorAccessorR input = make_input(allocator); + + SUBCASE("axis = 0") { + SplitAttrs attrs = SplitAttrs{ + /*splits=*/{1_p, 3_p}, + /*axis=*/ff_dim_t{0_n}, + }; + + std::vector outputs = { + create_2d_accessor_r_with_contents( + { + {1, 2, 3, 4, 5, 6}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {-1, -2, -3, -4, -5, -6}, + {0.5, 0.25, 0.125, 10, 20, 30}, + {7, 8, 9, 0.75, 0.5, 0.25}, + }, + allocator), + }; + + std::vector output_grads = { + create_2d_accessor_r_with_contents( + { + {-0.75, -0.5, -0.25, 0, 0.25, 0.5}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {9.25, 9.5, 9.75, 10, 10.25, 10.5}, + {10.75, 11, 11.25, 11.5, 11.75, 12}, + {12.25, 12.5, 12.75, 13, 13.25, 13.5}, + }, + allocator), + }; + + GenericTensorAccessorW input_grad = make_input_grad(allocator); + + split_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*outputs=*/outputs, + /*output_grads=*/output_grads, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {-3.75, -1.5, 0.75, 3, 5.25, 7.5}, + {18.25, 20.5, 22.75, 25, 27.25, 29.5}, + {31.75, 34, 36.25, 38.5, 40.75, 43}, + {45.25, 47.5, 49.75, 52, 54.25, 56.5}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + + SUBCASE("axis = 1") { + SplitAttrs attrs = SplitAttrs{ + /*splits=*/{2_p, 1_p, 3_p}, + /*axis=*/ff_dim_t{1_n}, + }; + + std::vector outputs = { + create_2d_accessor_r_with_contents( + { + {1, 2}, + {-1, -2}, + {0.5, 0.25}, + {7, 8}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {3}, + {-3}, + {0.125}, + {9}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {4, 5, 6}, + {-4, -5, -6}, + {10, 20, 30}, + {0.75, 0.5, 0.25}, + }, + allocator), + }; + + std::vector output_grads = { + create_2d_accessor_r_with_contents( + { + {-0.75, -0.5}, + {-0.25, 0}, + {0.25, 0.5}, + {0.75, 1}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {9.25}, + {9.5}, + {9.75}, + {10}, + }, + allocator), + create_2d_accessor_r_with_contents( + { + {19.25, 19.5, 19.75}, + {20, 20.25, 20.5}, + {20.75, 21, 21.25}, + {21.5, 21.75, 22}, + }, + allocator), + }; + + GenericTensorAccessorW input_grad = make_input_grad(allocator); + + split_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*outputs=*/outputs, + /*output_grads=*/output_grads, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {-3.75, -1.5, 10.25, 22.25, 24.5, 26.75}, + {8.75, 11, 22.5, 35, 37.25, 39.5}, + {21.25, 23.5, 34.75, 47.75, 50, 52.25}, + {33.75, 36, 47, 60.5, 62.75, 65}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + } +} diff --git a/lib/kernels/test/src/kernels/transpose_kernels_cpu.cc b/lib/kernels/test/src/kernels/transpose_kernels_cpu.cc new file mode 100644 index 0000000000..bdc64fbcf6 --- /dev/null +++ b/lib/kernels/test/src/kernels/transpose_kernels_cpu.cc @@ -0,0 +1,157 @@ +#include "kernels/transpose_kernels_cpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "kernels/local_cpu_allocator.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +// The permutation below maps each output dim to the input dim it comes from +// (i.e. `output[d] = input[permutation.at_l(d)]`), so it is equivalent to +// torch.permute(input, (1, 2, 0)). +static TransposeAttrs make_test_attrs() { + return TransposeAttrs{ + TensorDimPermutation{ + bidict{ + {ff_dim_t{0_n}, ff_dim_t{1_n}}, + {ff_dim_t{1_n}, ff_dim_t{2_n}}, + {ff_dim_t{2_n}, ff_dim_t{0_n}}, + }, + }, + }; +} + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("transpose_cpu_forward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + TransposeAttrs attrs = make_test_attrs(); + + GenericTensorAccessorR input = create_3d_accessor_r_with_contents( + { + { + {0.25, 0.5, 0.75, 1.0}, + {1.25, 1.5, 1.75, 2.0}, + {2.25, 2.5, 2.75, 3.0}, + }, + { + {3.25, 3.5, 3.75, 4.0}, + {4.25, 4.5, 4.75, 5.0}, + {5.25, 5.5, 5.75, 6.0}, + }, + }, + allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{3_p, 4_p, 2_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + transpose_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = create_3d_accessor_r_with_contents( + { + {{0.25, 3.25}, {0.5, 3.5}, {0.75, 3.75}, {1.0, 4.0}}, + {{1.25, 4.25}, {1.5, 4.5}, {1.75, 4.75}, {2.0, 5.0}}, + {{2.25, 5.25}, {2.5, 5.5}, {2.75, 5.75}, {3.0, 6.0}}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + TEST_CASE("transpose_cpu_backward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + TransposeAttrs attrs = make_test_attrs(); + + GenericTensorAccessorR input = create_3d_accessor_r_with_contents( + { + { + {0.25, 0.5, 0.75, 1.0}, + {1.25, 1.5, 1.75, 2.0}, + {2.25, 2.5, 2.75, 3.0}, + }, + { + {3.25, 3.5, 3.75, 4.0}, + {4.25, 4.5, 4.75, 5.0}, + {5.25, 5.5, 5.75, 6.0}, + }, + }, + allocator); + + GenericTensorAccessorR output = create_3d_accessor_r_with_contents( + { + {{0.25, 3.25}, {0.5, 3.5}, {0.75, 3.75}, {1.0, 4.0}}, + {{1.25, 4.25}, {1.5, 4.5}, {1.75, 4.75}, {2.0, 5.0}}, + {{2.25, 5.25}, {2.5, 5.5}, {2.75, 5.75}, {3.0, 6.0}}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_3d_accessor_r_with_contents( + { + {{-0.875, -0.75}, + {-0.625, -0.5}, + {-0.375, -0.25}, + {-0.125, 0.0}}, + {{0.125, 0.25}, {0.375, 0.5}, {0.625, 0.75}, {0.875, 1.0}}, + {{1.125, 1.25}, {1.375, 1.5}, {1.625, 1.75}, {1.875, 2.0}}, + }, + allocator); + + // input_grad is accumulated into, so it needs to start from a known value + GenericTensorAccessorW input_grad = + create_3d_accessor_w_with_contents( + { + { + {-2.5, -2.0, -1.5, -1.0}, + {-0.5, 0.0, 0.5, 1.0}, + {1.5, 2.0, 2.5, 3.0}, + }, + { + {3.5, 4.0, 4.5, 5.0}, + {5.5, 6.0, 6.5, 7.0}, + {7.5, 8.0, 8.5, 9.0}, + }, + }, + allocator); + + transpose_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = create_3d_accessor_r_with_contents( + { + { + {-3.375, -2.625, -1.875, -1.125}, + {-0.375, 0.375, 1.125, 1.875}, + {2.625, 3.375, 4.125, 4.875}, + }, + { + {2.75, 3.5, 4.25, 5.0}, + {5.75, 6.5, 7.25, 8.0}, + {8.75, 9.5, 10.25, 11.0}, + }, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } +} diff --git a/lib/kernels/test/src/kernels/transpose_kernels_gpu.cc b/lib/kernels/test/src/kernels/transpose_kernels_gpu.cc new file mode 100644 index 0000000000..313512902a --- /dev/null +++ b/lib/kernels/test/src/kernels/transpose_kernels_gpu.cc @@ -0,0 +1,168 @@ +#include "kernels/transpose_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +// The permutation below maps each output dim to the input dim it comes from +// (i.e. `output[d] = input[permutation.at_l(d)]`), so it is equivalent to +// torch.permute(input, (1, 2, 0)). +static TransposeAttrs make_test_attrs() { + return TransposeAttrs{ + TensorDimPermutation{ + bidict{ + {ff_dim_t{0_n}, ff_dim_t{1_n}}, + {ff_dim_t{1_n}, ff_dim_t{2_n}}, + {ff_dim_t{2_n}, ff_dim_t{0_n}}, + }, + }, + }; +} + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("transpose_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(); + + TransposeAttrs attrs = make_test_attrs(); + + GenericTensorAccessorR input = create_3d_accessor_r_with_contents( + { + { + {0.25, 0.5, 0.75, 1.0}, + {1.25, 1.5, 1.75, 2.0}, + {2.25, 2.5, 2.75, 3.0}, + }, + { + {3.25, 3.5, 3.75, 4.0}, + {4.25, 4.5, 4.75, 5.0}, + {5.25, 5.5, 5.75, 6.0}, + }, + }, + allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{3_p, 4_p, 2_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + transpose_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = create_3d_accessor_r_with_contents( + { + {{0.25, 3.25}, {0.5, 3.5}, {0.75, 3.75}, {1.0, 4.0}}, + {{1.25, 4.25}, {1.5, 4.5}, {1.75, 4.75}, {2.0, 5.0}}, + {{2.25, 5.25}, {2.5, 5.5}, {2.75, 5.75}, {3.0, 6.0}}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + TEST_CASE("transpose_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(); + + TransposeAttrs attrs = make_test_attrs(); + + GenericTensorAccessorR input = create_3d_accessor_r_with_contents( + { + { + {0.25, 0.5, 0.75, 1.0}, + {1.25, 1.5, 1.75, 2.0}, + {2.25, 2.5, 2.75, 3.0}, + }, + { + {3.25, 3.5, 3.75, 4.0}, + {4.25, 4.5, 4.75, 5.0}, + {5.25, 5.5, 5.75, 6.0}, + }, + }, + allocator); + + GenericTensorAccessorR output = create_3d_accessor_r_with_contents( + { + {{0.25, 3.25}, {0.5, 3.5}, {0.75, 3.75}, {1.0, 4.0}}, + {{1.25, 4.25}, {1.5, 4.5}, {1.75, 4.75}, {2.0, 5.0}}, + {{2.25, 5.25}, {2.5, 5.5}, {2.75, 5.75}, {3.0, 6.0}}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_3d_accessor_r_with_contents( + { + {{-0.875, -0.75}, + {-0.625, -0.5}, + {-0.375, -0.25}, + {-0.125, 0.0}}, + {{0.125, 0.25}, {0.375, 0.5}, {0.625, 0.75}, {0.875, 1.0}}, + {{1.125, 1.25}, {1.375, 1.5}, {1.625, 1.75}, {1.875, 2.0}}, + }, + allocator); + + // input_grad is accumulated into, so it needs to start from a known value + GenericTensorAccessorW input_grad = + create_3d_accessor_w_with_contents( + { + { + {-2.5, -2.0, -1.5, -1.0}, + {-0.5, 0.0, 0.5, 1.0}, + {1.5, 2.0, 2.5, 3.0}, + }, + { + {3.5, 4.0, 4.5, 5.0}, + {5.5, 6.0, 6.5, 7.0}, + {7.5, 8.0, 8.5, 9.0}, + }, + }, + allocator); + + transpose_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = create_3d_accessor_r_with_contents( + { + { + {-3.375, -2.625, -1.875, -1.125}, + {-0.375, 0.375, 1.125, 1.875}, + {2.625, 3.375, 4.125, 4.875}, + }, + { + {2.75, 3.5, 4.25, 5.0}, + {5.75, 6.5, 7.25, 8.0}, + {8.75, 9.5, 10.25, 11.0}, + }, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } +} diff --git a/lib/kernels/test/src/kernels/upsample_kernels_cpu.cc b/lib/kernels/test/src/kernels/upsample_kernels_cpu.cc new file mode 100644 index 0000000000..2e1e5edf61 --- /dev/null +++ b/lib/kernels/test/src/kernels/upsample_kernels_cpu.cc @@ -0,0 +1,179 @@ +#include "kernels/upsample_kernels_cpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "kernels/local_cpu_allocator.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("upsample_cpu_forward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + UpsampleAttrs attrs = UpsampleAttrs{ + /*scale_factor=*/int_ge_two{2}, + /*mode=*/UpsampleMode::NEAREST, + }; + + // NCHW + GenericTensorAccessorR input = create_4d_accessor_r_with_contents( + { + { + { + {1, 2, 3}, + {4, 5, 6}, + }, + { + {-1, -2, -3}, + {0.5, 0.25, 0.125}, + }, + }, + }, + allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{1_p, 2_p, 4_p, 6_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + upsample_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = create_4d_accessor_r_with_contents( + { + { + { + {1, 1, 2, 2, 3, 3}, + {1, 1, 2, 2, 3, 3}, + {4, 4, 5, 5, 6, 6}, + {4, 4, 5, 5, 6, 6}, + }, + { + {-1, -1, -2, -2, -3, -3}, + {-1, -1, -2, -2, -3, -3}, + {0.5, 0.5, 0.25, 0.25, 0.125, 0.125}, + {0.5, 0.5, 0.25, 0.25, 0.125, 0.125}, + }, + }, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + TEST_CASE("upsample_cpu_backward_kernel") { + Allocator allocator = create_local_cpu_memory_allocator(); + + UpsampleAttrs attrs = UpsampleAttrs{ + /*scale_factor=*/int_ge_two{2}, + /*mode=*/UpsampleMode::NEAREST, + }; + + GenericTensorAccessorR input = create_4d_accessor_r_with_contents( + { + { + { + {1, 2, 3}, + {4, 5, 6}, + }, + { + {-1, -2, -3}, + {0.5, 0.25, 0.125}, + }, + }, + }, + allocator); + + GenericTensorAccessorR output = create_4d_accessor_r_with_contents( + { + { + { + {1, 1, 2, 2, 3, 3}, + {1, 1, 2, 2, 3, 3}, + {4, 4, 5, 5, 6, 6}, + {4, 4, 5, 5, 6, 6}, + }, + { + {-1, -1, -2, -2, -3, -3}, + {-1, -1, -2, -2, -3, -3}, + {0.5, 0.5, 0.25, 0.25, 0.125, 0.125}, + {0.5, 0.5, 0.25, 0.25, 0.125, 0.125}, + }, + }, + }, + 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, 1.125, 1.25}, + {1.375, 1.5, 1.625, 1.75, 1.875, 2}, + }, + { + {2.125, 2.25, 2.375, 2.5, 2.625, 2.75}, + {2.875, 3, 3.125, 3.25, 3.375, 3.5}, + {3.625, 3.75, 3.875, 4, 4.125, 4.25}, + {4.375, 4.5, 4.625, 4.75, 4.875, 5}, + }, + }, + }, + allocator); + + // input_grad is accumulated into, so it needs to start from a known value + GenericTensorAccessorW input_grad = + create_4d_accessor_w_with_contents( + { + { + { + {-7, -4, -1}, + {2, 5, 8}, + }, + { + {11, 14, 17}, + {20, 23, 26}, + }, + }, + }, + allocator); + + upsample_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = create_4d_accessor_r_with_contents( + { + { + { + {-8.75, -4.75, -0.75}, + {6.25, 10.25, 14.25}, + }, + { + {21.25, 25.25, 29.25}, + {36.25, 40.25, 44.25}, + }, + }, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } +} diff --git a/lib/kernels/test/src/kernels/upsample_kernels_gpu.cc b/lib/kernels/test/src/kernels/upsample_kernels_gpu.cc new file mode 100644 index 0000000000..639d16e80f --- /dev/null +++ b/lib/kernels/test/src/kernels/upsample_kernels_gpu.cc @@ -0,0 +1,313 @@ +#include "kernels/upsample_kernels_gpu.h" +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +static UpsampleAttrs make_attrs(int scale_factor) { + return UpsampleAttrs{ + /*scale_factor=*/int_ge_two{scale_factor}, + /*mode=*/UpsampleMode::NEAREST, + }; +} + +// NCHW +static GenericTensorAccessorR make_input(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + { + {1, 2, 3}, + {4, 5, 6}, + }, + { + {-1, -2, -3}, + {0.5, 0.25, 0.125}, + }, + }, + }, + allocator); +} + +static GenericTensorAccessorR make_output(Allocator &allocator) { + return create_4d_accessor_r_with_contents( + { + { + { + {1, 1, 2, 2, 3, 3}, + {1, 1, 2, 2, 3, 3}, + {4, 4, 5, 5, 6, 6}, + {4, 4, 5, 5, 6, 6}, + }, + { + {-1, -1, -2, -2, -3, -3}, + {-1, -1, -2, -2, -3, -3}, + {0.5, 0.5, 0.25, 0.25, 0.125, 0.125}, + {0.5, 0.5, 0.25, 0.25, 0.125, 0.125}, + }, + }, + }, + allocator); +} + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("upsample_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(); + + UpsampleAttrs attrs = make_attrs(/*scale_factor=*/2); + + GenericTensorAccessorR input = make_input(allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{1_p, 2_p, 4_p, 6_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + upsample_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = make_output(allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + TEST_CASE("upsample_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(); + + UpsampleAttrs attrs = make_attrs(/*scale_factor=*/2); + + GenericTensorAccessorR input = make_input(allocator); + 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, 1.125, 1.25}, + {1.375, 1.5, 1.625, 1.75, 1.875, 2}, + }, + { + {2.125, 2.25, 2.375, 2.5, 2.625, 2.75}, + {2.875, 3, 3.125, 3.25, 3.375, 3.5}, + {3.625, 3.75, 3.875, 4, 4.125, 4.25}, + {4.375, 4.5, 4.625, 4.75, 4.875, 5}, + }, + }, + }, + allocator); + + // input_grad is accumulated into, so it needs to start from a known value + GenericTensorAccessorW input_grad = + create_4d_accessor_w_with_contents( + { + { + { + {-7, -4, -1}, + {2, 5, 8}, + }, + { + {11, 14, 17}, + {20, 23, 26}, + }, + }, + }, + allocator); + + upsample_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = create_4d_accessor_r_with_contents( + { + { + { + {-8.75, -4.75, -0.75}, + {6.25, 10.25, 14.25}, + }, + { + {21.25, 25.25, 29.25}, + {36.25, 40.25, 44.25}, + }, + }, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + + TEST_CASE("upsample_gpu_forward_kernel (scale_factor = 3)") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + UpsampleAttrs attrs = make_attrs(/*scale_factor=*/3); + + GenericTensorAccessorR input = create_4d_accessor_r_with_contents( + { + { + { + {1, 2}, + {-3, 0.5}, + }, + }, + }, + allocator); + + TensorShape output_shape = TensorShape{ + TensorDims{FFOrdered{1_p, 1_p, 6_p, 6_p}}, + DataType::FLOAT, + }; + + // Intentionally randomize this tensor so we can be confident we never read + // it + GenericTensorAccessorW output = + create_random_filled_accessor_w(output_shape, allocator); + + upsample_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = create_4d_accessor_r_with_contents( + { + { + { + {1, 1, 1, 2, 2, 2}, + {1, 1, 1, 2, 2, 2}, + {1, 1, 1, 2, 2, 2}, + {-3, -3, -3, 0.5, 0.5, 0.5}, + {-3, -3, -3, 0.5, 0.5, 0.5}, + {-3, -3, -3, 0.5, 0.5, 0.5}, + }, + }, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + TEST_CASE("upsample_gpu_backward_kernel (scale_factor = 3)") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + UpsampleAttrs attrs = make_attrs(/*scale_factor=*/3); + + GenericTensorAccessorR input = create_4d_accessor_r_with_contents( + { + { + { + {1, 2}, + {-3, 0.5}, + }, + }, + }, + allocator); + + GenericTensorAccessorR output = create_4d_accessor_r_with_contents( + { + { + { + {1, 1, 1, 2, 2, 2}, + {1, 1, 1, 2, 2, 2}, + {1, 1, 1, 2, 2, 2}, + {-3, -3, -3, 0.5, 0.5, 0.5}, + {-3, -3, -3, 0.5, 0.5, 0.5}, + {-3, -3, -3, 0.5, 0.5, 0.5}, + }, + }, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_4d_accessor_r_with_contents( + { + { + { + {-3.75, -3.5, -3.25, -3, -2.75, -2.5}, + {-2.25, -2, -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}, + {2.25, 2.5, 2.75, 3, 3.25, 3.5}, + {3.75, 4, 4.25, 4.5, 4.75, 5}, + }, + }, + }, + allocator); + + // input_grad is accumulated into, so it needs to start from a known value + GenericTensorAccessorW input_grad = + create_4d_accessor_w_with_contents( + { + { + { + {100, 200}, + {300, 400}, + }, + }, + }, + allocator); + + upsample_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = create_4d_accessor_r_with_contents( + { + { + { + {82, 188.75}, + {322.5, 429.25}, + }, + }, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + 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/kernels/test/src/test_concat_kernel.cc b/lib/kernels/test/src/test_concat_kernel.cc deleted file mode 100644 index e3fdd3ad61..0000000000 --- a/lib/kernels/test/src/test_concat_kernel.cc +++ /dev/null @@ -1,113 +0,0 @@ -#include "internal/test_utils.h" -#include "kernels/concat_kernels_gpu.h" -#include "utils/containers/repeat.h" -#include - -using namespace ::FlexFlow; -TEST_SUITE(FF_CUDA_TEST_SUITE) { - TEST_CASE("Test concat kernel forward and backward") { - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - ManagedFFStream managed_stream{}; - Allocator allocator = create_local_cuda_memory_allocator(); - - const positive_int num_inputs = 4_p; - - SUBCASE("gpu_forward_kernel") { - auto run_forward_test = [&](positive_int input_rows, - positive_int input_cols, - TensorShape output_shape, - ff_dim_t concat_axis) { - TensorShape input_shape = TensorShape{ - TensorDims{FFOrdered{input_rows, input_cols}}, - DataType::FLOAT, - }; - - std::vector input_accessors = - repeat(num_inputs.nonnegative_int_from_positive_int(), [&]() { - return create_random_filled_accessor_r(input_shape, allocator); - }); - - GenericTensorAccessorW output_accessor = - allocator.allocate_tensor(output_shape); - - Kernels::Concat::gpu_forward_kernel(managed_stream.raw_stream(), - output_accessor, - input_accessors, - concat_axis); - - CHECK(contains_non_zero(output_accessor)); - }; - - SUBCASE("test forward concat, axis = 0") { - positive_int input_rows = 2_p; - positive_int input_cols = 4_p; - TensorShape output_shape = TensorShape{ - TensorDims{FFOrdered{num_inputs * input_rows, input_cols}}, - DataType::FLOAT, - }; - run_forward_test(input_rows, input_cols, output_shape, ff_dim_t{0_n}); - } - - SUBCASE("test forward concat, axis = 1") { - positive_int input_rows = 4_p; - positive_int input_cols = 2_p; - TensorShape output_shape = TensorShape{ - TensorDims{FFOrdered{input_rows, num_inputs * input_cols}}, - DataType::FLOAT, - }; - run_forward_test(input_rows, input_cols, output_shape, ff_dim_t{1_n}); - } - } - - SUBCASE("gpu_backward_kernel") { - auto run_backward_test = [&](positive_int input_rows, - positive_int input_cols, - TensorShape output_shape, - ff_dim_t concat_axis) { - TensorShape input_shape = TensorShape{ - TensorDims{FFOrdered{input_rows, input_cols}}, - DataType::FLOAT, - }; - - GenericTensorAccessorR output_grad_accessor = - create_random_filled_accessor_r(output_shape, allocator); - - std::vector input_grad_accessors = - repeat(num_inputs.nonnegative_int_from_positive_int(), [&]() { - return create_zero_filled_accessor_w(input_shape, allocator); - }); - - Kernels::Concat::gpu_backward_kernel(managed_stream.raw_stream(), - output_grad_accessor, - input_grad_accessors, - concat_axis); - - for (auto &accessor : input_grad_accessors) { - CHECK(contains_non_zero(accessor)); - } - }; - - SUBCASE("test backward concat, axis = 0") { - positive_int input_rows = 2_p; - positive_int input_cols = 4_p; - TensorShape output_shape = TensorShape{ - TensorDims{FFOrdered{num_inputs * input_rows, input_cols}}, - DataType::FLOAT, - }; - run_backward_test(input_rows, input_cols, output_shape, ff_dim_t{0_n}); - } - - SUBCASE("test backward concat, axis = 1") { - positive_int input_rows = 4_p; - positive_int input_cols = 2_p; - TensorShape output_shape = TensorShape{ - TensorDims{FFOrdered{input_rows, num_inputs * input_cols}}, - DataType::FLOAT, - }; - run_backward_test(input_rows, input_cols, output_shape, ff_dim_t{1_n}); - } - } - } -} diff --git a/lib/kernels/test/src/test_pool_2d_kernels.cc b/lib/kernels/test/src/test_pool_2d_kernels.cc deleted file mode 100644 index 5aa32899cb..0000000000 --- a/lib/kernels/test/src/test_pool_2d_kernels.cc +++ /dev/null @@ -1,91 +0,0 @@ -#include "internal/test_utils.h" -#include "kernels/pool_2d_kernels_gpu.h" -#include "op-attrs/datatype_value.h" -#include - -using namespace ::FlexFlow; -TEST_SUITE(FF_CUDA_TEST_SUITE) { - TEST_CASE("Test Pool2D Forward and Backward Kernel") { - positive_int input_w = 10_p; - positive_int input_h = 10_p; - positive_int input_c = 3_p; - positive_int input_n = 1_p; - positive_int output_w = 5_p; - positive_int output_h = 5_p; - positive_int output_c = 3_p; - positive_int output_n = 1_p; - nonnegative_int pad_h = 0_n; - nonnegative_int pad_w = 0_n; - positive_int kernel_h = 2_p; - positive_int kernel_w = 2_p; - positive_int stride_h = 2_p; - positive_int stride_w = 2_p; - - PoolOp pool_type = PoolOp::MAX; - - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - ManagedFFStream managed_stream{}; - - Allocator allocator = create_local_cuda_memory_allocator(); - - Pool2DPerDeviceState state = Kernels::Pool2D::gpu_init_kernel( - /*handle=*/managed_handle.raw_handle(), - /*activation=*/std::nullopt, - /*input_w=*/input_w.int_from_positive_int(), - /*input_h=*/input_h.int_from_positive_int(), - /*input_c=*/input_c.int_from_positive_int(), - /*input_n=*/input_n.int_from_positive_int(), - /*output_w=*/output_w.int_from_positive_int(), - /*output_h=*/output_h.int_from_positive_int(), - /*output_c=*/output_c.int_from_positive_int(), - /*output_n=*/output_n.int_from_positive_int(), - /*pad_h=*/pad_h.unwrap_nonnegative(), - /*pad_w=*/pad_w.unwrap_nonnegative(), - /*kernel_h=*/kernel_h.int_from_positive_int(), - /*kernel_w=*/kernel_w.int_from_positive_int(), - /*stride_h=*/stride_h.int_from_positive_int(), - /*stride_w=*/stride_w.int_from_positive_int(), - /*pool_type=*/pool_type); - - TensorShape input_shape = TensorShape{ - TensorDims{FFOrdered{input_n, input_c, input_h, input_w}}, - DataType::FLOAT, - }; - TensorShape output_shape = TensorShape{ - TensorDims{FFOrdered{output_n, input_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); - - SUBCASE("gpu_forward_kernel") { - Kernels::Pool2D::gpu_forward_kernel(managed_stream.raw_stream(), - state, - input_accessor.ptr, - output_accessor.ptr); - - CHECK(contains_non_zero(output_accessor)); - } - - SUBCASE("gpu_backward_kernel") { - GenericTensorAccessorW output_grad_accessor = create_filled_accessor_w( - output_shape, allocator, make_float_data_type_value(1)); - GenericTensorAccessorW input_grad_accessor = - allocator.allocate_tensor(input_shape); - - Kernels::Pool2D::gpu_backward_kernel(managed_stream.raw_stream(), - state, - output_accessor.ptr, - output_grad_accessor.ptr, - input_accessor.ptr, - input_grad_accessor.ptr); - - CHECK(contains_non_zero(input_grad_accessor)); - } - } -} diff --git a/lib/kernels/test/src/test_reshape_kernel.cc b/lib/kernels/test/src/test_reshape_kernel.cc deleted file mode 100644 index ad598c9055..0000000000 --- a/lib/kernels/test/src/test_reshape_kernel.cc +++ /dev/null @@ -1,49 +0,0 @@ -#include "internal/test_utils.h" -#include "kernels/reshape_kernels_gpu.h" -#include - -using namespace ::FlexFlow; -TEST_SUITE(FF_CUDA_TEST_SUITE) { - TEST_CASE("Test Reshape Forward and Backward") { - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - ManagedFFStream managed_stream{}; - - Allocator allocator = create_local_cuda_memory_allocator(); - - TensorShape input_shape = TensorShape{ - TensorDims{FFOrdered{100_p}}, - DataType::FLOAT, - }; - TensorShape output_shape = TensorShape{ - TensorDims{FFOrdered{100_p}}, - DataType::INT32, - }; - - SUBCASE("gpu_forward_kernel") { - GenericTensorAccessorR input_accessor = - create_random_filled_accessor_r(input_shape, allocator); - GenericTensorAccessorW output_accessor = - allocator.allocate_tensor(output_shape); - - Kernels::Reshape::gpu_forward_kernel( - managed_stream.raw_stream(), input_accessor, output_accessor); - - CHECK(contains_non_zero(output_accessor)); - } - - SUBCASE("gpu_backward_kernel") { - GenericTensorAccessorR output_grad_accessor = - create_random_filled_accessor_r(output_shape, allocator); - GenericTensorAccessorW input_grad_accessor = - allocator.allocate_tensor(input_shape); - - Kernels::Reshape::gpu_backward_kernel(managed_stream.raw_stream(), - output_grad_accessor, - input_grad_accessor); - - CHECK(contains_non_zero(input_grad_accessor)); - } - } -} diff --git a/lib/kernels/test/src/test_split_kernel.cc b/lib/kernels/test/src/test_split_kernel.cc deleted file mode 100644 index 35866308ee..0000000000 --- a/lib/kernels/test/src/test_split_kernel.cc +++ /dev/null @@ -1,72 +0,0 @@ -#include "internal/test_utils.h" -#include "kernels/split_kernels_gpu.h" -#include "op-attrs/datatype_value.h" -#include "utils/containers/repeat.h" -#include - -using namespace ::FlexFlow; - -TEST_SUITE(FF_CUDA_TEST_SUITE) { - TEST_CASE("Test Split Forward and Backward Kernel") { - nonnegative_int num_outputs = 2_n; - int out_blk_sizes[] = {50, 50}; - int in_blk_size = 100; - int num_blks = 1; - - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - ManagedFFStream managed_stream{}; - - Allocator allocator = create_local_cuda_memory_allocator(); - - TensorShape input_shape = TensorShape{ - TensorDims{FFOrdered{100_p}}, - DataType::FLOAT, - }; - TensorShape output_shape = TensorShape{ - TensorDims{FFOrdered{50_p}}, - DataType::FLOAT, - }; - - SUBCASE("gpu_forward_kernel") { - GenericTensorAccessorW input_accessor = - create_random_filled_accessor_w(input_shape, allocator); - - std::vector output_ptrs = repeat(num_outputs, [&]() { - GenericTensorAccessorW output_accessor = - allocator.allocate_tensor(output_shape); - return output_accessor.get_float_ptr(); - }); - - Kernels::Split::gpu_forward_kernel(managed_stream.raw_stream(), - output_ptrs.data(), - input_accessor.get_float_ptr(), - out_blk_sizes, - in_blk_size, - num_blks, - num_outputs.unwrap_nonnegative()); - } - - SUBCASE("gpu_backward_kernel") { - std::vector output_grad_ptrs(num_outputs.unwrap_nonnegative()); - for (int i = 0; i < num_outputs; i++) { - GenericTensorAccessorW output_grad_accessor = - create_random_filled_accessor_w(output_shape, allocator); - output_grad_ptrs[i] = output_grad_accessor.get_float_ptr(); - } - - GenericTensorAccessorW input_grad_accessor = create_filled_accessor_w( - input_shape, allocator, make_float_data_type_value(0)); - - Kernels::Split::gpu_backward_kernel( - managed_stream.raw_stream(), - input_grad_accessor.get_float_ptr(), - (float const **)output_grad_ptrs.data(), - out_blk_sizes, - in_blk_size, - num_blks, - num_outputs.unwrap_nonnegative()); - } - } -} diff --git a/lib/kernels/test/src/test_transpose_kernel.cc b/lib/kernels/test/src/test_transpose_kernel.cc deleted file mode 100644 index d3aa9262f3..0000000000 --- a/lib/kernels/test/src/test_transpose_kernel.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "internal/test_utils.h" -#include "kernels/transpose_kernels_gpu.h" -#include - -using namespace ::FlexFlow; -TEST_SUITE(FF_CUDA_TEST_SUITE) { - TEST_CASE("Test Transpose Kernel Operations") { - TransposeAttrs attrs = TransposeAttrs{ - TensorDimPermutation{ - bidict{ - {ff_dim_t{1_n}, ff_dim_t{0_n}}, - {ff_dim_t{0_n}, ff_dim_t{1_n}}, - }, - }, - }; - - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - ManagedFFStream managed_stream{}; - - Allocator allocator = create_local_cuda_memory_allocator(); - - TensorShape input_shape = TensorShape{ - TensorDims{FFOrdered{10_p, 10_p}}, - DataType::FLOAT, - }; - TensorShape output_shape = input_shape; - - SUBCASE("gpu_forward_kernel") { - GenericTensorAccessorR input_accessor = - create_random_filled_accessor_r(input_shape, allocator); - GenericTensorAccessorW output_accessor = - allocator.allocate_tensor(output_shape); - - Kernels::Transpose::gpu_forward_kernel( - managed_stream.raw_stream(), attrs, input_accessor, output_accessor); - - CHECK(contains_non_zero(output_accessor)); - } - - SUBCASE("gpu_backward_kernel") { - GenericTensorAccessorR output_grad_accessor = - create_random_filled_accessor_r(output_shape, allocator); - GenericTensorAccessorW input_grad_accessor = - create_random_filled_accessor_w(input_shape, allocator); - - Kernels::Transpose::gpu_backward_kernel(managed_stream.raw_stream(), - attrs, - output_grad_accessor, - input_grad_accessor); - - CHECK(contains_non_zero(input_grad_accessor)); - } - } -} diff --git a/lib/task-spec/src/task-spec/ops/impl/batch_matmul.cc b/lib/task-spec/src/task-spec/ops/impl/batch_matmul.cc index 3445465ebd..7046d52587 100644 --- a/lib/task-spec/src/task-spec/ops/impl/batch_matmul.cc +++ b/lib/task-spec/src/task-spec/ops/impl/batch_matmul.cc @@ -9,6 +9,7 @@ static std::optional ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); + device_handle_t handle = acc.get_ff_handle(); auto lhs_input = acc.get_tensor(TensorSlotName::LHS_INPUT); auto rhs_input = acc.get_tensor(TensorSlotName::RHS_INPUT); auto output = acc.get_tensor(TensorSlotName::OUTPUT); @@ -17,6 +18,7 @@ static std::optional profiling, kernel_device_type, "[BatchMatmul] forward_time = {:.2lf}ms\n", + handle, lhs_input, rhs_input, output); @@ -26,6 +28,7 @@ static std::optional backward_task_impl(TaskArgumentAccessor const &acc) { ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); + device_handle_t handle = acc.get_ff_handle(); auto lhs_input = acc.get_tensor(TensorSlotName::LHS_INPUT); auto lhs_input_grad = @@ -43,6 +46,7 @@ static std::optional profiling, kernel_device_type, "[BatchMatmul] backward_time = {:.2lf}ms\n", + handle, output, output_grad, lhs_input, 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 e622f0bc4c..49a31a5325 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(); - ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - - auto output = acc.get_tensor(TensorSlotName::OUTPUT); BatchNormAttrs attrs = acc.get_op_attrs().require_batch_norm(); - positive_int output_w = dim_at_idx(output.shape.dims, legion_dim_t{0_n}); - positive_int output_h = dim_at_idx(output.shape.dims, legion_dim_t{1_n}); - positive_int output_c = dim_at_idx(output.shape.dims, legion_dim_t{2_n}); - positive_int output_n = dim_at_idx(output.shape.dims, legion_dim_t{3_n}); - - float *runningMean; + TensorShape input_shape = acc.get_tensor_shape(TensorSlotName::INPUT); + TensorShape output_shape = acc.get_tensor_shape(TensorSlotName::OUTPUT); - std::optional per_device_state = init_kernel( - /*device_type=*/kernel_device_type, - /*handle=*/handle, - /*allocator=*/allocator, - /*runningMean=*/runningMean, - /*output_n=*/output_n.int_from_positive_int(), - /*output_c=*/output_c.int_from_positive_int(), - /*output_h=*/output_h.int_from_positive_int(), - /*output_w=*/output_w.int_from_positive_int(), - /*relu=*/attrs.relu); + std::optional per_device_state = + batch_norm_init_kernel( + kernel_device_type, allocator, attrs, input_shape, output_shape); return DeviceSpecificPerDeviceOpState{ acc.make_device_specific(per_device_state), @@ -56,56 +39,73 @@ static DeviceSpecificPerDeviceOpState static std::optional forward_task_impl(TaskArgumentAccessor const &acc) { - auto per_device_state = - acc.get_per_device_op_state().require_batch_norm().value(); ProfilingSettings 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(); ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); + device_handle_t handle = acc.get_ff_handle(); + std::optional per_device_state = + acc.get_per_device_op_state().require_batch_norm(); + BatchNormAttrs attrs = acc.get_op_attrs().require_batch_norm(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto input_grad = acc.get_tensor_grad(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - auto output_grad = - acc.get_tensor_grad(TensorSlotName::OUTPUT); - auto scale = acc.get_tensor(TensorSlotName::SCALE); - auto scale_grad = acc.get_tensor_grad(TensorSlotName::SCALE); - auto bias_grad = acc.get_tensor_grad(TensorSlotName::BIAS); - - return profile(backward_kernel, + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorW input_grad = + acc.get_tensor_grad(TensorSlotName::INPUT); + GenericTensorAccessorR output = + acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR output_grad = + acc.get_tensor_grad(TensorSlotName::OUTPUT); + GenericTensorAccessorR gamma = + acc.get_tensor(TensorSlotName::GAMMA); + GenericTensorAccessorW gamma_grad = + acc.get_tensor_grad(TensorSlotName::GAMMA); + GenericTensorAccessorW beta_grad = + acc.get_tensor_grad(TensorSlotName::BETA); + + return profile(batch_norm_backward_kernel, profiling, kernel_device_type, "[BatchNorm] backward_time = {:.2lf}ms\n", + handle, per_device_state, - output.get_float_ptr(), - output_grad.get_float_ptr(), - input.get_float_ptr(), - input_grad.get_float_ptr(), - scale.get_float_ptr(), - scale_grad.get_float_ptr(), - bias_grad.get_float_ptr(), - get_num_elements(output.shape.dims).int_from_positive_int()); + attrs, + output, + output_grad, + input, + input_grad, + gamma, + gamma_grad, + beta_grad); } TaskImplFunction get_batch_norm_init_task_impl() { diff --git a/lib/task-spec/src/task-spec/ops/impl/concat.cc b/lib/task-spec/src/task-spec/ops/impl/concat.cc index 39f9806226..d5839a11cf 100644 --- a/lib/task-spec/src/task-spec/ops/impl/concat.cc +++ b/lib/task-spec/src/task-spec/ops/impl/concat.cc @@ -23,39 +23,38 @@ namespace FlexFlow { -using namespace FlexFlow::Kernels::Concat; - static std::vector get_input_slots(ConcatAttrs const &attrs) { return slice(get_variadic_inputs_slot_name_sequence(), 0, attrs.num_inputs.int_from_int_ge_two()); } +static std::vector + get_inputs(TaskArgumentAccessor const &acc, ConcatAttrs const &attrs) { + return transform( + get_input_slots(attrs), + [&](TensorSlotName input_slot_name) -> GenericTensorAccessorR { + return acc.get_tensor(input_slot_name); + }); +} + static std::optional forward_task_impl(TaskArgumentAccessor const &acc) { ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); ConcatAttrs attrs = acc.get_op_attrs().require_concat(); - std::vector input_slots = get_input_slots(attrs); - - std::vector inputs = - transform(input_slots, - [&](TensorSlotName input_slot_name) -> GenericTensorAccessorR { - return acc.get_tensor(input_slot_name); - }); - - ASSERT(inputs.size() <= MAX_NUM_INPUTS); + std::vector inputs = get_inputs(acc, attrs); + GenericTensorAccessorW output = + acc.get_tensor(TensorSlotName::OUTPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - - return profile(forward_kernel, + return profile(concat_forward_kernel, profiling, kernel_device_type, "[Concat] forward_time = {:.2lf}ms\n", - output, + attrs, inputs, - attrs.axis); + output); } static std::optional @@ -64,26 +63,28 @@ static std::optional DeviceType kernel_device_type = acc.get_kernel_device_type(); ConcatAttrs attrs = acc.get_op_attrs().require_concat(); - std::vector input_slots = get_input_slots(attrs); + std::vector inputs = get_inputs(acc, attrs); std::vector input_grads = - transform(input_slots, + transform(get_input_slots(attrs), [&](TensorSlotName input_slot_name) -> GenericTensorAccessorW { return acc.get_tensor_grad(input_slot_name); }); - ASSERT(input_grads.size() <= MAX_NUM_INPUTS); - - auto output_grad = + GenericTensorAccessorR output = + acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR output_grad = acc.get_tensor_grad(TensorSlotName::OUTPUT); - return profile(backward_kernel, + return profile(concat_backward_kernel, profiling, kernel_device_type, "[Concat] backward_time = {:.2lf}ms\n", + attrs, + output, output_grad, - input_grads, - attrs.axis); + inputs, + input_grads); } TaskImplFunction get_concat_fwd_task_impl() { diff --git a/lib/task-spec/src/task-spec/ops/impl/conv_2d.cc b/lib/task-spec/src/task-spec/ops/impl/conv_2d.cc index 28fe73c3fc..fb1233f882 100644 --- a/lib/task-spec/src/task-spec/ops/impl/conv_2d.cc +++ b/lib/task-spec/src/task-spec/ops/impl/conv_2d.cc @@ -4,35 +4,21 @@ namespace FlexFlow { -using namespace FlexFlow::Kernels::Conv2D; - static DeviceSpecificPerDeviceOpState init_task_impl(TaskArgumentAccessor const &acc) { - device_handle_t handle = acc.get_ff_handle(); DeviceType kernel_device_type = acc.get_kernel_device_type(); Conv2DAttrs attrs = acc.get_op_attrs().require_conv2d(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - auto filter = acc.get_tensor(TensorSlotName::FILTER); - auto filter_grad = - acc.get_tensor_grad(TensorSlotName::FILTER); - std::optional per_device_state = init_kernel( + TensorShape input_shape = acc.get_tensor_shape(TensorSlotName::INPUT); + TensorShape output_shape = acc.get_tensor_shape(TensorSlotName::OUTPUT); + + std::optional per_device_state = conv_2d_init_kernel( /*device_type=*/kernel_device_type, /*handle=*/handle, - /*activation=*/attrs.activation, - /*kernel_h=*/attrs.kernel_h.int_from_positive_int(), - /*kernel_w=*/attrs.kernel_w.int_from_positive_int(), - /*groups=*/attrs.groups.int_from_positive_int(), - /*padding_h=*/attrs.padding_h.unwrap_nonnegative(), - /*padding_w=*/attrs.padding_w.unwrap_nonnegative(), - /*stride_h=*/attrs.stride_h.int_from_positive_int(), - /*stride_w=*/attrs.stride_w.int_from_positive_int(), - /*input=*/input, - /*output=*/output, - /*filter_ptr=*/filter.get_float_ptr(), - /*filter_grad_ptr=*/filter_grad.get_float_ptr()); + /*attrs=*/attrs, + /*input_shape=*/input_shape, + /*output_shape=*/output_shape); return DeviceSpecificPerDeviceOpState{ acc.make_device_specific(per_device_state), @@ -43,59 +29,76 @@ static std::optional forward_task_impl(TaskArgumentAccessor const &acc) { ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - Conv2DPerDeviceState per_device_state = - acc.get_per_device_op_state().require_conv2d().value(); + device_handle_t handle = acc.get_ff_handle(); + std::optional per_device_state = + acc.get_per_device_op_state().require_conv2d(); Conv2DAttrs attrs = acc.get_op_attrs().require_conv2d(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto filter = acc.get_tensor(TensorSlotName::FILTER); - auto bias = acc.get_tensor(TensorSlotName::BIAS); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - - return profile(forward_kernel, + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorR filter = + acc.get_tensor(TensorSlotName::FILTER); + std::optional bias = + attrs.use_bias ? std::optional{acc.get_tensor< + Permissions::RO>(TensorSlotName::BIAS)} + : std::nullopt; + GenericTensorAccessorW output = + acc.get_tensor(TensorSlotName::OUTPUT); + + return profile(conv_2d_forward_kernel, profiling, kernel_device_type, - "[Conv2d] forward_time = {:.2lf}ms\n", + "[Conv2D] forward_time = {:.2lf}ms\n", + handle, per_device_state, - input.get_float_ptr(), - output.get_float_ptr(), - filter.get_float_ptr(), - bias.get_float_ptr(), - attrs.activation); + attrs, + input, + filter, + bias, + output); } static std::optional backward_task_impl(TaskArgumentAccessor const &acc) { ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - Conv2DPerDeviceState per_device_state = - acc.get_per_device_op_state().require_conv2d().value(); + device_handle_t handle = acc.get_ff_handle(); + std::optional per_device_state = + acc.get_per_device_op_state().require_conv2d(); Conv2DAttrs attrs = acc.get_op_attrs().require_conv2d(); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto filter = acc.get_tensor(TensorSlotName::FILTER); - - auto input_grad = acc.get_tensor_grad(TensorSlotName::INPUT); - auto output_grad = - acc.get_tensor_grad(TensorSlotName::OUTPUT); - auto filter_grad = + 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 filter = + acc.get_tensor(TensorSlotName::FILTER); + GenericTensorAccessorW filter_grad = acc.get_tensor_grad(TensorSlotName::FILTER); - auto bias_grad = acc.get_tensor_grad(TensorSlotName::BIAS); + std::optional bias_grad = + attrs.use_bias + ? std::optional{acc.get_tensor_grad< + Permissions::RW>(TensorSlotName::BIAS)} + : std::nullopt; - return profile(backward_kernel, + return profile(conv_2d_backward_kernel, profiling, kernel_device_type, - "[Conv2d] backward_time = {:.2lf}ms\n", + "[Conv2D] 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(), - filter.get_float_ptr(), - filter_grad.get_float_ptr(), - bias_grad.get_float_ptr(), - attrs.activation); + attrs, + output, + output_grad, + input, + input_grad, + filter, + filter_grad, + bias_grad); } TaskImplFunction get_conv_2d_init_task_impl() { diff --git a/lib/task-spec/src/task-spec/ops/impl/pool_2d.cc b/lib/task-spec/src/task-spec/ops/impl/pool_2d.cc index ba2d984115..ce90f33486 100644 --- a/lib/task-spec/src/task-spec/ops/impl/pool_2d.cc +++ b/lib/task-spec/src/task-spec/ops/impl/pool_2d.cc @@ -2,64 +2,19 @@ #include "kernels/pool_2d_kernels.h" #include "op-attrs/ops/pool_2d.h" #include "task-spec/profiling.h" -#include "utils/exception.h" -#include "utils/hash-utils.h" - -using namespace FlexFlow::Kernels::Pool2D; namespace FlexFlow { -static nonnegative_int calculate_padding(nonnegative_int output_size, - nonnegative_int stride, - nonnegative_int kernel_size, - nonnegative_int input_size) { - int o = output_size.unwrap_nonnegative(); - int s = stride.unwrap_nonnegative(); - int k = kernel_size.unwrap_nonnegative(); - int i = kernel_size.unwrap_nonnegative(); - - return nonnegative_int{ - ((o - 1) * s + k - i + 1) / 2, - }; -} - static DeviceSpecificPerDeviceOpState init_task_impl(TaskArgumentAccessor const &acc) { Pool2DAttrs attrs = acc.get_op_attrs().require_pool2d(); - device_handle_t handle = acc.get_ff_handle(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - - positive_int input_w = dim_at_idx(input.shape.dims, ff_dim_t{0_n}); - positive_int input_h = dim_at_idx(input.shape.dims, ff_dim_t{1_n}); - positive_int input_c = dim_at_idx(input.shape.dims, ff_dim_t{2_n}); - positive_int input_n = dim_at_idx(input.shape.dims, ff_dim_t{3_n}); - positive_int output_w = dim_at_idx(output.shape.dims, ff_dim_t{0_n}); - positive_int output_h = dim_at_idx(output.shape.dims, ff_dim_t{1_n}); - positive_int output_c = dim_at_idx(output.shape.dims, ff_dim_t{2_n}); - positive_int output_n = dim_at_idx(output.shape.dims, ff_dim_t{3_n}); + 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(kernel_device_type, - handle, - attrs.activation, - input_w.int_from_positive_int(), - input_h.int_from_positive_int(), - input_c.int_from_positive_int(), - input_n.int_from_positive_int(), - output_w.int_from_positive_int(), - output_h.int_from_positive_int(), - output_c.int_from_positive_int(), - output_n.int_from_positive_int(), - attrs.padding_h.unwrap_nonnegative(), - attrs.padding_w.unwrap_nonnegative(), - attrs.kernel_h.int_from_positive_int(), - attrs.kernel_w.int_from_positive_int(), - attrs.stride_h.int_from_positive_int(), - attrs.stride_w.int_from_positive_int(), - attrs.pool_type); + pool_2d_init_kernel(kernel_device_type, attrs, input_shape, output_shape); return DeviceSpecificPerDeviceOpState{ acc.make_device_specific(per_device_state), @@ -70,42 +25,56 @@ static std::optional forward_task_impl(TaskArgumentAccessor const &acc) { ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - Pool2DPerDeviceState state = - acc.get_per_device_op_state().require_pool_2d().value(); + device_handle_t handle = acc.get_ff_handle(); + std::optional per_device_state = + acc.get_per_device_op_state().require_pool_2d(); + Pool2DAttrs attrs = acc.get_op_attrs().require_pool2d(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorW output = + acc.get_tensor(TensorSlotName::OUTPUT); - return profile(forward_kernel, + return profile(pool_2d_forward_kernel, profiling, kernel_device_type, "[Pool2D] forward_time = {:.2lf}ms\n", - state, - input.get_float_ptr(), - output.get_float_ptr()); + handle, + per_device_state, + attrs, + input, + output); } static std::optional backward_task_impl(TaskArgumentAccessor const &acc) { ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - Pool2DPerDeviceState state = - acc.get_per_device_op_state().require_pool_2d().value(); + device_handle_t handle = acc.get_ff_handle(); + std::optional per_device_state = + acc.get_per_device_op_state().require_pool_2d(); + Pool2DAttrs attrs = acc.get_op_attrs().require_pool2d(); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - auto output_grad = acc.get_tensor(TensorSlotName::OUTPUT); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto input_grad = acc.get_tensor(TensorSlotName::INPUT); + 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); - return profile(backward_kernel, + return profile(pool_2d_backward_kernel, profiling, kernel_device_type, "[Pool2D] backward_time = {:.2lf}ms\n", - state, - output.get_float_ptr(), - output_grad.get_float_ptr(), - input.get_float_ptr(), - input_grad.get_float_ptr()); + handle, + per_device_state, + attrs, + output, + output_grad, + input, + input_grad); } TaskImplFunction get_pool_2d_init_task_impl() { diff --git a/lib/task-spec/src/task-spec/ops/impl/reshape.cc b/lib/task-spec/src/task-spec/ops/impl/reshape.cc index 98470004c2..602cf2375d 100644 --- a/lib/task-spec/src/task-spec/ops/impl/reshape.cc +++ b/lib/task-spec/src/task-spec/ops/impl/reshape.cc @@ -19,18 +19,17 @@ namespace FlexFlow { -using namespace FlexFlow::Kernels::Reshape; - static std::optional forward_task_impl(TaskArgumentAccessor const &acc) { ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - ReshapeAttrs attrs = acc.get_op_attrs().require_reshape(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorW output = + acc.get_tensor(TensorSlotName::OUTPUT); - return profile(forward_kernel, + return profile(reshape_forward_kernel, profiling, kernel_device_type, "[Reshape] forward time = {:.2lf}ms\n", @@ -42,17 +41,23 @@ static std::optional backward_task_impl(TaskArgumentAccessor const &acc) { ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - ReshapeAttrs attrs = acc.get_op_attrs().require_reshape(); - auto input_grad = acc.get_tensor_grad(TensorSlotName::INPUT); - auto output_grad = + 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); - return profile(backward_kernel, + return profile(reshape_backward_kernel, profiling, kernel_device_type, "[Reshape] backward time = {:.2lf}ms\n", + output, output_grad, + input, input_grad); } diff --git a/lib/task-spec/src/task-spec/ops/impl/split.cc b/lib/task-spec/src/task-spec/ops/impl/split.cc index 1e3d3dde92..a9be5b1b9c 100644 --- a/lib/task-spec/src/task-spec/ops/impl/split.cc +++ b/lib/task-spec/src/task-spec/ops/impl/split.cc @@ -15,29 +15,25 @@ #include "task-spec/ops/impl/split.h" #include "kernels/split_kernels.h" +#include "op-attrs/tensor_slot_name.h" #include "task-spec/profiling.h" -#include "utils/exception.h" -#include "utils/hash-utils.h" -#include "utils/nonnegative_int/nonnegative_range.h" +#include "utils/containers/slice.h" +#include "utils/containers/transform.h" namespace FlexFlow { -using namespace FlexFlow::Kernels::Split; - -static std::pair - calc_block_size(TensorShape const &tensor_shape, ff_dim_t axis) { - positive_int num_blocks = 1_p; - positive_int block_size = 1_p; - for (nonnegative_int d : - nonnegative_range(get_num_elements(tensor_shape.dims) - .nonnegative_int_from_positive_int())) { - if (d <= axis.value) { - block_size *= dim_at_idx(tensor_shape.dims, legion_dim_t{d}); - } else { - num_blocks *= dim_at_idx(tensor_shape.dims, legion_dim_t{d}); - } - } - return {num_blocks, block_size}; +static std::vector get_output_slots(SplitAttrs const &attrs) { + return slice( + get_variadic_outputs_slot_name_sequence(), 0, attrs.splits.size()); +} + +static std::vector + get_outputs(TaskArgumentAccessor const &acc, SplitAttrs const &attrs) { + return transform( + get_output_slots(attrs), + [&](TensorSlotName output_slot_name) -> GenericTensorAccessorR { + return acc.get_tensor(output_slot_name); + }); } static std::optional @@ -46,27 +42,22 @@ static std::optional DeviceType kernel_device_type = acc.get_kernel_device_type(); SplitAttrs attrs = acc.get_op_attrs().require_split(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); - int out_block_sizes[MAX_NUM_OUTPUTS]; - auto [num_blocks, in_block_size] = calc_block_size(input.shape, attrs.axis); + std::vector outputs = + transform(get_output_slots(attrs), + [&](TensorSlotName output_slot_name) -> GenericTensorAccessorW { + return acc.get_tensor(output_slot_name); + }); - for (int i = 0; i < attrs.splits.size(); i++) { - auto [_, out_block_size] = calc_block_size(output.shape, attrs.axis); - out_block_sizes[i] = out_block_size.int_from_positive_int(); - } - float *output_float_ptr = output.get_float_ptr(); - return profile(forward_kernel, + return profile(split_forward_kernel, profiling, kernel_device_type, "[Split] forward_time = {:.2lf}ms\n", - &output_float_ptr, - input.get_float_ptr(), - out_block_sizes, - in_block_size.int_from_positive_int(), - num_blocks.int_from_positive_int(), - attrs.splits.size()); + attrs, + input, + outputs); } static std::optional @@ -75,30 +66,28 @@ static std::optional DeviceType kernel_device_type = acc.get_kernel_device_type(); SplitAttrs attrs = acc.get_op_attrs().require_split(); - auto input_grad = acc.get_tensor_grad(TensorSlotName::INPUT); - auto output_grad = - acc.get_tensor_grad(TensorSlotName::OUTPUT); - - int out_block_sizes[MAX_NUM_OUTPUTS]; - auto [num_blocks, in_block_size] = - calc_block_size(input_grad.shape, attrs.axis); - - for (int i = 0; i < attrs.splits.size(); i++) { - int out_num_blocks; - auto [_, out_block_size] = calc_block_size(output_grad.shape, attrs.axis); - out_block_sizes[i] = out_block_size.int_from_positive_int(); - } - float const *output_grad_ptr = output_grad.get_float_ptr(); - return profile(backward_kernel, + std::vector outputs = get_outputs(acc, attrs); + + std::vector output_grads = + transform(get_output_slots(attrs), + [&](TensorSlotName output_slot_name) -> GenericTensorAccessorR { + return acc.get_tensor_grad(output_slot_name); + }); + + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorW input_grad = + acc.get_tensor_grad(TensorSlotName::INPUT); + + return profile(split_backward_kernel, profiling, kernel_device_type, "[Split] backward_time = {:.2lf}ms\n", - input_grad.get_float_ptr(), - &output_grad_ptr, - out_block_sizes, - in_block_size.int_from_positive_int(), - num_blocks.int_from_positive_int(), - attrs.splits.size()); + attrs, + outputs, + output_grads, + input, + input_grad); } TaskImplFunction get_split_fwd_task_impl() { diff --git a/lib/task-spec/src/task-spec/ops/impl/transpose.cc b/lib/task-spec/src/task-spec/ops/impl/transpose.cc index 6b0c1d6d44..59f88de113 100644 --- a/lib/task-spec/src/task-spec/ops/impl/transpose.cc +++ b/lib/task-spec/src/task-spec/ops/impl/transpose.cc @@ -17,9 +17,6 @@ #include "kernels/transpose_kernels.h" #include "op-attrs/ops/transpose.h" #include "task-spec/profiling.h" -#include "utils/integer_conversions.h" - -using namespace FlexFlow::Kernels::Transpose; namespace FlexFlow { @@ -29,13 +26,15 @@ static std::optional TransposeAttrs attrs = acc.get_op_attrs().require_transpose(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorW output = + acc.get_tensor(TensorSlotName::OUTPUT); - return profile(forward_kernel, + return profile(transpose_forward_kernel, profiling, kernel_device_type, - "[Transpose] Forward_time = {:.2lf} [ms]", + "[Transpose] forward_time = {:.2lf}ms\n", attrs, input, output); @@ -47,16 +46,23 @@ static std::optional TransposeAttrs attrs = acc.get_op_attrs().require_transpose(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - auto input_grad = acc.get_tensor_grad(TensorSlotName::INPUT); - auto output_grad = + 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); - return profile(backward_kernel, + return profile(transpose_backward_kernel, profiling, kernel_device_type, - "[Transpose] Backward_time = {:.2lf} [ms]", + "[Transpose] backward_time = {:.2lf}ms\n", attrs, + output, output_grad, + input, input_grad); } From e7d87188ecb51572d47bb585794d02a664b4b76e Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 3 Sep 2026 16:20:05 -0700 Subject: [PATCH 2/2] Update batch_matmul kernels to support > 3D. --- .../src/cuda/ops/batch_matmul_kernels.cu | 42 +++++++++++++------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/lib/kernels/src/cuda/ops/batch_matmul_kernels.cu b/lib/kernels/src/cuda/ops/batch_matmul_kernels.cu index e8e4363dd8..291a9809a4 100644 --- a/lib/kernels/src/cuda/ops/batch_matmul_kernels.cu +++ b/lib/kernels/src/cuda/ops/batch_matmul_kernels.cu @@ -1,13 +1,21 @@ #include "internal/device.h" #include "kernels/batch_matmul_kernels_gpu.h" +#include "op-attrs/ff_ordered/ff_ordered_slice.h" #include "op-attrs/tensor_dims.h" +#include "utils/containers/product.h" #include "utils/containers/require_same.h" +#include "utils/containers/vector_of.h" namespace FlexFlow { // The dimensions of a batched matmul `output[b] = input_lhs[b] * input_rhs[b]`, // where input_lhs is (batch_size, lhs_rows, inner), input_rhs is (batch_size, // inner, rhs_cols) and output is (batch_size, lhs_rows, rhs_cols). +// +// Inputs may have more than one leading (batch) dimension, matching +// @ref batch_matmul_get_output_shape. Since the tensors are contiguous and +// row-major, all leading dimensions fold into a single cuBLAS batch count +// without any data movement. struct BatchMatmulDims { int batch_size; int lhs_rows; @@ -15,11 +23,18 @@ struct BatchMatmulDims { int rhs_cols; }; +// The dimensions preceding the trailing (rows, cols) pair, which together make +// up the batch of matrices being multiplied. +static FFOrdered get_leading_dims(TensorShape const &shape) { + return ff_ordered_slice( + shape.dims.ff_ordered, relative_ff_dim_t{0}, relative_ff_dim_t{-2}); +} + static BatchMatmulDims get_batch_matmul_dims(TensorShape const &lhs_shape, TensorShape const &rhs_shape, TensorShape const &output_shape) { - ASSERT(get_num_dims(lhs_shape.dims) == num_tensor_dims_t{3_n}, - "BatchMatmul expects 3-dimensional tensors", + ASSERT(get_num_dims(lhs_shape.dims) >= num_tensor_dims_t{3_n}, + "BatchMatmul expects tensors of at least 3 dimensions", lhs_shape); require_same(lhs_shape.data_type, rhs_shape.data_type); @@ -29,21 +44,24 @@ static BatchMatmulDims get_batch_matmul_dims(TensorShape const &lhs_shape, "If you need this feature, please create an issue.", lhs_shape.data_type); - positive_int batch_size = - require_same(dim_at_idx(lhs_shape.dims, ff_dim_t{0_n}), - dim_at_idx(rhs_shape.dims, ff_dim_t{0_n}), - dim_at_idx(output_shape.dims, ff_dim_t{0_n})); + FFOrdered leading_dims = + require_same(get_leading_dims(lhs_shape), + get_leading_dims(rhs_shape), + get_leading_dims(output_shape)); + + positive_int batch_size = product(vector_of(leading_dims)); positive_int lhs_rows = - require_same(dim_at_idx(lhs_shape.dims, ff_dim_t{1_n}), - dim_at_idx(output_shape.dims, ff_dim_t{1_n})); + require_same(dim_at_idx(lhs_shape.dims, relative_ff_dim_t{-2}), + dim_at_idx(output_shape.dims, relative_ff_dim_t{-2})); - positive_int inner = require_same(dim_at_idx(lhs_shape.dims, ff_dim_t{2_n}), - dim_at_idx(rhs_shape.dims, ff_dim_t{1_n})); + positive_int inner = + require_same(dim_at_idx(lhs_shape.dims, relative_ff_dim_t{-1}), + dim_at_idx(rhs_shape.dims, relative_ff_dim_t{-2})); positive_int rhs_cols = - require_same(dim_at_idx(rhs_shape.dims, ff_dim_t{2_n}), - dim_at_idx(output_shape.dims, ff_dim_t{2_n})); + require_same(dim_at_idx(rhs_shape.dims, relative_ff_dim_t{-1}), + dim_at_idx(output_shape.dims, relative_ff_dim_t{-1})); return BatchMatmulDims{ /*batch_size=*/batch_size.int_from_positive_int(),