flexflow · lockshaw · Jul 10, 2024 · May 6, 2024 · May 17, 2024 · May 30, 2024
diff --git a/.proj.toml b/.proj.toml
@@ -11,8 +11,9 @@ build_targets = [
   # "substitutions",
   # "compiler",
   "substitution-generator",
-  "local-execution",
+  "local-execution", 
 ]
+
 test_targets = [
   "utils-tests",
   "op-attrs-tests",

diff --git a/flake.lock b/flake.lock
diff --git a/flake.nix b/flake.nix
@@ -152,4 +152,4 @@
       };
     }
   );
-}
+}
diff --git a/lib/kernels/CMakeLists.txt b/lib/kernels/CMakeLists.txt
@@ -7,6 +7,7 @@ file(GLOB_RECURSE SRC
      CONFIGURE_DEPENDS
      LIST_DIRECTORIES False
      src/*.cc
+     src/cuda/cuda_helper.cu
      src/cuda/ops/*.cu
      )
 
@@ -28,6 +29,7 @@ target_link_libraries(
   cuda
   cudnn
   nccl
+  utils
 )
 
 define_ff_vars(${project_target})
@@ -37,3 +39,5 @@ set_target_properties(
   PROPERTIES 
   CUDA_STANDARD 17
 )
+
+add_subdirectory(test)
diff --git a/lib/kernels/include/kernels/accessor.h b/lib/kernels/include/kernels/accessor.h
@@ -142,6 +142,9 @@ std::vector<real_type<DT> const *>
   return out;
 }
 
+GenericTensorAccessorR read_only_accessor_from_write_accessor(
+    GenericTensorAccessorW const &write_accessor);
+
 } // namespace FlexFlow
 
 namespace FlexFlow {

diff --git a/lib/kernels/include/kernels/array_shape.h b/lib/kernels/include/kernels/array_shape.h
@@ -41,8 +41,6 @@ struct ArrayShape {
 
   std::optional<std::size_t> at_maybe(std::size_t) const;
 
-  ArrayShape reversed_dim_order() const;
-
   ArrayShape
       sub_shape(std::optional<std::variant<ff_dim_t, legion_dim_t>> start,
                 std::optional<std::variant<ff_dim_t, legion_dim_t>> end) const;

diff --git a/lib/kernels/include/kernels/attention_kernels.h b/lib/kernels/include/kernels/attention_kernels.h
@@ -3,6 +3,7 @@
 
 #include "device.h"
 #include "kernels/allocation.h"
+#include "kernels/device.h"
 #include "kernels/ff_handle.h"
 #include "op-attrs/ops/attention.h"
 #include <memory>

diff --git a/lib/kernels/include/kernels/conv_2d_kernels.h b/lib/kernels/include/kernels/conv_2d_kernels.h
@@ -46,7 +46,7 @@ Conv2DPerDeviceState init_kernel(PerDeviceFFHandle handle,
                                  int padding_w,
                                  int stride_h,
                                  int stride_w,
-                                 GenericTensorAccessorR const &input,
+                                 GenericTensorAccessorW const &input,
                                  GenericTensorAccessorW const &output,
                                  float const *filter_ptr,
                                  float *filter_grad_ptr);

diff --git a/lib/kernels/include/kernels/device.h b/lib/kernels/include/kernels/device.h
@@ -26,9 +26,12 @@
 #include <iostream>
 #include <sstream>
 
+namespace FlexFlow {
+cudaError_t get_legion_stream(cudaStream_t *stream);
+} // namespace FlexFlow
+
 #if defined(FF_USE_CUDA) || defined(FF_USE_HIP_CUDA)
 typedef cudaStream_t ffStream_t;
-cudaError_t get_legion_stream(cudaStream_t *stream);
 typedef cudnnTensorDescriptor_t ffTensorDescriptor_t;
 typedef cudnnActivationDescriptor_t ffActivationDescriptor_t;
 typedef cudnnPoolingDescriptor_t ffPoolingDescriptor_t;
@@ -96,7 +99,8 @@ using coord_t = long long;
   do {                                                                         \
     std::stringstream _error;                                                  \
     if (status != 0) {                                                         \
-      _error << "Cuda failure: " << status;                                    \
+      _error << "CUDA failure: " << cudaGetErrorString(status) << " ("         \
+             << status << ")";                                                 \
       FatalError(_error.str());                                                \
     }                                                                          \
   } while (0)

diff --git a/lib/kernels/include/kernels/element_unary_kernels.h b/lib/kernels/include/kernels/element_unary_kernels.h
@@ -29,14 +29,14 @@ ElementUnaryPerDeviceState init_kernel(ArrayShape const &input_shape,
 void forward_kernel(ffStream_t stream,
                     ElementUnaryPerDeviceState const &device_state,
                     ElementUnaryAttrs const &attrs,
-                    PerDeviceFFHandle &handle,
+                    PerDeviceFFHandle const &handle,
                     GenericTensorAccessorR const &input,
                     GenericTensorAccessorW const &output);
 
 void backward_kernel(ffStream_t stream,
                      ElementUnaryPerDeviceState const &device_state,
                      ElementUnaryAttrs const &attrs,
-                     PerDeviceFFHandle &handle,
+                     PerDeviceFFHandle const &handle,
                      GenericTensorAccessorR const &input,
                      GenericTensorAccessorW const &input_grad,
                      GenericTensorAccessorR const &output,

diff --git a/lib/kernels/include/kernels/layer_norm_kernels.h b/lib/kernels/include/kernels/layer_norm_kernels.h
@@ -34,8 +34,8 @@ namespace Kernels {
 namespace LayerNorm {
 
 // todo: this may have some problem.
-LayerNormPerDeviceState init_kernel(PerDeviceFFHandle const &,
-                                    Allocator const &,
+LayerNormPerDeviceState init_kernel(PerDeviceFFHandle const &handle,
+                                    Allocator &allocator,
                                     bool elementwise_affine,
                                     int64_t effective_batch_size,
                                     int64_t effective_num_elements,

diff --git a/lib/kernels/include/kernels/legion_dim.h b/lib/kernels/include/kernels/legion_dim.h
@@ -6,7 +6,7 @@
 
 namespace FlexFlow {
 
-legion_dim_t add_to_legion_dim(legion_dim_t, int);
+legion_dim_t add_to_legion_dim(legion_dim_t legion_dim, int value);
 
 legion_dim_t legion_dim_from_ff_dim(ff_dim_t, int num_dimensions);
 

diff --git a/lib/kernels/include/kernels/linear_kernels.h b/lib/kernels/include/kernels/linear_kernels.h
@@ -38,6 +38,7 @@ namespace Linear {
 
 LinearPerDeviceState init_kernel(PerDeviceFFHandle handle,
                                  float *one_ptr,
+                                 std::optional<Activation> activation,
                                  std::optional<RegularizerAttrs> regularizer,
                                  bool use_bias,
                                  DataType input_type,
@@ -57,6 +58,7 @@ void forward_kernel(ffStream_t stream,
                     int in_dim,
                     int out_dim,
                     int batch_size);
+
 void backward_kernel(ffStream_t stream,
                      LinearPerDeviceState const &m,
                      void const *input_ptr,

diff --git a/lib/kernels/include/kernels/local_cuda_allocator.h b/lib/kernels/include/kernels/local_cuda_allocator.h
@@ -0,0 +1,22 @@
+#include "kernels/allocation.h"
+#include <unordered_set>
+
+namespace FlexFlow {
+
+struct LocalCudaAllocator : public IAllocator {
+  LocalCudaAllocator() = default;
+  LocalCudaAllocator(LocalCudaAllocator const &) = delete;
+  LocalCudaAllocator(LocalCudaAllocator &&) = delete;
+  ~LocalCudaAllocator() override;
+
+  void *allocate(size_t) override;
+  void deallocate(void *) override;
+
+private:
+  std::unordered_set<void *> ptrs;
+};
+CHECK_RC_COPY_VIRTUAL_COMPLIANT(LocalCudaAllocator);
+
+Allocator create_local_cuda_memory_allocator();
+
+} // namespace FlexFlow
diff --git a/lib/kernels/include/kernels/managed_ff_stream.h b/lib/kernels/include/kernels/managed_ff_stream.h
@@ -0,0 +1,28 @@
+#ifndef _FLEXFLOW_KERNELS_MANAGED_FF_STREAM_H
+#define _FLEXFLOW_KERNELS_MANAGED_FF_STREAM_H
+
+#include "device.h"
+
+namespace FlexFlow {
+
+struct ManagedFFStream {
+public:
+  ManagedFFStream();
+
+  ManagedFFStream(ManagedFFStream const &) = delete;
+  ManagedFFStream &operator=(ManagedFFStream const &) = delete;
+
+  ManagedFFStream(ManagedFFStream &&other) noexcept;
+  ManagedFFStream &operator=(ManagedFFStream &&other) noexcept;
+
+  ~ManagedFFStream();
+
+  ffStream_t const &raw_stream() const;
+
+private:
+  ffStream_t *stream;
+};
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/kernels/include/kernels/managed_per_device_ff_handle.h b/lib/kernels/include/kernels/managed_per_device_ff_handle.h
@@ -0,0 +1,30 @@
+#ifndef _FLEXFLOW_KERNELS_MANAGED_HANDLE_H
+#define _FLEXFLOW_KERNELS_MANAGED_HANDLE_H
+
+#include "kernels/ff_handle.h"
+
+namespace FlexFlow {
+
+struct ManagedPerDeviceFFHandle {
+public:
+  ManagedPerDeviceFFHandle();
+
+  ManagedPerDeviceFFHandle(ManagedPerDeviceFFHandle const &) = delete;
+  ManagedPerDeviceFFHandle &
+      operator=(ManagedPerDeviceFFHandle const &) = delete;
+
+  ManagedPerDeviceFFHandle(ManagedPerDeviceFFHandle &&other) noexcept;
+  ManagedPerDeviceFFHandle &
+      operator=(ManagedPerDeviceFFHandle &&other) noexcept;
+
+  ~ManagedPerDeviceFFHandle();
+
+  PerDeviceFFHandle const &raw_handle() const;
+
+private:
+  PerDeviceFFHandle *handle;
+};
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/kernels/include/kernels/reduce_kernels.h b/lib/kernels/include/kernels/reduce_kernels.h
@@ -31,8 +31,8 @@ namespace Reduce {
 ReducePerDeviceState init_kernel(PerDeviceFFHandle const &,
                                  OperatorType const &,
                                  size_t const &,
-                                 ArrayShape input_shape,
-                                 ArrayShape output_shape);
+                                 ArrayShape const &input_shape,
+                                 ArrayShape const &output_shape);
 
 void forward_kernel(ffStream_t stream,
                     ReducePerDeviceState const &m,

diff --git a/lib/kernels/include/kernels/replicate_kernels.h b/lib/kernels/include/kernels/replicate_kernels.h
@@ -13,8 +13,8 @@ void forward_kernel(ffStream_t stream,
                     GenericTensorAccessorW const &output);
 
 void backward_kernel(ffStream_t stream,
-                     GenericTensorAccessorR const &input,
-                     GenericTensorAccessorW const &output,
+                     GenericTensorAccessorW const &input,
+                     GenericTensorAccessorR const &output,
                      size_t num_replicas);
 
 } // namespace Replicate

diff --git a/lib/kernels/include/kernels/softmax_kernels.h b/lib/kernels/include/kernels/softmax_kernels.h
@@ -18,12 +18,18 @@ FF_VISITABLE_STRUCT(SoftmaxPerDeviceState, handle, inputTensor, dim);
 namespace Kernels {
 namespace Softmax {
 
-SoftmaxPerDeviceState init_kernel(PerDeviceFFHandle const &, int);
+SoftmaxPerDeviceState init_kernel(PerDeviceFFHandle const &handle,
+                                  int dim,
+                                  int input_n,
+                                  int input_c,
+                                  int input_h,
+                                  int input_w);
 
 void forward_kernel(ffStream_t stream,
                     SoftmaxPerDeviceState const &m,
                     float const *input_ptr,
                     float *output_ptr);
+
 void backward_kernel(ffStream_t stream,
                      float *input_grad_ptr,
                      float const *output_grad_ptr,

diff --git a/lib/kernels/include/kernels/transpose_kernels.h b/lib/kernels/include/kernels/transpose_kernels.h
@@ -2,6 +2,7 @@
 #define _FLEXFLOW_OPS_KERNELS_TRANSPOSE_KERNELS_H
 
 #include "device.h"
+#include "kernels/accessor.h"
 #include <vector>
 
 namespace FlexFlow {

diff --git a/lib/kernels/src/accessor.cc b/lib/kernels/src/accessor.cc
@@ -2,6 +2,46 @@
 
 namespace FlexFlow {
 
+int32_t *GenericTensorAccessorW::get_int32_ptr() const {
+  return this->get<DataType::INT32>();
+}
+
+int64_t *GenericTensorAccessorW::get_int64_ptr() const {
+  return this->get<DataType::INT64>();
+}
+
+float *GenericTensorAccessorW::get_float_ptr() const {
+  return this->get<DataType::FLOAT>();
+}
+
+double *GenericTensorAccessorW::get_double_ptr() const {
+  return this->get<DataType::DOUBLE>();
+}
+
+half *GenericTensorAccessorW::get_half_ptr() const {
+  return this->get<DataType::HALF>();
+}
+
+int32_t const *GenericTensorAccessorR::get_int32_ptr() const {
+  return this->get<DataType::INT32>();
+}
+
+int64_t const *GenericTensorAccessorR::get_int64_ptr() const {
+  return this->get<DataType::INT64>();
+}
+
+float const *GenericTensorAccessorR::get_float_ptr() const {
+  return this->get<DataType::FLOAT>();
+}
+
+double const *GenericTensorAccessorR::get_double_ptr() const {
+  return this->get<DataType::DOUBLE>();
+}
+
+half const *GenericTensorAccessorR::get_half_ptr() const {
+  return get<DataType::HALF>();
+}
+
 int32_t *get_int32_ptr(GenericTensorAccessorW const &a) {
   return get<DataType::INT32>(a);
 }
@@ -92,4 +132,10 @@ std::vector<half const *>
   return get<DataType::HALF>(a);
 }
 
+GenericTensorAccessorR read_only_accessor_from_write_accessor(
+    GenericTensorAccessorW const &writable) {
+  return GenericTensorAccessorR{
+      writable.data_type, writable.shape, req<void const *>(writable.ptr)};
+}
+
 } // namespace FlexFlow
-Original file line number
+Diff line change
@@ Expand Up / @@ -152,4 +152,4 @@ @@
           };
         }
       );
-    }
+    }