Skip to content

Commit ed48111

Browse files
elliottslaughterelliottslaughter
authored
Add function to convert from ComputationGraph to DynamicOpenDataflowGraph (#1624)
* Add function to convert from ComputationGraph to DynamicOpenDataflowGraph. * Fixup unit tests. * Switch role to nullopt. * Apply PR feedback. * Format.
1 parent f4ace86 commit ed48111

File tree

10 files changed

+168
-34
lines changed

10 files changed

+168
-34
lines changed

‎lib/task-spec/include/task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.toml‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
namespace = "FlexFlow"
2+
name = "dynamic_layer_guid_t"
3+
type = "variant"
4+
features = [
5+
"eq",
6+
"hash",
7+
"fmt",
8+
]
9+
10+
includes = [
11+
"pcg/layer_guid_t.dtg.h",
12+
"pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
13+
]
14+
15+
[[values]]
16+
name = "layer_guid"
17+
type = "::FlexFlow::layer_guid_t"
18+
19+
[[values]]
20+
name = "pcg_layer_guid"
21+
type = "::FlexFlow::parallel_layer_guid_t"

‎lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.dtg.toml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ includes = [
1313
"pcg/machine_space_coordinate.dtg.h",
1414
"pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h",
1515
"op-attrs/pcg_operator_attrs.dtg.h",
16-
"pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
16+
"task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h",
1717
"task-spec/device_specific_per_device_op_state.dtg.h",
1818
]
1919

@@ -38,8 +38,8 @@ name = "op_attrs"
3838
type = "std::optional<::FlexFlow::PCGOperatorAttrs>"
3939

4040
[[fields]]
41-
name = "pcg_layer_guid"
42-
type = "::FlexFlow::parallel_layer_guid_t"
41+
name = "layer_guid"
42+
type = "::FlexFlow::dynamic_layer_guid_t"
4343

4444
[[fields]]
4545
name = "per_device_op_state"

‎lib/task-spec/include/task-spec/dynamic_graph/dynamic_tensor_guid_t.dtg.toml‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
namespace = "FlexFlow"
2+
name = "dynamic_tensor_guid_t"
3+
type = "variant"
4+
features = [
5+
"eq",
6+
"hash",
7+
"fmt",
8+
]
9+
10+
includes = [
11+
"pcg/tensor_guid_t.dtg.h",
12+
"pcg/parallel_computation_graph/parallel_tensor_guid_t.dtg.h",
13+
]
14+
15+
[[values]]
16+
name = "tensor_guid"
17+
type = "::FlexFlow::tensor_guid_t"
18+
19+
[[values]]
20+
name = "pcg_tensor_guid"
21+
type = "::FlexFlow::parallel_tensor_guid_t"

‎lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ features = [
99

1010
includes = [
1111
"<optional>",
12+
"task-spec/dynamic_graph/dynamic_tensor_guid_t.dtg.h",
1213
"op-attrs/parallel_tensor_shape.dtg.h",
1314
"kernels/accessor.h",
1415
"pcg/parallel_computation_graph/parallel_tensor_guid_t.dtg.h",
@@ -21,8 +22,8 @@ src_includes = [
2122
]
2223

2324
[[fields]]
24-
name = "pcg_tensor_guid"
25-
type = "::FlexFlow::parallel_tensor_guid_t"
25+
name = "tensor_guid"
26+
type = "::FlexFlow::dynamic_tensor_guid_t"
2627

2728
[[fields]]
2829
name = "parallel_tensor_shape"

‎lib/task-spec/include/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_cg.h‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_OPEN_DATAFLOW_GRAPH_FROM_CG_H
2+
#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_OPEN_DATAFLOW_GRAPH_FROM_CG_H
3+
4+
#include "pcg/computation_graph.dtg.h"
5+
#include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.dtg.h"
6+
7+
namespace FlexFlow {
8+
9+
DynamicOpenDataflowGraph
10+
make_dynamic_open_dataflow_graph_from_cg(ComputationGraph const &);
11+
12+
} // namespace FlexFlow
13+
14+
#endif

‎lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_cg.cc‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
#include "task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_cg.h"
2+
#include "op-attrs/pcg_operator_attrs.h"
3+
#include "pcg/computation_graph.h"
4+
#include "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h"
5+
#include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h"
6+
#include "task-spec/dynamic_graph/dynamic_tensor_role.h"
7+
#include "utils/containers/generate_map.h"
8+
#include <optional>
9+
#include <unordered_map>
10+
#include <utility>
11+
12+
namespace FlexFlow {
13+
14+
DynamicOpenDataflowGraph
15+
make_dynamic_open_dataflow_graph_from_cg(ComputationGraph const &cg) {
16+
DynamicOpenDataflowGraph result = make_empty_dynamic_open_dataflow_graph();
17+
18+
for (auto const &[layer, attrs] : get_layer_attrs_mapping(cg)) {
19+
DynamicNodeAttrs result_attrs{
20+
/*task_type=*/std::nullopt,
21+
/*device_coord=*/std::nullopt,
22+
/*mapping=*/std::nullopt,
23+
/*op_attrs=*/pcg_op_attrs_from_compgraph_op_attrs(attrs.op_attrs),
24+
/*pcg_layer_guid=*/dynamic_layer_guid_t{layer},
25+
/*per_device_op_state=*/std::nullopt,
26+
};
27+
28+
std::unordered_map<DynamicTensorSlot, DynamicValueAttrs> result_inputs =
29+
transform(
30+
get_incoming_tensors(cg, layer),
31+
[](TensorSlotName const &slot_name, tensor_guid_t const &tensor) {
32+
return std::pair<DynamicTensorSlot, DynamicValueAttrs>{
33+
DynamicTensorSlot{
34+
/*slot_name=*/slot_name,
35+
/*slot_tensor_role=*/std::nullopt,
36+
},
37+
DynamicValueAttrs{
38+
/*tensor_guid=*/dynamic_tensor_guid_t{tensor},
39+
/*parallel_tensor_shape=*/std::nullopt,
40+
/*shard_coord=*/std::nullopt,
41+
/*accessor=*/std::nullopt,
42+
/*role=*/std::nullopt,
43+
},
44+
};
45+
});
46+
std::unordered_map<DynamicTensorSlot, DynamicValueAttrs> result_outputs =
47+
transform(
48+
get_outgoing_tensors(cg, layer),
49+
[](TensorSlotName const &slot_name, tensor_guid_t const &tensor) {
50+
return std::pair<DynamicTensorSlot, DynamicValueAttrs>{
51+
DynamicTensorSlot{
52+
/*slot_name=*/slot_name,
53+
/*slot_tensor_role=*/std::nullopt,
54+
},
55+
DynamicValueAttrs{
56+
/*tensor_guid=*/dynamic_tensor_guid_t{tensor},
57+
/*parallel_tensor_shape=*/std::nullopt,
58+
/*shard_coord=*/std::nullopt,
59+
/*accessor=*/std::nullopt,
60+
/*role=*/std::nullopt,
61+
},
62+
};
63+
});
64+
65+
result.invocations.emplace(result_outputs, result_attrs, result_outputs);
66+
}
67+
68+
return result;
69+
}
70+
71+
} // namespace FlexFlow

‎lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,38 +7,38 @@ using namespace ::FlexFlow;
77
TEST_SUITE(FF_TEST_SUITE) {
88
TEST_CASE("dynamic_op_dataflow_graph_from_invocation_set") {
99
DynamicValueAttrs value_1 = DynamicValueAttrs{
10-
/*pcg_tensor_guid=*/parallel_tensor_guid_t{
10+
/*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
1111
KwargDataflowOutput{
1212
Node{1},
1313
TensorSlotName::OUTPUT,
1414
},
15-
},
15+
}},
1616
/*parallel_tensor_shape=*/std::nullopt,
1717
/*shard_coord=*/std::nullopt,
1818
/*accessor=*/std::nullopt,
1919
/*tensor_type=*/std::nullopt,
2020
};
2121

2222
DynamicValueAttrs value_2 = DynamicValueAttrs{
23-
/*pcg_tensor_guid=*/parallel_tensor_guid_t{
23+
/*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
2424
KwargDataflowOutput{
2525
Node{2},
2626
TensorSlotName::OUTPUT,
2727
},
28-
},
28+
}},
2929
/*parallel_tensor_shape=*/std::nullopt,
3030
/*shard_coord=*/std::nullopt,
3131
/*accessor=*/std::nullopt,
3232
/*tensor_type=*/std::nullopt,
3333
};
3434

3535
DynamicValueAttrs value_3 = DynamicValueAttrs{
36-
/*pcg_tensor_guid=*/parallel_tensor_guid_t{
36+
/*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
3737
KwargDataflowOutput{
3838
Node{3},
3939
TensorSlotName::OUTPUT,
4040
},
41-
},
41+
}},
4242
/*parallel_tensor_shape=*/std::nullopt,
4343
/*shard_coord=*/std::nullopt,
4444
/*accessor=*/std::nullopt,
@@ -50,7 +50,7 @@ TEST_SUITE(FF_TEST_SUITE) {
5050
/*device_coord=*/std::nullopt,
5151
/*mapping=*/std::nullopt,
5252
/*op_attrs=*/std::nullopt,
53-
/*pcg_layer_guid=*/parallel_layer_guid_t{Node{4}},
53+
/*layer_guid=*/dynamic_layer_guid_t{parallel_layer_guid_t{Node{4}}},
5454
/*per_device_op_state=*/std::nullopt,
5555
};
5656

‎lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -67,12 +67,12 @@ TEST_SUITE(FF_TEST_SUITE) {
6767
std::optional<ParallelTensorSpaceCoordinate> const &shard_coord)
6868
-> DynamicValueAttrs {
6969
return DynamicValueAttrs{
70-
/*pcg_tensor_guid=*/parallel_tensor_guid_t{
70+
/*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
7171
KwargDataflowOutput<TensorSlotName>{
7272
Node{src_node_id},
7373
src_slot_name,
7474
},
75-
},
75+
}},
7676
/*parallel_tensor_shape=*/std::nullopt,
7777
/*shard_coord=*/shard_coord,
7878
/*accessor=*/std::nullopt,
@@ -114,7 +114,8 @@ TEST_SUITE(FF_TEST_SUITE) {
114114
/*device_coord=*/mc2,
115115
/*mapping=*/std::nullopt,
116116
/*op_attrs=*/std::nullopt,
117-
/*pcg_layer_guid=*/parallel_layer_guid_t{Node{invocation1_id}},
117+
/*layer_guid=*/
118+
dynamic_layer_guid_t{parallel_layer_guid_t{Node{invocation1_id}}},
118119
/*per_device_op_state=*/std::nullopt,
119120
},
120121
/*outputs=*/
@@ -140,7 +141,8 @@ TEST_SUITE(FF_TEST_SUITE) {
140141
/*device_coord=*/mc1,
141142
/*mapping=*/std::nullopt,
142143
/*op_attrs=*/std::nullopt,
143-
/*pcg_layer_guid=*/parallel_layer_guid_t{Node{invocation2_id}},
144+
/*layer_guid=*/
145+
dynamic_layer_guid_t{parallel_layer_guid_t{Node{invocation2_id}}},
144146
/*per_device_op_state=*/std::nullopt,
145147
},
146148
/*outputs=*/
@@ -173,7 +175,8 @@ TEST_SUITE(FF_TEST_SUITE) {
173175
/*device_coord=*/mc2,
174176
/*mapping=*/std::nullopt,
175177
/*op_attrs=*/std::nullopt,
176-
/*pcg_layer_guid=*/parallel_layer_guid_t{Node{invocation3_id}},
178+
/*layer_guid=*/
179+
dynamic_layer_guid_t{parallel_layer_guid_t{Node{invocation3_id}}},
177180
/*per_device_op_state=*/std::nullopt,
178181
},
179182
/*outputs=*/

‎lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,12 @@ TEST_SUITE(FF_TEST_SUITE) {
1111
[](size_t node_id, std::optional<DynamicTensorRole> const &tensor_role)
1212
-> DynamicValueAttrs {
1313
return DynamicValueAttrs{
14-
/*pcg_tensor_guid=*/parallel_tensor_guid_t{
14+
/*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
1515
KwargDataflowOutput<TensorSlotName>{
1616
Node{node_id},
1717
TensorSlotName::OUTPUT,
1818
},
19-
},
19+
}},
2020
/*parallel_tensor_shape=*/std::nullopt,
2121
/*shard_coord=*/std::nullopt,
2222
/*accessor=*/std::nullopt,
@@ -33,7 +33,7 @@ TEST_SUITE(FF_TEST_SUITE) {
3333
};
3434
};
3535

36-
parallel_layer_guid_t pcg_layer_guid = parallel_layer_guid_t{Node{20}};
36+
dynamic_layer_guid_t layer_guid{parallel_layer_guid_t{Node{20}}};
3737

3838
DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation {
3939
DynamicValueAttrs v1 = mk_value_attrs(0, std::nullopt);
@@ -53,7 +53,7 @@ TEST_SUITE(FF_TEST_SUITE) {
5353
/*device_coord=*/std::nullopt,
5454
/*mapping=*/std::nullopt,
5555
/*op_attrs=*/std::nullopt,
56-
/*pcg_layer_guid=*/pcg_layer_guid,
56+
/*layer_guid=*/layer_guid,
5757
/*per_device_op_state=*/std::nullopt,
5858
},
5959
/*outputs=*/
@@ -86,7 +86,7 @@ TEST_SUITE(FF_TEST_SUITE) {
8686
/*device_coord=*/std::nullopt,
8787
/*mapping=*/std::nullopt,
8888
/*op_attrs=*/std::nullopt,
89-
/*pcg_layer_guid=*/pcg_layer_guid,
89+
/*layer_guid=*/layer_guid,
9090
/*per_device_op_state=*/std::nullopt,
9191
},
9292
/*outputs=*/
@@ -104,12 +104,12 @@ TEST_SUITE(FF_TEST_SUITE) {
104104
[](size_t node_id, std::optional<DynamicTensorRole> const &tensor_role)
105105
-> DynamicValueAttrs {
106106
return DynamicValueAttrs{
107-
/*pcg_tensor_guid=*/parallel_tensor_guid_t{
107+
/*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
108108
KwargDataflowOutput<TensorSlotName>{
109109
Node{node_id},
110110
TensorSlotName::OUTPUT,
111111
},
112-
},
112+
}},
113113
/*parallel_tensor_shape=*/std::nullopt,
114114
/*shard_coord=*/std::nullopt,
115115
/*accessor=*/std::nullopt,
@@ -126,7 +126,7 @@ TEST_SUITE(FF_TEST_SUITE) {
126126
};
127127
};
128128

129-
parallel_layer_guid_t pcg_layer_guid = parallel_layer_guid_t{Node{20}};
129+
dynamic_layer_guid_t layer_guid{parallel_layer_guid_t{Node{20}}};
130130

131131
DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation {
132132
DynamicValueAttrs v1 = mk_value_attrs(0, std::nullopt);
@@ -146,7 +146,7 @@ TEST_SUITE(FF_TEST_SUITE) {
146146
/*device_coord=*/std::nullopt,
147147
/*mapping=*/std::nullopt,
148148
/*op_attrs=*/std::nullopt,
149-
/*pcg_layer_guid=*/pcg_layer_guid,
149+
/*layer_guid=*/layer_guid,
150150
/*per_device_op_state=*/std::nullopt,
151151
},
152152
/*outputs=*/
@@ -185,7 +185,7 @@ TEST_SUITE(FF_TEST_SUITE) {
185185
/*device_coord=*/std::nullopt,
186186
/*mapping=*/std::nullopt,
187187
/*op_attrs=*/std::nullopt,
188-
/*pcg_layer_guid=*/pcg_layer_guid,
188+
/*layer_guid=*/layer_guid,
189189
/*per_device_op_state=*/std::nullopt,
190190
},
191191
/*outputs=*/
@@ -210,7 +210,8 @@ TEST_SUITE(FF_TEST_SUITE) {
210210
/*device_coord=*/std::nullopt,
211211
/*mapping=*/std::nullopt,
212212
/*op_attrs=*/std::nullopt,
213-
/*pcg_layer_guid=*/parallel_layer_guid_t{Node{layer_id}},
213+
/*layer_guid=*/
214+
dynamic_layer_guid_t{parallel_layer_guid_t{Node{layer_id}}},
214215
/*per_device_op_state=*/std::nullopt,
215216
};
216217
};
@@ -219,12 +220,12 @@ TEST_SUITE(FF_TEST_SUITE) {
219220
[](size_t node_id, std::optional<DynamicTensorRole> const &tensor_type)
220221
-> DynamicValueAttrs {
221222
return DynamicValueAttrs{
222-
/*pcg_tensor_guid=*/parallel_tensor_guid_t{
223+
/*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
223224
KwargDataflowOutput{
224225
Node{node_id},
225226
TensorSlotName::OUTPUT,
226227
},
227-
},
228+
}},
228229
/*parallel_tensor_shape=*/std::nullopt,
229230
/*shard_coord=*/std::nullopt,
230231
/*accessor=*/std::nullopt,

0 commit comments

Comments
 (0)