Skip to content

Update kernels for conv_2d, batch_norm, pool_2d, concat, reshape, split, batch_matmul, transpose and upsample #1676

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

Draft
elliottslaughter wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

Update kernels for conv_2d, batch_norm, pool_2d, concat, reshape, split, batch_matmul, transpose and upsample #1676

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cf328b4
Update kernels for conv_2d, batch_norm, pool_2d, concat, reshape, spl…
elliottslaughter Sep 3, 2026
e7d8718
Update batch_matmul kernels to support > 3D.
elliottslaughter Sep 3, 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
3 changes: 3 additions & 0 deletions lib/kernels/include/kernels/batch_matmul_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions lib/kernels/include/kernels/batch_matmul_kernels_gpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
71 changes: 38 additions & 33 deletions lib/kernels/include/kernels/batch_norm_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<BatchNormPerDeviceState>
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<BatchNormPerDeviceState> 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<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_cleanup_kernel(
DeviceType device_type,
Allocator &allocator,
std::optional<BatchNormPerDeviceState> const &per_device_state);
std::optional<BatchNormPerDeviceState> &per_device_state);

} // namespace FlexFlow

} // namespace FlexFlow::Kernels::BatchNorm
#endif
36 changes: 17 additions & 19 deletions lib/kernels/include/kernels/batch_norm_kernels_cpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
73 changes: 38 additions & 35 deletions lib/kernels/include/kernels/batch_norm_kernels_gpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
29 changes: 0 additions & 29 deletions lib/kernels/include/kernels/batch_norm_per_device_state.dtg.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -24,10 +19,6 @@ type = "ffTensorDescriptor_t"
name = "biasTensor"
type = "ffTensorDescriptor_t"

[[fields]]
name = "actiDesc"
type = "ffActivationDescriptor_t"

[[fields]]
name = "mode"
type = "ffBatchNormMode_t"
Expand All @@ -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"
24 changes: 14 additions & 10 deletions lib/kernels/include/kernels/concat_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<GenericTensorAccessorR> const &inputs,
ff_dim_t axis);
void concat_forward_kernel(device_stream_t const &stream,
ConcatAttrs const &attrs,
std::vector<GenericTensorAccessorR> const &inputs,
GenericTensorAccessorW const &output);

void backward_kernel(device_stream_t const &stream,
GenericTensorAccessorR const &output_grad,
std::vector<GenericTensorAccessorW> 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<GenericTensorAccessorR> const &inputs,
std::vector<GenericTensorAccessorW> const &input_grads);

} // namespace FlexFlow::Kernels::Concat
} // namespace FlexFlow

#endif
22 changes: 13 additions & 9 deletions lib/kernels/include/kernels/concat_kernels_cpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<GenericTensorAccessorR> const &inputs,
ff_dim_t axis);
void concat_cpu_forward_kernel(
ConcatAttrs const &attrs,
std::vector<GenericTensorAccessorR> const &inputs,
GenericTensorAccessorW const &output);

void cpu_backward_kernel(GenericTensorAccessorR const &output_grad,
std::vector<GenericTensorAccessorW> const &input_grads,
ff_dim_t axis);
void concat_cpu_backward_kernel(
ConcatAttrs const &attrs,
GenericTensorAccessorR const &output,
GenericTensorAccessorR const &output_grad,
std::vector<GenericTensorAccessorR> const &inputs,
std::vector<GenericTensorAccessorW> const &input_grads);

} // namespace FlexFlow::Kernels::Concat
} // namespace FlexFlow

#endif
Loading
Loading