Skip to content

Update and sync element unary/binary kernels #1667

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
elliottslaughter wants to merge 6 commits into
base: master
Choose a base branch
from
Open

Update and sync element unary/binary kernels #1667

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
17d86fd
Update and sync element unary/binary kernels.
elliottslaughter Aug 27, 2026
614d760
Sort GPU tests into correct test suite.
elliottslaughter Aug 27, 2026
c559098
Fix dimensioning of cuDNN tensors.
elliottslaughter Aug 28, 2026
393efd7
Intentionally randomize tensors we expect to override and zero-initia…
elliottslaughter Aug 28, 2026
f07e7c7
Cleanup.
elliottslaughter Aug 28, 2026
ce3bf40
Return to exact testing.
elliottslaughter Aug 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 28 additions & 35 deletions lib/kernels/include/kernels/element_binary_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<ElementBinaryPerDeviceState>
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<ElementBinaryPerDeviceState> 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<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,
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<ElementBinaryPerDeviceState> const &per_device_state);

} // namespace FlexFlow::Kernels::ElementBinary
} // namespace FlexFlow

#endif
32 changes: 16 additions & 16 deletions lib/kernels/include/kernels/element_binary_kernels_cpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
74 changes: 37 additions & 37 deletions lib/kernels/include/kernels/element_binary_kernels_gpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
5 changes: 0 additions & 5 deletions lib/kernels/include/kernels/element_binary_per_device_state.dtg.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
15 changes: 8 additions & 7 deletions lib/kernels/include/kernels/element_unary_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<ElementUnaryPerDeviceState>
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<ElementUnaryPerDeviceState> const &device_state,
ElementUnaryAttrs const &attrs,
device_handle_t const &handle,
std::optional<ElementUnaryPerDeviceState> 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<ElementUnaryPerDeviceState> const &device_state,
ElementUnaryAttrs const &attrs,
device_handle_t const &handle,
std::optional<ElementUnaryPerDeviceState> const &per_device_state,
ElementUnaryAttrs const &attrs,
GenericTensorAccessorR const &output,
GenericTensorAccessorR const &output_grad,
GenericTensorAccessorR const &input,
Expand Down
1 change: 0 additions & 1 deletion lib/kernels/include/kernels/element_unary_kernels_cpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
11 changes: 5 additions & 6 deletions lib/kernels/include/kernels/element_unary_kernels_gpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion lib/kernels/include/kernels/map_tensor_accessors.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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});
Expand Down
11 changes: 8 additions & 3 deletions lib/kernels/src/cuda/cuda_helper.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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})
Expand Down
Loading
Loading