feat: intial implementation of realm-backend

lib/realm-backend/include/realm-backend/allocated_tensors.h

@@ -0,0 +1,30 @@
+#ifndef _FLEXFLOW_LOCAL_EXECUTION_ALLOCATED_TENSORS_H
+#define _FLEXFLOW_LOCAL_EXECUTION_ALLOCATED_TENSORS_H
+
+#include "realm-backend/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/realm-backend/include/realm-backend/allocated_tensors.struct.toml

@@ -0,0 +1,32 @@
+namespace = "FlexFlow"
+name = "AllocatedTensors"
+features = [
+  "eq",
+  "fmt",
+  "hash",
+]
+
+includes = [
+  "task-spec/tensor_type_t.dtg.h",
+  "kernels/accessor.h",
+  "realm-backend/realm_allocator.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, std::pair<::FlexFlow::RealmRegion,::FlexFlow::TensorShape>>"
+
+[[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/realm-backend/include/realm-backend/model_training_instance.h

@@ -12,26 +12,23 @@ using PerLayerElapsedTime =
     std::unordered_map<layer_guid_t, std::optional<float>>;
 
 struct ModelTrainingInstance {
-  ModelTrainingInstance(ComputationGraph const &,
-                        RuntimeArgConfig const &,
-                        LossAttrs const &,
+  ModelTrainingInstance(RealmTrainingBacking const &,
                         tensor_guid_t const &logit_tensor,
-                        loss_tensor_t const &label_tensor,
+                        TensorShape const &label_tensor_shape,
+                        LossAttrs const &,
                         OptimizerAttrs const &);
 
-  void execute_init();
-  PerLayerElapsedTime execute_forward();
-  PerLayerElapsedTime execute_backward();
-  void execute_update();
-
-  ComputationGraph computation_graph;
   RealmTrainingBacking training_backing;
-  LossAttrs loss_attrs;
   tensor_guid_t logit_tensor;
   loss_tensor_t label_tensor;
+  LossAttrs loss_attrs;
   OptimizerAttrs optimizer_attrs;
 };
 
+PerLayerElapsedTime forward(ModelTrainingInstance &);
+PerLayerElapsedTime backward(ModelTrainingInstance &);
+void update(ModelTrainingInstance &);
+
 } // namespace FlexFlow
 
 #endif

lib/realm-backend/include/realm-backend/realm_args_backing.h

@@ -1,38 +1,38 @@
 #ifndef _FLEXFLOW_REALM_BACKEND_REALM_ARGS_BACKING_H
 #define _FLEXFLOW_REALM_BACKEND_REALM_ARGS_BACKING_H
 
-#include "local-execution/op_task_invocation.h"
-#include "local-execution/per_device_op_state.h"
-#include "local-execution/runtime_arg_config.h"
-#include "local-execution/task_invocation.dtg.h"
 #include "pcg/computation_graph.h"
 #include "pcg/layer_guid_t.dtg.h"
 #include "realm-backend/realm_task_argument_accessor.h"
 #include "realm-backend/task_result.h"
+#include "task-spec/op_task_invocation.h"
+#include "task-spec/per_device_op_state.h"
+#include "task-spec/runtime_arg_config.h"
+#include "task-spec/task_invocation.dtg.h"
 
 namespace FlexFlow {
 
 struct RealmArgsBacking {
-  RealmArgsBacking(RuntimeArgConfig const &);
-
-public:
-  void add_per_device_op_state(layer_guid_t const &,
-                               Future<DeviceSpecificDeviceStates> &&);
-
-  ArgSlotsBacking construct_arg_slots_backing(TaskBinding const &) const;
-
-  ConcreteArgSpec lower_to_concrete_arg_spec(RuntimeArgRefSpec const &) const;
-  ConcreteArgSpec lower_to_concrete_arg_spec(OpArgRefSpec const &,
-                                             ComputationGraph const &,
-                                             layer_guid_t const &) const;
+  RealmArgsBacking(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;
 };
 
+RealmArgsBacking
+make_args_backing_with_empty_device_states(RuntimeArgConfig const &);
+
+std::optional<DeviceSpecificDeviceStates>
+get_per_device_op_state_if_exists(RealmArgsBacking const &,
+                                  layer_guid_t const &);
+
+ArgSlotsBacking construct_arg_slots_backing(TaskBinding const &,
+                                            RuntimeArgConfig const &);
+
 } // namespace FlexFlow
 
 #endif

lib/realm-backend/include/realm-backend/realm_task_argument_accessor.h

@@ -1,8 +1,9 @@
 #ifndef _FLEXFLOW_REALM_BACKEND_REALM_TASK_ARGUMENT_ACCESSOR_H
 #define _FLEXFLOW_REALM_BACKEND_REALM_TASK_ARGUMENT_ACCESSOR_H
 
-#include "local-execution/slot_tensor_type_id.dtg.h"
 #include "local-execution/task_argument_accessor.h"
+#include "realm-backend/realm_allocator.h"
+#include "task-spec/slot_tensor_type_id.dtg.h"
 #include <unordered_map>
 #include <variant>
 
@@ -14,7 +15,7 @@ using TensorSlotsBacking = std::unordered_map<
 using ArgSlotsBacking = std::unordered_map<slot_id_t, ConcreteArgSpec>;
 
 struct RealmTaskArgumentAccessor : public ITaskArgumentAccessor {
-  RealmTaskArgumentAccessor(Allocator const &allocator,
+  RealmTaskArgumentAccessor(RealmAllocator const &allocator,
                             TensorSlotsBacking const &tensor_slots_backing,
                             ArgSlotsBacking const &arg_slots_backing);
 
@@ -23,18 +24,18 @@ struct RealmTaskArgumentAccessor : public ITaskArgumentAccessor {
 
   ConcreteArgSpec const &get_concrete_arg(slot_id_t) const override;
 
-  GenericTensorAccessor get_tensor(slot_id_t slot,
-                                   Permissions priv,
+  GenericTensorAccessor get_tensor(slot_id_t slot, Permissions priv,
                                    TensorType tensor_type) const override;
-  VariadicGenericTensorAccessor get_variadic_tensor(
-      slot_id_t slot, Permissions priv, TensorType tensor_type) const override;
+  VariadicGenericTensorAccessor
+  get_variadic_tensor(slot_id_t slot, Permissions priv,
+                      TensorType tensor_type) const override;
 
   Allocator get_allocator() const override;
 
   size_t get_device_idx() const override;
 
 private:
-  Allocator allocator;
+  RealmAllocator allocator;
   TensorSlotsBacking tensor_slots_backing;
   ArgSlotsBacking arg_slots_backing;
 };
@@ -45,8 +46,8 @@ using TensorSlotsBackingWithoutAddresses = std::unordered_map<
                  std::vector<std::pair<ArrayShape, DataType>>>>;
 
 TensorSlotsBackingWithoutAddresses
-    get_slots_backing_without_tensor_allocation_addresses(
-        TensorSlotsBacking const &);
+get_slots_backing_without_tensor_allocation_addresses(
+    TensorSlotsBacking const &);
 
 CHECK_RC_COPY_VIRTUAL_COMPLIANT(RealmTaskArgumentAccessor);
 

lib/realm-backend/include/realm-backend/realm_tensor_backing.h

@@ -3,58 +3,70 @@
 #define _FLEXFLOW_REALM_BACKEND_REALM_TENSOR_BACKING_H
 
 #include "kernels/accessor.h"
-#include "realm-backend/realm_task_argument_accessor.h"
-#include "realm-backend/realm_allocator.h"
-#include "local-execution/task_invocation.dtg.h"
-#include "local-execution/tensor_role.dtg.h"
-#include "local-execution/lowered_tensor_t.dtg.h"
+#include "local-execution/gradient_tensor_source.h"
+#include "local-execution/loss_tensor_source.h"
 #include "local-execution/lowered_tensor_source.h"
-#include "local-execution/optimizer_tensor_t.dtg.h"
-#include "local-execution/loss_tensor_t.dtg.h"
+#include "local-execution/optimizer_tensor_source.h"
 #include "pcg/computation_graph.dtg.h"
-#include "pcg/tensor_guid_t.dtg.h"
 #include "pcg/layer_guid_t.dtg.h"
+#include "pcg/optimizer_attrs.dtg.h"
+#include "realm-backend/allocated_tensors.dtg.h"
+#include "realm-backend/realm_allocator.h"
+#include "realm-backend/realm_task_argument_accessor.h"
+#include "realm-backend/unallocated_tensors.dtg.h"
+#include "task-spec/lowered_tensor_t.dtg.h"
+#include "task-spec/task_invocation.dtg.h"
+#include "task-spec/tensor_role.dtg.h"
 
 namespace FlexFlow {
 
-using TensorRegionMap =
-    std::unordered_map<lowered_tensor_t, RealmRegion>;
-using TensorShapeMap =
-    std::unordered_map<lowered_tensor_t, TensorShape>;
+using TensorBackingMap = std::unordered_map<lowered_tensor_t, std::pair<RealmRegion, TensorShape>>;
 
 struct RealmTensorBacking {
-  RealmTensorBacking();
+  RealmTensorBacking(AllocatedTensors const &, UnallocatedTensors const &,
+                     RealmAllocator const &);
 
 public:
-  void allocate_layer_tensors(layer_guid_t const &,
-                              ComputationGraph const &,
-                              RealmAllocator &);
-  void allocate_tensors_by_role(TensorRole const &,
-                                layer_guid_t const &,
-                                ComputationGraph const &,
-                                RealmAllocator &);
-  void allocate_optimizer_tensors(tensor_guid_t const &,
-                                  std::vector<optimizer_tensor_t> const &,
-                                  RealmAllocator &);
-  TensorSlotsBacking
-      construct_tensor_slots_backing(TaskBinding const &) const;
-
-  GenericTensorAccessorW const &
-      get_tensor_backing(lowered_tensor_t const &) const;
-
-  bool is_tensor_allocated(lowered_tensor_t const &) const;
+  GenericTensorAccessorW get_tensor(TensorTypeVariant const &) const;
 
 public:
   // tensors
-  TensorRegionMap tensor_regions;
-  TensorShapeMap tensor_shapes;
+  TensorBackingMap tensor_backings;
+
   std::unordered_map<tensor_guid_t, lowered_tensor_t> tensor_lowering_mapping;
-  std::unordered_map<tensor_guid_t, lowered_tensor_t> gradient_tensor_lowering_mapping;
-  std::unordered_map<optimizer_tensor_t, lowered_tensor_t> optimizer_tensor_lowering_mapping;
-  std::unordered_map<loss_tensor_t, lowered_tensor_t> loss_tensor_lowering_mapping;
+  std::unordered_map<gradient_tensor_t, lowered_tensor_t>
+      gradient_tensor_lowering_mapping;
+  std::unordered_map<optimizer_tensor_t, lowered_tensor_t>
+      optimizer_tensor_lowering_mapping;
+  std::unordered_map<loss_tensor_t, lowered_tensor_t>
+      loss_tensor_lowering_mapping;
+
+  std::unordered_map<tensor_guid_t, gradient_tensor_t> tensor_gradient_mapping;
+  std::unordered_map<tensor_guid_t, std::vector<optimizer_tensor_t>>
+      tensor_optimizer_mapping;
+
+  RealmAllocator allocator;
+
+private:
+  lowered_tensor_t insert_tensor(TensorTypeVariant const &);
   LoweredTensorSource lowered_tensor_source;
 };
 
+GenericTensorAccessorW wrappup_tensor_accessor(std::pair<RealmRegion, TensorShape> const &);
+
+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(RealmTensorBacking const &,
+                                                  TaskBinding const &);
+
 } // namespace FlexFlow
 
 #endif