Add v1 file format for MappedParallelComputationGraph and from_v1 conversion functions

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

@@ -6,7 +6,14 @@
 #include "utils/bidict/algorithms/bidict_from_enumerating.h"
 #include "utils/containers/map_values.h"
 #include "utils/containers/transform.h"
+#include "utils/graph/digraph/algorithms/get_topological_ordering.h"
+#include "utils/graph/digraph/digraph.h"
+#include "utils/graph/digraph/directed_edge.dtg.h"
+#include "utils/graph/instances/adjacency_digraph.h"
+#include "utils/graph/instances/unordered_set_labelled_open_kwarg_dataflow_graph.h"
 #include "utils/graph/kwarg_dataflow_graph/algorithms/get_outgoing_kwarg_dataflow_outputs_for_node.h"
+#include "utils/graph/kwarg_dataflow_graph/kwarg_node_added_result.dtg.h"
+#include "utils/graph/labelled_kwarg_dataflow_graph/labelled_kwarg_dataflow_graph.h"
 #include "utils/graph/labelled_kwarg_dataflow_graph/labelled_kwarg_dataflow_graph_view.h"
 #include "utils/graph/node/algorithms.h"
 
@@ -50,6 +57,56 @@ V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> to_v1(
   return to_v1_including_node_numbering(g).first;
 }
 
+template <typename NodeLabel, typename OutputLabel, typename SlotName>
+LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> from_v1(
+    V1LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName> const &v1) {
+  // Build incoming-edge map
+  std::unordered_map<nonnegative_int, std::vector<V1GraphEdge<SlotName>>>
+      incoming;
+  for (nonnegative_int const &n : v1.graph.nodes) {
+    incoming[n] = {};
+  }
+  for (V1GraphEdge<SlotName> const &e : v1.graph.edges) {
+    incoming[e.dstNode].push_back(e);
+  }
+
+  // Build a DiGraph with V1 indices as Node raw_uids to get topological order
+  DiGraph dg = DiGraph::create<AdjacencyDiGraph>();
+  for (nonnegative_int const &n : v1.graph.nodes) {
+    dg.add_node_unsafe(Node{static_cast<size_t>(n.unwrap_nonnegative())});
+  }
+  for (V1GraphEdge<SlotName> const &e : v1.graph.edges) {
+    dg.add_edge(DirectedEdge{
+        Node{static_cast<size_t>(e.srcNode.unwrap_nonnegative())},
+        Node{static_cast<size_t>(e.dstNode.unwrap_nonnegative())}});
+  }
+
+  auto g = LabelledKwargDataflowGraph<NodeLabel, OutputLabel, SlotName>::
+      template create<UnorderedSetLabelledOpenKwargDataflowGraph<NodeLabel,
+                                                                 OutputLabel,
+                                                                 int,
+                                                                 SlotName>>();
+
+  std::unordered_map<nonnegative_int, Node> node_map;
+  for (Node const &topo_node : get_topological_ordering(dg)) {
+    nonnegative_int v1_idx{static_cast<size_t>(topo_node.raw_uid)};
+
+    std::unordered_map<SlotName, KwargDataflowOutput<SlotName>> inputs;
+    for (V1GraphEdge<SlotName> const &e : incoming.at(v1_idx)) {
+      inputs.emplace(
+          e.dstSlot,
+          KwargDataflowOutput<SlotName>{node_map.at(e.srcNode), e.srcSlot});
+    }
+
+    KwargNodeAddedResult<SlotName> result = g.add_node(
+        v1.node_labels.at(v1_idx), inputs, v1.output_labels.at(v1_idx));
+
+    node_map.emplace(v1_idx, result.node);
+  }
+
+  return g;
+}
+
 } // namespace FlexFlow
 
 #endif

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

@@ -8,6 +8,7 @@
 namespace FlexFlow {
 
 V1ComputationGraph to_v1(ComputationGraph const &);
+ComputationGraph from_v1(V1ComputationGraph const &);
 
 std::pair<V1ComputationGraph, bidict<nonnegative_int, layer_guid_t>>
     to_v1_including_node_numbering(ComputationGraph const &);

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

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

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

@@ -6,6 +6,7 @@ features = [
   "ord",
   "hash",
   "fmt",
+  "json",
 ]
 
 includes = [

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

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

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

@@ -10,6 +10,12 @@ V1ComputationGraph to_v1(ComputationGraph const &g) {
   };
 }
 
+ComputationGraph from_v1(V1ComputationGraph const &v1) {
+  return ComputationGraph{
+      from_v1(v1.raw_graph),
+  };
+}
+
 std::pair<V1ComputationGraph, bidict<nonnegative_int, layer_guid_t>>
     to_v1_including_node_numbering(ComputationGraph const &cg) {
   std::pair<

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

@@ -10,4 +10,10 @@ V1ParallelComputationGraph to_v1(ParallelComputationGraph const &g) {
   };
 }
 
+ParallelComputationGraph from_v1(V1ParallelComputationGraph const &v1) {
+  return ParallelComputationGraph{
+      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));
+    }
   }
 }

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

@@ -0,0 +1,81 @@
+#include "pcg/file_format/v1/v1_mapped_parallel_computation_graph.h"
+#include "op-attrs/parallel_tensor_space_coordinate.dtg.h"
+#include "op-attrs/tensor_slot_name.dtg.h"
+#include "pcg/device_type.dtg.h"
+#include "pcg/file_format/v1/v1_mapped_operator_task_group.h"
+#include "pcg/machine_space_coordinate.dtg.h"
+#include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h"
+#include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.dtg.h"
+#include "pcg/mapped_parallel_computation_graph/operator_atomic_task_shard_binding.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
+#include "pcg/parallel_computation_graph/parallel_layer_added_result.dtg.h"
+#include "utils/bidict/bidict.h"
+#include <doctest/doctest.h>
+#include <nlohmann/json.hpp>
+
+using namespace ::FlexFlow;
+
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("V1MappedParallelComputationGraph") {
+    ParallelComputationGraph pcg = empty_parallel_computation_graph();
+
+    TensorShape input_shape = TensorShape{
+        TensorDims{
+            FFOrdered{
+                12_p,
+                16_p,
+            },
+        },
+        DataType::FLOAT,
+    };
+
+    ParallelLayerAddedResult result = pcg_add_input_layer(pcg, input_shape);
+    parallel_layer_guid_t layer = result.parallel_layer;
+
+    MachineSpaceCoordinate coord = MachineSpaceCoordinate{
+        /*node_idx=*/0_n,
+        /*device_idx=*/0_n,
+        /*device_type=*/DeviceType::GPU,
+    };
+
+    OperatorAtomicTaskShardBinding binding = OperatorAtomicTaskShardBinding{
+        /*tensor_coords=*/{
+            {
+                TensorSlotName::OUTPUT,
+                ParallelTensorSpaceCoordinate{
+                    /*sum_component=*/0_n,
+                    /*discard_copy_component=*/0_n,
+                    /*shard_components=*/FFOrdered<nonnegative_int>{0_n, 0_n},
+                },
+            },
+        },
+    };
+
+    MappedOperatorTaskGroup task_group = MappedOperatorTaskGroup{
+        bidict<MachineSpaceCoordinate, OperatorAtomicTaskShardBinding>{
+            {coord, binding},
+        },
+    };
+
+    MappedParallelComputationGraph mpcg = MappedParallelComputationGraph{
+        /*pcg=*/pcg,
+        /*mapped_tasks=*/{{layer, task_group}},
+    };
+
+    V1MappedParallelComputationGraph v1_mpcg = to_v1(mpcg);
+
+    SUBCASE("serializes to JSON") {
+      nlohmann::json j = v1_mpcg;
+    }
+
+    SUBCASE("MappedOperatorTaskGroup round-trips via from_v1") {
+      MappedOperatorTaskGroup result = from_v1(to_v1(task_group));
+      CHECK(result == task_group);
+    }
+
+    SUBCASE("MappedParallelComputationGraph round-trips via from_v1") {
+      MappedParallelComputationGraph result = from_v1(v1_mpcg);
+      CHECK(pcgs_are_isomorphic(mpcg.pcg, result.pcg));
+    }
+  }
+}

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

@@ -1,6 +1,8 @@
 #include "pcg/file_format/v1/v1_parallel_computation_graph.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph_builder.h"
 #include <doctest/doctest.h>
+#include <nlohmann/json.hpp>
 
 using namespace ::FlexFlow;
 
@@ -29,6 +31,14 @@ TEST_SUITE(FF_TEST_SUITE) {
     }();
 
     V1ParallelComputationGraph v1_pcg = to_v1(pcg);
-    nlohmann::json j = v1_pcg;
+
+    SUBCASE("serializes to JSON") {
+      nlohmann::json j = v1_pcg;
+    }
+
+    SUBCASE("round-trips via from_v1") {
+      ParallelComputationGraph result = from_v1(v1_pcg);
+      CHECK(pcgs_are_isomorphic(pcg, result));
+    }
   }
 }