Extend kwarg dataflow graph changes into substitutions

Author: lockshaw

lib/pcg/include/pcg/parallel_computation_graph/parallel_tensor_use_t.dtg.toml

@@ -9,9 +9,10 @@ features = [
 ]
 
 includes = [
-  "utils/graph/dataflow_graph/dataflow_input.dtg.h",
+  "utils/graph/kwarg_dataflow_graph/kwarg_dataflow_input.dtg.h",
+  "op-attrs/tensor_slot_name.dtg.h"
 ]
 
 [[fields]]
 name = "raw_dataflow_input"
-type = "::FlexFlow::DataflowInput"
+type = "::FlexFlow::KwargDataflowInput<::FlexFlow::TensorSlotName>"

lib/substitutions/include/substitutions/apply_substitution/perform_shape_inference.h

@@ -2,8 +2,10 @@
 #define _FLEXFLOW_LIB_SUBSTITUTIONS_INCLUDE_SUBSTITUTIONS_APPLY_SUBSTITUTION_PERFORM_SHAPE_INFERENCE_H
 
 #include "op-attrs/parallel_tensor_shape.dtg.h"
+#include "op-attrs/tensor_slot_name.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_layer_attrs.dtg.h"
-#include "utils/graph/labelled_open_dataflow_graph/labelled_open_dataflow_graph_view.h"
+#include "utils/graph/labelled_open_kwarg_dataflow_graph/labelled_open_kwarg_dataflow_graph_view.h"
+#include "utils/graph/open_kwarg_dataflow_graph/kwarg_dataflow_graph_input.dtg.h"
 
 namespace FlexFlow {
 
@@ -22,11 +24,11 @@ namespace FlexFlow {
  * Exists only to enable apply_substitution(SubParallelComputationGraph const &,
  * Substitution const &, PCGPatternMatch const &)
  */
-LabelledOpenDataflowGraphView<ParallelLayerAttrs, ParallelTensorShape>
+LabelledOpenKwargDataflowGraphView<ParallelLayerAttrs, ParallelTensorShape, int, TensorSlotName>
     perform_shape_inference(
-        LabelledOpenDataflowGraphView<ParallelLayerAttrs, std::monostate> const
+        LabelledOpenKwargDataflowGraphView<ParallelLayerAttrs, std::monostate, int, TensorSlotName> const
             &g,
-        std::unordered_map<DataflowGraphInput, ParallelTensorShape> const
+        std::unordered_map<KwargDataflowGraphInput<int>, ParallelTensorShape> const
             &input_shapes);
 
 } // namespace FlexFlow

lib/substitutions/include/substitutions/input_parallel_tensor_guid_t.dtg.toml

@@ -9,9 +9,9 @@ features = [
 ]
 
 includes = [
-  "utils/graph/open_dataflow_graph/dataflow_graph_input.dtg.h",
+  "utils/graph/open_kwarg_dataflow_graph/kwarg_dataflow_graph_input.dtg.h",
 ]
 
 [[fields]]
 name = "raw_dataflow_graph_input"
-type = "::FlexFlow::DataflowGraphInput"
+type = "::FlexFlow::KwargDataflowGraphInput<int>"

lib/substitutions/include/substitutions/open_parallel_tensor_guid_t.dtg.toml

@@ -9,9 +9,10 @@ features = [
 ]
 
 includes = [
-  "utils/graph/open_dataflow_graph/open_dataflow_value.dtg.h"
+  "utils/graph/open_kwarg_dataflow_graph/open_kwarg_dataflow_value.dtg.h",
+  "op-attrs/tensor_slot_name.dtg.h",
 ]
 
 [[fields]]
 name = "raw_open_dataflow_value"
-type = "::FlexFlow::OpenDataflowValue"
+type = "::FlexFlow::OpenKwargDataflowValue<int, ::FlexFlow::TensorSlotName>"

lib/substitutions/include/substitutions/open_parallel_tensor_guid_t.h

@@ -17,8 +17,8 @@ template <typename F,
           typename Ret = std::invoke_result_t<F, parallel_tensor_guid_t>>
 Ret visit_open_parallel_tensor_guid(open_parallel_tensor_guid_t t, F f) {
   return t.raw_open_dataflow_value.visit<Ret>(overload{
-      [&](DataflowOutput const &o) { return f(parallel_tensor_guid_t{o}); },
-      [&](DataflowGraphInput const &i) {
+      [&](KwargDataflowOutput<TensorSlotName> const &o) { return f(parallel_tensor_guid_t{o}); },
+      [&](KwargDataflowGraphInput<int> const &i) {
         return f(input_parallel_tensor_guid_t{i});
       },
   });

lib/substitutions/include/substitutions/output_graph/output_graph_expr.dtg.toml

@@ -4,11 +4,12 @@ type = "struct"
 features = []
 
 includes = [
-  "utils/graph/labelled_open_dataflow_graph/labelled_open_dataflow_graph.h",
+  "utils/graph/labelled_open_kwarg_dataflow_graph/labelled_open_kwarg_dataflow_graph.h",
   "substitutions/output_graph/output_operator_attrs_assignment.dtg.h",
   "<variant>",
+  "op-attrs/tensor_slot_name.dtg.h",
 ]
 
 [[fields]]
 name = "raw_graph"
-type = "::FlexFlow::LabelledOpenDataflowGraph<::FlexFlow::OutputOperatorAttrsAssignment, std::monostate>"
+type = "::FlexFlow::LabelledOpenKwargDataflowGraph<::FlexFlow::OutputOperatorAttrsAssignment, std::monostate, int, ::FlexFlow::TensorSlotName>"

lib/substitutions/include/substitutions/output_graph/output_graph_expr.h

@@ -10,7 +10,7 @@ namespace FlexFlow {
 
 std::unordered_set<OutputGraphExprNode> get_nodes(OutputGraphExpr const &);
 
-std::vector<OutputGraphExprNodeOutput>
+std::unordered_map<TensorSlotName, OutputGraphExprNodeOutput>
     get_node_outputs(OutputGraphExpr const &, OutputGraphExprNode const &);
 
 std::unordered_set<OutputGraphExprInput> get_inputs(OutputGraphExpr const &);

lib/substitutions/include/substitutions/output_graph/output_graph_expr_input.dtg.toml

@@ -9,9 +9,9 @@ features = [
 ]
 
 includes = [
-  "utils/graph/open_dataflow_graph/dataflow_graph_input.dtg.h",
+  "utils/graph/open_kwarg_dataflow_graph/kwarg_dataflow_graph_input.dtg.h",
 ]
 
 [[fields]]
 name = "raw_dataflow_graph_input"
-type = "::FlexFlow::DataflowGraphInput"
+type = "::FlexFlow::KwargDataflowGraphInput<int>"

lib/substitutions/include/substitutions/output_graph/output_graph_expr_node_output.dtg.toml

@@ -9,9 +9,10 @@ features = [
 ]
 
 includes = [
-  "utils/graph/dataflow_graph/dataflow_output.dtg.h",
+  "utils/graph/kwarg_dataflow_graph/kwarg_dataflow_output.dtg.h",
+  "op-attrs/tensor_slot_name.dtg.h",
 ]
 
 [[fields]]
 name = "raw_dataflow_output"
-type = "::FlexFlow::DataflowOutput"
+type = "::FlexFlow::KwargDataflowOutput<::FlexFlow::TensorSlotName>"

lib/substitutions/include/substitutions/output_graph/output_graph_expr_value.h

@@ -2,14 +2,14 @@
 #define _FLEXFLOW_LIB_SUBSTITUTIONS_INCLUDE_SUBSTITUTIONS_OUTPUT_GRAPH_OUTPUT_GRAPH_EXPR_VALUE_H
 
 #include "substitutions/output_graph/output_graph_expr_value.dtg.h"
-#include "utils/graph/open_dataflow_graph/open_dataflow_value.dtg.h"
+#include "utils/graph/open_kwarg_dataflow_graph/open_kwarg_dataflow_value.dtg.h"
 
 namespace FlexFlow {
 
-OpenDataflowValue raw_open_dataflow_value_from_output_graph_expr_value(
+OpenKwargDataflowValue<int, TensorSlotName> raw_open_dataflow_value_from_output_graph_expr_value(
     OutputGraphExprValue const &);
 OutputGraphExprValue output_graph_expr_value_from_raw_open_dataflow_value(
-    OpenDataflowValue const &);
+    OpenKwargDataflowValue<int, TensorSlotName> const &);
 
 } // namespace FlexFlow