More unit test fixes and cpu kernel implementations

Author: lockshaw

lib/kernels/include/kernels/accessor.h

@@ -242,7 +242,7 @@ std::pair<ArrayShape, DataType>
 TensorShape get_tensor_shape_for_accessor_r(GenericTensorAccessorR const &);
 TensorShape get_tensor_shape_for_accessor_w(GenericTensorAccessorW const &);
 
-void copy_accessor_data_to_l_from_r(GenericTensorAccessorW &dst_accessor,
+void copy_accessor_data_to_l_from_r(GenericTensorAccessorW const &dst_accessor,
                                     GenericTensorAccessorR const &src_accessor);
 
 template <DataType DT>

lib/kernels/include/kernels/batch_matmul_kernels.h

@@ -3,11 +3,12 @@
 
 #include "kernels/device_stream_t.dtg.h"
 #include "kernels/ff_handle.h"
+#include "kernels/device_handle_t.dtg.h"
 
 namespace FlexFlow::Kernels::BatchMatmul {
 
 void forward_kernel(device_stream_t const &stream,
-                    PerDeviceFFHandle const &handle,
+                    device_handle_t const &handle,
                     float *output_ptr,
                     float const *a_input_ptr,
                     float const *b_input_ptr,
@@ -20,7 +21,7 @@ void forward_kernel(device_stream_t const &stream,
                     int b_seq_length_dim);
 
 void backward_kernel(device_stream_t const &stream,
-                     PerDeviceFFHandle const &handle,
+                     device_handle_t const &handle,
                      float const *o_ptr,
                      float const *o_grad_ptr,
                      float const *a_ptr,

lib/kernels/include/kernels/batch_matmul_kernels_cpu.h

@@ -2,12 +2,10 @@
 #define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_BATCH_MATMUL_KERNELS_CPU_H
 
 #include "kernels/allocation.h"
-#include "kernels/ff_handle.h"
 
 namespace FlexFlow::Kernels::BatchMatmul {
 
-void cpu_forward_kernel(PerDeviceFFHandle const &handle,
-                        float *output_ptr,
+void cpu_forward_kernel(float *output_ptr,
                         float const *a_input_ptr,
                         float const *b_input_ptr,
                         int m,
@@ -18,8 +16,7 @@ void cpu_forward_kernel(PerDeviceFFHandle const &handle,
                         int a_seq_length_dim,
                         int b_seq_length_dim);
 
-void cpu_backward_kernel(PerDeviceFFHandle const &handle,
-                         float const *o_ptr,
+void cpu_backward_kernel(float const *o_ptr,
                          float const *o_grad_ptr,
                          float const *a_ptr,
                          float *a_grad_ptr,

lib/kernels/include/kernels/batch_norm_kernels.h

@@ -5,12 +5,13 @@
 #include "kernels/batch_norm_per_device_state.dtg.h"
 #include "kernels/device_stream_t.dtg.h"
 #include "kernels/ff_handle.h"
+#include "kernels/device_handle_t.dtg.h"
 
 namespace FlexFlow::Kernels::BatchNorm {
 
 std::optional<BatchNormPerDeviceState>
     init_kernel(DeviceType device_type,
-                PerDeviceFFHandle const &handle,
+                device_handle_t const &handle,
                 Allocator &allocator,
                 float *runningMean,
                 int output_n,

lib/kernels/include/kernels/conv_2d_kernels.h

@@ -6,12 +6,13 @@
 #include "kernels/device_stream_t.dtg.h"
 #include "kernels/ff_handle.h"
 #include "op-attrs/activation.dtg.h"
+#include "kernels/device_handle_t.dtg.h"
 
 namespace FlexFlow::Kernels::Conv2D {
 
 std::optional<Conv2DPerDeviceState>
     init_kernel(DeviceType device_type,
-                PerDeviceFFHandle handle,
+                device_handle_t const &handle,
                 std::optional<Activation> activation,
                 int kernel_h,
                 int kernel_w,

lib/kernels/include/kernels/dropout_kernels.h

@@ -5,14 +5,14 @@
 #include "kernels/array_shape.h"
 #include "kernels/device_stream_t.dtg.h"
 #include "kernels/dropout_per_device_state.dtg.h"
-#include "kernels/ff_handle.h"
 #include <cstddef>
+#include "kernels/device_handle_t.dtg.h"
 
 namespace FlexFlow::Kernels::Dropout {
 
 std::optional<DropoutPerDeviceState>
     init_kernel(DeviceType device_type,
-                PerDeviceFFHandle const &handle,
+                device_handle_t const &handle,
                 float rate,
                 unsigned long long seed,
                 ArrayShape const &output_domain,

lib/kernels/include/kernels/element_binary_kernels.h

@@ -9,12 +9,13 @@
 #include "op-attrs/datatype.h"
 #include "op-attrs/operator_type.h"
 #include "pcg/device_type.dtg.h"
+#include "kernels/device_handle_t.dtg.h"
 
 namespace FlexFlow::Kernels::ElementBinary {
 
 std::optional<ElementBinaryPerDeviceState>
     init_kernel(DeviceType device_type,
-                PerDeviceFFHandle handle,
+                device_handle_t const &handle,
                 OperatorType op_type,
                 bool should_broadcast_lhs,
                 bool should_broadcast_rhs,
@@ -30,7 +31,7 @@ void forward_kernel(
     float *out_ptr,
     OperatorType op_type,
     bool broadcast_inputLHS,
-    PerDeviceFFHandle handle);
+    device_handle_t const &handle);
 
 void backward_kernel(
     device_stream_t const &stream,
@@ -43,7 +44,7 @@ void backward_kernel(
     OperatorType op_type,
     bool broadcast_inputLHS,
     bool broadcast_inputRHS,
-    PerDeviceFFHandle handle);
+    device_handle_t const &handle);
 
 void cleanup_kernel(
     DeviceType device_type,

lib/kernels/include/kernels/element_unary_kernels.h

@@ -6,6 +6,7 @@
 #include "kernels/element_unary_per_device_state.dtg.h"
 #include "kernels/ff_handle.h"
 #include "op-attrs/ops/element_unary_attrs.dtg.h"
+#include "kernels/device_handle_t.dtg.h"
 
 namespace FlexFlow::Kernels::ElementUnary {
 
@@ -19,15 +20,15 @@ void forward_kernel(
     device_stream_t const &stream,
     std::optional<ElementUnaryPerDeviceState> const &device_state,
     ElementUnaryAttrs const &attrs,
-    PerDeviceFFHandle const &handle,
+    device_handle_t const &handle,
     GenericTensorAccessorR const &input,
     GenericTensorAccessorW const &output);
 
 void backward_kernel(
     device_stream_t const &stream,
     std::optional<ElementUnaryPerDeviceState> const &device_state,
     ElementUnaryAttrs const &attrs,
-    PerDeviceFFHandle const &handle,
+    device_handle_t const &handle,
     GenericTensorAccessorR const &output,
     GenericTensorAccessorR const &output_grad,
     GenericTensorAccessorR const &input,

lib/kernels/include/kernels/element_unary_kernels_cpu.h

@@ -8,12 +8,10 @@
 namespace FlexFlow::Kernels::ElementUnary {
 
 void cpu_forward_kernel(ElementUnaryAttrs const &attrs,
-                        PerDeviceFFHandle const &handle,
                         GenericTensorAccessorR const &input,
                         GenericTensorAccessorW const &output);
 
 void cpu_backward_kernel(ElementUnaryAttrs const &attrs,
-                         PerDeviceFFHandle const &handle,
                          GenericTensorAccessorR const &output,
                          GenericTensorAccessorR const &output_grad,
                          GenericTensorAccessorR const &input,

lib/kernels/include/kernels/fill_tensor_accessor.h

@@ -7,7 +7,7 @@
 
 namespace FlexFlow {
 
-void fill_tensor_accessor(GenericTensorAccessorW &, DataTypeValue val);
+void fill_tensor_accessor(GenericTensorAccessorW const &, DataTypeValue val);
 
 GenericTensorAccessorW create_accessor_w_filled_with(
     TensorShape const &shape, DataTypeValue val, Allocator const &allocator);