From 17d86fd5f96d37f1218e8bdeec04737ba7f627a3 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 27 Aug 2026 12:07:06 -0700 Subject: [PATCH 1/6] Update and sync element unary/binary kernels. --- .../include/kernels/element_binary_kernels.h | 63 ++-- .../kernels/element_binary_kernels_cpu.h | 32 +- .../kernels/element_binary_kernels_gpu.h | 74 ++--- .../element_binary_per_device_state.dtg.toml | 5 - .../include/kernels/element_unary_kernels.h | 15 +- .../kernels/element_unary_kernels_cpu.h | 1 - .../kernels/element_unary_kernels_gpu.h | 11 +- .../include/kernels/map_tensor_accessors.h | 3 +- .../src/cuda/ops/element_binary_kernels.cu | 288 +++++++++--------- .../src/cuda/ops/element_unary_kernels.cu | 40 +-- .../src/kernels/element_binary_kernels.cc | 133 ++++---- .../src/kernels/element_binary_kernels_cpu.cc | 132 ++++++-- .../src/kernels/element_unary_kernels.cc | 27 +- .../src/kernels/element_unary_kernels_cpu.cc | 3 +- lib/kernels/test/src/internal/test_utils.cc | 17 ++ lib/kernels/test/src/internal/test_utils.h | 4 + .../src/kernels/element_binary_kernels_cpu.cc | 159 ++++++++++ .../src/kernels/element_binary_kernels_gpu.cc | 187 ++++++++++++ .../src/task-spec/ops/impl/element_binary.cc | 81 +++-- .../src/task-spec/ops/impl/element_unary.cc | 23 +- 20 files changed, 868 insertions(+), 430 deletions(-) create mode 100644 lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc create mode 100644 lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc diff --git a/lib/kernels/include/kernels/element_binary_kernels.h b/lib/kernels/include/kernels/element_binary_kernels.h index 8c9a405e6f..924d5bcea0 100644 --- a/lib/kernels/include/kernels/element_binary_kernels.h +++ b/lib/kernels/include/kernels/element_binary_kernels.h @@ -1,55 +1,48 @@ #ifndef _FLEXFLOW_OPS_KERNELS_ELEMENT_BINARY_KERNELS_H #define _FLEXFLOW_OPS_KERNELS_ELEMENT_BINARY_KERNELS_H -#include "kernels/device.h" +#include "kernels/accessor.h" #include "kernels/device_handle_t.dtg.h" #include "kernels/device_stream_t.dtg.h" #include "kernels/element_binary_per_device_state.dtg.h" -#include "kernels/ff_handle.h" -#include "op-attrs/datatype.h" -#include "op-attrs/operator_type.h" +#include "op-attrs/ops/element_binary_attrs.dtg.h" #include "op-attrs/tensor_shape.dtg.h" #include "pcg/device_type.dtg.h" -namespace FlexFlow::Kernels::ElementBinary { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - OperatorType op_type, - bool should_broadcast_lhs, - bool should_broadcast_rhs, - TensorShape const &lhs_shape, - TensorShape const &rhs_shape, - TensorShape const &output_shape); - -void forward_kernel( + element_binary_init_kernel(DeviceType device_type, + ElementBinaryAttrs const &attrs, + TensorShape const &lhs_shape, + TensorShape const &rhs_shape, + TensorShape const &output_shape); + +void element_binary_forward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - float const *lhs_ptr, - float const *rhs_ptr, - float *out_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - device_handle_t const &handle); - -void backward_kernel( + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &output); + +void element_binary_backward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - float const *out_grad_ptr, - float const *lhs_ptr, - float const *rhs_ptr, - float *lhs_grad_ptr, - float *rhs_grad_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - bool broadcast_inputRHS, - device_handle_t const &handle); - -void cleanup_kernel( + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorW const &lhs_grad, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &rhs_grad); + +void element_binary_cleanup_kernel( DeviceType device_type, std::optional const &per_device_state); -} // namespace FlexFlow::Kernels::ElementBinary +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/element_binary_kernels_cpu.h b/lib/kernels/include/kernels/element_binary_kernels_cpu.h index c53920764c..85e70657f3 100644 --- a/lib/kernels/include/kernels/element_binary_kernels_cpu.h +++ b/lib/kernels/include/kernels/element_binary_kernels_cpu.h @@ -1,25 +1,25 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_ELEMENT_BINARY_KERNELS_CPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_ELEMENT_BINARY_KERNELS_CPU_H -#include "op-attrs/operator_type.dtg.h" +#include "kernels/accessor.h" +#include "op-attrs/ops/element_binary_attrs.dtg.h" -namespace FlexFlow::Kernels::ElementBinary { +namespace FlexFlow { -void cpu_forward_kernel(float const *lhs_ptr, - float const *rhs_ptr, - float *out_ptr, - OperatorType op_type, - bool broadcast_inputLHS); +void element_binary_cpu_forward_kernel(ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &output); -void cpu_backward_kernel(float const *out_grad_ptr, - float const *lhs_ptr, - float const *rhs_ptr, - float *lhs_grad_ptr, - float *rhs_grad_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - bool broadcast_inputRHS); +void element_binary_cpu_backward_kernel( + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorW const &lhs_grad, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &rhs_grad); -} // namespace FlexFlow::Kernels::ElementBinary +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/element_binary_kernels_gpu.h b/lib/kernels/include/kernels/element_binary_kernels_gpu.h index 58a06edb4d..4a2430be6a 100644 --- a/lib/kernels/include/kernels/element_binary_kernels_gpu.h +++ b/lib/kernels/include/kernels/element_binary_kernels_gpu.h @@ -1,43 +1,43 @@ #ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_ELEMENT_BINARY_KERNELS_GPU_H #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_ELEMENT_BINARY_KERNELS_GPU_H +#include "kernels/accessor.h" +#include "kernels/device.h" #include "kernels/element_binary_per_device_state.dtg.h" -#include "op-attrs/operator_type.h" -#include "op-attrs/tensor_shape.dtg.h" - -namespace FlexFlow::Kernels::ElementBinary { - -ElementBinaryPerDeviceState gpu_init_kernel(PerDeviceFFHandle handle, - OperatorType op_type, - bool should_broadcast_lhs, - bool should_broadcast_rhs, - TensorShape const &lhs_shape, - TensorShape const &rhs_shape, - TensorShape const &output_shape); - -void gpu_forward_kernel(ffStream_t stream, - ElementBinaryPerDeviceState const &per_device_state, - float const *lhs_ptr, - float const *rhs_ptr, - float *out_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - PerDeviceFFHandle handle); - -void gpu_backward_kernel(ffStream_t stream, - ElementBinaryPerDeviceState const &per_device_state, - float const *out_grad_ptr, - float const *lhs_ptr, - float const *rhs_ptr, - float *lhs_grad_ptr, - float *rhs_grad_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - bool broadcast_inputRHS, - PerDeviceFFHandle handle); - -void gpu_cleanup_kernel(ElementBinaryPerDeviceState const &per_device_state); - -} // namespace FlexFlow::Kernels::ElementBinary +#include "op-attrs/ops/element_binary_attrs.dtg.h" + +namespace FlexFlow { + +ElementBinaryPerDeviceState + element_binary_gpu_init_kernel(ElementBinaryAttrs const &attrs, + TensorShape const &lhs_shape, + TensorShape const &rhs_shape, + TensorShape const &output_shape); + +void element_binary_gpu_forward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + ElementBinaryPerDeviceState const &per_device_state, + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &output); + +void element_binary_gpu_backward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + ElementBinaryPerDeviceState const &per_device_state, + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorW const &lhs_grad, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &rhs_grad); + +void element_binary_gpu_cleanup_kernel( + ElementBinaryPerDeviceState const &per_device_state); + +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/element_binary_per_device_state.dtg.toml b/lib/kernels/include/kernels/element_binary_per_device_state.dtg.toml index d0fb28dab4..b2729280a2 100644 --- a/lib/kernels/include/kernels/element_binary_per_device_state.dtg.toml +++ b/lib/kernels/include/kernels/element_binary_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 = "inputLHSTensor" type = "ffTensorDescriptor_t" diff --git a/lib/kernels/include/kernels/element_unary_kernels.h b/lib/kernels/include/kernels/element_unary_kernels.h index 6fd5a4b095..e88bace5fb 100644 --- a/lib/kernels/include/kernels/element_unary_kernels.h +++ b/lib/kernels/include/kernels/element_unary_kernels.h @@ -5,30 +5,31 @@ #include "kernels/device_handle_t.dtg.h" #include "kernels/device_stream_t.dtg.h" #include "kernels/element_unary_per_device_state.dtg.h" -#include "kernels/ff_handle.h" #include "op-attrs/ops/element_unary_attrs.dtg.h" +#include "op-attrs/tensor_shape.dtg.h" +#include "pcg/device_type.dtg.h" namespace FlexFlow { std::optional element_unary_init_kernel(DeviceType device_type, + ElementUnaryAttrs const &attrs, TensorShape const &input_shape, - TensorShape const &output_shape, - ElementUnaryAttrs const &attrs); + TensorShape const &output_shape); void element_unary_forward_kernel( device_stream_t const &stream, - std::optional const &device_state, - ElementUnaryAttrs const &attrs, device_handle_t const &handle, + std::optional const &per_device_state, + ElementUnaryAttrs const &attrs, GenericTensorAccessorR const &input, GenericTensorAccessorW const &output); void element_unary_backward_kernel( device_stream_t const &stream, - std::optional const &device_state, - ElementUnaryAttrs const &attrs, device_handle_t const &handle, + std::optional const &per_device_state, + ElementUnaryAttrs const &attrs, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input, diff --git a/lib/kernels/include/kernels/element_unary_kernels_cpu.h b/lib/kernels/include/kernels/element_unary_kernels_cpu.h index 83119176b1..d7c7a7e419 100644 --- a/lib/kernels/include/kernels/element_unary_kernels_cpu.h +++ b/lib/kernels/include/kernels/element_unary_kernels_cpu.h @@ -2,7 +2,6 @@ #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_ELEMENT_UNARY_KERNELS_CPU_H #include "kernels/accessor.h" -#include "kernels/ff_handle.h" #include "op-attrs/ops/element_unary_attrs.dtg.h" namespace FlexFlow { diff --git a/lib/kernels/include/kernels/element_unary_kernels_gpu.h b/lib/kernels/include/kernels/element_unary_kernels_gpu.h index 9292448f54..e9776a48f3 100644 --- a/lib/kernels/include/kernels/element_unary_kernels_gpu.h +++ b/lib/kernels/include/kernels/element_unary_kernels_gpu.h @@ -4,29 +4,28 @@ #include "kernels/accessor.h" #include "kernels/device.h" #include "kernels/element_unary_per_device_state.dtg.h" -#include "kernels/ff_handle.h" #include "op-attrs/ops/element_unary_attrs.dtg.h" namespace FlexFlow { ElementUnaryPerDeviceState - element_unary_gpu_init_kernel(TensorShape const &input_shape, - TensorShape const &output_shape, - ElementUnaryAttrs const &attrs); + element_unary_gpu_init_kernel(ElementUnaryAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape); void element_unary_gpu_forward_kernel( ffStream_t stream, + PerDeviceFFHandle const &handle, ElementUnaryPerDeviceState const &per_device_state, ElementUnaryAttrs const &attrs, - PerDeviceFFHandle const &handle, GenericTensorAccessorR const &input, GenericTensorAccessorW const &output); void element_unary_gpu_backward_kernel( ffStream_t stream, + PerDeviceFFHandle const &handle, ElementUnaryPerDeviceState const &per_device_state, ElementUnaryAttrs const &attrs, - PerDeviceFFHandle const &handle, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input, diff --git a/lib/kernels/include/kernels/map_tensor_accessors.h b/lib/kernels/include/kernels/map_tensor_accessors.h index 2e21907c46..940b2c630f 100644 --- a/lib/kernels/include/kernels/map_tensor_accessors.h +++ b/lib/kernels/include/kernels/map_tensor_accessors.h @@ -229,7 +229,8 @@ GenericTensorAccessorW map_tensor_accessors3(GenericTensorAccessorR const &lhs, DataType output_data_type, F &&f, Allocator &output_allocator) { - TensorDims output_dims = require_same(lhs.shape.dims, rhs.shape.dims); + TensorDims output_dims = + require_same(lhs.shape.dims, chs.shape.dims, rhs.shape.dims); GenericTensorAccessorW output = output_allocator.allocate_tensor( TensorShape{output_dims, output_data_type}); diff --git a/lib/kernels/src/cuda/ops/element_binary_kernels.cu b/lib/kernels/src/cuda/ops/element_binary_kernels.cu index e4c698ad02..d8c3379049 100644 --- a/lib/kernels/src/cuda/ops/element_binary_kernels.cu +++ b/lib/kernels/src/cuda/ops/element_binary_kernels.cu @@ -17,12 +17,10 @@ #include "kernels/element_binary_kernels_gpu.h" #include "kernels/ff_handle.h" #include "op-attrs/datatype.h" -#include "op-attrs/operator_type.h" +#include "op-attrs/get_op_type.h" #include "utils/exception.h" namespace FlexFlow { -namespace Kernels { -namespace ElementBinary { __global__ void elewise_binary_backward_kernel(size_t volume, float const alpha, @@ -80,13 +78,11 @@ __global__ void elewise_binary_backward_kernel(size_t volume, } } -ElementBinaryPerDeviceState gpu_init_kernel(PerDeviceFFHandle handle, - OperatorType op_type, - bool should_broadcast_lhs, - bool should_broadcast_rhs, - TensorShape const &lhs_shape, - TensorShape const &rhs_shape, - TensorShape const &output_shape) { +ElementBinaryPerDeviceState + element_binary_gpu_init_kernel(ElementBinaryAttrs const &attrs, + TensorShape const &lhs_shape, + TensorShape const &rhs_shape, + TensorShape const &output_shape) { ffTensorDescriptor_t inputLHSTensor; ffTensorDescriptor_t inputRHSTensor; ffTensorDescriptor_t outputTensor; @@ -100,7 +96,7 @@ ElementBinaryPerDeviceState gpu_init_kernel(PerDeviceFFHandle handle, checkCUDNN(cudnnCreateOpTensorDescriptor(&opDesc)); checkCUDNN(cudnnCreateReduceTensorDescriptor(&reduceAddDesc)); - switch (op_type) { + switch (get_op_type(attrs)) { case OperatorType::EW_ADD: case OperatorType::EW_SUB: mode = CUDNN_OP_TENSOR_ADD; @@ -133,7 +129,6 @@ ElementBinaryPerDeviceState gpu_init_kernel(PerDeviceFFHandle handle, cudnnSetTensorDescriptorFromTensorShape(outputTensor, output_shape)); ElementBinaryPerDeviceState per_device_state = ElementBinaryPerDeviceState{ - /*handle=*/handle, /*inputLHSTensor=*/inputLHSTensor, /*inputRHSTensor=*/inputRHSTensor, /*outputTensor=*/outputTensor, @@ -143,16 +138,17 @@ ElementBinaryPerDeviceState gpu_init_kernel(PerDeviceFFHandle handle, return per_device_state; } -void gpu_forward_kernel(cudaStream_t stream, - ElementBinaryPerDeviceState const &m, - float const *lhs_ptr, - float const *rhs_ptr, - float *out_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - PerDeviceFFHandle handle) { +void element_binary_gpu_forward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + ElementBinaryPerDeviceState const &per_device_state, + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &output) { checkCUBLAS(cublasSetStream(handle.blas, stream)); checkCUDNN(cudnnSetStream(handle.dnn, stream)); + OperatorType op_type = get_op_type(attrs); float alpha1 = 1.0f, alpha2 = 1.0f, beta = 0.0f; switch (op_type) { case OperatorType::EW_SUB: @@ -168,232 +164,235 @@ void gpu_forward_kernel(cudaStream_t stream, } // cudnn currently does not support broadcasting the first input in // cudnnOpTensor - if (broadcast_inputLHS) { + if (attrs.should_broadcast_lhs) { // currently only handle add and sub assert(op_type == OperatorType::EW_SUB || op_type == OperatorType::EW_ADD || op_type == OperatorType::EW_MUL); if (op_type == OperatorType::EW_SUB || op_type == OperatorType::EW_ADD) { // output = (beta*output + alpha1*input1) + beta*output = input1 checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &beta, - m.outputTensor, - out_ptr, + per_device_state.outputTensor, + output.get_float_ptr(), &alpha1, - m.inputLHSTensor, - lhs_ptr, + per_device_state.inputLHSTensor, + lhs.get_float_ptr(), &beta, - m.outputTensor, - out_ptr)); + per_device_state.outputTensor, + output.get_float_ptr())); // output = (beta*output + alpha2*input2) + alpha1*output = alpha2*input2 // + alpha1*input1 checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &beta, - m.outputTensor, - out_ptr, + per_device_state.outputTensor, + output.get_float_ptr(), &alpha2, - m.inputRHSTensor, - rhs_ptr, + per_device_state.inputRHSTensor, + rhs.get_float_ptr(), &alpha1, - m.outputTensor, - out_ptr)); + per_device_state.outputTensor, + output.get_float_ptr())); } else if (op_type == OperatorType::EW_MUL) { - checkCUDNN(cudnnSetOpTensorDescriptor(m.opDesc, + checkCUDNN(cudnnSetOpTensorDescriptor(per_device_state.opDesc, CUDNN_OP_TENSOR_ADD, CUDNN_DATA_FLOAT, CUDNN_PROPAGATE_NAN)); // output = (beta*output + alpha1*input1) + beta*output = input1 checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &beta, - m.outputTensor, - out_ptr, + per_device_state.outputTensor, + output.get_float_ptr(), &alpha1, - m.inputLHSTensor, - lhs_ptr, + per_device_state.inputLHSTensor, + lhs.get_float_ptr(), &beta, - m.outputTensor, - out_ptr)); - checkCUDNN(cudnnSetOpTensorDescriptor(m.opDesc, + per_device_state.outputTensor, + output.get_float_ptr())); + checkCUDNN(cudnnSetOpTensorDescriptor(per_device_state.opDesc, CUDNN_OP_TENSOR_MUL, CUDNN_DATA_FLOAT, CUDNN_PROPAGATE_NAN)); // output = (alpha1*output * alpha2*input2) + beta*output checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &alpha1, - m.outputTensor, - out_ptr, + per_device_state.outputTensor, + output.get_float_ptr(), &alpha2, - m.inputRHSTensor, - rhs_ptr, + per_device_state.inputRHSTensor, + rhs.get_float_ptr(), &beta, - m.outputTensor, - out_ptr)); + per_device_state.outputTensor, + output.get_float_ptr())); } } else { checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &alpha1, - m.inputLHSTensor, - lhs_ptr, + per_device_state.inputLHSTensor, + lhs.get_float_ptr(), &alpha2, - m.inputRHSTensor, - rhs_ptr, + per_device_state.inputRHSTensor, + rhs.get_float_ptr(), &beta, - m.outputTensor, - out_ptr)); + per_device_state.outputTensor, + output.get_float_ptr())); } } -void gpu_backward_kernel(cudaStream_t stream, - ElementBinaryPerDeviceState const &m, - float const *out_grad_ptr, - float const *lhs_ptr, - float const *rhs_ptr, - float *lhs_grad_ptr, - float *rhs_grad_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - bool broadcast_inputRHS, - PerDeviceFFHandle handle) { +void element_binary_gpu_backward_kernel( + ffStream_t stream, + PerDeviceFFHandle const &handle, + ElementBinaryPerDeviceState const &per_device_state, + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorW const &lhs_grad, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &rhs_grad) { checkCUBLAS(cublasSetStream(handle.blas, stream)); checkCUDNN(cudnnSetStream(handle.dnn, stream)); + OperatorType op_type = get_op_type(attrs); if (op_type == OperatorType::EW_ADD || op_type == OperatorType::EW_SUB) { float alpha = 1.0f, beta = 1.0f; - if (lhs_grad_ptr != nullptr) { - if (broadcast_inputLHS) { + if (lhs_grad.get_float_ptr() != nullptr) { + if (attrs.should_broadcast_lhs) { checkCUDNN(cudnnReduceTensor(handle.dnn, - m.reduceAddDesc, + per_device_state.reduceAddDesc, nullptr /*indices*/, 0 /*indicesSizeInBytes*/, handle.workSpace, handle.workSpaceSize, &alpha, - m.outputTensor, - out_grad_ptr, + per_device_state.outputTensor, + output_grad.get_float_ptr(), &beta, - m.inputLHSTensor, - lhs_grad_ptr)); + per_device_state.inputLHSTensor, + lhs_grad.get_float_ptr())); } else { checkCUDNN(cudnnAddTensor(handle.dnn, &alpha, - m.outputTensor, - out_grad_ptr, + per_device_state.outputTensor, + output_grad.get_float_ptr(), &beta, - m.inputLHSTensor, - lhs_grad_ptr)); + per_device_state.inputLHSTensor, + lhs_grad.get_float_ptr())); } } if (op_type == OperatorType::EW_SUB) { alpha = -1.0f; } - if (rhs_grad_ptr != nullptr) { - if (broadcast_inputRHS) { + if (rhs_grad.get_float_ptr() != nullptr) { + if (attrs.should_broadcast_rhs) { checkCUDNN(cudnnReduceTensor(handle.dnn, - m.reduceAddDesc, + per_device_state.reduceAddDesc, nullptr /*indices*/, 0 /*indicesSizeInBytes*/, handle.workSpace, handle.workSpaceSize, &alpha, - m.outputTensor, - out_grad_ptr, + per_device_state.outputTensor, + output_grad.get_float_ptr(), &beta, - m.inputRHSTensor, - rhs_grad_ptr)); + per_device_state.inputRHSTensor, + rhs_grad.get_float_ptr())); } else { checkCUDNN(cudnnAddTensor(handle.dnn, &alpha, - m.outputTensor, - out_grad_ptr, + per_device_state.outputTensor, + output_grad.get_float_ptr(), &beta, - m.inputRHSTensor, - rhs_grad_ptr)); + per_device_state.inputRHSTensor, + rhs_grad.get_float_ptr())); } } } else if (op_type == OperatorType::EW_MUL) { float alpha1 = 1.0f, alpha2 = 1.0f, beta = 1.0f, zero = 0.0f; - if (lhs_grad_ptr != nullptr) { - if (broadcast_inputLHS) { + if (lhs_grad.get_float_ptr() != nullptr) { + if (attrs.should_broadcast_lhs) { checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &alpha1, - m.outputTensor, - out_grad_ptr, + per_device_state.outputTensor, + output_grad.get_float_ptr(), &alpha2, - m.inputRHSTensor, - rhs_ptr, + per_device_state.inputRHSTensor, + rhs.get_float_ptr(), &zero, - m.outputTensor, + per_device_state.outputTensor, handle.workSpace)); checkCUDNN(cudnnReduceTensor( handle.dnn, - m.reduceAddDesc, + per_device_state.reduceAddDesc, nullptr /*indices*/, 0 /*indicesSizeInBytes*/, - (void *)((char *)handle.workSpace + sizeof(*out_grad_ptr)), - handle.workSpaceSize - sizeof(*out_grad_ptr), + (void *)((char *)handle.workSpace + + sizeof(*output_grad.get_float_ptr())), + handle.workSpaceSize - sizeof(*output_grad.get_float_ptr()), &alpha1, - m.outputTensor, + per_device_state.outputTensor, handle.workSpace, &beta, - m.inputLHSTensor, - lhs_grad_ptr)); + per_device_state.inputLHSTensor, + lhs_grad.get_float_ptr())); } else { checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &alpha1, - m.outputTensor, - out_grad_ptr, + per_device_state.outputTensor, + output_grad.get_float_ptr(), &alpha2, - m.inputRHSTensor, - rhs_ptr, + per_device_state.inputRHSTensor, + rhs.get_float_ptr(), &beta, - m.inputLHSTensor, - lhs_grad_ptr)); + per_device_state.inputLHSTensor, + lhs_grad.get_float_ptr())); } } - if (rhs_grad_ptr != nullptr) { - if (broadcast_inputRHS) { + if (rhs_grad.get_float_ptr() != nullptr) { + if (attrs.should_broadcast_rhs) { checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &alpha1, - m.outputTensor, - out_grad_ptr, + per_device_state.outputTensor, + output_grad.get_float_ptr(), &alpha2, - m.inputLHSTensor, - lhs_ptr, + per_device_state.inputLHSTensor, + lhs.get_float_ptr(), &zero, - m.outputTensor, + per_device_state.outputTensor, handle.workSpace)); checkCUDNN(cudnnReduceTensor( handle.dnn, - m.reduceAddDesc, + per_device_state.reduceAddDesc, nullptr /*indices*/, 0 /*indicesSizeInBytes*/, - (void *)((char *)handle.workSpace + sizeof(*out_grad_ptr)), - handle.workSpaceSize - sizeof(*out_grad_ptr), + (void *)((char *)handle.workSpace + + sizeof(*output_grad.get_float_ptr())), + handle.workSpaceSize - sizeof(*output_grad.get_float_ptr()), &alpha1, - m.outputTensor, + per_device_state.outputTensor, handle.workSpace, &beta, - m.inputRHSTensor, - rhs_grad_ptr)); + per_device_state.inputRHSTensor, + rhs_grad.get_float_ptr())); } else { checkCUDNN(cudnnOpTensor(handle.dnn, - m.opDesc, + per_device_state.opDesc, &alpha1, - m.outputTensor, - out_grad_ptr, + per_device_state.outputTensor, + output_grad.get_float_ptr(), &alpha2, - m.inputLHSTensor, - lhs_ptr, + per_device_state.inputLHSTensor, + lhs.get_float_ptr(), &beta, - m.inputRHSTensor, - rhs_grad_ptr)); + per_device_state.inputRHSTensor, + rhs_grad.get_float_ptr())); } } } else if (op_type == OperatorType::EW_MIN || @@ -403,8 +402,12 @@ void gpu_backward_kernel(cudaStream_t stream, int n; int dims[MAX_TENSOR_DIM]; int strides[MAX_TENSOR_DIM]; - checkCUDNN(cudnnGetTensorNdDescriptor( - m.outputTensor, MAX_TENSOR_DIM, &dataType, &n, dims, strides)); + checkCUDNN(cudnnGetTensorNdDescriptor(per_device_state.outputTensor, + MAX_TENSOR_DIM, + &dataType, + &n, + dims, + strides)); size_t volume = 1; for (int i = 0; i < n; i++) { volume *= dims[i]; @@ -416,20 +419,19 @@ void gpu_backward_kernel(cudaStream_t stream, alpha, beta, op_type, - out_grad_ptr, - lhs_ptr, - rhs_ptr, - lhs_grad_ptr, - rhs_grad_ptr); + output_grad.get_float_ptr(), + lhs.get_float_ptr(), + rhs.get_float_ptr(), + lhs_grad.get_float_ptr(), + rhs_grad.get_float_ptr()); } else { assert(false && "Unsupported ElementWise Binary Type"); } } -void gpu_cleanup_kernel(ElementBinaryPerDeviceState const &per_device_state) { +void element_binary_gpu_cleanup_kernel( + ElementBinaryPerDeviceState const &per_device_state) { NOT_IMPLEMENTED(); } -} // namespace ElementBinary -} // 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 7042ff664e..99ce5ba888 100644 --- a/lib/kernels/src/cuda/ops/element_unary_kernels.cu +++ b/lib/kernels/src/cuda/ops/element_unary_kernels.cu @@ -46,10 +46,10 @@ static bool use_scalar(OperatorType op_type) { } } -static ElementUnaryPerDeviceState - element_unary_gpu_init_kernel(TensorShape const &input_shape, - TensorShape const &output_shape, - OperatorType op_type) { +ElementUnaryPerDeviceState + element_unary_gpu_init_kernel(ElementUnaryAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { ffTensorDescriptor_t inputTensor; ffTensorDescriptor_t outputTensor; @@ -59,9 +59,9 @@ static ElementUnaryPerDeviceState checkCUDNN(cudnnCreateTensorDescriptor(&outputTensor)); checkCUDNN(cudnnCreateActivationDescriptor(&actiDesc)); - if (use_cudnn(op_type)) { + if (use_cudnn(get_op_type(attrs))) { cudnnActivationMode_t mode; - switch (op_type) { + switch (get_op_type(attrs)) { case OperatorType::SIGMOID: mode = CUDNN_ACTIVATION_SIGMOID; break; @@ -92,14 +92,6 @@ static ElementUnaryPerDeviceState }; } -ElementUnaryPerDeviceState - element_unary_gpu_init_kernel(TensorShape const &input_shape, - TensorShape const &output_shape, - ElementUnaryAttrs const &attrs) { - return element_unary_gpu_init_kernel( - input_shape, output_shape, get_op_type(attrs)); -} - template __global__ void elewise_scalar_unary_forward_kernel( coord_t volume, T scalar, OperatorType type, T const *in, T *out) { @@ -252,10 +244,10 @@ __global__ void elewise_unary_backward_kernel(coord_t volume, template struct ForwardKernel { void operator()(ffStream_t stream, + PerDeviceFFHandle const &handle, ElementUnaryPerDeviceState const &m, OperatorType op_type, std::optional scalar, - PerDeviceFFHandle const &handle, GenericTensorAccessorR const &input, GenericTensorAccessorW const &output) const { checkCUDNN(cudnnSetStream(handle.dnn, stream)); @@ -293,10 +285,10 @@ struct ForwardKernel { template struct BackwardKernel { void operator()(ffStream_t stream, + PerDeviceFFHandle const &handle, ElementUnaryPerDeviceState const &m, OperatorType op_type, std::optional scalar, - PerDeviceFFHandle const &handle, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input, @@ -345,36 +337,36 @@ struct BackwardKernel { void element_unary_gpu_forward_kernel( ffStream_t stream, - ElementUnaryPerDeviceState const &device_state, - ElementUnaryAttrs const &attrs, PerDeviceFFHandle const &handle, + ElementUnaryPerDeviceState const &per_device_state, + ElementUnaryAttrs const &attrs, GenericTensorAccessorR const &input, GenericTensorAccessorW const &output) { DataTypeDispatch1{}(input.shape.data_type, stream, - device_state, + handle, + per_device_state, get_op_type(attrs), attrs.scalar, - handle, input, output); } void element_unary_gpu_backward_kernel( ffStream_t stream, - ElementUnaryPerDeviceState const &device_state, - ElementUnaryAttrs const &attrs, PerDeviceFFHandle const &handle, + ElementUnaryPerDeviceState const &per_device_state, + ElementUnaryAttrs const &attrs, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input, GenericTensorAccessorW const &input_grad) { DataTypeDispatch1{}(input.shape.data_type, stream, - device_state, + handle, + per_device_state, get_op_type(attrs), attrs.scalar, - handle, output, output_grad, input, diff --git a/lib/kernels/src/kernels/element_binary_kernels.cc b/lib/kernels/src/kernels/element_binary_kernels.cc index 1d8fbaaf77..03b1ea43e9 100644 --- a/lib/kernels/src/kernels/element_binary_kernels.cc +++ b/lib/kernels/src/kernels/element_binary_kernels.cc @@ -1,117 +1,104 @@ #include "kernels/element_binary_kernels.h" #include "kernels/element_binary_kernels_cpu.h" #include "kernels/element_binary_kernels_gpu.h" -#include +#include "utils/optional.h" -namespace FlexFlow::Kernels::ElementBinary { +namespace FlexFlow { std::optional - init_kernel(DeviceType device_type, - device_handle_t const &handle, - OperatorType op_type, - bool should_broadcast_lhs, - bool should_broadcast_rhs, - TensorShape const &lhs_shape, - TensorShape const &rhs_shape, - TensorShape const &output_shape) { + element_binary_init_kernel(DeviceType device_type, + ElementBinaryAttrs const &attrs, + TensorShape const &lhs_shape, + TensorShape const &rhs_shape, + TensorShape const &output_shape) { if (device_type == DeviceType::GPU) { - return gpu_init_kernel( - /*handle=*/handle.require_for_gpu(), - /*op_type=*/op_type, - /*should_broadcast_lhs=*/should_broadcast_lhs, - /*should_broadcast_rhs=*/should_broadcast_rhs, + return element_binary_gpu_init_kernel( + /*attrs=*/attrs, /*lhs_shape=*/lhs_shape, /*rhs_shape=*/rhs_shape, /*output_shape=*/output_shape); } else { ASSERT(device_type == DeviceType::CPU); - ASSERT(handle.is_for_cpu()); return std::nullopt; } } -void forward_kernel( +void element_binary_forward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - float const *lhs_ptr, - float const *rhs_ptr, - float *out_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - device_handle_t const &handle) { + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + element_binary_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*lhs_ptr=*/lhs_ptr, - /*rhs_ptr=*/rhs_ptr, - /*out_ptr=*/out_ptr, - /*op_type=*/op_type, - /*broadcast_inputLHS=*/broadcast_inputLHS, - /*handle=*/handle.require_for_gpu()); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*lhs=*/lhs, + /*rhs=*/rhs, + /*output=*/output); } else { ASSERT(stream.is_cpu()); - ASSERT(per_device_state == std::nullopt); ASSERT(handle.is_for_cpu()); - cpu_forward_kernel( - /*lhs_ptr=*/lhs_ptr, - /*rhs_ptr=*/rhs_ptr, - /*out_ptr=*/out_ptr, - /*op_type=*/op_type, - /*broadcast_inputLHS=*/broadcast_inputLHS); + ASSERT(!per_device_state.has_value()); + element_binary_cpu_forward_kernel( + /*attrs=*/attrs, + /*lhs=*/lhs, + /*lhs=*/rhs, + /*output=*/output); } } -void backward_kernel( +void element_binary_backward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - float const *out_grad_ptr, - float const *lhs_ptr, - float const *rhs_ptr, - float *lhs_grad_ptr, - float *rhs_grad_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - bool broadcast_inputRHS, - device_handle_t const &handle) { + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorW const &lhs_grad, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &rhs_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + element_binary_gpu_backward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*out_grad_ptr=*/out_grad_ptr, - /*lhs_ptr=*/lhs_ptr, - /*rhs_ptr=*/rhs_ptr, - /*lhs_grad_ptr=*/lhs_grad_ptr, - /*rhs_grad_ptr=*/rhs_grad_ptr, - /*op_type=*/op_type, - /*broadcast_inputLHS=*/broadcast_inputLHS, - /*broadcast_inputRHS=*/broadcast_inputRHS, - /*handle=*/handle.require_for_gpu()); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*lhs=*/lhs, + /*lhs_grad=*/lhs_grad, + /*rhs=*/rhs, + /*rhs_grad=*/rhs_grad); } else { ASSERT(stream.is_cpu()); - ASSERT(per_device_state == std::nullopt); ASSERT(handle.is_for_cpu()); - cpu_backward_kernel( - /*out_grad_ptr=*/out_grad_ptr, - /*lhs_ptr=*/lhs_ptr, - /*rhs_ptr=*/rhs_ptr, - /*lhs_grad_ptr=*/lhs_grad_ptr, - /*rhs_grad_ptr=*/rhs_grad_ptr, - /*op_type=*/op_type, - /*broadcast_inputLHS=*/broadcast_inputLHS, - /*broadcast_inputRHS=*/broadcast_inputRHS); + ASSERT(!per_device_state.has_value()); + element_binary_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*lhs=*/lhs, + /*lhs_grad=*/lhs_grad, + /*rhs=*/rhs, + /*rhs_grad=*/rhs_grad); } } -void cleanup_kernel( +void element_binary_cleanup_kernel( DeviceType device_type, std::optional const &per_device_state) { if (device_type == DeviceType::GPU) { - gpu_cleanup_kernel(per_device_state.value()); + element_binary_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::ElementBinary +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/element_binary_kernels_cpu.cc b/lib/kernels/src/kernels/element_binary_kernels_cpu.cc index cbcd98dc7e..eaaca9f1d8 100644 --- a/lib/kernels/src/kernels/element_binary_kernels_cpu.cc +++ b/lib/kernels/src/kernels/element_binary_kernels_cpu.cc @@ -1,25 +1,123 @@ #include "kernels/element_binary_kernels_cpu.h" +#include "kernels/map_tensor_accessors.h" +#include "kernels/tensor_accessor_binary_ops.h" #include "utils/exception.h" -namespace FlexFlow::Kernels::ElementBinary { +namespace FlexFlow { -void cpu_forward_kernel(float const *lhs_ptr, - float const *rhs_ptr, - float *out_ptr, - OperatorType op_type, - bool broadcast_inputLHS) { - NOT_IMPLEMENTED(); +void element_binary_cpu_forward_kernel(ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &output) { + std::function element_function = + [&]() -> std::function { + switch (attrs.type) { + case OperatorType::EW_ADD: + return [](float x, float y) -> float { return x + y; }; + case OperatorType::EW_SUB: + return [](float x, float y) -> float { return x - y; }; + case OperatorType::EW_MUL: + return [](float x, float y) -> float { return x * y; }; + case OperatorType::EW_DIV: + return [](float x, float y) -> float { return x / y; }; + case OperatorType::EW_MAX: + return [](float x, float y) -> float { return std::max(x, y); }; + case OperatorType::EW_MIN: + return [](float x, float y) -> float { return std::min(x, y); }; + default: + PANIC("Unhandled OperatorType {}", attrs.type); + } + }(); + + map_tensor_accessors2_to( + /*lhs=*/lhs, + /*rhs=*/rhs, + /*output_data_type=*/ + require_same(lhs.shape.data_type, rhs.shape.data_type), + /*f=*/element_function, + /*output=*/output); } -void cpu_backward_kernel(float const *out_grad_ptr, - float const *lhs_ptr, - float const *rhs_ptr, - float *lhs_grad_ptr, - float *rhs_grad_ptr, - OperatorType op_type, - bool broadcast_inputLHS, - bool broadcast_inputRHS) { - NOT_IMPLEMENTED(); +void element_binary_cpu_backward_kernel( + ElementBinaryAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &lhs, + GenericTensorAccessorW const &lhs_grad, + GenericTensorAccessorR const &rhs, + GenericTensorAccessorW const &rhs_grad) { + + std::function lhs_grad_update_function = + [&]() -> std::function { + switch (attrs.type) { + case OperatorType::EW_ADD: + return [](float og, float l, float r) -> float { return og; }; + case OperatorType::EW_SUB: + return [](float og, float l, float r) -> float { return og; }; + case OperatorType::EW_MUL: + return [](float og, float l, float r) -> float { return og * r; }; + case OperatorType::EW_DIV: + return [](float og, float l, float r) -> float { return og / r; }; + case OperatorType::EW_MAX: + NOT_IMPLEMENTED(); + case OperatorType::EW_MIN: + NOT_IMPLEMENTED(); + default: + PANIC("Unhandled OperatorType {}", attrs.type); + } + }(); + + std::function rhs_grad_update_function = + [&]() -> std::function { + switch (attrs.type) { + case OperatorType::EW_ADD: + return [](float og, float l, float r) -> float { return og; }; + case OperatorType::EW_SUB: + return [](float og, float l, float r) -> float { return -og; }; + case OperatorType::EW_MUL: + return [](float og, float l, float r) -> float { return og * l; }; + case OperatorType::EW_DIV: + return [](float og, float l, float r) -> float { + return -og / l * (r * r); + }; + case OperatorType::EW_MAX: + NOT_IMPLEMENTED(); + case OperatorType::EW_MIN: + NOT_IMPLEMENTED(); + default: + PANIC("Unhandled OperatorType {}", attrs.type); + } + }(); + + Allocator cpu_allocator = create_local_cpu_memory_allocator(); + + GenericTensorAccessorR lhs_grad_update = + read_only_accessor_from_write_accessor(map_tensor_accessors3( + /*lhs=*/output_grad, + /*chs=*/lhs, + /*rhs=*/rhs, + /*output_data_type=*/ + require_same(output_grad.shape.data_type, + lhs.shape.data_type, + rhs.shape.data_type), + /*f=*/lhs_grad_update_function, + /*output_allocator=*/cpu_allocator)); + + GenericTensorAccessorR rhs_grad_update = + read_only_accessor_from_write_accessor(map_tensor_accessors3( + /*lhs=*/output_grad, + /*chs=*/lhs, + /*rhs=*/rhs, + /*output_data_type=*/ + require_same(output_grad.shape.data_type, + lhs.shape.data_type, + rhs.shape.data_type), + /*f=*/rhs_grad_update_function, + /*output_allocator=*/cpu_allocator)); + + tensor_accessor_elementwise_add_to(lhs_grad, lhs_grad_update, lhs_grad); + + tensor_accessor_elementwise_add_to(rhs_grad, rhs_grad_update, rhs_grad); } -} // namespace FlexFlow::Kernels::ElementBinary +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/element_unary_kernels.cc b/lib/kernels/src/kernels/element_unary_kernels.cc index 070eb39e1b..3ff94e878a 100644 --- a/lib/kernels/src/kernels/element_unary_kernels.cc +++ b/lib/kernels/src/kernels/element_unary_kernels.cc @@ -1,19 +1,20 @@ #include "kernels/element_unary_kernels.h" #include "kernels/element_unary_kernels_cpu.h" #include "kernels/element_unary_kernels_gpu.h" +#include "utils/optional.h" namespace FlexFlow { std::optional element_unary_init_kernel(DeviceType device_type, + ElementUnaryAttrs const &attrs, TensorShape const &input_shape, - TensorShape const &output_shape, - ElementUnaryAttrs const &attrs) { + TensorShape const &output_shape) { if (device_type == DeviceType::GPU) { return element_unary_gpu_init_kernel( + /*attrs=*/attrs, /*input_shape=*/input_shape, - /*output_shape=*/output_shape, - /*attrs=*/attrs); + /*output_shape=*/output_shape); } else { ASSERT(device_type == DeviceType::CPU); return std::nullopt; @@ -22,23 +23,23 @@ std::optional void element_unary_forward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, ElementUnaryAttrs const &attrs, - device_handle_t const &handle, GenericTensorAccessorR const &input, GenericTensorAccessorW const &output) { if (stream.is_gpu()) { element_unary_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*attrs=*/attrs, /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, /*input=*/input, /*output=*/output); } else { ASSERT(stream.is_cpu()); - ASSERT(per_device_state == std::nullopt); ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); element_unary_cpu_forward_kernel( /*attrs=*/attrs, /*input=*/input, @@ -48,9 +49,9 @@ void element_unary_forward_kernel( void element_unary_backward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, ElementUnaryAttrs const &attrs, - device_handle_t const &handle, GenericTensorAccessorR const &output, GenericTensorAccessorR const &output_grad, GenericTensorAccessorR const &input, @@ -58,17 +59,17 @@ void element_unary_backward_kernel( if (stream.is_gpu()) { element_unary_gpu_backward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*attrs=*/attrs, /*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()); - ASSERT(per_device_state == std::nullopt); ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); element_unary_cpu_backward_kernel( /*attrs=*/attrs, /*output=*/output, @@ -85,7 +86,7 @@ void element_unary_cleanup_kernel( element_unary_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()); } } diff --git a/lib/kernels/src/kernels/element_unary_kernels_cpu.cc b/lib/kernels/src/kernels/element_unary_kernels_cpu.cc index 40c6376072..183cdfd4cf 100644 --- a/lib/kernels/src/kernels/element_unary_kernels_cpu.cc +++ b/lib/kernels/src/kernels/element_unary_kernels_cpu.cc @@ -163,8 +163,7 @@ void element_unary_cpu_backward_kernel( /*f=*/element_function, /*output_allocator=*/cpu_allocator)); - return tensor_accessor_elementwise_multiply_to( - output_grad, df_dx, input_grad); + tensor_accessor_elementwise_multiply_to(output_grad, df_dx, input_grad); } } // namespace FlexFlow diff --git a/lib/kernels/test/src/internal/test_utils.cc b/lib/kernels/test/src/internal/test_utils.cc index d7baaeef48..41c161e3f3 100644 --- a/lib/kernels/test/src/internal/test_utils.cc +++ b/lib/kernels/test/src/internal/test_utils.cc @@ -1,5 +1,8 @@ #include "internal/test_utils.h" #include "kernels/fill_tensor_accessor.h" +#include "kernels/map_tensor_accessors.h" +#include "kernels/tensor_accessor_binary_ops.h" +#include "kernels/tensor_accessor_reductions.h" #include "op-attrs/tensor_shape.h" #include "utils/containers/require_all_same1.h" #include "utils/join_strings.h" @@ -132,6 +135,20 @@ bool accessors_are_equal(GenericTensorAccessorR const &accessor_a, accessor_a.shape.data_type, accessor_a, accessor_b); } +bool accessors_within_epsilon(GenericTensorAccessorR const &accessor_a, + GenericTensorAccessorR const &accessor_b, + float epsilon) { + ASSERT(accessor_a.shape == accessor_b.shape, + "accessors_are_equal expects accessors to have the same shape"); + + Allocator cpu_allocator = create_local_cpu_memory_allocator(); + GenericTensorAccessorW diff = tensor_accessor_elementwise_subtract( + accessor_a, accessor_b, cpu_allocator); + GenericTensorAccessorW within_epsilon = map_tensor_accessor( + diff, [&](float x) { return std::abs(x) < epsilon; }, cpu_allocator); + return tensor_accessor_all(within_epsilon); +} + template struct CreateFilledAccessorW { GenericTensorAccessorW operator()(TensorShape const &shape, diff --git a/lib/kernels/test/src/internal/test_utils.h b/lib/kernels/test/src/internal/test_utils.h index 3a2c9b773c..ab24142fc0 100644 --- a/lib/kernels/test/src/internal/test_utils.h +++ b/lib/kernels/test/src/internal/test_utils.h @@ -37,6 +37,10 @@ void print_2d_tensor_accessor_contents(GenericTensorAccessorR const &accessor, bool accessors_are_equal(GenericTensorAccessorR const &accessor_a, GenericTensorAccessorR const &accessor_b); +bool accessors_within_epsilon(GenericTensorAccessorR const &accessor_a, + GenericTensorAccessorR const &accessor_b, + float epsilon); + GenericTensorAccessorW create_filled_accessor_w(TensorShape const &shape, Allocator &allocator, DataTypeValue val); diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc new file mode 100644 index 0000000000..b52d847178 --- /dev/null +++ b/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc @@ -0,0 +1,159 @@ +#include "kernels/element_binary_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 "op-attrs/ops/element_binary.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("element_binary_cpu_forward_kernel") { + SUBCASE("add") { + Allocator cpu_allocator = create_local_cpu_memory_allocator(); + + ElementBinaryAttrs attrs{ + /*type=*/OperatorType::EW_ADD, + /*compute_type=*/DataType::FLOAT, + /*should_broadcast_lhs=*/false, + /*should_broadcast_rhs=*/false, + }; + + GenericTensorAccessorR lhs = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {7, 8, 9}, + {-1, -2, -3}, + }, + cpu_allocator); + + GenericTensorAccessorR rhs = create_2d_accessor_r_with_contents( + { + {10, 50, -10}, + {20, 60, -20}, + {30, 70, -30}, + {40, 80, -40}, + }, + cpu_allocator); + + GenericTensorAccessorW result = + create_zero_filled_accessor_w(lhs.shape, cpu_allocator); + + element_binary_cpu_forward_kernel( + /*attrs=*/attrs, + /*lhs=*/lhs, + /*rhs=*/rhs, + /*output=*/result); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {11, 52, -7}, + {24, 65, -14}, + {37, 78, -21}, + {39, 78, -43}, + }, + cpu_allocator); + + CHECK_MESSAGE(accessors_are_equal(result, correct), + check_kv("result", format_accessor_w_contents(result))); + } + } + + TEST_CASE("element_binary_cpu_backward_kernel") { + SUBCASE("add") { + Allocator cpu_allocator = create_local_cpu_memory_allocator(); + + ElementBinaryAttrs attrs{ + /*type=*/OperatorType::EW_ADD, + /*compute_type=*/DataType::FLOAT, + /*should_broadcast_lhs=*/false, + /*should_broadcast_rhs=*/false, + }; + + GenericTensorAccessorR lhs = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {7, 8, 9}, + {-1, -2, -3}, + }, + cpu_allocator); + + GenericTensorAccessorR rhs = create_2d_accessor_r_with_contents( + { + {10, 50, -10}, + {20, 60, -20}, + {30, 70, -30}, + {40, 80, -40}, + }, + cpu_allocator); + + GenericTensorAccessorW lhs_grad = + cpu_allocator.allocate_tensor(get_tensor_shape_for_accessor_r(lhs)); + + GenericTensorAccessorW rhs_grad = + cpu_allocator.allocate_tensor(get_tensor_shape_for_accessor_r(rhs)); + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {11, 52, -7}, + {24, 65, -14}, + {73, 78, -21}, + {39, 78, -43}, + }, + cpu_allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, -1}, + {6, 4, 2}, + {0.5, 0.1, -2}, + {0, 0.5, 0}, + }, + cpu_allocator); + + element_binary_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*lhs=*/lhs, + /*lhs_grad=*/lhs_grad, + /*rhs=*/rhs, + /*rhs_grad=*/rhs_grad); + + GenericTensorAccessorR correct_lhs_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, -1}, + {6, 4, 2}, + {0.5, 0.1, -2}, + {0, 0.5, 0}, + }, + cpu_allocator); + + GenericTensorAccessorR correct_rhs_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, -1}, + {6, 4, 2}, + {0.5, 0.1, -2}, + {0, 0.5, 0}, + }, + cpu_allocator); + + float epsilon = 1e-20; + CHECK_MESSAGE( + accessors_within_epsilon(lhs_grad, correct_lhs_grad, epsilon), + check_kv("lhs_grad", format_accessor_w_contents(lhs_grad))); + + CHECK_MESSAGE( + accessors_within_epsilon(rhs_grad, correct_rhs_grad, epsilon), + check_kv("rhs_grad", format_accessor_w_contents(rhs_grad))); + } + } +} diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc new file mode 100644 index 0000000000..96b0d589b3 --- /dev/null +++ b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc @@ -0,0 +1,187 @@ +#include "kernels/element_binary_kernels_gpu.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 "op-attrs/ops/element_binary.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("element_binary_gpu_forward_kernel") { + SUBCASE("add") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + ElementBinaryAttrs attrs{ + /*type=*/OperatorType::EW_ADD, + /*compute_type=*/DataType::FLOAT, + /*should_broadcast_lhs=*/false, + /*should_broadcast_rhs=*/false, + }; + + GenericTensorAccessorR lhs = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {7, 8, 9}, + {-1, -2, -3}, + }, + allocator); + + GenericTensorAccessorR rhs = create_2d_accessor_r_with_contents( + { + {10, 50, -10}, + {20, 60, -20}, + {30, 70, -30}, + {40, 80, -40}, + }, + allocator); + + GenericTensorAccessorW result = + create_zero_filled_accessor_w(lhs.shape, allocator); + + ElementBinaryPerDeviceState per_device_state = + element_binary_gpu_init_kernel(/*attrs=*/attrs, + /*lhs_shape=*/lhs.shape, + /*rhs_shape=*/rhs.shape, + /*output_shape=*/result.shape); + + element_binary_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*lhs=*/lhs, + /*rhs=*/rhs, + /*output=*/result); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {11, 52, -7}, + {24, 65, -14}, + {37, 78, -21}, + {39, 78, -43}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(result, correct), + check_kv("result", format_accessor_w_contents(result))); + } + } + + TEST_CASE("element_binary_gpu_backward_kernel") { + SUBCASE("add") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + ElementBinaryAttrs attrs{ + /*type=*/OperatorType::EW_ADD, + /*compute_type=*/DataType::FLOAT, + /*should_broadcast_lhs=*/false, + /*should_broadcast_rhs=*/false, + }; + + GenericTensorAccessorR lhs = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {7, 8, 9}, + {-1, -2, -3}, + }, + allocator); + + GenericTensorAccessorR rhs = create_2d_accessor_r_with_contents( + { + {10, 50, -10}, + {20, 60, -20}, + {30, 70, -30}, + {40, 80, -40}, + }, + allocator); + + GenericTensorAccessorW lhs_grad = + allocator.allocate_tensor(get_tensor_shape_for_accessor_r(lhs)); + + GenericTensorAccessorW rhs_grad = + allocator.allocate_tensor(get_tensor_shape_for_accessor_r(rhs)); + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {11, 52, -7}, + {24, 65, -14}, + {73, 78, -21}, + {39, 78, -43}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, -1}, + {6, 4, 2}, + {0.5, 0.1, -2}, + {0, 0.5, 0}, + }, + allocator); + + ElementBinaryPerDeviceState per_device_state = + element_binary_gpu_init_kernel(/*attrs=*/attrs, + /*lhs_shape=*/lhs.shape, + /*rhs_shape=*/rhs.shape, + /*output_shape=*/output.shape); + + element_binary_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, + /*lhs=*/lhs, + /*lhs_grad=*/lhs_grad, + /*rhs=*/rhs, + /*rhs_grad=*/rhs_grad); + + GenericTensorAccessorR correct_lhs_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, -1}, + {6, 4, 2}, + {0.5, 0.1, -2}, + {0, 0.5, 0}, + }, + allocator); + + GenericTensorAccessorR correct_rhs_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, -1}, + {6, 4, 2}, + {0.5, 0.1, -2}, + {0, 0.5, 0}, + }, + allocator); + + float epsilon = 1e-20; + CHECK_MESSAGE( + accessors_within_epsilon(lhs_grad, correct_lhs_grad, epsilon), + check_kv("lhs_grad", format_accessor_w_contents(lhs_grad))); + + CHECK_MESSAGE( + accessors_within_epsilon(rhs_grad, correct_rhs_grad, epsilon), + check_kv("rhs_grad", format_accessor_w_contents(rhs_grad))); + } + } +} diff --git a/lib/task-spec/src/task-spec/ops/impl/element_binary.cc b/lib/task-spec/src/task-spec/ops/impl/element_binary.cc index c8460af538..ad38eb3b5d 100644 --- a/lib/task-spec/src/task-spec/ops/impl/element_binary.cc +++ b/lib/task-spec/src/task-spec/ops/impl/element_binary.cc @@ -5,27 +5,21 @@ namespace FlexFlow { -using namespace FlexFlow::Kernels::ElementBinary; - static DeviceSpecificPerDeviceOpState init_task_impl(TaskArgumentAccessor const &acc) { - auto input_lhs = acc.get_tensor(TensorSlotName::LHS_INPUT); - auto input_rhs = acc.get_tensor(TensorSlotName::RHS_INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - - device_handle_t handle = acc.get_ff_handle(); - DeviceType kernel_device_type = acc.get_kernel_device_type(); ElementBinaryAttrs attrs = acc.get_op_attrs().require_element_binary(); + DeviceType kernel_device_type = acc.get_kernel_device_type(); + + TensorShape input_lhs_shape = acc.get_tensor_shape(TensorSlotName::LHS_INPUT); + TensorShape input_rhs_shape = acc.get_tensor_shape(TensorSlotName::RHS_INPUT); + TensorShape output_shape = acc.get_tensor_shape(TensorSlotName::OUTPUT); std::optional per_device_state = - init_kernel(kernel_device_type, - handle, - attrs.type, - attrs.should_broadcast_lhs, - attrs.should_broadcast_rhs, - input_lhs.shape, - input_rhs.shape, - output.shape); + element_binary_init_kernel(kernel_device_type, + attrs, + input_lhs_shape, + input_rhs_shape, + output_shape); return DeviceSpecificPerDeviceOpState{ acc.make_device_specific(per_device_state), @@ -41,21 +35,23 @@ static std::optional ElementBinaryAttrs attrs = acc.get_op_attrs().require_element_binary(); device_handle_t handle = acc.get_ff_handle(); - auto input_lhs = acc.get_tensor(TensorSlotName::LHS_INPUT); - auto input_rhs = acc.get_tensor(TensorSlotName::RHS_INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR input_lhs = + acc.get_tensor(TensorSlotName::LHS_INPUT); + GenericTensorAccessorR input_rhs = + acc.get_tensor(TensorSlotName::RHS_INPUT); + GenericTensorAccessorW output = + acc.get_tensor(TensorSlotName::OUTPUT); - return profile(forward_kernel, + return profile(element_binary_forward_kernel, profiling, kernel_device_type, "[ElementBinary] forward_time = {:.2lf}ms\n", + handle, per_device_state, - input_lhs.get_float_ptr(), - input_rhs.get_float_ptr(), - output.get_float_ptr(), - attrs.type, - attrs.should_broadcast_lhs, - handle); + attrs, + input_lhs, + input_rhs, + output); } static std::optional @@ -67,30 +63,33 @@ static std::optional ElementBinaryAttrs attrs = acc.get_op_attrs().require_element_binary(); device_handle_t handle = acc.get_ff_handle(); - auto input_lhs = acc.get_tensor(TensorSlotName::LHS_INPUT); - auto input_rhs = acc.get_tensor(TensorSlotName::RHS_INPUT); + GenericTensorAccessorR input_lhs = + acc.get_tensor(TensorSlotName::LHS_INPUT); + GenericTensorAccessorR input_rhs = + acc.get_tensor(TensorSlotName::RHS_INPUT); + GenericTensorAccessorR output = + acc.get_tensor(TensorSlotName::OUTPUT); - auto output_grad = + GenericTensorAccessorR output_grad = acc.get_tensor_grad(TensorSlotName::OUTPUT); - auto input_lhs_grad = + GenericTensorAccessorW input_lhs_grad = acc.get_tensor_grad(TensorSlotName::LHS_INPUT); - auto input_rhs_grad = + GenericTensorAccessorW input_rhs_grad = acc.get_tensor_grad(TensorSlotName::RHS_INPUT); - return profile(backward_kernel, + return profile(element_binary_backward_kernel, profiling, kernel_device_type, "[ElementBinary] backward_time = {:.2lf}ms\n", + handle, per_device_state, - output_grad.get_float_ptr(), - input_lhs.get_float_ptr(), - input_rhs.get_float_ptr(), - input_lhs_grad.get_float_ptr(), - input_rhs_grad.get_float_ptr(), - attrs.type, - attrs.should_broadcast_lhs, - attrs.should_broadcast_rhs, - handle); + attrs, + output, + output_grad, + input_lhs, + input_lhs_grad, + input_rhs, + input_rhs_grad); } TaskImplFunction get_element_binary_init_task_impl() { diff --git a/lib/task-spec/src/task-spec/ops/impl/element_unary.cc b/lib/task-spec/src/task-spec/ops/impl/element_unary.cc index e51d268c53..d721f5c41d 100644 --- a/lib/task-spec/src/task-spec/ops/impl/element_unary.cc +++ b/lib/task-spec/src/task-spec/ops/impl/element_unary.cc @@ -17,7 +17,7 @@ static DeviceSpecificPerDeviceOpState std::optional per_device_state = element_unary_init_kernel( - kernel_device_type, input_shape, output_shape, attrs); + kernel_device_type, attrs, input_shape, output_shape); return DeviceSpecificPerDeviceOpState{ acc.make_device_specific(per_device_state), @@ -26,8 +26,10 @@ static DeviceSpecificPerDeviceOpState static std::optional forward_task_impl(TaskArgumentAccessor const &acc) { - 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); ElementUnaryAttrs attrs = acc.get_op_attrs().require_element_unary(); device_handle_t handle = acc.get_ff_handle(); @@ -41,19 +43,22 @@ static std::optional profiling, kernel_device_type, "[ElementUnary] forward_time = {:.2lf}ms\n", + handle, per_device_state, attrs, - handle, input, output); } static std::optional backward_task_impl(TaskArgumentAccessor const &acc) { - 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 = + 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); ElementUnaryAttrs attrs = acc.get_op_attrs().require_element_unary(); @@ -68,9 +73,9 @@ static std::optional profiling, kernel_device_type, "[ElementUnary] backward_time = {:.2lf}ms\n", + handle, per_device_state, attrs, - handle, output, output_grad, input, From 614d760fe90e0c6c4d18efb24df976b9db584987 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 27 Aug 2026 14:10:23 -0700 Subject: [PATCH 2/6] Sort GPU tests into correct test suite. --- lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc index 96b0d589b3..04772b2469 100644 --- a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc +++ b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc @@ -9,7 +9,7 @@ using namespace ::FlexFlow; -TEST_SUITE(FF_TEST_SUITE) { +TEST_SUITE(FF_CUDA_TEST_SUITE) { TEST_CASE("element_binary_gpu_forward_kernel") { SUBCASE("add") { ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( From c55909868f1d371007e119fbc23bbd18883093f6 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 28 Aug 2026 09:33:32 -0700 Subject: [PATCH 3/6] Fix dimensioning of cuDNN tensors. --- lib/kernels/src/cuda/cuda_helper.cu | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/lib/kernels/src/cuda/cuda_helper.cu b/lib/kernels/src/cuda/cuda_helper.cu index cd89945579..d720c9cf92 100644 --- a/lib/kernels/src/cuda/cuda_helper.cu +++ b/lib/kernels/src/cuda/cuda_helper.cu @@ -240,17 +240,22 @@ ffCudnnDataType_t ff_to_cudnn_datatype(DataType flexflow_data_type) { ffStatus_t cudnnSetTensorDescriptorFromTensorShape(cudnnTensorDescriptor_t tensor, TensorShape const &shape) { + ASSERT(get_num_dims(shape.dims) <= 4, + "cudnnSetTensorDescriptorFromTensorShape only supports tensors of up " + "to 4 dimensions", + shape); + return cudnnSetTensor4dDescriptor( tensor, CUDNN_TENSOR_NCHW, ff_to_cudnn_datatype(shape.data_type), - try_dim_at_idx(shape.dims, relative_ff_dim_t{3}) + try_dim_at_idx(shape.dims, relative_ff_dim_t{0}) .value_or(1_p) .int_from_positive_int(), - try_dim_at_idx(shape.dims, relative_ff_dim_t{3}) + try_dim_at_idx(shape.dims, relative_ff_dim_t{1}) .value_or(1_p) .int_from_positive_int(), - try_dim_at_idx(shape.dims, relative_ff_dim_t{3}) + try_dim_at_idx(shape.dims, relative_ff_dim_t{2}) .value_or(1_p) .int_from_positive_int(), try_dim_at_idx(shape.dims, relative_ff_dim_t{3}) From 393efd7f0274e49d6a8ac7a4a16c4bdeb620dcdb Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 28 Aug 2026 09:58:05 -0700 Subject: [PATCH 4/6] Intentionally randomize tensors we expect to override and zero-initialize the rest. --- .../src/kernels/element_binary_kernels_cpu.cc | 11 +- .../src/kernels/element_binary_kernels_gpu.cc | 11 +- .../src/kernels/element_unary_kernels_cpu.cc | 8 +- .../src/kernels/element_unary_kernels_gpu.cc | 122 ++++++++++++++++++ 4 files changed, 139 insertions(+), 13 deletions(-) create mode 100644 lib/kernels/test/src/kernels/element_unary_kernels_gpu.cc diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc index b52d847178..10aab827db 100644 --- a/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc +++ b/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc @@ -39,8 +39,9 @@ TEST_SUITE(FF_TEST_SUITE) { }, cpu_allocator); + // Intentionally randomize this tensor so we can be confident we never read it GenericTensorAccessorW result = - create_zero_filled_accessor_w(lhs.shape, cpu_allocator); + create_random_filled_accessor_w(lhs.shape, cpu_allocator); element_binary_cpu_forward_kernel( /*attrs=*/attrs, @@ -92,11 +93,11 @@ TEST_SUITE(FF_TEST_SUITE) { }, cpu_allocator); - GenericTensorAccessorW lhs_grad = - cpu_allocator.allocate_tensor(get_tensor_shape_for_accessor_r(lhs)); + GenericTensorAccessorW lhs_grad = create_zero_filled_accessor_w( + get_tensor_shape_for_accessor_r(lhs), cpu_allocator); - GenericTensorAccessorW rhs_grad = - cpu_allocator.allocate_tensor(get_tensor_shape_for_accessor_r(rhs)); + GenericTensorAccessorW rhs_grad = create_zero_filled_accessor_w( + get_tensor_shape_for_accessor_r(rhs), cpu_allocator); GenericTensorAccessorR output = create_2d_accessor_r_with_contents( { diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc index 04772b2469..fabd2da47c 100644 --- a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc +++ b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc @@ -44,8 +44,9 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { }, allocator); + // Intentionally randomize this tensor so we can be confident we never read it GenericTensorAccessorW result = - create_zero_filled_accessor_w(lhs.shape, allocator); + create_random_filled_accessor_w(lhs.shape, allocator); ElementBinaryPerDeviceState per_device_state = element_binary_gpu_init_kernel(/*attrs=*/attrs, @@ -111,11 +112,11 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { }, allocator); - GenericTensorAccessorW lhs_grad = - allocator.allocate_tensor(get_tensor_shape_for_accessor_r(lhs)); + GenericTensorAccessorW lhs_grad = create_zero_filled_accessor_w( + get_tensor_shape_for_accessor_r(lhs), allocator); - GenericTensorAccessorW rhs_grad = - allocator.allocate_tensor(get_tensor_shape_for_accessor_r(rhs)); + GenericTensorAccessorW rhs_grad = create_zero_filled_accessor_w( + get_tensor_shape_for_accessor_r(rhs), allocator); GenericTensorAccessorR output = create_2d_accessor_r_with_contents( { diff --git a/lib/kernels/test/src/kernels/element_unary_kernels_cpu.cc b/lib/kernels/test/src/kernels/element_unary_kernels_cpu.cc index 799e8ee74e..c797931a95 100644 --- a/lib/kernels/test/src/kernels/element_unary_kernels_cpu.cc +++ b/lib/kernels/test/src/kernels/element_unary_kernels_cpu.cc @@ -25,8 +25,9 @@ TEST_SUITE(FF_TEST_SUITE) { }, cpu_allocator); + // Intentionally randomize this tensor so we can be confident we never read it GenericTensorAccessorW result = - create_zero_filled_accessor_w(input.shape, cpu_allocator); + create_random_filled_accessor_w(input.shape, cpu_allocator); element_unary_cpu_forward_kernel( /*attrs=*/attrs, @@ -63,8 +64,9 @@ TEST_SUITE(FF_TEST_SUITE) { }, cpu_allocator); - GenericTensorAccessorW input_grad = - cpu_allocator.allocate_tensor(get_tensor_shape_for_accessor_r(input)); + // 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), cpu_allocator); GenericTensorAccessorR output = create_2d_accessor_r_with_contents( { diff --git a/lib/kernels/test/src/kernels/element_unary_kernels_gpu.cc b/lib/kernels/test/src/kernels/element_unary_kernels_gpu.cc new file mode 100644 index 0000000000..f7e21b7e1f --- /dev/null +++ b/lib/kernels/test/src/kernels/element_unary_kernels_gpu.cc @@ -0,0 +1,122 @@ +#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; + +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{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + ElementUnaryAttrs attrs = make_relu_attrs(); + + GenericTensorAccessorR input = create_2d_accessor_r_with_contents( + { + {3, -3, 6}, + {0, 1, 5}, + {1, 2, -2}, + {-8, 0.5, -3}, + }, + allocator); + + // 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( + /*attrs=*/attrs, + /*input=*/input, + /*output=*/result); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {3, 0, 6}, + {0, 1, 5}, + {1, 2, 0}, + {0, 0.5, 0}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(result, correct), + 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{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + ElementUnaryAttrs attrs = make_relu_attrs(); + + GenericTensorAccessorR input = create_2d_accessor_r_with_contents( + { + {3, -3, 6}, + {0, 1, 5}, + {1, 2, -2}, + {-8, 0.5, -3}, + }, + 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); + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {3, 0, 6}, + {0, 1, 5}, + {1, 2, -2}, + {0, 0.5, 0}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, -1}, + {6, 4, 2}, + {0.5, 0.1, -2}, + {0, 0.5, 0}, + }, + allocator); + + element_unary_cpu_backward_kernel( + /*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.0f, 0.0f, -1.0f}, + {0.0f, 4.0f, 2.0f}, + {0.5f, 0.1f, 0.0f}, + {0.0f, 0.5f, 0.0f}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct_input_grad), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + } +} From f07e7c73f757cdb5e90081745abee55b6322aa7b Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 28 Aug 2026 11:16:28 -0700 Subject: [PATCH 5/6] Cleanup. --- lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc | 1 - lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc | 1 - 2 files changed, 2 deletions(-) diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc index 10aab827db..e30073b80e 100644 --- a/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc +++ b/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc @@ -3,7 +3,6 @@ #include "kernels/create_accessor_with_contents.h" #include "kernels/format_accessor_contents.h" #include "kernels/local_cpu_allocator.h" -#include "op-attrs/ops/element_binary.h" #include "test/utils/doctest/check_kv.h" #include diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc index fabd2da47c..be0e056302 100644 --- a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc +++ b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc @@ -3,7 +3,6 @@ #include "kernels/create_accessor_with_contents.h" #include "kernels/format_accessor_contents.h" #include "kernels/local_cpu_allocator.h" -#include "op-attrs/ops/element_binary.h" #include "test/utils/doctest/check_kv.h" #include From ce3bf405073b21fc8e70abde8ea24c8047451e10 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 28 Aug 2026 12:01:41 -0700 Subject: [PATCH 6/6] Return to exact testing. --- .../test/src/kernels/element_binary_kernels_cpu.cc | 13 +++++-------- .../test/src/kernels/element_binary_kernels_gpu.cc | 13 +++++-------- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc index e30073b80e..c30cdd1045 100644 --- a/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc +++ b/lib/kernels/test/src/kernels/element_binary_kernels_cpu.cc @@ -146,14 +146,11 @@ TEST_SUITE(FF_TEST_SUITE) { }, cpu_allocator); - float epsilon = 1e-20; - CHECK_MESSAGE( - accessors_within_epsilon(lhs_grad, correct_lhs_grad, epsilon), - check_kv("lhs_grad", format_accessor_w_contents(lhs_grad))); - - CHECK_MESSAGE( - accessors_within_epsilon(rhs_grad, correct_rhs_grad, epsilon), - check_kv("rhs_grad", format_accessor_w_contents(rhs_grad))); + CHECK_MESSAGE(accessors_are_equal(lhs_grad, correct_lhs_grad), + check_kv("lhs_grad", format_accessor_w_contents(lhs_grad))); + + CHECK_MESSAGE(accessors_are_equal(rhs_grad, correct_rhs_grad), + check_kv("rhs_grad", format_accessor_w_contents(rhs_grad))); } } } diff --git a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc index be0e056302..bcc770ca87 100644 --- a/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc +++ b/lib/kernels/test/src/kernels/element_binary_kernels_gpu.cc @@ -174,14 +174,11 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) { }, allocator); - float epsilon = 1e-20; - CHECK_MESSAGE( - accessors_within_epsilon(lhs_grad, correct_lhs_grad, epsilon), - check_kv("lhs_grad", format_accessor_w_contents(lhs_grad))); - - CHECK_MESSAGE( - accessors_within_epsilon(rhs_grad, correct_rhs_grad, epsilon), - check_kv("rhs_grad", format_accessor_w_contents(rhs_grad))); + CHECK_MESSAGE(accessors_are_equal(lhs_grad, correct_lhs_grad), + check_kv("lhs_grad", format_accessor_w_contents(lhs_grad))); + + CHECK_MESSAGE(accessors_are_equal(rhs_grad, correct_rhs_grad), + check_kv("rhs_grad", format_accessor_w_contents(rhs_grad))); } } }