From 7602aaf437b80a21c2279544b923e54c26d9cf7b Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 6 Mar 2026 15:55:50 -0800 Subject: [PATCH 01/11] Add V1MappedParallelComputationGraph. --- .../v1/v1_mapped_operator_task_group.dtg.toml | 19 +++++ .../v1/v1_mapped_operator_task_group.h | 13 ++++ ...mapped_parallel_computation_graph.dtg.toml | 29 ++++++++ .../v1/v1_mapped_parallel_computation_graph.h | 13 ++++ .../v1/v1_mapped_operator_task_group.cc | 9 +++ .../v1_mapped_parallel_computation_graph.cc | 17 +++++ .../v1_mapped_parallel_computation_graph.cc | 69 +++++++++++++++++++ 7 files changed, 169 insertions(+) create mode 100644 lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.dtg.toml create mode 100644 lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h create mode 100644 lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.dtg.toml create mode 100644 lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h create mode 100644 lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc create mode 100644 lib/pcg/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc create mode 100644 lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc diff --git a/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.dtg.toml b/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.dtg.toml new file mode 100644 index 0000000000..2e4300745d --- /dev/null +++ b/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>" diff --git a/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h b/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h new file mode 100644 index 0000000000..7c1788b9d0 --- /dev/null +++ b/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h @@ -0,0 +1,13 @@ +#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 &); + +} // namespace FlexFlow + +#endif diff --git a/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.dtg.toml b/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.dtg.toml new file mode 100644 index 0000000000..8dc336e4ea --- /dev/null +++ b/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 = [ + "", + "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>" diff --git a/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h b/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h new file mode 100644 index 0000000000..5b9d18ccc4 --- /dev/null +++ b/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h @@ -0,0 +1,13 @@ +#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 &); + +} // namespace FlexFlow + +#endif diff --git a/lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc b/lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc new file mode 100644 index 0000000000..480ea7197a --- /dev/null +++ b/lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc @@ -0,0 +1,9 @@ +#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()}; +} + +} // namespace FlexFlow diff --git a/lib/pcg/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc new file mode 100644 index 0000000000..96429e06a6 --- /dev/null +++ b/lib/pcg/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc @@ -0,0 +1,17 @@ +#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); }), + }; +} + +} // namespace FlexFlow diff --git a/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc b/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc new file mode 100644 index 0000000000..dd2cdb35ed --- /dev/null +++ b/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc @@ -0,0 +1,69 @@ +#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/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 +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("V1MappedParallelComputationGraph") { + MappedParallelComputationGraph mpcg = [] { + 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{0_n, 0_n}, + }, + }, + }, + }; + + MappedOperatorTaskGroup task_group = MappedOperatorTaskGroup{ + bidict{ + {coord, binding}, + }, + }; + + return MappedParallelComputationGraph{ + /*pcg=*/pcg, + /*mapped_tasks=*/{{layer, task_group}}, + }; + }(); + + V1MappedParallelComputationGraph v1_mpcg = to_v1(mpcg); + nlohmann::json j = v1_mpcg; + } +} From 0ecdb3ed8eb40f7083975cc7421d15ec8d380b28 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 6 Mar 2026 16:44:43 -0800 Subject: [PATCH 02/11] Add from_v1. --- .../graphs/v1_labelled_kwarg_dataflow_graph.h | 57 +++++++++++ .../pcg/file_format/v1/v1_computation_graph.h | 1 + .../v1/v1_mapped_operator_task_group.h | 1 + .../v1/v1_mapped_parallel_computation_graph.h | 2 + .../v1/v1_parallel_computation_graph.h | 1 + .../v1_labelled_kwarg_dataflow_graph.cc | 3 + .../file_format/v1/v1_computation_graph.cc | 6 ++ .../v1/v1_mapped_operator_task_group.cc | 4 + .../v1_mapped_parallel_computation_graph.cc | 9 ++ .../v1/v1_parallel_computation_graph.cc | 6 ++ .../file_format/v1/v1_computation_graph.cc | 12 ++- .../v1_mapped_parallel_computation_graph.cc | 94 +++++++++++-------- .../v1/v1_parallel_computation_graph.cc | 12 ++- 13 files changed, 165 insertions(+), 43 deletions(-) diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h index dbe660c3a6..1c1bc70c88 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h +++ b/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 to_v1( return to_v1_including_node_numbering(g).first; } +template +LabelledKwargDataflowGraph from_v1( + V1LabelledKwargDataflowGraph const &v1) { + // Build incoming-edge map + std::unordered_map>> + incoming; + for (nonnegative_int const &n : v1.graph.nodes) { + incoming[n] = {}; + } + for (V1GraphEdge 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(); + for (nonnegative_int const &n : v1.graph.nodes) { + dg.add_node_unsafe(Node{static_cast(n.unwrap_nonnegative())}); + } + for (V1GraphEdge const &e : v1.graph.edges) { + dg.add_edge(DirectedEdge{ + Node{static_cast(e.srcNode.unwrap_nonnegative())}, + Node{static_cast(e.dstNode.unwrap_nonnegative())}}); + } + + auto g = LabelledKwargDataflowGraph:: + template create>(); + + std::unordered_map node_map; + for (Node const &topo_node : get_topological_ordering(dg)) { + nonnegative_int v1_idx{static_cast(topo_node.raw_uid)}; + + std::unordered_map> inputs; + for (V1GraphEdge const &e : incoming.at(v1_idx)) { + inputs.emplace( + e.dstSlot, + KwargDataflowOutput{node_map.at(e.srcNode), e.srcSlot}); + } + + KwargNodeAddedResult 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 diff --git a/lib/pcg/include/pcg/file_format/v1/v1_computation_graph.h b/lib/pcg/include/pcg/file_format/v1/v1_computation_graph.h index c0e9966425..8b6128d603 100644 --- a/lib/pcg/include/pcg/file_format/v1/v1_computation_graph.h +++ b/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> to_v1_including_node_numbering(ComputationGraph const &); diff --git a/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h b/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h index 7c1788b9d0..8e386e156f 100644 --- a/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h +++ b/lib/pcg/include/pcg/file_format/v1/v1_mapped_operator_task_group.h @@ -7,6 +7,7 @@ namespace FlexFlow { V1MappedOperatorTaskGroup to_v1(MappedOperatorTaskGroup const &); +MappedOperatorTaskGroup from_v1(V1MappedOperatorTaskGroup const &); } // namespace FlexFlow diff --git a/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h b/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h index 5b9d18ccc4..f78efc4591 100644 --- a/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/v1_mapped_parallel_computation_graph.h @@ -7,6 +7,8 @@ namespace FlexFlow { V1MappedParallelComputationGraph to_v1(MappedParallelComputationGraph const &); +MappedParallelComputationGraph + from_v1(V1MappedParallelComputationGraph const &); } // namespace FlexFlow diff --git a/lib/pcg/include/pcg/file_format/v1/v1_parallel_computation_graph.h b/lib/pcg/include/pcg/file_format/v1/v1_parallel_computation_graph.h index aceb59f5af..d481096d49 100644 --- a/lib/pcg/include/pcg/file_format/v1/v1_parallel_computation_graph.h +++ b/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 diff --git a/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc b/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc index 4e7b9b651f..a2953e8fb3 100644 --- a/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc @@ -18,4 +18,7 @@ template std::pair< template V1LabelledKwargDataflowGraph to_v1( LabelledKwargDataflowGraphView const &); +template LabelledKwargDataflowGraph from_v1( + V1LabelledKwargDataflowGraph const &); + } // namespace FlexFlow diff --git a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc index 852ca73a36..71fc105711 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc +++ b/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> to_v1_including_node_numbering(ComputationGraph const &cg) { std::pair< diff --git a/lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc b/lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc index 480ea7197a..465dd01fb6 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_mapped_operator_task_group.cc @@ -6,4 +6,8 @@ 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 diff --git a/lib/pcg/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc index 96429e06a6..0236a8834c 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc @@ -14,4 +14,13 @@ V1MappedParallelComputationGraph }; } +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 diff --git a/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc index e14d15d66a..a169abe4c1 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc +++ b/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 diff --git a/lib/pcg/test/src/pcg/file_format/v1/v1_computation_graph.cc b/lib/pcg/test/src/pcg/file_format/v1/v1_computation_graph.cc index 7af3f648d9..2ae643bd0f 100644 --- a/lib/pcg/test/src/pcg/file_format/v1/v1_computation_graph.cc +++ b/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 +#include 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)); + } } } diff --git a/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc b/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc index dd2cdb35ed..78da5430b7 100644 --- a/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc +++ b/lib/pcg/test/src/pcg/file_format/v1/v1_mapped_parallel_computation_graph.cc @@ -2,6 +2,7 @@ #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" @@ -16,54 +17,65 @@ using namespace ::FlexFlow; TEST_SUITE(FF_TEST_SUITE) { TEST_CASE("V1MappedParallelComputationGraph") { - MappedParallelComputationGraph mpcg = [] { - ParallelComputationGraph pcg = empty_parallel_computation_graph(); + ParallelComputationGraph pcg = empty_parallel_computation_graph(); - TensorShape input_shape = TensorShape{ - TensorDims{ - FFOrdered{ - 12_p, - 16_p, - }, - }, - DataType::FLOAT, - }; + 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; + 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, - }; + 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{0_n, 0_n}, - }, - }, - }, - }; + OperatorAtomicTaskShardBinding binding = OperatorAtomicTaskShardBinding{ + /*tensor_coords=*/{ + { + TensorSlotName::OUTPUT, + ParallelTensorSpaceCoordinate{ + /*sum_component=*/0_n, + /*discard_copy_component=*/0_n, + /*shard_components=*/FFOrdered{0_n, 0_n}, + }, + }, + }, + }; - MappedOperatorTaskGroup task_group = MappedOperatorTaskGroup{ - bidict{ - {coord, binding}, - }, - }; + MappedOperatorTaskGroup task_group = MappedOperatorTaskGroup{ + bidict{ + {coord, binding}, + }, + }; - return MappedParallelComputationGraph{ - /*pcg=*/pcg, - /*mapped_tasks=*/{{layer, task_group}}, - }; - }(); + MappedParallelComputationGraph mpcg = MappedParallelComputationGraph{ + /*pcg=*/pcg, + /*mapped_tasks=*/{{layer, task_group}}, + }; V1MappedParallelComputationGraph v1_mpcg = to_v1(mpcg); - nlohmann::json j = 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)); + } } } diff --git a/lib/pcg/test/src/pcg/file_format/v1/v1_parallel_computation_graph.cc b/lib/pcg/test/src/pcg/file_format/v1/v1_parallel_computation_graph.cc index ec6a4ab006..033626ab5c 100644 --- a/lib/pcg/test/src/pcg/file_format/v1/v1_parallel_computation_graph.cc +++ b/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 +#include 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)); + } } } From 337f77d1e0b0167f3d55cc70abd546ed02cb2ea4 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Mon, 23 Mar 2026 11:29:14 -0700 Subject: [PATCH 03/11] Refactoring for PR feedback. --- .../graphs/v1_labelled_kwarg_dataflow_graph.h | 9 ++++----- .../utils/nonnegative_int/nonnegative_int.h | 4 ++++ .../utils/nonnegative_int/nonnegative_int.cc | 12 ++++++++++++ .../utils/nonnegative_int/nonnegative_int.cc | 18 ++++++++++++++++++ 4 files changed, 38 insertions(+), 5 deletions(-) diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h index 1c1bc70c88..eb2c78aa09 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h @@ -73,12 +73,11 @@ LabelledKwargDataflowGraph from_v1( // Build a DiGraph with V1 indices as Node raw_uids to get topological order DiGraph dg = DiGraph::create(); for (nonnegative_int const &n : v1.graph.nodes) { - dg.add_node_unsafe(Node{static_cast(n.unwrap_nonnegative())}); + dg.add_node_unsafe(Node{n.size_t_from_nonnegative_int()}); } for (V1GraphEdge const &e : v1.graph.edges) { - dg.add_edge(DirectedEdge{ - Node{static_cast(e.srcNode.unwrap_nonnegative())}, - Node{static_cast(e.dstNode.unwrap_nonnegative())}}); + dg.add_edge(DirectedEdge{Node{e.srcNode.size_t_from_nonnegative_int()}, + Node{e.dstNode.size_t_from_nonnegative_int()}}); } auto g = LabelledKwargDataflowGraph:: @@ -89,7 +88,7 @@ LabelledKwargDataflowGraph from_v1( std::unordered_map node_map; for (Node const &topo_node : get_topological_ordering(dg)) { - nonnegative_int v1_idx{static_cast(topo_node.raw_uid)}; + nonnegative_int v1_idx{topo_node.raw_uid}; std::unordered_map> inputs; for (V1GraphEdge const &e : incoming.at(v1_idx)) { diff --git a/lib/utils/include/utils/nonnegative_int/nonnegative_int.h b/lib/utils/include/utils/nonnegative_int/nonnegative_int.h index c775cfc9ed..1972cc6786 100644 --- a/lib/utils/include/utils/nonnegative_int/nonnegative_int.h +++ b/lib/utils/include/utils/nonnegative_int/nonnegative_int.h @@ -17,6 +17,7 @@ class nonnegative_int { explicit nonnegative_int(unsigned long long int value); explicit operator int() const noexcept; + explicit operator size_t() const noexcept; bool operator<(nonnegative_int const &other) const; bool operator==(nonnegative_int const &other) const; @@ -56,6 +57,9 @@ class nonnegative_int { nonnegative_int operator%(nonnegative_int const &other) const; nonnegative_int &operator%=(nonnegative_int const &other); + int int_from_nonnegative_int() const; + size_t size_t_from_nonnegative_int() const; + friend std::ostream &operator<<(std::ostream &os, nonnegative_int const &n); friend int format_as(nonnegative_int const &); diff --git a/lib/utils/src/utils/nonnegative_int/nonnegative_int.cc b/lib/utils/src/utils/nonnegative_int/nonnegative_int.cc index 7593a8e9ec..8d1c4383a9 100644 --- a/lib/utils/src/utils/nonnegative_int/nonnegative_int.cc +++ b/lib/utils/src/utils/nonnegative_int/nonnegative_int.cc @@ -24,6 +24,10 @@ nonnegative_int::operator int() const noexcept { return this->value_; } +nonnegative_int::operator size_t() const noexcept { + return static_cast(this->value_); +} + bool nonnegative_int::operator<(nonnegative_int const &other) const { return this->value_ < other.value_; } @@ -151,6 +155,14 @@ nonnegative_int &nonnegative_int::operator%=(nonnegative_int const &other) { return *this; } +int nonnegative_int::int_from_nonnegative_int() const { + return this->value_; +} + +size_t nonnegative_int::size_t_from_nonnegative_int() const { + return static_cast(this->value_); +} + std::ostream &operator<<(std::ostream &os, nonnegative_int const &n) { os << n.value_; return os; diff --git a/lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc b/lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc index 58fb151313..8c5ecd3e2c 100644 --- a/lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc +++ b/lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc @@ -315,6 +315,24 @@ TEST_SUITE(FF_TEST_SUITE) { CHECK(result == correct); } + TEST_CASE("nonnegative_int::int_from_nonnegative_int()") { + nonnegative_int input = nonnegative_int{3}; + + int result = input.int_from_nonnegative_int(); + int correct = 3; + + CHECK(result == correct); + } + + TEST_CASE("nonnegative_int::size_t_from_nonnegative_int()") { + nonnegative_int input = nonnegative_int{3}; + + size_t result = input.size_t_from_nonnegative_int(); + size_t correct = 3; + + CHECK(result == correct); + } + TEST_CASE("adl_serializer") { SUBCASE("to_json") { nonnegative_int input = nonnegative_int{5}; From b60cc29a899a140ab040d2733059d9209528aa96 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Mon, 23 Mar 2026 11:33:40 -0700 Subject: [PATCH 04/11] Switch to insert. --- .../v1/graphs/v1_labelled_kwarg_dataflow_graph.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h index eb2c78aa09..20f61f6580 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h @@ -92,15 +92,15 @@ LabelledKwargDataflowGraph from_v1( std::unordered_map> inputs; for (V1GraphEdge const &e : incoming.at(v1_idx)) { - inputs.emplace( + inputs.insert(std::pair{ e.dstSlot, - KwargDataflowOutput{node_map.at(e.srcNode), e.srcSlot}); + KwargDataflowOutput{node_map.at(e.srcNode), e.srcSlot}}); } KwargNodeAddedResult result = g.add_node( v1.node_labels.at(v1_idx), inputs, v1.output_labels.at(v1_idx)); - node_map.emplace(v1_idx, result.node); + node_map.insert(std::pair{v1_idx, result.node}); } return g; From 1a4445e64e4e69905266bd992764a361eec4ec76 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Mon, 23 Mar 2026 14:24:38 -0700 Subject: [PATCH 05/11] Refactor to use container APIs. --- .../v1/graphs/v1_labelled_kwarg_dataflow_graph.h | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h index 20f61f6580..1d5356f783 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h @@ -6,6 +6,7 @@ #include "utils/bidict/algorithms/bidict_from_enumerating.h" #include "utils/containers/map_values.h" #include "utils/containers/transform.h" +#include "utils/containers/unordered_map_from_pairs.h" #include "utils/graph/digraph/algorithms/get_topological_ordering.h" #include "utils/graph/digraph/digraph.h" #include "utils/graph/digraph/directed_edge.dtg.h" @@ -90,12 +91,13 @@ LabelledKwargDataflowGraph from_v1( for (Node const &topo_node : get_topological_ordering(dg)) { nonnegative_int v1_idx{topo_node.raw_uid}; - std::unordered_map> inputs; - for (V1GraphEdge const &e : incoming.at(v1_idx)) { - inputs.insert(std::pair{ - e.dstSlot, - KwargDataflowOutput{node_map.at(e.srcNode), e.srcSlot}}); - } + std::unordered_map> inputs = + unordered_map_from_pairs( + transform(incoming.at(v1_idx), [&](V1GraphEdge const &e) { + return std::pair{e.dstSlot, + KwargDataflowOutput{ + node_map.at(e.srcNode), e.srcSlot}}; + })); KwargNodeAddedResult result = g.add_node( v1.node_labels.at(v1_idx), inputs, v1.output_labels.at(v1_idx)); From 3d2df8bc03211f60493242e24f9ca5e15d783aa0 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 24 Mar 2026 14:39:48 -0700 Subject: [PATCH 06/11] Implement from_v1 for KwargDataflowGraph. --- .../v1/graphs/v1_kwarg_dataflow_graph.h | 33 +++++++++++++++++++ .../v1/graphs/v1_kwarg_dataflow_graph.cc | 3 ++ 2 files changed, 36 insertions(+) diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h index a923d18ce6..beb1670d6e 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h @@ -1,14 +1,19 @@ #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/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 +46,34 @@ V1KwargDataflowGraph }; } +template +KwargDataflowGraphView + from_v1(V1KwargDataflowGraph const &v1_g) { + std::unordered_set graph_nodes = + unordered_set_of(transform(v1_g.nodes, [](nonnegative_int n) { + return Node{n.size_t_from_nonnegative_int()}; + })); + + std::unordered_set> graph_edges = + transform(v1_g.edges, [](V1GraphEdge const &e) { + Node srcNode = Node{e.srcNode.size_t_from_nonnegative_int()}; + Node dstNode = Node{e.dstNode.size_t_from_nonnegative_int()}; + return OpenKwargDataflowEdge{KwargDataflowEdge{ + /*src=*/KwargDataflowOutput{srcNode, e.srcSlot}, + /*dst=*/KwargDataflowInput{dstNode, e.dstSlot}, + }}; + }); + + OpenKwargDataflowGraphData graph_data = + OpenKwargDataflowGraphData{ + /*nodes=*/graph_nodes, + /*edges=*/graph_edges, + /*inputs=*/{}, + /*outputs=*/{}, + }; + return view_from_open_kwarg_dataflow_graph_data(graph_data); +} + } // namespace FlexFlow #endif diff --git a/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc b/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc index 9e4a46b87a..eb4801ce89 100644 --- a/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc @@ -12,4 +12,7 @@ template V1KwargDataflowGraph to_v1(KwargDataflowGraphView const &, std::unordered_map const &); +template KwargDataflowGraphView + from_v1(V1KwargDataflowGraph const &); + } // namespace FlexFlow From 25b004124e43706a8b9e71c8675516345b38a9ad Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 24 Mar 2026 14:44:44 -0700 Subject: [PATCH 07/11] Cherry pick kwarg_dataflow_graph_view_with_labelling from lockshaw/data-parallelism. --- ...kwarg_dataflow_graph_view_with_labelling.h | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 lib/utils/include/utils/graph/labelled_kwarg_dataflow_graph/algorithms/kwarg_dataflow_graph_view_with_labelling.h diff --git a/lib/utils/include/utils/graph/labelled_kwarg_dataflow_graph/algorithms/kwarg_dataflow_graph_view_with_labelling.h b/lib/utils/include/utils/graph/labelled_kwarg_dataflow_graph/algorithms/kwarg_dataflow_graph_view_with_labelling.h new file mode 100644 index 0000000000..782e63889b --- /dev/null +++ b/lib/utils/include/utils/graph/labelled_kwarg_dataflow_graph/algorithms/kwarg_dataflow_graph_view_with_labelling.h @@ -0,0 +1,73 @@ +#ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_GRAPH_LABELLED_KWARG_DATAFLOW_GRAPH_ALGORITHMS_KWARG_DATAFLOW_GRAPH_VIEW_WITH_LABELLING_H +#define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_GRAPH_LABELLED_KWARG_DATAFLOW_GRAPH_ALGORITHMS_KWARG_DATAFLOW_GRAPH_VIEW_WITH_LABELLING_H + +#include "utils/graph/kwarg_dataflow_graph/kwarg_dataflow_graph_view.h" +#include "utils/graph/labelled_kwarg_dataflow_graph/i_labelled_kwarg_dataflow_graph_view.h" +#include "utils/graph/labelled_kwarg_dataflow_graph/labelled_kwarg_dataflow_graph_view.h" + +namespace FlexFlow { + +template +struct KwargDataflowGraphLabellingWrapper final + : public ILabelledKwargDataflowGraphView { +public: + KwargDataflowGraphLabellingWrapper() = delete; + KwargDataflowGraphLabellingWrapper( + KwargDataflowGraphView const &unlabelled, + std::unordered_map const &node_labels, + std::unordered_map, OutputLabel> const + &output_labels) + : unlabelled(unlabelled), node_labels(node_labels), + output_labels(output_labels) {} + + std::unordered_set query_nodes(NodeQuery const &q) const override { + return this->unlabelled.query_nodes(q); + } + + std::unordered_set> + query_edges(KwargDataflowEdgeQuery const &q) const override { + return this->unlabelled.query_edges(q); + } + + std::unordered_set> query_outputs( + KwargDataflowOutputQuery const &q) const override { + return this->unlabelled.query_outputs(q); + } + + NodeLabel at(Node const &n) const override { + return this->node_labels.at(n); + } + + OutputLabel at(KwargDataflowOutput const &v) const override { + return this->output_labels.at(v); + } + + KwargDataflowGraphLabellingWrapper *clone() const override { + return new KwargDataflowGraphLabellingWrapper{ + this->unlabelled, + this->node_labels, + this->output_labels, + }; + } + +private: + KwargDataflowGraphView unlabelled; + std::unordered_map node_labels; + std::unordered_map, OutputLabel> output_labels; +}; + +template +LabelledKwargDataflowGraphView + kwarg_dataflow_graph_view_with_labelling( + KwargDataflowGraphView const &g, + std::unordered_map const &node_labels, + std::unordered_map, OutputLabel> const + &value_labels) { + return LabelledKwargDataflowGraphView:: + template create< + KwargDataflowGraphLabellingWrapper>( + g, node_labels, value_labels); +} +} // namespace FlexFlow + +#endif From c277800a2f4c645fdbce668dae804784302d3f58 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 24 Mar 2026 15:25:09 -0700 Subject: [PATCH 08/11] Finish from_v1 for LabelledKwargDataflowGraph. --- .../v1/graphs/v1_kwarg_dataflow_graph.h | 24 ++++--- .../graphs/v1_labelled_kwarg_dataflow_graph.h | 64 +++++-------------- .../v1/graphs/v1_kwarg_dataflow_graph.cc | 3 +- .../v1_labelled_kwarg_dataflow_graph.cc | 7 +- .../file_format/v1/v1_computation_graph.cc | 12 +++- .../v1/v1_parallel_computation_graph.cc | 15 ++++- 6 files changed, 58 insertions(+), 67 deletions(-) diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h index beb1670d6e..7a81f98b46 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h @@ -5,6 +5,7 @@ #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" @@ -47,15 +48,17 @@ V1KwargDataflowGraph } template -KwargDataflowGraphView - from_v1(V1KwargDataflowGraph const &v1_g) { - std::unordered_set graph_nodes = - unordered_set_of(transform(v1_g.nodes, [](nonnegative_int n) { +std::pair, + std::unordered_map> + from_v1(V1KwargDataflowGraph const &v1) { + std::unordered_map node_map = + generate_map(v1.nodes, [](nonnegative_int n) { return Node{n.size_t_from_nonnegative_int()}; - })); + }); + std::unordered_set node_set = unordered_set_of(values(node_map)); - std::unordered_set> graph_edges = - transform(v1_g.edges, [](V1GraphEdge const &e) { + std::unordered_set> edges = + transform(v1.edges, [](V1GraphEdge const &e) { Node srcNode = Node{e.srcNode.size_t_from_nonnegative_int()}; Node dstNode = Node{e.dstNode.size_t_from_nonnegative_int()}; return OpenKwargDataflowEdge{KwargDataflowEdge{ @@ -66,12 +69,13 @@ KwargDataflowGraphView OpenKwargDataflowGraphData graph_data = OpenKwargDataflowGraphData{ - /*nodes=*/graph_nodes, - /*edges=*/graph_edges, + /*nodes=*/node_set, + /*edges=*/edges, /*inputs=*/{}, /*outputs=*/{}, }; - return view_from_open_kwarg_dataflow_graph_data(graph_data); + return std::pair{view_from_open_kwarg_dataflow_graph_data(graph_data), + node_map}; } } // namespace FlexFlow diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h index 1d5356f783..7fc9737a70 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h @@ -4,6 +4,7 @@ #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/containers/unordered_map_from_pairs.h" @@ -14,9 +15,10 @@ #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/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" +#include "utils/nonnegative_int/nonnegative_int.h" namespace FlexFlow { @@ -59,53 +61,19 @@ V1LabelledKwargDataflowGraph to_v1( } template -LabelledKwargDataflowGraph from_v1( - V1LabelledKwargDataflowGraph const &v1) { - // Build incoming-edge map - std::unordered_map>> - incoming; - for (nonnegative_int const &n : v1.graph.nodes) { - incoming[n] = {}; - } - for (V1GraphEdge 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(); - for (nonnegative_int const &n : v1.graph.nodes) { - dg.add_node_unsafe(Node{n.size_t_from_nonnegative_int()}); - } - for (V1GraphEdge const &e : v1.graph.edges) { - dg.add_edge(DirectedEdge{Node{e.srcNode.size_t_from_nonnegative_int()}, - Node{e.dstNode.size_t_from_nonnegative_int()}}); - } - - auto g = LabelledKwargDataflowGraph:: - template create>(); - - std::unordered_map node_map; - for (Node const &topo_node : get_topological_ordering(dg)) { - nonnegative_int v1_idx{topo_node.raw_uid}; - - std::unordered_map> inputs = - unordered_map_from_pairs( - transform(incoming.at(v1_idx), [&](V1GraphEdge const &e) { - return std::pair{e.dstSlot, - KwargDataflowOutput{ - node_map.at(e.srcNode), e.srcSlot}}; - })); - - KwargNodeAddedResult result = g.add_node( - v1.node_labels.at(v1_idx), inputs, v1.output_labels.at(v1_idx)); - - node_map.insert(std::pair{v1_idx, result.node}); - } - - return g; +std::pair, + std::unordered_map> + from_v1(V1LabelledKwargDataflowGraph const + &v1) { + auto [graph_view, node_map] = from_v1(v1.graph); + + std::unordered_map node_labels = map_keys( + v1.node_labels, [&](nonnegative_int n) { return node_map.at(n); }); + std::unordered_map, OutputLabel> value_labels; + + return std::pair{kwarg_dataflow_graph_view_with_labelling( + graph_view, node_labels, value_labels), + node_map}; } } // namespace FlexFlow diff --git a/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc b/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc index eb4801ce89..6dcfe4ce91 100644 --- a/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc @@ -12,7 +12,8 @@ template V1KwargDataflowGraph to_v1(KwargDataflowGraphView const &, std::unordered_map const &); -template KwargDataflowGraphView +template std::pair, + std::unordered_map> from_v1(V1KwargDataflowGraph const &); } // namespace FlexFlow diff --git a/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc b/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc index a2953e8fb3..7c8c2a77dc 100644 --- a/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc @@ -18,7 +18,10 @@ template std::pair< template V1LabelledKwargDataflowGraph to_v1( LabelledKwargDataflowGraphView const &); -template LabelledKwargDataflowGraph from_v1( - V1LabelledKwargDataflowGraph const &); +template std::pair< + LabelledKwargDataflowGraphView, + std::unordered_map> + from_v1( + V1LabelledKwargDataflowGraph const &); } // namespace FlexFlow diff --git a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc index 71fc105711..230c48a368 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc @@ -1,6 +1,7 @@ #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/labelled_kwarg_dataflow_graph/labelled_kwarg_dataflow_graph.h" namespace FlexFlow { @@ -11,9 +12,14 @@ V1ComputationGraph to_v1(ComputationGraph const &g) { } ComputationGraph from_v1(V1ComputationGraph const &v1) { - return ComputationGraph{ - from_v1(v1.raw_graph), - }; + LabelledKwargDataflowGraph + raw_graph = + LabelledKwargDataflowGraph:: + create_copy_of>( + from_v1(v1.raw_graph).first); + return ComputationGraph{raw_graph}; } std::pair> diff --git a/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc index a169abe4c1..c82eaca960 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc @@ -1,5 +1,6 @@ #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/labelled_kwarg_dataflow_graph/labelled_kwarg_dataflow_graph.h" namespace FlexFlow { @@ -11,9 +12,17 @@ V1ParallelComputationGraph to_v1(ParallelComputationGraph const &g) { } ParallelComputationGraph from_v1(V1ParallelComputationGraph const &v1) { - return ParallelComputationGraph{ - from_v1(v1.raw_graph), - }; + LabelledKwargDataflowGraph + raw_graph = LabelledKwargDataflowGraph:: + create_copy_of>( + from_v1(v1.raw_graph).first); + return ParallelComputationGraph{raw_graph}; } } // namespace FlexFlow From 8628fd378357ec9999a65ff8a955605d014c95ef Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 24 Mar 2026 15:33:54 -0700 Subject: [PATCH 09/11] Fix build. --- .../file_format/v1/v1_computation_graph.cc | 17 +++++++------- .../v1/v1_parallel_computation_graph.cc | 22 +++++++++---------- 2 files changed, 20 insertions(+), 19 deletions(-) diff --git a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc index 230c48a368..d4cf7612bd 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc @@ -1,6 +1,7 @@ #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 { @@ -12,14 +13,14 @@ V1ComputationGraph to_v1(ComputationGraph const &g) { } ComputationGraph from_v1(V1ComputationGraph const &v1) { - LabelledKwargDataflowGraph - raw_graph = - LabelledKwargDataflowGraph:: - create_copy_of>( - from_v1(v1.raw_graph).first); - return ComputationGraph{raw_graph}; + return ComputationGraph{ + LabelledKwargDataflowGraph:: + create_copy_of< + UnorderedSetLabelledOpenKwargDataflowGraph>( + from_v1(v1.raw_graph).first)}; } std::pair> diff --git a/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc index c82eaca960..c1a394e4c5 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc @@ -1,5 +1,6 @@ #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 { @@ -12,17 +13,16 @@ V1ParallelComputationGraph to_v1(ParallelComputationGraph const &g) { } ParallelComputationGraph from_v1(V1ParallelComputationGraph const &v1) { - LabelledKwargDataflowGraph - raw_graph = LabelledKwargDataflowGraph:: - create_copy_of>( - from_v1(v1.raw_graph).first); - return ParallelComputationGraph{raw_graph}; + return ParallelComputationGraph{ + LabelledKwargDataflowGraph:: + create_copy_of< + UnorderedSetLabelledOpenKwargDataflowGraph>( + from_v1(v1.raw_graph).first)}; } } // namespace FlexFlow From 87b6a73a5aa6598d686e304a0c888237c8a5b920 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 24 Mar 2026 15:40:14 -0700 Subject: [PATCH 10/11] Various cleanup. --- .../graphs/v1_labelled_kwarg_dataflow_graph.h | 8 ------- .../pcg/file_format/v1/v1_computation_graph.h | 3 ++- .../v1/v1_parallel_computation_graph.h | 1 + .../file_format/v1/v1_computation_graph.cc | 22 +++++++++---------- 4 files changed, 14 insertions(+), 20 deletions(-) diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h index 7fc9737a70..b58de15cfd 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h @@ -7,18 +7,10 @@ #include "utils/containers/map_keys.h" #include "utils/containers/map_values.h" #include "utils/containers/transform.h" -#include "utils/containers/unordered_map_from_pairs.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/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" -#include "utils/nonnegative_int/nonnegative_int.h" namespace FlexFlow { diff --git a/lib/pcg/include/pcg/file_format/v1/v1_computation_graph.h b/lib/pcg/include/pcg/file_format/v1/v1_computation_graph.h index 8b6128d603..a3ddbdf7be 100644 --- a/lib/pcg/include/pcg/file_format/v1/v1_computation_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/v1_computation_graph.h @@ -8,11 +8,12 @@ namespace FlexFlow { V1ComputationGraph to_v1(ComputationGraph const &); -ComputationGraph from_v1(V1ComputationGraph const &); std::pair> to_v1_including_node_numbering(ComputationGraph const &); +ComputationGraph from_v1(V1ComputationGraph const &); + } // namespace FlexFlow #endif diff --git a/lib/pcg/include/pcg/file_format/v1/v1_parallel_computation_graph.h b/lib/pcg/include/pcg/file_format/v1/v1_parallel_computation_graph.h index d481096d49..1ec9ee0e8c 100644 --- a/lib/pcg/include/pcg/file_format/v1/v1_parallel_computation_graph.h +++ b/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 diff --git a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc index d4cf7612bd..674a87651b 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc @@ -12,17 +12,6 @@ V1ComputationGraph to_v1(ComputationGraph const &g) { }; } -ComputationGraph from_v1(V1ComputationGraph const &v1) { - return ComputationGraph{ - LabelledKwargDataflowGraph:: - create_copy_of< - UnorderedSetLabelledOpenKwargDataflowGraph>( - from_v1(v1.raw_graph).first)}; -} - std::pair> to_v1_including_node_numbering(ComputationGraph const &cg) { std::pair< @@ -38,4 +27,15 @@ std::pair> return {v1_cg, v1_node_ids}; } +ComputationGraph from_v1(V1ComputationGraph const &v1) { + return ComputationGraph{ + LabelledKwargDataflowGraph:: + create_copy_of< + UnorderedSetLabelledOpenKwargDataflowGraph>( + from_v1(v1.raw_graph).first)}; +} + } // namespace FlexFlow From 9327f4a121baf65110143dcef21443028e75288a Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 24 Mar 2026 15:46:52 -0700 Subject: [PATCH 11/11] Split the with/without node numbering APIs. --- .../file_format/v1/graphs/v1_kwarg_dataflow_graph.h | 8 +++++++- .../v1/graphs/v1_labelled_kwarg_dataflow_graph.h | 13 ++++++++++--- .../v1/graphs/v1_kwarg_dataflow_graph.cc | 3 +++ .../v1/graphs/v1_labelled_kwarg_dataflow_graph.cc | 4 ++++ .../src/pcg/file_format/v1/v1_computation_graph.cc | 2 +- .../file_format/v1/v1_parallel_computation_graph.cc | 2 +- 6 files changed, 26 insertions(+), 6 deletions(-) diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h index 7a81f98b46..f8cc29c900 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h @@ -50,7 +50,7 @@ V1KwargDataflowGraph template std::pair, std::unordered_map> - from_v1(V1KwargDataflowGraph const &v1) { + from_v1_including_node_numbering(V1KwargDataflowGraph const &v1) { std::unordered_map node_map = generate_map(v1.nodes, [](nonnegative_int n) { return Node{n.size_t_from_nonnegative_int()}; @@ -78,6 +78,12 @@ std::pair, node_map}; } +template +KwargDataflowGraphView + from_v1(V1KwargDataflowGraph const &v1) { + return from_v1_including_node_numbering(v1).first; +} + } // namespace FlexFlow #endif diff --git a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h index b58de15cfd..9fe2d53db1 100644 --- a/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h +++ b/lib/pcg/include/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.h @@ -55,9 +55,10 @@ V1LabelledKwargDataflowGraph to_v1( template std::pair, std::unordered_map> - from_v1(V1LabelledKwargDataflowGraph const - &v1) { - auto [graph_view, node_map] = from_v1(v1.graph); + from_v1_including_node_numbering( + V1LabelledKwargDataflowGraph const + &v1) { + auto [graph_view, node_map] = from_v1_including_node_numbering(v1.graph); std::unordered_map node_labels = map_keys( v1.node_labels, [&](nonnegative_int n) { return node_map.at(n); }); @@ -68,6 +69,12 @@ std::pair, node_map}; } +template +LabelledKwargDataflowGraphView from_v1( + V1LabelledKwargDataflowGraph const &v1) { + return from_v1_including_node_numbering(v1).first; +} + } // namespace FlexFlow #endif diff --git a/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc b/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc index 6dcfe4ce91..cc10bbf4cb 100644 --- a/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.cc @@ -14,6 +14,9 @@ template V1KwargDataflowGraph template std::pair, std::unordered_map> + from_v1_including_node_numbering(V1KwargDataflowGraph const &); + +template KwargDataflowGraphView from_v1(V1KwargDataflowGraph const &); } // namespace FlexFlow diff --git a/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc b/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc index 7c8c2a77dc..4e50949e3f 100644 --- a/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/graphs/v1_labelled_kwarg_dataflow_graph.cc @@ -21,6 +21,10 @@ template V1LabelledKwargDataflowGraph to_v1( template std::pair< LabelledKwargDataflowGraphView, std::unordered_map> + from_v1_including_node_numbering( + V1LabelledKwargDataflowGraph const &); + +template LabelledKwargDataflowGraphView from_v1( V1LabelledKwargDataflowGraph const &); diff --git a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc index 674a87651b..e52b5708e5 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_computation_graph.cc @@ -35,7 +35,7 @@ ComputationGraph from_v1(V1ComputationGraph const &v1) { TensorAttrs, int, TensorSlotName>>( - from_v1(v1.raw_graph).first)}; + from_v1(v1.raw_graph))}; } } // namespace FlexFlow diff --git a/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc b/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc index c1a394e4c5..a5afa3ebdc 100644 --- a/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc +++ b/lib/pcg/src/pcg/file_format/v1/v1_parallel_computation_graph.cc @@ -22,7 +22,7 @@ ParallelComputationGraph from_v1(V1ParallelComputationGraph const &v1) { ParallelTensorAttrs, int, TensorSlotName>>( - from_v1(v1.raw_graph).first)}; + from_v1(v1.raw_graph))}; } } // namespace FlexFlow