CPU Kernel Tests

lib/kernels/include/kernels/cast_kernels.h

@@ -3,8 +3,6 @@
 
 #include "device.h"
 #include "kernels/accessor.h"
-#include "kernels/ff_handle.h"
-#include "op-attrs/activation.dtg.h"
 
 namespace FlexFlow {
 namespace Kernels {

lib/kernels/include/kernels/cast_kernels_cpu.h

@@ -0,0 +1,27 @@
+#ifndef _FLEXFLOW_OPS_KERNELS_CAST_KERNELS_CPU_H
+#define _FLEXFLOW_OPS_KERNELS_CAST_KERNELS_CPU_H
+
+#include "device.h"
+#include "kernels/accessor.h"
+
+namespace FlexFlow {
+namespace Kernels {
+namespace Cast {
+namespace CPU {
+
+void forward_kernel(GenericTensorAccessorR const &input,
+                    GenericTensorAccessorW const &output,
+                    DataType input_type,
+                    DataType output_type);
+
+void backward_kernel(GenericTensorAccessorR const &input,
+                     GenericTensorAccessorW const &output,
+                     DataType input_type,
+                     DataType output_type);
+
+} // namespace CPU
+} // namespace Cast
+} // namespace Kernels
+} // namespace FlexFlow
+
+#endif

lib/kernels/include/kernels/combine_kernels_cpu.h

@@ -0,0 +1,23 @@
+#ifndef _FLEXFLOW_OPS_KERNELS_COMBINE_KERNELS_CPU_H
+#define _FLEXFLOW_OPS_KERNELS_COMBINE_KERNELS_CPU_H
+
+#include "device.h"
+#include "kernels/accessor.h"
+
+namespace FlexFlow {
+namespace Kernels {
+namespace Combine {
+namespace CPU {
+
+void forward_kernel(GenericTensorAccessorR const &input,
+                    GenericTensorAccessorW const &output);
+
+void backward_kernel(GenericTensorAccessorR const &output_grad,
+                     GenericTensorAccessorW const &input_grad);
+
+} // namespace CPU
+} // namespace Combine
+} // namespace Kernels
+} // namespace FlexFlow
+
+#endif // _FLEXFLOW_OPS_KERNELS_COMBINE_KERNELS_CPU_H

lib/kernels/include/kernels/local_cpu_allocator.h

@@ -0,0 +1,22 @@
+#include "kernels/allocation.h"
+#include <unordered_set>
+
+namespace FlexFlow {
+
+struct LocalCPUAllocator : public IAllocator {
+  LocalCPUAllocator() = default;
+  LocalCPUAllocator(LocalCPUAllocator const &) = delete;
+  LocalCPUAllocator(LocalCPUAllocator &&) = delete;
+  ~LocalCPUAllocator() override;
+
+  void *allocate(size_t) override;
+  void deallocate(void *) override;
+
+private:
+  std::unordered_set<void *> ptrs;
+};
+CHECK_RC_COPY_VIRTUAL_COMPLIANT(LocalCPUAllocator);
+
+Allocator create_local_cpu_memory_allocator();
+
+} // namespace FlexFlow

lib/kernels/include/kernels/replicate_kernels_cpu.h

@@ -0,0 +1,24 @@
+#ifndef _FLEXFLOW_OPS_KERNELS_REPLICATE_KERNELS_CPU_H
+#define _FLEXFLOW_OPS_KERNELS_REPLICATE_KERNELS_CPU_H
+
+#include "device.h"
+#include "kernels/accessor.h"
+
+namespace FlexFlow {
+namespace Kernels {
+namespace Replicate {
+namespace CPU {
+
+void forward_kernel(GenericTensorAccessorR const &input,
+                    GenericTensorAccessorW const &output);
+
+void backward_kernel(GenericTensorAccessorW const &input,
+                     GenericTensorAccessorR const &output,
+                     size_t num_replicas);
+
+} // namespace CPU
+} // namespace Replicate
+} // namespace Kernels
+} // namespace FlexFlow
+
+#endif // _FLEXFLOW_OPS_KERNELS_REPLICATE_KERNELS_CPU_H

lib/kernels/include/kernels/reverse_kernels_cpu.h

@@ -0,0 +1,29 @@
+#ifndef _FLEXFLOW_OPS_KERNELS_REVERSE_KERNELS_CPU_H
+#define _FLEXFLOW_OPS_KERNELS_REVERSE_KERNELS_CPU_H
+
+#include "device.h"
+
+namespace FlexFlow {
+namespace Kernels {
+namespace Reverse {
+namespace CPU {
+
+void forward_kernel(float const *in_ptr,
+                    float *out_ptr,
+                    coord_t num_out_blks,
+                    coord_t reverse_dim_size,
+                    coord_t in_blk_size,
+                    coord_t output_size);
+
+void backward_kernel(float const *out_grad_ptr,
+                     float *in_grad_ptr,
+                     coord_t num_out_blks,
+                     coord_t reverse_dim_size,
+                     coord_t in_blk_size,
+                     coord_t input_size);
+} // namespace CPU
+} // namespace Reverse
+} // namespace Kernels
+} // namespace FlexFlow
+
+#endif // _FLEXFLOW_OPS_KERNELS_REVERSE_KERNELS_CPU_H

lib/kernels/src/cpu/cast_kernels.cc

@@ -0,0 +1,59 @@
+#include "kernels/cast_kernels_cpu.h"
+#include "kernels/datatype_dispatch.h"
+
+namespace FlexFlow {
+namespace Kernels {
+namespace Cast {
+namespace CPU {
+
+template <typename IDT, typename ODT>
+void cast_forward(IDT const *input, ODT *output, size_t volume) {
+  for (size_t i = 0; i < volume; ++i) {
+    output[i] = static_cast<ODT>(input[i]);
+  }
+}
+
+template <typename IDT, typename ODT>
+void cast_backward(IDT const *input, ODT *output, size_t volume, ODT beta) {
+  for (size_t i = 0; i < volume; i++) {
+    output[i] = static_cast<ODT>(input[i]) + beta * output[i];
+  }
+}
+
+template <DataType IDT, DataType ODT>
+struct ForwardKernel {
+  void operator()(GenericTensorAccessorR const &input,
+                  GenericTensorAccessorW const &output) {
+    size_t volume = input.shape.get_volume();
+    cast_forward(input.get<IDT>(), output.get<ODT>(), volume);
+  }
+};
+
+template <DataType IDT, DataType ODT>
+struct BackwardKernel {
+  void operator()(GenericTensorAccessorR const &input,
+                  GenericTensorAccessorW const &output) {
+    size_t volume = input.shape.get_volume();
+    cast_backward(
+        input.get<IDT>(), output.get<ODT>(), volume, cast_to<ODT>(1.0f));
+  }
+};
+
+void forward_kernel(GenericTensorAccessorR const &input,
+                    GenericTensorAccessorW const &output,
+                    DataType input_type,
+                    DataType output_type) {
+  DataTypeDispatch2<ForwardKernel>{}(input_type, output_type, input, output);
+}
+
+void backward_kernel(GenericTensorAccessorR const &input,
+                     GenericTensorAccessorW const &output,
+                     DataType input_type,
+                     DataType output_type) {
+  DataTypeDispatch2<BackwardKernel>{}(input_type, output_type, input, output);
+}
+
+} // namespace CPU
+} // namespace Cast
+} // namespace Kernels
+} // namespace FlexFlow

lib/kernels/src/cpu/combine_kernels.cc

@@ -0,0 +1,44 @@
+#include "kernels/combine_kernels_cpu.h"
+#include "kernels/datatype_dispatch.h"
+
+namespace FlexFlow {
+namespace Kernels {
+namespace Combine {
+namespace CPU {
+
+template <DataType DT>
+struct ForwardKernel {
+  void operator()(GenericTensorAccessorR const &input,
+                  GenericTensorAccessorW const &output) {
+    memcpy(output.get<DT>(),
+           input.get<DT>(),
+           input.shape.get_volume() * size_of_datatype(DT));
+  }
+};
+
+template <DataType DT>
+struct BackwardKernel {
+  void operator()(GenericTensorAccessorR const &output_grad,
+                  GenericTensorAccessorW const &input_grad) {
+    size_t num_elements = output_grad.shape.get_volume();
+    for (int i = 0; i < num_elements; ++i) {
+      input_grad.get<DT>()[i] += output_grad.get<DT>()[i];
+    }
+  }
+};
+
+void forward_kernel(GenericTensorAccessorR const &input,
+                    GenericTensorAccessorW const &output) {
+  DataTypeDispatch1<ForwardKernel>{}(input.data_type, input, output);
+}
+
+void backward_kernel(GenericTensorAccessorR const &output_grad,
+                     GenericTensorAccessorW const &input_grad) {
+  DataTypeDispatch1<BackwardKernel>{}(
+      input_grad.data_type, output_grad, input_grad);
+}
+
+} // namespace CPU
+} // namespace Combine
+} // namespace Kernels
+} // namespace FlexFlow

lib/kernels/src/cpu/replicate_kernels.cc

@@ -0,0 +1,61 @@
+#include "kernels/datatype_dispatch.h"
+#include "kernels/replicate_kernels_cpu.h"
+
+namespace FlexFlow {
+namespace Kernels {
+namespace Replicate {
+namespace CPU {
+
+template <typename T>
+void replicate_backward_kernel(T *input,
+                               T const *output,
+                               size_t num_elements,
+                               size_t num_replicas) {
+  for (size_t i = 0; i < num_elements; ++i) {
+    T sum = 0;
+    for (size_t j = 0; j < num_replicas; ++j) {
+      sum += output[i + j * num_elements];
+    }
+    input[i] = sum;
+  }
+}
+
+// Why does replicate forward seem to only transfer memory? Shouldn't it also
+// handle the replication?
+template <DataType T>
+struct ForwardKernel {
+  void operator()(GenericTensorAccessorR const &input,
+                  GenericTensorAccessorW const &output) {
+    memcpy(output.get<T>(),
+           input.get<T>(),
+           input.shape.num_elements() * size_of_datatype(T));
+  }
+};
+
+template <DataType T>
+struct BackwardKernel {
+  void operator()(GenericTensorAccessorW const &input,
+                  GenericTensorAccessorR const &output,
+                  size_t num_replicas) {
+    size_t total_elements = input.shape.num_elements() * num_replicas;
+    replicate_backward_kernel(
+        input.get<T>(), output.get<T>(), total_elements, num_replicas);
+  }
+};
+
+void forward_kernel(GenericTensorAccessorR const &input,
+                    GenericTensorAccessorW const &output) {
+  DataTypeDispatch1<ForwardKernel>{}(input.data_type, input, output);
+}
+
+void backward_kernel(GenericTensorAccessorW const &input,
+                     GenericTensorAccessorR const &output,
+                     size_t num_replicas) {
+  DataTypeDispatch1<BackwardKernel>{}(
+      input.data_type, input, output, num_replicas);
+}
+
+} // namespace CPU
+} // namespace Replicate
+} // namespace Kernels
+} // namespace FlexFlow

lib/kernels/src/cpu/reverse_kernels.cc

@@ -0,0 +1,48 @@
+#include "kernels/reverse_kernels_cpu.h"
+
+namespace FlexFlow {
+namespace Kernels {
+namespace Reverse {
+namespace CPU {
+
+void reverse_forward_kernel(float const *in_ptr,
+                            float *out_ptr,
+                            coord_t num_out_blks,
+                            coord_t reverse_dim_size,
+                            coord_t in_blk_size) {
+  coord_t total_elements = num_out_blks * reverse_dim_size * in_blk_size;
+  for (coord_t i = 0; i < total_elements; ++i) {
+    coord_t blk_idx = i / (reverse_dim_size * in_blk_size);
+    coord_t offset = i - blk_idx * (reverse_dim_size * in_blk_size);
+    coord_t reverse_dim_idx = offset / in_blk_size;
+    coord_t in_idx = blk_idx * (reverse_dim_size * in_blk_size) +
+                     (reverse_dim_size - 1 - reverse_dim_idx) * in_blk_size +
+                     (offset % in_blk_size);
+    out_ptr[i] = in_ptr[in_idx];
+  }
+}
+
+void forward_kernel(float const *in_ptr,
+                    float *out_ptr,
+                    coord_t num_out_blks,
+                    coord_t reverse_dim_size,
+                    coord_t in_blk_size,
+                    coord_t output_size) {
+  reverse_forward_kernel(
+      in_ptr, out_ptr, num_out_blks, reverse_dim_size, in_blk_size);
+}
+
+void backward_kernel(float const *out_grad_ptr,
+                     float *in_grad_ptr,
+                     coord_t num_out_blks,
+                     coord_t reverse_dim_size,
+                     coord_t in_blk_size,
+                     coord_t input_size) {
+  reverse_forward_kernel(
+      out_grad_ptr, in_grad_ptr, num_out_blks, reverse_dim_size, in_blk_size);
+}
+
+} // namespace CPU
+} // namespace Reverse
+} // namespace Kernels
+} // namespace FlexFlow