PR todo cleanup

Author: lockshaw

lib/local-execution/include/local-execution/atomic_task_invocation/atomic_task_binding.dtg.toml

@@ -14,9 +14,16 @@ namespace FlexFlow {
 struct LocalTaskArgumentAccessor : public ITaskArgumentAccessor {
   explicit LocalTaskArgumentAccessor(
       Allocator const &allocator,
-      std::unordered_map<TaskTensorParameter, TensorSlotBacking> const
+      std::unordered_map<TaskTensorParameter, GenericTensorAccessorW> const
           &tensor_slots_backing,
-      std::unordered_map<arg_slot_id_t, ConcreteArgSpec> const &arg_slots_backing,
+      ProfilingSettings const &profiling_settings,
+      device_handle_t const &ff_handle,
+      DeviceType kernel_device_type,
+      PCGOperatorAttrs const &op_attrs,
+      std::optional<LossAttrs> const &loss_attrs,
+      std::optional<PerDeviceOpState> const &per_device_op_state,
+      FFIterationConfig const &iteration_config,
+      std::optional<OptimizerAttrs> const &optimizer_attrs,
       size_t device_idx);
 
   LocalTaskArgumentAccessor(LocalTaskArgumentAccessor const &) = delete;
@@ -26,18 +33,33 @@ struct LocalTaskArgumentAccessor : public ITaskArgumentAccessor {
 
   GenericTensorAccessor get_tensor(TaskTensorParameter slot,
                                    Permissions priv) const override;
-  VariadicGenericTensorAccessor get_variadic_tensor(
-      TaskTensorParameter, Permissions priv) const override;
+
+  ProfilingSettings get_profiling_settings() const override;
+  device_handle_t get_ff_handle() const override;
+  DeviceType get_kernel_device_type() const override;
+  PCGOperatorAttrs get_op_attrs() const override;
+  LossAttrs get_loss_attrs() const override;
+  PerDeviceOpState get_per_device_op_state() const override;
+  FFIterationConfig get_iteration_config() const override;
+  OptimizerAttrs get_optimizer_attrs() const override;
 
   Allocator get_allocator() const override;
 
   device_id_t get_device_idx() const override;
 
 private:
   Allocator allocator;
-  std::unordered_map<TaskTensorParameter, TensorSlotBacking>
-      tensor_slots_backing;
-  std::unordered_map<arg_slot_id_t, ConcreteArgSpec> arg_slots_backing;
+  std::unordered_map<TaskTensorParameter, GenericTensorAccessorW> tensor_slots_backing;
+
+  ProfilingSettings profiling_settings;
+  device_handle_t ff_handle;
+  DeviceType kernel_device_type;
+  PCGOperatorAttrs op_attrs;
+  std::optional<LossAttrs> loss_attrs;
+  std::optional<PerDeviceOpState> per_device_op_state;
+  FFIterationConfig iteration_config;
+  std::optional<OptimizerAttrs> optimizer_attrs;
+
   device_id_t device_idx; 
 };
 

lib/local-execution/include/local-execution/atomic_task_invocation/atomic_task_graph.dtg.toml

@@ -8,7 +8,7 @@ features = [
 ]
 
 includes = [
-  "task-spec/task_signature_impl.dtg.h",
+  "task-spec/task_impl_function.dtg.h",
 ]
 
 src_includes = [
@@ -18,4 +18,4 @@ src_includes = [
 
 [[fields]]
 name = "task_mapping"
-type = "std::unordered_map<::FlexFlow::task_id_t, ::FlexFlow::TaskSignatureAndImpl>"
+type = "std::unordered_map<::FlexFlow::task_id_t, ::FlexFlow::TaskImplFunction>"

lib/local-execution/include/local-execution/atomic_task_invocation/atomic_task_invocation.dtg.toml

@@ -0,0 +1,13 @@
+#ifndef _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_PER_DEVICE_OP_STATE_INITIALIZATION_H
+#define _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_PER_DEVICE_OP_STATE_INITIALIZATION_H
+
+#include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h"
+namespace FlexFlow {
+
+DynamicOpenDataflowGraph
+  perform_per_device_op_state_initialization(
+    DynamicOpenDataflowGraph const &);
+
+} // namespace FlexFlow
+
+#endif

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

@@ -0,0 +1,36 @@
+#ifndef _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_TASK_EXECUTION_H
+#define _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_TASK_EXECUTION_H
+
+#include "kernels/profiling_settings.dtg.h"
+#include "op-attrs/ops/loss_functions/loss_attrs.dtg.h"
+#include "task-spec/per_device_op_state.dtg.h"
+#include "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h"
+#include "task-spec/task_argument_accessor/task_argument_accessor.h"
+
+namespace FlexFlow {
+
+TaskArgumentAccessor
+  make_task_argument_accessor_for_invocation(
+    DynamicNodeInvocation const &invocation,
+    ProfilingSettings const &profiling_settings,
+    DeviceType kernel_device_type,
+    PCGOperatorAttrs op_attrs,
+    std::optional<LossAttrs> const &loss_attrs,
+    std::optional<PerDeviceOpState> const &per_device_op_state,
+    FFIterationConfig iteration_config,
+    std::optional<OptimizerAttrs> const &optimizer_attrs);
+
+void
+  execute_dynamic_node_invocation(
+    DynamicNodeInvocation const &invocation,
+    ProfilingSettings const &profiling_settings,
+    DeviceType kernel_device_type,
+    PCGOperatorAttrs op_attrs,
+    std::optional<LossAttrs> const &loss_attrs,
+    std::optional<PerDeviceOpState> const &per_device_op_state,
+    FFIterationConfig iteration_config,
+    std::optional<OptimizerAttrs> const &optimizer_attrs);
+
+} // namespace FlexFlow
+
+#endif

lib/local-execution/include/local-execution/local_task_registry.dtg.toml

@@ -0,0 +1,23 @@
+#ifndef _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_TENSOR_ALLOCATION_H
+#define _FLEXFLOW_LIB_LOCAL_EXECUTION_INCLUDE_LOCAL_EXECUTION_TENSOR_ALLOCATION_H
+
+#include "kernels/allocation.h"
+#include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h"
+
+namespace FlexFlow {
+
+bool no_tensors_are_allocated(DynamicOpenDataflowGraph const &);
+bool all_tensors_are_allocated(DynamicOpenDataflowGraph const &);
+
+DynamicValueAttrs
+  perform_tensor_allocation_for_value(DynamicValueAttrs const &, 
+                                      Allocator &);
+
+DynamicOpenDataflowGraph
+  perform_tensor_allocation(DynamicOpenDataflowGraph const &, 
+                            std::unordered_map<DynamicValueAttrs, GenericTensorAccessorW> const &preallocated,
+                            Allocator &);
+
+} // namespace FlexFlow
+
+#endif