flexflow · KateUnger · Aug 22, 2023 · Aug 22, 2023 · Aug 23, 2023 · Aug 23, 2023
diff --git a/lib/kernels/include/kernels/batch_matmul_kernels.h b/lib/kernels/include/kernels/batch_matmul_kernels.h
@@ -6,38 +6,43 @@
 
 namespace FlexFlow {
 
-class BatchMatmulPerDeviceState : public PerDeviceOpState {
-public:
-  BatchMatmulPerDeviceState(FFHandler handler);
-  int a_seq_length_dim, b_seq_length_dim;
+struct BMMPerDeviceState {
+  PerDeviceFFHandle handle;
+  Allocator allocator;
+  int a_seq_length_dim;
+  req<int> b_seq_length_dim;
 };
 
+FF_VISITABLE_STRUCT_NO_EQ(
+    BMMPerDeviceState, handle, allocator, a_seq_length_dim, b_seq_length_dim);
+
 namespace Kernels {
 namespace BatchMatmul {
 
+BMMPerDeviceState init_kernel(PerDeviceFFHandle const &handle,
+                              Allocator const &allocator,
+                              int a_seq_length_dim,
+                              int b_seq_length_dim);
+
 void forward_kernel(ffStream_t stream,
-                    BatchMatmulPerDeviceState const *,
-                    float *o_ptr,
-                    float const *a_ptr,
-                    float const *b_ptr,
-                    float const *c_ptr,
+                    BMMPerDeviceState const &meta,
+                    float *output_ptr,
+                    float const *lhs_input_ptr,
+                    float const *rhs_input_ptr,
                     int m,
                     int n,
                     int k,
                     int batch,
-                    int a_seq_length_dim = -1,
-                    int b_seq_length_dim = -1,
                     int seq_length = -1);
 
 void backward_kernel(ffStream_t stream,
-                     BatchMatmulPerDeviceState const *,
+                     BMMPerDeviceState const &meta,
                      float const *o_ptr,
                      float const *o_grad_ptr,
                      float const *a_ptr,
                      float *a_grad_ptr,
                      float const *b_ptr,
                      float *b_grad_ptr,
-                     float *c_grad_ptr,
                      int m,
                      int n,
                      int k,

diff --git a/lib/kernels/include/kernels/batch_norm_kernels.h b/lib/kernels/include/kernels/batch_norm_kernels.h
@@ -8,30 +8,66 @@
 
 namespace FlexFlow {
 
-class BatchNormPerDeviceState : public PerDeviceOpState {
-public:
-  BatchNormPerDeviceState(FFHandler handle,
-                          std::unique_ptr<IAllocator> allocator,
-                          int output_n,
-                          int output_c,
-                          int output_h,
-                          int output_w,
-                          bool relu,
-                          bool profiling);
-  ~BatchNormPerDeviceState(void);
-
-  ffTensorDescriptor_t inputTensor, outputTensor, biasTensor;
+struct BatchNormPerDeviceState {
+  PerDeviceFFHandle handle;
+  Allocator allocator;
+  ffTensorDescriptor_t inputTensor;
+  ffTensorDescriptor_t outputTensor;
+  ffTensorDescriptor_t biasTensor;
   ffActivationDescriptor_t actiDesc;
   ffBatchNormMode_t mode;
-  float *runningMean, *runningVar, *saveMean, *saveVar;
-  bool relu;
-  bool profiling;
-  std::unique_ptr<IAllocator> allocator;
+  float *runningMean;
+  float *runningVar;
+  float *saveMean;
+  float *saveVar;
+  int output_n;
+  int output_c;
+  int output_h;
+  int output_w;
+  ProfilingSettings profiling;
+  req<bool> relu;
 };
 
+FF_VISITABLE_STRUCT_NO_EQ(BatchNormPerDeviceState,
+                          handle,
+                          allocator,
+                          inputTensor,
+                          outputTensor,
+                          biasTensor,
+                          actiDesc,
+                          mode,
+                          runningMean,
+                          runningVar,
+                          saveMean,
+                          saveVar,
+                          output_n,
+                          output_c,
+                          output_h,
+                          output_w,
+                          profiling,
+                          relu);
+
 namespace Kernels {
 namespace BatchNorm {
 
+BatchNormPerDeviceState init_kernel(PerDeviceFFHandle handle,
+                                    Allocator allocator,
+                                    ffTensorDescriptor_t inputTensor,
+                                    ffTensorDescriptor_t outputTensor,
+                                    ffTensorDescriptor_t biasTensor,
+                                    ffActivationDescriptor_t actiDesc,
+                                    ffBatchNormMode_t mode,
+                                    float *runningMean,
+                                    float *runningVar,
+                                    float *saveMean,
+                                    float *saveVar,
+                                    int output_n,
+                                    int output_c,
+                                    int output_h,
+                                    int output_w,
+                                    ProfilingSettings profiling,
+                                    bool relu);
+
 void forward_kernel(ffStream_t stream,
                     BatchNormPerDeviceState *m,
                     float const *input_ptr,

diff --git a/lib/kernels/include/kernels/cast_kernels.h b/lib/kernels/include/kernels/cast_kernels.h
@@ -3,19 +3,26 @@
 
 #include "kernels/accessor.h"
 #include "kernels/device.h"
-#include "op-attrs/ffconst.h"
 
 namespace FlexFlow {
 
-class CastPerDeviceState : public PerDeviceOpState {
-public:
-  CastPerDeviceState(FFHandler handle);
-  DataType input_data_type, output_data_type;
+struct CastPerDeviceState {
+  PerDeviceFFHandle handle;
+  DataType input_data_type;
+  req<DataType> output_data_type;
 };
 
+FF_VISITABLE_STRUCT_NO_EQ(CastPerDeviceState,
+                          handle,
+                          input_data_type,
+                          output_data_type);
+
 namespace Kernels {
 namespace Cast {
 
+CastPerDeviceState
+    init_kernel(PerDeviceFFHandle const &, DataType input, DataType output);
+
 void forward_kernel(ffStream_t stream,
                     CastPerDeviceState const *,
                     GenericTensorAccessorR const &input,

diff --git a/lib/kernels/include/kernels/combine_kernels.h b/lib/kernels/include/kernels/combine_kernels.h
@@ -6,15 +6,17 @@
 
 namespace FlexFlow {
 
-class CombinePerDeviceState : public PerDeviceOpState {
-public:
-  CombinePerDeviceState(FFHandler handle);
-  DataType data_type;
+struct CombinePerDeviceState {
+  req<DataType> data_type;
 };
 
+FF_VISITABLE_STRUCT_NO_EQ(CombinePerDeviceState, data_type);
+
 namespace Kernels {
 namespace Combine {
 
+CombinePerDeviceState init_kernel(DataType data_type);
+
 void forward_kernel(ffStream_t stream,
                     CombinePerDeviceState const *m,
                     GenericTensorAccessorR const &input,

diff --git a/lib/kernels/include/kernels/concat_kernels.h b/lib/kernels/include/kernels/concat_kernels.h
@@ -6,29 +6,29 @@
 
 namespace FlexFlow {
 
-class ConcatPerDeviceState : public PerDeviceOpState {
-public:
-  ConcatPerDeviceState(FFHandler handle) : PerDeviceOpState(handle){};
-  int legion_axis;
-  char op_name[MAX_OPNAME];
+struct ConcatPerDeviceState {
+  req<ff_dim_t> legion_axis;
 };
 
+FF_VISITABLE_STRUCT_NONSTANDARD_CONSTRUCTION(ConcatPerDeviceState, legion_axis);
+
 namespace Kernels {
 namespace Concat {
 
-void init_meta(ConcatPerDeviceState *meta, int legion_axis);
+ConcatPerDeviceState init_kernel(ff_dim_t legion_axis);
 
 void forward_kernel(ffStream_t stream,
                     ConcatPerDeviceState const *m,
                     GenericTensorAccessorW const &output,
-                    GenericTensorAccessorR const *inputs,
+                    std::vector<FlexFlow::GenericTensorAccessorR> const &inputs,
                     int num_inputs);
 
-void backward_kernel(ffStream_t stream,
-                     ConcatPerDeviceState const *m,
-                     GenericTensorAccessorR const &output_grad,
-                     GenericTensorAccessorW const *input_grads,
-                     int num_inputs);
+void backward_kernel(
+    ffStream_t stream,
+    ConcatPerDeviceState const *m,
+    GenericTensorAccessorR const &output_grad,
+    std::vector<FlexFlow::GenericTensorAccessorW> const &input_grads,
+    int num_inputs);
 
 } // namespace Concat
 } // namespace Kernels

diff --git a/lib/kernels/include/kernels/conv_2d_kernels.h b/lib/kernels/include/kernels/conv_2d_kernels.h
@@ -5,45 +5,50 @@
 
 namespace FlexFlow {
 
-class Conv2DPerDeviceState : public PerDeviceOpState {
-public:
-  Conv2DPerDeviceState(FFHandler handler);
-  ffTensorDescriptor_t inputTensor, biasTensor, outputTensor;
+struct Conv2DPerDeviceState {
+  PerDeviceFFHandle handle;
+  ffTensorDescriptor_t inputTensor;
+  ffTensorDescriptor_t biasTensor;
+  ffTensorDescriptor_t outputTensor;
   ffFilterDescriptor_t filterDesc;
   ffActivationDescriptor_t actiDesc;
   ffConvolutionDescriptor_t convDesc;
   ffConvolutionFwdAlgo_t fwdAlgo;
   ffConvolutionBwdFilterAlgo_t bwdFilterAlgo;
   ffConvolutionBwdDataAlgo_t bwdDataAlgo;
-  bool relu, use_bias;
-  char op_name[MAX_OPNAME];
+  req<optional<Activation>> activation;
+  req<bool> use_bias;
 };
 
+FF_VISITABLE_STRUCT_NO_EQ(Conv2DPerDeviceState,
+                          handle,
+                          inputTensor,
+                          biasTensor,
+                          outputTensor,
+                          filterDesc,
+                          actiDesc,
+                          convDesc,
+                          fwdAlgo,
+                          bwdFilterAlgo,
+                          bwdDataAlgo,
+                          activation,
+                          use_bias);
+
 namespace Kernels {
 namespace Conv2D {
 
-void init_kernel(Conv2DPerDeviceState *m,
-                 int input_w,
-                 int input_h,
-                 int input_c,
-                 int input_n,
-                 int output_w,
-                 int output_h,
-                 int output_c,
-                 int output_n,
-                 int kernel_h,
-                 int kernel_w,
-                 int groups,
-                 int stride_h,
-                 int stride_w,
-                 int pad_h,
-                 int pad_w,
-                 float const *input_ptr,
-                 float *output_ptr,
-                 float const *kernel_ptr,
-                 float *kernel_grad_ptr,
-                 float *forward_time = nullptr,
-                 float *backward_time = nullptr);
+Conv2DPerDeviceState init_kernel(PerDeviceFFHandle handle,
+                                 ffTensorDescriptor_t inputTensor,
+                                 ffTensorDescriptor_t biasTensor,
+                                 ffTensorDescriptor_t outputTensor,
+                                 ffFilterDescriptor_t filterDesc,
+                                 ffActivationDescriptor_t actiDesc,
+                                 ffConvolutionDescriptor_t convDesc,
+                                 ffConvolutionFwdAlgo_t fwdAlgo,
+                                 ffConvolutionBwdFilterAlgo_t bwdFilterAlgo,
+                                 ffConvolutionBwdDataAlgo_t bwdDataAlgo,
+                                 req<optional<Activation>> relu,
+                                 bool use_bias);
 
 void forward_kernel(ffStream_t stream,
                     Conv2DPerDeviceState const *m,
@@ -58,8 +63,8 @@ void backward_kernel(ffStream_t stream,
                      float *input_grad_ptr,
                      float const *output_ptr,
                      float *output_grad_ptr,
-                     float const *kernel_ptr,
-                     float *kernel_grad_ptr,
+                     float const *filter_ptr,
+                     float *filter_grad_ptr,
                      float *bias_grad_ptr);
 
 } // namespace Conv2D

diff --git a/lib/kernels/src/cuda/batch_matmul_kernels.cu b/lib/kernels/src/cuda/batch_matmul_kernels.cu
@@ -18,9 +18,6 @@
 
 namespace FlexFlow {
 
-BatchMatmulPerDeviceState::BatchMatmulPerDeviceState(FFHandler handler)
-    : PerDeviceOpState(handler) {}
-
 namespace Kernels {
 namespace BatchMatmul {
 
@@ -124,7 +121,7 @@ O = A * B
 */
 
 void forward_kernel(cudaStream_t stream,
-                    BatchMatmulPerDeviceState const *meta,
+                    BatchMatmulPerDeviceState const &meta,
                     float *o_ptr,
                     float const *a_ptr,
                     float const *b_ptr,

diff --git a/lib/kernels/src/hip/batch_matmul_kernels.cpp b/lib/kernels/src/hip/batch_matmul_kernels.cpp
@@ -19,9 +19,6 @@
 
 namespace FlexFlow {
 
-BatchMatmulPerDeviceState::BatchMatmulPerDeviceState(FFHandler handler)
-    : PerDeviceOpState(handler) {}
-
 namespace Kernels {
 namespace BatchMatmul {
 
@@ -32,7 +29,7 @@ O: (batch, n, m)
 O = A * B
 */
 void forward_kernel(hipStream_t stream,
-                    BatchMatmulPerDeviceState const *meta,
+                    BatchMatmulPerDeviceState const &meta,
                     float *o_ptr,
                     float const *a_ptr,
                     float const *b_ptr,
@@ -42,9 +39,9 @@ void forward_kernel(hipStream_t stream,
                     int k,
                     int batch,
                     hipStream_t stream,
-                    int a_seq_length_dim,
-                    int b_seq_length_dim,
                     int seq_length) {
+  int a_seq_length_dim = meta->a_seq_length_dim;
+  int b_seq_length_dim = meta->b_seq_length_dim;
   checkCUDA(hipblasSetStream(meta->handle.blas, stream));
   checkCUDNN(miopenSetStream(meta->handle.dnn, stream));
 

diff --git a/lib/kernels/src/hip/concat_kernels.cpp b/lib/kernels/src/hip/concat_kernels.cpp
@@ -26,10 +26,6 @@ using Legion::Rect;
 namespace Kernels {
 namespace Concat {
 
-void init_meta(ConcatPerDeviceState *m, int legion_axis) {
-  m->legion_axis = legion_axis;
-}
-
 template <int N>
 void calc_blk_size(coord_t &num_blocks,
                    coord_t &blk_size,

diff --git a/lib/op-attrs/include/op-attrs/get_op_type.h b/lib/op-attrs/include/op-attrs/get_op_type.h
@@ -12,6 +12,7 @@ OperatorType get_op_type(BatchMatmulAttrs const &);
 OperatorType get_op_type(BatchNormAttrs const &);
 OperatorType get_op_type(BroadcastAttrs const &);
 OperatorType get_op_type(CastAttrs const &);
+OperatorType get_op_type(CombineAttrs const &);
 OperatorType get_op_type(ConcatAttrs const &);
 OperatorType get_op_type(Conv2DAttrs const &);
 OperatorType get_op_type(DropoutAttrs const &);