1 // Copyright 2020 The Marl Authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 #include "marl/dag.h"
16
17 #include "marl_test.h"
18
19 using namespace testing;
20
21 namespace {
22
23 struct Data {
24 std::mutex mutex;
25 std::vector<std::string> order;
26
push__anon45c983720111::Data27 void push(std::string&& s) {
28 std::unique_lock<std::mutex> lock(mutex);
29 order.emplace_back(std::move(s));
30 }
31 };
32
33 template <typename T>
slice(const std::vector<T> & in,size_t from,size_t to)34 std::vector<T> slice(const std::vector<T>& in, size_t from, size_t to) {
35 return {in.begin() + from, in.begin() + to};
36 }
37
38 } // namespace
39
40 // [A] --> [B] --> [C] |
TEST_P(WithBoundScheduler,DAGChainNoArg)41 TEST_P(WithBoundScheduler, DAGChainNoArg) {
42 marl::DAG<>::Builder builder;
43
44 Data data;
45 builder.root()
46 .then([&] { data.push("A"); })
47 .then([&] { data.push("B"); })
48 .then([&] { data.push("C"); });
49
50 auto dag = builder.build();
51 dag->run();
52
53 ASSERT_THAT(data.order, ElementsAre("A", "B", "C"));
54 }
55
56 // [A] --> [B] --> [C] |
TEST_P(WithBoundScheduler,DAGChain)57 TEST_P(WithBoundScheduler, DAGChain) {
58 marl::DAG<Data&>::Builder builder;
59
60 builder.root()
61 .then([](Data& data) { data.push("A"); })
62 .then([](Data& data) { data.push("B"); })
63 .then([](Data& data) { data.push("C"); });
64
65 auto dag = builder.build();
66
67 Data data;
68 dag->run(data);
69
70 ASSERT_THAT(data.order, ElementsAre("A", "B", "C"));
71 }
72
73 // [A] --> [B] --> [C] |
TEST_P(WithBoundScheduler,DAGRunRepeat)74 TEST_P(WithBoundScheduler, DAGRunRepeat) {
75 marl::DAG<Data&>::Builder builder;
76
77 builder.root()
78 .then([](Data& data) { data.push("A"); })
79 .then([](Data& data) { data.push("B"); })
80 .then([](Data& data) { data.push("C"); });
81
82 auto dag = builder.build();
83
84 Data dataA, dataB;
85 dag->run(dataA);
86 dag->run(dataB);
87 dag->run(dataA);
88
89 ASSERT_THAT(dataA.order, ElementsAre("A", "B", "C", "A", "B", "C"));
90 ASSERT_THAT(dataB.order, ElementsAre("A", "B", "C"));
91 }
92
93 // /--> [A] |
94 // [root] --|--> [B] |
95 // \--> [C] |
TEST_P(WithBoundScheduler,DAGFanOutFromRoot)96 TEST_P(WithBoundScheduler, DAGFanOutFromRoot) {
97 marl::DAG<Data&>::Builder builder;
98
99 auto root = builder.root();
100 root.then([](Data& data) { data.push("A"); });
101 root.then([](Data& data) { data.push("B"); });
102 root.then([](Data& data) { data.push("C"); });
103
104 auto dag = builder.build();
105
106 Data data;
107 dag->run(data);
108
109 ASSERT_THAT(data.order, UnorderedElementsAre("A", "B", "C"));
110 }
111
112 // /--> [A] |
113 // [root] -->[N]--|--> [B] |
114 // \--> [C] |
TEST_P(WithBoundScheduler,DAGFanOutFromNonRoot)115 TEST_P(WithBoundScheduler, DAGFanOutFromNonRoot) {
116 marl::DAG<Data&>::Builder builder;
117
118 auto root = builder.root();
119 auto node = root.then([](Data& data) { data.push("N"); });
120 node.then([](Data& data) { data.push("A"); });
121 node.then([](Data& data) { data.push("B"); });
122 node.then([](Data& data) { data.push("C"); });
123
124 auto dag = builder.build();
125
126 Data data;
127 dag->run(data);
128
129 ASSERT_THAT(data.order, UnorderedElementsAre("N", "A", "B", "C"));
130 ASSERT_EQ(data.order[0], "N");
131 ASSERT_THAT(slice(data.order, 1, 4), UnorderedElementsAre("A", "B", "C"));
132 }
133
134 // /--> [A0] --\ /--> [C0] --\ /--> [E0] --\ |
135 // [root] --|--> [A1] --|-->[B]--|--> [C1] --|-->[D]--|--> [E1] --|-->[F] |
136 // \--> [C2] --/ |--> [E2] --| |
137 // \--> [E3] --/ |
TEST_P(WithBoundScheduler,DAGFanOutFanIn)138 TEST_P(WithBoundScheduler, DAGFanOutFanIn) {
139 marl::DAG<Data&>::Builder builder;
140
141 auto root = builder.root();
142 auto a0 = root.then([](Data& data) { data.push("A0"); });
143 auto a1 = root.then([](Data& data) { data.push("A1"); });
144
145 auto b = builder.node([](Data& data) { data.push("B"); }, {a0, a1});
146
147 auto c0 = b.then([](Data& data) { data.push("C0"); });
148 auto c1 = b.then([](Data& data) { data.push("C1"); });
149 auto c2 = b.then([](Data& data) { data.push("C2"); });
150
151 auto d = builder.node([](Data& data) { data.push("D"); }, {c0, c1, c2});
152
153 auto e0 = d.then([](Data& data) { data.push("E0"); });
154 auto e1 = d.then([](Data& data) { data.push("E1"); });
155 auto e2 = d.then([](Data& data) { data.push("E2"); });
156 auto e3 = d.then([](Data& data) { data.push("E3"); });
157
158 builder.node([](Data& data) { data.push("F"); }, {e0, e1, e2, e3});
159
160 auto dag = builder.build();
161
162 Data data;
163 dag->run(data);
164
165 ASSERT_THAT(data.order,
166 UnorderedElementsAre("A0", "A1", "B", "C0", "C1", "C2", "D", "E0",
167 "E1", "E2", "E3", "F"));
168 ASSERT_THAT(slice(data.order, 0, 2), UnorderedElementsAre("A0", "A1"));
169 ASSERT_THAT(data.order[2], "B");
170 ASSERT_THAT(slice(data.order, 3, 6), UnorderedElementsAre("C0", "C1", "C2"));
171 ASSERT_THAT(data.order[6], "D");
172 ASSERT_THAT(slice(data.order, 7, 11),
173 UnorderedElementsAre("E0", "E1", "E2", "E3"));
174 ASSERT_THAT(data.order[11], "F");
175 }
176