Skip to content

Commit 80f4efb

Browse files
elliottslaughterelliottslaughter
committed
Add V1MappedParallelComputationGraph.
1 parent 9fbfbef commit 80f4efb

8 files changed

Lines changed: 170 additions & 0 deletions

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
namespace = "FlexFlow"
2+
name = "V1MappedOperatorTaskGroup"
3+
type = "struct"
4+
features = [
5+
"eq",
6+
"hash",
7+
"fmt",
8+
"json",
9+
]
10+
11+
includes = [
12+
"pcg/machine_space_coordinate.dtg.h",
13+
"pcg/mapped_parallel_computation_graph/operator_atomic_task_shard_binding.dtg.h",
14+
"utils/bidict/bidict.h",
15+
]
16+
17+
[[fields]]
18+
name = "shard_bindings"
19+
type = "::FlexFlow::bidict<::FlexFlow::MachineSpaceCoordinate, ::FlexFlow::OperatorAtomicTaskShardBinding>"

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
#ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_V1_MAPPED_OPERATOR_TASK_GROUP_H
2+
#define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_V1_MAPPED_OPERATOR_TASK_GROUP_H
3+
4+
#include "pcg/file_format/v1/v1_mapped_operator_task_group.dtg.h"
5+
#include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h"
6+
7+
namespace FlexFlow {
8+
9+
V1MappedOperatorTaskGroup to_v1(MappedOperatorTaskGroup const &);
10+
11+
} // namespace FlexFlow
12+
13+
#endif

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
namespace = "FlexFlow"
2+
name = "V1MappedParallelComputationGraph"
3+
type = "struct"
4+
features = [
5+
"eq",
6+
"hash",
7+
"fmt",
8+
"json",
9+
]
10+
11+
includes = [
12+
"<unordered_map>",
13+
"pcg/file_format/v1/v1_parallel_computation_graph.dtg.h",
14+
"pcg/file_format/v1/v1_mapped_operator_task_group.dtg.h",
15+
"pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
16+
]
17+
18+
src_includes = [
19+
"utils/hash/unordered_map.h",
20+
"utils/fmt/unordered_map.h",
21+
]
22+
23+
[[fields]]
24+
name = "pcg"
25+
type = "::FlexFlow::V1ParallelComputationGraph"
26+
27+
[[fields]]
28+
name = "mapped_tasks"
29+
type = "std::unordered_map<::FlexFlow::parallel_layer_guid_t, ::FlexFlow::V1MappedOperatorTaskGroup>"

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
#ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_V1_MAPPED_PARALLEL_COMPUTATION_GRAPH_H
2+
#define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_V1_MAPPED_PARALLEL_COMPUTATION_GRAPH_H
3+
4+
#include "pcg/file_format/v1/v1_mapped_parallel_computation_graph.dtg.h"
5+
#include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.dtg.h"
6+
7+
namespace FlexFlow {
8+
9+
V1MappedParallelComputationGraph to_v1(MappedParallelComputationGraph const &);
10+
11+
} // namespace FlexFlow
12+
13+
#endif

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ features = [
66
"ord",
77
"hash",
88
"fmt",
9+
"json",
910
]
1011

1112
includes = [

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#include "pcg/file_format/v1/v1_mapped_operator_task_group.h"
2+
3+
namespace FlexFlow {
4+
5+
V1MappedOperatorTaskGroup to_v1(MappedOperatorTaskGroup const &g) {
6+
return V1MappedOperatorTaskGroup{g.get_shard_bindings()};
7+
}
8+
9+
} // namespace FlexFlow

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#include "pcg/file_format/v1/v1_mapped_parallel_computation_graph.h"
2+
#include "pcg/file_format/v1/v1_mapped_operator_task_group.h"
3+
#include "pcg/file_format/v1/v1_parallel_computation_graph.h"
4+
#include "utils/containers/map_values.h"
5+
6+
namespace FlexFlow {
7+
8+
V1MappedParallelComputationGraph
9+
to_v1(MappedParallelComputationGraph const &mpcg) {
10+
return V1MappedParallelComputationGraph{
11+
to_v1(mpcg.pcg),
12+
map_values(mpcg.mapped_tasks,
13+
[](MappedOperatorTaskGroup const &g) { return to_v1(g); }),
14+
};
15+
}
16+
17+
} // namespace FlexFlow

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
#include "pcg/file_format/v1/v1_mapped_parallel_computation_graph.h"
2+
#include "op-attrs/parallel_tensor_space_coordinate.dtg.h"
3+
#include "op-attrs/tensor_slot_name.dtg.h"
4+
#include "pcg/device_type.dtg.h"
5+
#include "pcg/machine_space_coordinate.dtg.h"
6+
#include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h"
7+
#include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.dtg.h"
8+
#include "pcg/mapped_parallel_computation_graph/operator_atomic_task_shard_binding.dtg.h"
9+
#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
10+
#include "pcg/parallel_computation_graph/parallel_layer_added_result.dtg.h"
11+
#include "utils/bidict/bidict.h"
12+
#include <doctest/doctest.h>
13+
#include <nlohmann/json.hpp>
14+
15+
using namespace ::FlexFlow;
16+
17+
TEST_SUITE(FF_TEST_SUITE) {
18+
TEST_CASE("V1MappedParallelComputationGraph") {
19+
MappedParallelComputationGraph mpcg = [] {
20+
ParallelComputationGraph pcg = empty_parallel_computation_graph();
21+
22+
TensorShape input_shape = TensorShape{
23+
TensorDims{
24+
FFOrdered{
25+
12_p,
26+
16_p,
27+
},
28+
},
29+
DataType::FLOAT,
30+
};
31+
32+
ParallelLayerAddedResult result = pcg_add_input_layer(pcg, input_shape);
33+
parallel_layer_guid_t layer = result.parallel_layer;
34+
35+
MachineSpaceCoordinate coord = MachineSpaceCoordinate{
36+
/*node_idx=*/0_n,
37+
/*device_idx=*/0_n,
38+
/*device_type=*/DeviceType::GPU,
39+
};
40+
41+
OperatorAtomicTaskShardBinding binding = OperatorAtomicTaskShardBinding{
42+
/*tensor_coords=*/{
43+
{
44+
TensorSlotName::OUTPUT,
45+
ParallelTensorSpaceCoordinate{
46+
/*sum_component=*/0_n,
47+
/*discard_copy_component=*/0_n,
48+
/*shard_components=*/FFOrdered<nonnegative_int>{0_n, 0_n},
49+
},
50+
},
51+
},
52+
};
53+
54+
MappedOperatorTaskGroup task_group = MappedOperatorTaskGroup{
55+
bidict<MachineSpaceCoordinate, OperatorAtomicTaskShardBinding>{
56+
{coord, binding},
57+
},
58+
};
59+
60+
return MappedParallelComputationGraph{
61+
/*pcg=*/pcg,
62+
/*mapped_tasks=*/{{layer, task_group}},
63+
};
64+
}();
65+
66+
V1MappedParallelComputationGraph v1_mpcg = to_v1(mpcg);
67+
nlohmann::json j = v1_mpcg;
68+
}
69+
}

0 commit comments

Comments
 (0)