Add v1 file format for MappedParallelComputationGraph and from_v1 conversion functions

lib/pcg/include/pcg/file_format/v1/graphs/v1_graph_output.dtg.toml

@@ -10,10 +10,18 @@ features = [
   "fmt",
 ]
 
+template_params = [
+  "SlotName",
+]
+
+includes = [
+  "utils/nonnegative_int/nonnegative_int.h",
+]
+
 [[fields]]
-name = "srcNode"
-type = "size_t"
+name = "node"
+type = "::FlexFlow::nonnegative_int"
 
 [[fields]]
-name = "srcIdx"
-type = "size_t"
+name = "slot_name"
+type = "SlotName"

lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.dtg.toml

@@ -18,6 +18,7 @@ includes = [
   "<vector>",
   "<unordered_set>",
   "pcg/file_format/v1/graphs/v1_graph_edge.dtg.h",
+  "pcg/file_format/v1/graphs/v1_graph_output.dtg.h",
   "utils/nonnegative_int/nonnegative_int.h",
 ]
 
@@ -35,3 +36,7 @@ type = "std::vector<::FlexFlow::nonnegative_int>"
 [[fields]]
 name = "edges"
 type = "std::unordered_set<::FlexFlow::V1GraphEdge<SlotName>>"
+
+[[fields]]
+name = "outputs"
+type = "std::unordered_set<::FlexFlow::V1GraphOutput<SlotName>>"

lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h

@@ -1,14 +1,22 @@
 #ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_GRAPHS_V1_KWARG_DATAFLOW_GRAPH_H
 #define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_GRAPHS_V1_KWARG_DATAFLOW_GRAPH_H
 
+#include "pcg/file_format/v1/graphs/v1_graph_edge.dtg.h"
+#include "pcg/file_format/v1/graphs/v1_graph_output.dtg.h"
 #include "pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.dtg.h"
 #include "utils/bidict/algorithms/bidict_from_enumerating.h"
 #include "utils/containers/enumerate.h"
+#include "utils/containers/generate_map.h"
 #include "utils/containers/sorted.h"
+#include "utils/containers/transform.h"
+#include "utils/containers/unordered_set_of.h"
 #include "utils/containers/values.h"
 #include "utils/graph/kwarg_dataflow_graph/algorithms/get_all_kwarg_dataflow_edges.h"
+#include "utils/graph/kwarg_dataflow_graph/algorithms/get_all_kwarg_dataflow_outputs.h"
 #include "utils/graph/kwarg_dataflow_graph/kwarg_dataflow_graph_view.h"
 #include "utils/graph/node/algorithms.h"
+#include "utils/graph/open_kwarg_dataflow_graph/algorithms/open_kwarg_dataflow_graph_data.dtg.h"
+#include "utils/graph/open_kwarg_dataflow_graph/algorithms/view_from_open_kwarg_dataflow_graph_data.h"
 #include "utils/integer_conversions.h"
 
 namespace FlexFlow {
@@ -27,20 +35,71 @@ template <typename SlotName>
 V1KwargDataflowGraph<SlotName>
     to_v1(KwargDataflowGraphView<SlotName> const &g,
           std::unordered_map<Node, nonnegative_int> const &nodes) {
-  std::unordered_set<V1GraphEdge<SlotName>> edges;
-  for (KwargDataflowEdge<SlotName> const &e : get_all_kwarg_dataflow_edges(g)) {
-    edges.insert(V1GraphEdge{nodes.at(e.src.node),
-                             e.src.slot_name,
-                             nodes.at(e.dst.node),
-                             e.dst.slot_name});
-  }
+  std::unordered_set<V1GraphEdge<SlotName>> edges =
+      transform(get_all_kwarg_dataflow_edges(g),
+                [&](KwargDataflowEdge<SlotName> const &e) {
+                  return V1GraphEdge{nodes.at(e.src.node),
+                                     e.src.slot_name,
+                                     nodes.at(e.dst.node),
+                                     e.dst.slot_name};
+                });
+
+  std::unordered_set<V1GraphOutput<SlotName>> outputs =
+      transform(get_all_kwarg_dataflow_outputs(g),
+                [&](KwargDataflowOutput<SlotName> const &o) {
+                  return V1GraphOutput{nodes.at(o.node), o.slot_name};
+                });
 
   return V1KwargDataflowGraph<SlotName>{
       sorted(values(nodes)),
       edges,
+      outputs,
   };
 }
 
+template <typename SlotName>
+std::pair<KwargDataflowGraphView<SlotName>,
+          std::unordered_map<nonnegative_int, Node>>
+    from_v1_including_node_numbering(V1KwargDataflowGraph<SlotName> const &v1) {
+  std::unordered_map<nonnegative_int, Node> node_map =
+      generate_map(v1.nodes, [](nonnegative_int n) {
+        return Node{n.size_t_from_nonnegative_int()};
+      });
+  std::unordered_set<Node> node_set = unordered_set_of(values(node_map));
+
+  std::unordered_set<OpenKwargDataflowEdge<int, SlotName>> edges =
+      transform(v1.edges, [](V1GraphEdge<SlotName> const &e) {
+        Node srcNode = Node{e.srcNode.size_t_from_nonnegative_int()};
+        Node dstNode = Node{e.dstNode.size_t_from_nonnegative_int()};
+        return OpenKwargDataflowEdge<int, SlotName>{KwargDataflowEdge<SlotName>{
+            /*src=*/KwargDataflowOutput<SlotName>{srcNode, e.srcSlot},
+            /*dst=*/KwargDataflowInput<SlotName>{dstNode, e.dstSlot},
+        }};
+      });
+
+  std::unordered_set<KwargDataflowOutput<SlotName>> outputs =
+      transform(v1.outputs, [](V1GraphOutput<SlotName> const &o) {
+        Node n = Node{o.node.size_t_from_nonnegative_int()};
+        return KwargDataflowOutput<SlotName>{n, o.slot_name};
+      });
+
+  OpenKwargDataflowGraphData<int, SlotName> graph_data =
+      OpenKwargDataflowGraphData<int, SlotName>{
+          /*nodes=*/node_set,
+          /*edges=*/edges,
+          /*inputs=*/{},
+          /*outputs=*/outputs,
+      };
+  return std::pair{view_from_open_kwarg_dataflow_graph_data(graph_data),
+                   node_map};
+}
+
+template <typename SlotName>
+KwargDataflowGraphView<SlotName>
+    from_v1(V1KwargDataflowGraph<SlotName> const &v1) {
+  return from_v1_including_node_numbering(v1).first;
+}
+
 } // namespace FlexFlow
 
 #endif

lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.dtg.toml

@@ -36,7 +36,7 @@ type = "std::unordered_map<::FlexFlow::nonnegative_int, NodeLabel>"
 
 [[fields]]
 name = "output_labels"
-type = "std::unordered_map<::FlexFlow::nonnegative_int, std::unordered_map<SlotName, OutputLabel>>"
+type = "std::unordered_map<::FlexFlow::V1GraphOutput<SlotName>, OutputLabel>"
 
 [[fields]]
 name = "graph"

lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h

@@ -4,9 +4,12 @@
 #include "pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h"
 #include "pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.dtg.h"
 #include "utils/bidict/algorithms/bidict_from_enumerating.h"
+#include "utils/containers/map_keys.h"
 #include "utils/containers/map_values.h"
 #include "utils/containers/transform.h"
-#include "utils/graph/kwarg_dataflow_graph/algorithms/get_outgoing_kwarg_dataflow_outputs_for_node.h"
+#include "utils/containers/unordered_map_from_pairs.h"
+#include "utils/graph/kwarg_dataflow_graph/algorithms/get_all_kwarg_dataflow_outputs.h"
+#include "utils/graph/labelled_kwarg_dataflow_graph/algorithms/kwarg_dataflow_graph_view_with_labelling.h"
 #include "utils/graph/labelled_kwarg_dataflow_graph/labelled_kwarg_dataflow_graph_view.h"
 #include "utils/graph/node/algorithms.h"
 
@@ -26,16 +29,13 @@ std::pair<V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName>,
   std::unordered_map<nonnegative_int, NodeLabel> node_labels = map_values(
       nodes.as_unordered_map(), [&](Node const &n) { return g.at(n); });
 
-  std::unordered_map<nonnegative_int, std::unordered_map<SlotName, OutputLabel>>
-      output_labels = map_values(
-          nodes.as_unordered_map(),
-          [&](Node const &n) -> std::unordered_map<SlotName, OutputLabel> {
-            return map_values(
-                get_outgoing_kwarg_dataflow_outputs_for_node(g, n),
-                [&](KwargDataflowOutput<SlotName> const &o) {
-                  return g.at(o);
-                });
-          });
+  std::unordered_map<V1GraphOutput<SlotName>, OutputLabel> output_labels =
+      unordered_map_from_pairs(transform(
+          get_all_kwarg_dataflow_outputs(g),
+          [&](KwargDataflowOutput<SlotName> const &o) {
+            return std::pair{V1GraphOutput{nodes.at_r(o.node), o.slot_name},
+                             g.at(o)};
+          }));
 
   return {
       V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName>{
@@ -50,6 +50,33 @@ V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> to_v1(
   return to_v1_including_node_numbering(g).first;
 }
 
+template <typename NodeLabel, typename OutputLabel, typename SlotName>
+std::pair<LabelledKwargDataflowGraphView<NodeLabel, OutputLabel, SlotName>,
+          std::unordered_map<nonnegative_int, Node>>
+    from_v1_including_node_numbering(
+        V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> const
+            &v1) {
+  auto [graph_view, node_map] = from_v1_including_node_numbering(v1.graph);
+
+  std::unordered_map<Node, NodeLabel> node_labels = map_keys(
+      v1.node_labels, [&](nonnegative_int n) { return node_map.at(n); });
+
+  std::unordered_map<KwargDataflowOutput<SlotName>, OutputLabel> value_labels =
+      map_keys(v1.output_labels, [&](V1GraphOutput<SlotName> const &o) {
+        return KwargDataflowOutput<SlotName>{node_map.at(o.node), o.slot_name};
+      });
+
+  return std::pair{kwarg_dataflow_graph_view_with_labelling(
+                       graph_view, node_labels, value_labels),
+                   node_map};
+}
+
+template <typename NodeLabel, typename OutputLabel, typename SlotName>
+LabelledKwargDataflowGraphView<NodeLabel, OutputLabel, SlotName> from_v1(
+    V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> const &v1) {
+  return from_v1_including_node_numbering(v1).first;
+}
+
 } // namespace FlexFlow
 
 #endif

lib/pcg/include/pcg/file_format/v1/v1_computation_graph.h

@@ -12,6 +12,8 @@ V1ComputationGraph to_v1(ComputationGraph const &);
 std::pair<V1ComputationGraph, bidict<nonnegative_int, layer_guid_t>>
     to_v1_including_node_numbering(ComputationGraph const &);
 
+ComputationGraph from_v1(V1ComputationGraph const &);
+
 } // namespace FlexFlow
 
 #endif

lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.dtg.toml

@@ -0,0 +1,19 @@
+namespace = "FlexFlow"
+name = "V1MappedOperatorTaskGroup"
+type = "struct"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+  "json",
+]
+
+includes = [
+  "pcg/machine_space_coordinate.dtg.h",
+  "pcg/mapped_parallel_computation_graph/operator_atomic_task_shard_binding.dtg.h",
+  "utils/bidict/bidict.h",
+]
+
+[[fields]]
+name = "shard_bindings"
+type = "::FlexFlow::bidict<::FlexFlow::MachineSpaceCoordinate, ::FlexFlow::OperatorAtomicTaskShardBinding>"

lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h

@@ -0,0 +1,14 @@
+#ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_V1_MAPPED_OPERATOR_TASK_GROUP_H
+#define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_V1_MAPPED_OPERATOR_TASK_GROUP_H
+
+#include "pcg/file_format/v1/v1_mapped_operator_task_group.dtg.h"
+#include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h"
+
+namespace FlexFlow {
+
+V1MappedOperatorTaskGroup to_v1(MappedOperatorTaskGroup const &);
+MappedOperatorTaskGroup from_v1(V1MappedOperatorTaskGroup const &);
+
+} // namespace FlexFlow
+
+#endif

lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.dtg.toml

@@ -0,0 +1,29 @@
+namespace = "FlexFlow"
+name = "V1MappedParallelComputationGraph"
+type = "struct"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+  "json",
+]
+
+includes = [
+  "<unordered_map>",
+  "pcg/file_format/v1/v1_parallel_computation_graph.dtg.h",
+  "pcg/file_format/v1/v1_mapped_operator_task_group.dtg.h",
+  "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
+]
+
+src_includes = [
+  "utils/hash/unordered_map.h",
+  "utils/fmt/unordered_map.h",
+]
+
+[[fields]]
+name = "pcg"
+type = "::FlexFlow::V1ParallelComputationGraph"
+
+[[fields]]
+name = "mapped_tasks"
+type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, ::FlexFlow::V1MappedOperatorTaskGroup>"

lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h

@@ -0,0 +1,15 @@
+#ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_V1_MAPPED_PARALLEL_COMPUTATION_GRAPH_H
+#define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_V1_MAPPED_PARALLEL_COMPUTATION_GRAPH_H
+
+#include "pcg/file_format/v1/v1_mapped_parallel_computation_graph.dtg.h"
+#include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.dtg.h"
+
+namespace FlexFlow {
+
+V1MappedParallelComputationGraph to_v1(MappedParallelComputationGraph const &);
+MappedParallelComputationGraph
+    from_v1(V1MappedParallelComputationGraph const &);
+
+} // namespace FlexFlow
+
+#endif

lib/pcg/include/pcg/file_format/v1/v1_parallel_computation_graph.h

@@ -8,6 +8,8 @@ namespace FlexFlow {
 
 V1ParallelComputationGraph to_v1(ParallelComputationGraph const &);
 
+ParallelComputationGraph from_v1(V1ParallelComputationGraph const &);
+
 } // namespace FlexFlow
 
 #endif

lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc

@@ -12,4 +12,11 @@ template V1KwargDataflowGraph<SlotName>
     to_v1(KwargDataflowGraphView<SlotName> const &,
           std::unordered_map<Node, nonnegative_int> const &);
 
+template std::pair<KwargDataflowGraphView<SlotName>,
+                   std::unordered_map<nonnegative_int, Node>>
+    from_v1_including_node_numbering(V1KwargDataflowGraph<SlotName> const &);
+
+template KwargDataflowGraphView<SlotName>
+    from_v1(V1KwargDataflowGraph<SlotName> const &);
+
 } // namespace FlexFlow

lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc

@@ -18,4 +18,14 @@ template std::pair<
 template V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> to_v1(
     LabelledKwargDataflowGraphView<NodeLabel, OutputLabel, SlotName> const &);
 
+template std::pair<
+    LabelledKwargDataflowGraphView<NodeLabel, OutputLabel, SlotName>,
+    std::unordered_map<nonnegative_int, Node>>
+    from_v1_including_node_numbering(
+        V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> const &);
+
+template LabelledKwargDataflowGraphView<NodeLabel, OutputLabel, SlotName>
+    from_v1(
+        V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> const &);
+
 } // namespace FlexFlow

lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc

@@ -1,6 +1,8 @@
 #include "pcg/file_format/v1/v1_computation_graph.h"
 #include "pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h"
 #include "utils/bidict/algorithms/transform_values.h"
+#include "utils/graph/instances/unordered_set_labelled_open_kwarg_dataflow_graph.h"
+#include "utils/graph/labelled_kwarg_dataflow_graph/labelled_kwarg_dataflow_graph.h"
 
 namespace FlexFlow {
 
@@ -25,4 +27,15 @@ std::pair<V1ComputationGraph, bidict<nonnegative_int, layer_guid_t>>
   return {v1_cg, v1_node_ids};
 }
 
+ComputationGraph from_v1(V1ComputationGraph const &v1) {
+  return ComputationGraph{
+      LabelledKwargDataflowGraph<LayerAttrs, TensorAttrs, TensorSlotName>::
+          create_copy_of<
+              UnorderedSetLabelledOpenKwargDataflowGraph<LayerAttrs,
+                                                         TensorAttrs,
+                                                         int,
+                                                         TensorSlotName>>(
+              from_v1(v1.raw_graph))};
+}
+
 } // namespace FlexFlow

lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc

@@ -0,0 +1,13 @@
+#include "pcg/file_format/v1/v1_mapped_operator_task_group.h"
+
+namespace FlexFlow {
+
+V1MappedOperatorTaskGroup to_v1(MappedOperatorTaskGroup const &g) {
+  return V1MappedOperatorTaskGroup{g.get_shard_bindings()};
+}
+
+MappedOperatorTaskGroup from_v1(V1MappedOperatorTaskGroup const &v1) {
+  return MappedOperatorTaskGroup{v1.shard_bindings};
+}
+
+} // namespace FlexFlow