Add v1 file format for MappedParallelComputationGraph and from_v1 conversion functions

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

@@ -1,14 +1,20 @@
 #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_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/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 {
@@ -41,6 +47,43 @@ V1KwargDataflowGraph<SlotName>
   };
 }
 
+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},
+        }};
+      });
+
+  OpenKwargDataflowGraphData<int, SlotName> graph_data =
+      OpenKwargDataflowGraphData<int, SlotName>{
+          /*nodes=*/node_set,
+          /*edges=*/edges,
+          /*inputs=*/{},
+          /*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.h

@@ -4,9 +4,11 @@
 #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/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"
 
@@ -50,6 +52,29 @@ 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;
+
+  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

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

@@ -0,0 +1,26 @@
+#include "pcg/file_format/v1/v1_mapped_parallel_computation_graph.h"
+#include "pcg/file_format/v1/v1_mapped_operator_task_group.h"
+#include "pcg/file_format/v1/v1_parallel_computation_graph.h"
+#include "utils/containers/map_values.h"
+
+namespace FlexFlow {
+
+V1MappedParallelComputationGraph
+    to_v1(MappedParallelComputationGraph const &mpcg) {
+  return V1MappedParallelComputationGraph{
+      to_v1(mpcg.pcg),
+      map_values(mpcg.mapped_tasks,
+                 [](MappedOperatorTaskGroup const &g) { return to_v1(g); }),
+  };
+}
+
+MappedParallelComputationGraph
+    from_v1(V1MappedParallelComputationGraph const &v1) {
+  return MappedParallelComputationGraph{
+      from_v1(v1.pcg),
+      map_values(v1.mapped_tasks,
+                 [](V1MappedOperatorTaskGroup const &g) { return from_v1(g); }),
+  };
+}
+
+} // namespace FlexFlow

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

@@ -1,5 +1,7 @@
 #include "pcg/file_format/v1/v1_parallel_computation_graph.h"
 #include "pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.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 {
 
@@ -10,4 +12,17 @@ V1ParallelComputationGraph to_v1(ParallelComputationGraph const &g) {
   };
 }
 
+ParallelComputationGraph from_v1(V1ParallelComputationGraph const &v1) {
+  return ParallelComputationGraph{
+      LabelledKwargDataflowGraph<ParallelLayerAttrs,
+                                 ParallelTensorAttrs,
+                                 TensorSlotName>::
+          create_copy_of<
+              UnorderedSetLabelledOpenKwargDataflowGraph<ParallelLayerAttrs,
+                                                         ParallelTensorAttrs,
+                                                         int,
+                                                         TensorSlotName>>(
+              from_v1(v1.raw_graph))};
+}
+
 } // namespace FlexFlow

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

@@ -1,6 +1,8 @@
 #include "pcg/file_format/v1/v1_computation_graph.h"
+#include "pcg/computation_graph.h"
 #include "pcg/computation_graph_builder.h"
 #include <doctest/doctest.h>
+#include <nlohmann/json.hpp>
 
 using namespace ::FlexFlow;
 
@@ -25,6 +27,14 @@ TEST_SUITE(FF_TEST_SUITE) {
     }();
 
     V1ComputationGraph v1_cg = to_v1(cg);
-    nlohmann::json j = v1_cg;
+
+    SUBCASE("serializes to JSON") {
+      nlohmann::json j = v1_cg;
+    }
+
+    SUBCASE("round-trips via from_v1") {
+      ComputationGraph result = from_v1(v1_cg);
+      CHECK(computation_graphs_are_isomorphic(cg, result));
+    }
   }
 }