Realm backend

lib/kernels/include/kernels/array_shape.h

@@ -69,4 +69,11 @@ std::ostream &operator<<(std::ostream &, ArrayShape const &);
 
 } // namespace FlexFlow
 
+namespace std {
+template <>
+struct hash<::FlexFlow::ArrayShape> {
+  size_t operator()(::FlexFlow::ArrayShape const &) const;
+};
+} // namespace std
+
 #endif

lib/kernels/src/array_shape.cc

@@ -128,3 +128,14 @@ std::ostream &operator<<(std::ostream &s, ArrayShape const &x) {
 }
 
 } // namespace FlexFlow
+
+namespace std {
+size_t hash<FlexFlow::ArrayShape>::operator()(
+    ::FlexFlow::ArrayShape const &x) const {
+  size_t result = 0;
+  result ^= std::hash<::FlexFlow::LegionOrdered<::FlexFlow::nonnegative_int>>{}(
+                x.dims) +
+            0x9e3779b9 + (result << 6) + (result >> 2);
+  return result;
+}
+} // namespace std

lib/kernels/src/legion_dim.cc

@@ -19,4 +19,10 @@ ff_dim_t legion_dim_from_ff_dim(legion_dim_t legion_dim,
                                   legion_dim.value.unwrap_nonnegative() - 1}};
 }
 
+ff_dim_t ff_dim_from_legion_dim(legion_dim_t legion_dim,
+                                nonnegative_int num_dimensions) {
+  return ff_dim_t{nonnegative_int{num_dimensions.unwrap_nonnegative() -
+                                  legion_dim.value.unwrap_nonnegative() - 1}};
+}
+
 } // namespace FlexFlow

lib/local-execution/include/local-execution/allocated_tensors.h

@@ -0,0 +1,30 @@
+#ifndef _FLEXFLOW_LOCAL_EXECUTION_ALLOCATED_TENSORS_H
+#define _FLEXFLOW_LOCAL_EXECUTION_ALLOCATED_TENSORS_H
+
+#include "local-execution/allocated_tensors.dtg.h"
+#include "pcg/computation_graph.h"
+
+namespace FlexFlow {
+
+bool are_allocated_forward_tensors_valid(
+    AllocatedTensors const &,
+    std::unordered_map<tensor_guid_t, TensorAttrs> const &);
+bool are_allocated_gradient_tensors_valid(
+    AllocatedTensors const &,
+    std::unordered_map<tensor_guid_t, TensorAttrs> const &);
+bool are_allocated_optimizer_tensors_valid(
+    AllocatedTensors const &,
+    std::unordered_map<tensor_guid_t, TensorAttrs> const &);
+
+bool are_allocated_tensors_valid(
+    AllocatedTensors const &,
+    std::unordered_map<tensor_guid_t, TensorAttrs> const &);
+
+bool is_allocated_tensor_backing_valid(
+    TensorTypeVariant const &,
+    std::unordered_map<TensorTypeVariant, GenericTensorAccessorW> const &,
+    ArrayShape const &);
+
+} // namespace FlexFlow
+
+#endif

lib/local-execution/include/local-execution/allocated_tensors.struct.toml

@@ -0,0 +1,31 @@
+namespace = "FlexFlow"
+name = "AllocatedTensors"
+features = [
+  "eq",
+  "fmt",
+  "hash",
+]
+
+includes = [
+  "task-spec/tensor_type_t.dtg.h",
+  "kernels/accessor.h"
+]
+
+src_includes = [
+  "utils/hash/unordered_map.h",
+  "utils/fmt/unordered_map.h",
+  "utils/hash/vector.h",
+  "utils/fmt/vector.h"
+]
+
+[[fields]]
+name = "tensor_type_backings"
+type = "std::unordered_map<::FlexFlow::TensorTypeVariant, ::FlexFlow::GenericTensorAccessorW>"
+
+[[fields]]
+name = "gradient_mapping"
+type = "std::unordered_map<::FlexFlow::tensor_guid_t, ::FlexFlow::gradient_tensor_t>"
+
+[[fields]]
+name = "optimizer_mapping"
+type = "std::unordered_map<::FlexFlow::tensor_guid_t, std::vector<::FlexFlow::optimizer_tensor_t>>"

lib/local-execution/include/local-execution/gradient_tensor_source.h

@@ -11,6 +11,8 @@ struct GradientTensorSource {
 
   gradient_tensor_t new_gradient_tensor();
 
+  void reset();
+
 private:
   static size_t next_available_gradient_tensor_id;
 };

lib/local-execution/include/local-execution/local_args_backing.h

@@ -12,18 +12,19 @@
 namespace FlexFlow {
 
 struct LocalArgsBacking {
-  LocalArgsBacking(RuntimeArgConfig const &);
+  LocalArgsBacking(
+      RuntimeArgConfig const &,
+      std::unordered_map<layer_guid_t, DeviceSpecificDeviceStates> const &);
 
 public:
   // arguments
+  RuntimeArgConfig runtime_arg_config;
   std::unordered_map<layer_guid_t, DeviceSpecificDeviceStates>
       per_device_op_states;
-  RuntimeArgConfig runtime_arg_config;
 };
 
-void add_per_device_op_state(LocalArgsBacking &,
-                             layer_guid_t const &,
-                             DeviceSpecificDeviceStates const &);
+LocalArgsBacking
+    make_args_backing_with_empty_device_states(RuntimeArgConfig const &);
 
 std::optional<DeviceSpecificDeviceStates>
     get_per_device_op_state_if_exists(LocalArgsBacking const &,

lib/local-execution/include/local-execution/local_tensor_backing.h

@@ -3,41 +3,32 @@
 #define _FLEXFLOW_LOCAL_EXECUTION_LOCAL_TENSOR_BACKING_H
 
 #include "kernels/accessor.h"
+#include "local-execution/allocated_tensors.dtg.h"
 #include "local-execution/gradient_tensor_source.h"
 #include "local-execution/local_task_argument_accessor.h"
 #include "local-execution/loss_tensor_source.h"
 #include "local-execution/lowered_tensor_source.h"
 #include "local-execution/optimizer_tensor_source.h"
-#include "op-attrs/tensor_shape.dtg.h"
+#include "local-execution/unallocated_tensors.dtg.h"
 #include "pcg/computation_graph.dtg.h"
 #include "pcg/layer_guid_t.dtg.h"
 #include "pcg/optimizer_attrs.dtg.h"
-#include "pcg/tensor_guid_t.dtg.h"
-#include "task-spec/loss_tensor_t.dtg.h"
 #include "task-spec/lowered_tensor_t.dtg.h"
-#include "task-spec/optimizer_tensor_t.dtg.h"
 #include "task-spec/task_invocation.dtg.h"
 #include "task-spec/tensor_role.dtg.h"
-#include "task-spec/tensor_type_t.dtg.h"
 
 namespace FlexFlow {
 
 using TensorBackingMap =
     std::unordered_map<lowered_tensor_t, GenericTensorAccessorW>;
 
 struct LocalTensorBacking {
-  LocalTensorBacking() = default;
-  LocalTensorBacking(
-      std::unordered_map<TensorTypeVariant, GenericTensorAccessorW> const
-          &allocated_tensor_backings,
-      std::unordered_set<tensor_guid_t> const &allocated_tensor_guids,
-      std::unordered_map<tensor_guid_t, gradient_tensor_t> const
-          &allocated_gradient_mapping,
-      std::unordered_map<tensor_guid_t, std::vector<optimizer_tensor_t>> const
-          &allocated_optimizer_mapping,
-      std::unordered_set<loss_tensor_t> const &allocated_loss_tensors);
-
-  lowered_tensor_t allocate_tensor(TensorShape const &, Allocator &);
+  LocalTensorBacking(AllocatedTensors const &,
+                     UnallocatedTensors const &,
+                     Allocator const &);
+
+public:
+  GenericTensorAccessorW get_tensor(TensorTypeVariant const &) const;
 
 public:
   // tensors
@@ -55,39 +46,24 @@ struct LocalTensorBacking {
   std::unordered_map<tensor_guid_t, std::vector<optimizer_tensor_t>>
       tensor_optimizer_mapping;
 
+  Allocator allocator;
+
 private:
-  lowered_tensor_t insert_tensor(GenericTensorAccessorW const &);
+  lowered_tensor_t insert_tensor(TensorTypeVariant const &);
   LoweredTensorSource lowered_tensor_source;
 };
 
-void allocate_tensor_guid(LocalTensorBacking &,
-                          tensor_guid_t const &,
-                          TensorShape const &,
-                          Allocator &);
-void allocate_gradient_tensor(LocalTensorBacking &,
-                              gradient_tensor_t const &,
-                              tensor_guid_t const &,
-                              TensorShape const &,
-                              Allocator &);
-void allocate_optimizer_tensors(LocalTensorBacking &,
-                                std::vector<optimizer_tensor_t> const &,
-                                tensor_guid_t const &,
-                                TensorShape const &,
-                                Allocator &);
-
-void allocate_all_computation_graph_tensors(LocalTensorBacking &,
-                                            GradientTensorSource &,
-                                            ComputationGraph const &,
-                                            Allocator &);
-void allocate_all_optimizer_tensors(LocalTensorBacking &,
-                                    OptimizerTensorSource &,
-                                    ComputationGraph const &,
-                                    Allocator &,
-                                    OptimizerAttrs const &);
-loss_tensor_t allocate_loss_tensor(LocalTensorBacking &,
-                                   LossTensorSource const &,
-                                   TensorShape const &,
-                                   Allocator &);
+UnallocatedTensors generate_unallocated_tensors(
+    AllocatedTensors const &,
+    std::unordered_map<tensor_guid_t, TensorAttrs> const &,
+    GradientTensorSource &);
+
+UnallocatedTensors generate_unallocated_tensors_with_optimizer(
+    AllocatedTensors const &,
+    std::unordered_map<tensor_guid_t, TensorAttrs> const &,
+    GradientTensorSource &,
+    OptimizerTensorSource &,
+    OptimizerAttrs const &);
 
 TensorSlotsBacking construct_tensor_slots_backing(LocalTensorBacking const &,
                                                   TaskBinding const &);

lib/local-execution/include/local-execution/local_training_backing.h

@@ -1,6 +1,7 @@
 #ifndef _FLEXFLOW_LOCAL_EXECUTION_LOCAL_TRAINING_BACKING_H
 #define _FLEXFLOW_LOCAL_EXECUTION_LOCAL_TRAINING_BACKING_H
 
+#include "local-execution/allocated_tensors.dtg.h"
 #include "local-execution/local_args_backing.h"
 #include "local-execution/local_tensor_backing.h"
 #include "local-execution/optimizer_tensor_source.h"
@@ -13,46 +14,51 @@ namespace FlexFlow {
 
 struct LocalTrainingBacking {
   LocalTrainingBacking(Allocator const &,
+                       AllocatedTensors const &,
                        ComputationGraph const &,
-                       LocalTensorBacking const &,
-                       LocalArgsBacking const &);
+                       RuntimeArgConfig const &);
+
+  LocalTrainingBacking(Allocator const &,
+                       AllocatedTensors const &,
+                       ComputationGraph const &,
+                       RuntimeArgConfig const &,
+                       OptimizerAttrs const &);
 
 public:
   LocalTensorBacking local_tensor_backing;
   LocalArgsBacking local_args_backing;
 
-  Allocator allocator;
   ComputationGraph computation_graph;
   TaskRegistry task_registry;
 
   GradientTensorSource gradient_tensor_source;
+  OptimizerTensorSource optimizer_tensor_source;
 };
 
-DeviceSpecificDeviceStates call_init_task_impl(TaskRegistry const &,
-                                               task_id_t task_id,
-                                               TaskArgumentAccessor const &acc);
+LocalArgsBacking initialize_args_backing(TaskRegistry const &,
+                                         ComputationGraph const &,
+                                         RuntimeArgConfig const &,
+                                         LocalTensorBacking const &);
 
 std::optional<float> call_task_impl(TaskRegistry const &,
-                                    task_id_t task_id,
-                                    TaskArgumentAccessor acc);
+                                    task_id_t const &task_id,
+                                    TaskArgumentAccessor const &acc);
 
-void execute_init(LocalTrainingBacking &, layer_guid_t const &);
-std::optional<float> execute_forward(LocalTrainingBacking &,
+std::optional<float> execute_forward(LocalTrainingBacking const &,
                                      layer_guid_t const &);
-std::optional<float> execute_backward(LocalTrainingBacking &,
+std::optional<float> execute_backward(LocalTrainingBacking const &,
                                       layer_guid_t const &);
-void compute_loss(LocalTrainingBacking &,
+void compute_loss(LocalTrainingBacking const &,
                   LossAttrs const &,
                   tensor_guid_t const &logit_tensor,
                   loss_tensor_t const &label_tensor);
-void execute_update(LocalTrainingBacking &,
+void execute_update(LocalTrainingBacking const &,
                     layer_guid_t const &,
                     OptimizerAttrs const &);
 
 TaskArgumentAccessor get_task_arg_accessor(LocalTensorBacking const &,
                                            LocalArgsBacking const &,
-                                           TaskInvocation const &,
-                                           Allocator &);
+                                           TaskInvocation const &);
 
 } // namespace FlexFlow
 

lib/local-execution/include/local-execution/model_training_instance.h

@@ -14,24 +14,19 @@ using PerLayerElapsedTime =
 struct ModelTrainingInstance {
   ModelTrainingInstance(LocalTrainingBacking const &,
                         tensor_guid_t const &logit_tensor,
-                        TensorShape const &label_tensor_shape,
+                        loss_tensor_t const &label_tensor,
                         LossAttrs const &,
                         OptimizerAttrs const &);
 
   LocalTrainingBacking training_backing;
-  LossAttrs loss_attrs;
-  OptimizerAttrs optimizer_attrs;
   tensor_guid_t logit_tensor;
   loss_tensor_t label_tensor;
-
-private:
-  OptimizerTensorSource optimizer_tensor_source;
-  LossTensorSource loss_tensor_source;
+  LossAttrs loss_attrs;
+  OptimizerAttrs optimizer_attrs;
 };
 
-void init(ModelTrainingInstance &);
-PerLayerElapsedTime forward(ModelTrainingInstance &);
-PerLayerElapsedTime backward(ModelTrainingInstance &);
+PerLayerElapsedTime forward(ModelTrainingInstance const &);
+PerLayerElapsedTime backward(ModelTrainingInstance const &);
 void update(ModelTrainingInstance &);
 
 } // namespace FlexFlow

lib/local-execution/include/local-execution/optimizer_tensor_source.h

@@ -11,6 +11,8 @@ struct OptimizerTensorSource {
 
   optimizer_tensor_t new_optimizer_tensor();
 
+  void reset();
+
 private:
   static size_t next_available_optimizer_tensor_id;
 };

lib/local-execution/include/local-execution/task_registry.h

@@ -9,19 +9,13 @@
 
 namespace FlexFlow {
 
-TaskRegistry empty_task_registry();
-
-void register_tasks_for_layer(TaskRegistry &,
-                              layer_guid_t const &,
-                              ComputationGraphOpAttrs const &attrs);
+TaskRegistry construct_task_registry(
+    std::unordered_map<layer_guid_t, LayerAttrs> const &);
 
 bool registry_contains_task_for_layer(TaskRegistry const &,
                                       layer_guid_t const &,
                                       OpTaskType const &);
 
-void register_all_computation_graph_tasks(TaskRegistry &,
-                                          ComputationGraph const &);
-
 } // namespace FlexFlow
 
 #endif