schedule_test.cpp 25.3 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
5
#include <migraphx/schedule.hpp>
#include <migraphx/operators.hpp>
#include <migraphx/generate.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/iterator_for.hpp>
Paul's avatar
Paul committed
6
#include <migraphx/ranges.hpp>
Paul's avatar
Paul committed
7
#include <migraphx/dfor.hpp>
Paul's avatar
Paul committed
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
#include <basic_ops.hpp>
#include <test.hpp>

struct unary_op
{
    std::string name() const { return "unary"; }
    migraphx::argument
    compute(migraphx::context&, const migraphx::shape&, std::vector<migraphx::argument> args) const
    {
        if(args.empty())
            return {};
        return args.front();
    }

    migraphx::shape compute_shape(std::vector<migraphx::shape> inputs) const
    {
        if(inputs.empty())
            return {};
        return inputs.front();
    }
    int output_alias(const std::vector<migraphx::shape>&) const { return 0; }
};

Paul's avatar
Paul committed
31
struct nary_op
Paul's avatar
Paul committed
32
{
Paul's avatar
Paul committed
33
34
35
36
37
38
    std::string comment = "";
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::pack(f(self.comment, "comment"));
    }
Paul's avatar
Paul committed
39
    std::string name() const { return "nary"; }
Paul's avatar
Paul committed
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    migraphx::argument
    compute(migraphx::context&, const migraphx::shape&, std::vector<migraphx::argument> args) const
    {
        if(args.empty())
            return {};
        return args.front();
    }

    migraphx::shape compute_shape(std::vector<migraphx::shape> inputs) const
    {
        if(inputs.empty())
            return {};
        return inputs.front();
    }
};

Paul's avatar
Paul committed
56
57
struct wait_event
{
Paul's avatar
Paul committed
58
59
    std::shared_ptr<std::vector<std::size_t>> wait_for =
        std::make_shared<std::vector<std::size_t>>();
Paul's avatar
Paul committed
60
61
62
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
63
        return migraphx::pack(f(*self.wait_for, "wait_for"));
Paul's avatar
Paul committed
64
65
66
67
    }
    std::string name() const { return "wait_event"; }
    migraphx::shape compute_shape(const std::vector<migraphx::shape>&) const { return {}; }

Paul's avatar
Paul committed
68
69
70
    migraphx::argument compute(migraphx::context&,
                               const migraphx::shape&,
                               const std::vector<migraphx::argument>&) const
Paul's avatar
Paul committed
71
    {
Paul's avatar
Paul committed
72
73
        assert(wait_for != nullptr);
        assert(not wait_for->empty());
Paul's avatar
Paul committed
74
75
76
77
        return {};
    }
};

Paul's avatar
Paul committed
78
using instruction_map = std::unordered_map<migraphx::instruction_ref, std::size_t>;
Paul's avatar
Paul committed
79
80
using wait_map =
    std::unordered_map<migraphx::instruction_ref, std::shared_ptr<std::vector<std::size_t>>>;
Paul's avatar
Paul committed
81
82
83

struct schedule_model_test
{
Paul's avatar
Paul committed
84
    std::shared_ptr<instruction_map> ins2stream = std::make_shared<instruction_map>();
Paul's avatar
Paul committed
85
86
    std::shared_ptr<std::unordered_map<std::size_t, std::size_t>> wait2stream =
        std::make_shared<std::unordered_map<std::size_t, std::size_t>>();
Paul's avatar
Paul committed
87
    std::shared_ptr<wait_map> ins2wait_for = std::make_shared<wait_map>();
Paul's avatar
Paul committed
88
    std::size_t concurrency() const { return 4; }
Paul's avatar
Paul committed
89
    void sched(migraphx::program&, migraphx::instruction_ref ins, std::size_t n) const
Paul's avatar
Paul committed
90
91
92
    {
        (*ins2stream)[ins] = n;
    }
Paul's avatar
Paul committed
93
    void wait(migraphx::program& p, migraphx::instruction_ref ins, std::size_t wait_id) const
Paul's avatar
Paul committed
94
    {
Paul's avatar
Paul committed
95
        if(ins2wait_for->count(ins) == 0)
Paul's avatar
Paul committed
96
97
98
99
100
101
102
        {
            auto event = wait_event{};
            p.insert_instruction(ins, event);
            (*ins2wait_for)[ins] = event.wait_for;
        }
        (*ins2wait_for)[ins]->push_back(wait2stream->at(wait_id));
    }
Paul's avatar
Paul committed
103
    void record(migraphx::program&, migraphx::instruction_ref ins, std::size_t wait_id) const
Paul's avatar
Paul committed
104
105
    {
        (*wait2stream)[wait_id] = ins2stream->at(ins);
Paul's avatar
Paul committed
106
107
108
    }
    std::size_t weight(const migraphx::operation& op) const
    {
Paul's avatar
Paul committed
109
        if(op.name() == "binary" or op.name() == "unary")
Paul's avatar
Paul committed
110
111
112
113
114
115
116
117
            return 4;
        else
            return 1;
    }
};

struct schedule_target
{
Paul's avatar
Paul committed
118
    schedule_model_test model{};
Paul's avatar
Paul committed
119
120
121
    std::string name() const { return "schedule"; }
    std::vector<migraphx::pass> get_passes(migraphx::context&) const
    {
Paul's avatar
Paul committed
122
        return {migraphx::schedule{model}};
Paul's avatar
Paul committed
123
124
    }
    migraphx::context get_context() const { return {}; }
Paul's avatar
Paul committed
125

Paul's avatar
Paul committed
126
    std::size_t get_stream(migraphx::instruction_ref ins) { return model.ins2stream->at(ins); }
Paul's avatar
Paul committed
127

Paul's avatar
Paul committed
128
129
130
131
132
133
134
135
136
    std::vector<std::size_t> get_streams(std::vector<migraphx::instruction_ref> inss)
    {
        std::vector<std::size_t> result;
        std::transform(inss.begin(), inss.end(), std::back_inserter(result), [&](auto ins) {
            return this->get_stream(ins);
        });
        return result;
    }

Paul's avatar
Paul committed
137
    bool has_stream(migraphx::instruction_ref ins) { return model.ins2stream->count(ins) > 0; }
Paul's avatar
Paul committed
138
139
140
141
};

bool check_conflicts(migraphx::program& p, migraphx::instruction_ref x, migraphx::instruction_ref y)
{
Paul's avatar
Paul committed
142
    for(auto ins : migraphx::iterator_for(p))
Paul's avatar
Paul committed
143
    {
Paul's avatar
Paul committed
144
        if(ins->name() != "identity")
Paul's avatar
Paul committed
145
            continue;
Paul's avatar
Paul committed
146
        if(not migraphx::contains(ins->inputs(), x))
Paul's avatar
Paul committed
147
            continue;
Paul's avatar
Paul committed
148
        if(not migraphx::contains(ins->inputs(), y))
Paul's avatar
Paul committed
149
150
            continue;
        return true;
Paul's avatar
Paul committed
151
152
153
154
    }
    return false;
}

Paul's avatar
Paul committed
155
void check_conflicts(migraphx::program& p,
Paul's avatar
Paul committed
156
157
                     std::vector<std::vector<migraphx::instruction_ref>> conflicts,
                     bool result = true)
Paul's avatar
Paul committed
158
159
{
    migraphx::dfor(conflicts.size(), conflicts.size())([&](auto i, auto j) {
Paul's avatar
Paul committed
160
        if(i == j)
Paul's avatar
Paul committed
161
            return;
Paul's avatar
Paul committed
162
163
        for(auto ins1 : conflicts[i])
            for(auto ins2 : conflicts[j])
Paul's avatar
Paul committed
164
                CHECK(check_conflicts(p, ins1, ins2) == result);
Paul's avatar
Paul committed
165
166
167
    });
}

Paul's avatar
Paul committed
168
template <class T>
Paul's avatar
Paul committed
169
170
171
172
173
174
std::vector<T> sorted(std::vector<T> x)
{
    std::sort(x.begin(), x.end());
    return x;
}

Paul's avatar
Paul committed
175
template <class T>
Paul's avatar
Paul committed
176
177
178
179
180
181
182
183
184
185
186
187
188
std::vector<T> unique(std::vector<T> x)
{
    std::sort(x.begin(), x.end());
    x.erase(std::unique(x.begin(), x.end()), x.end());
    return x;
}

std::vector<std::size_t> get_wait_for(std::vector<std::size_t> wait_for)
{
    std::sort(wait_for.begin(), wait_for.end());
    return wait_for;
}

Paul's avatar
Paul committed
189
190
191
192
193
194
195
196
197
198
std::vector<std::size_t> get_wait_for(std::size_t wait_on, std::vector<std::size_t> wait_for)
{
    wait_for.erase(std::find(wait_for.begin(), wait_for.end(), wait_on));
    std::sort(wait_for.begin(), wait_for.end());
    return wait_for;
}

std::vector<std::size_t> get_wait_for(migraphx::instruction_ref ins)
{
    auto wait_ins = std::prev(ins);
Paul's avatar
Paul committed
199
200
201
    // Skip identity operators
    while(wait_ins->name() == "identity")
        wait_ins = std::prev(wait_ins);
Paul's avatar
Paul committed
202
    if(wait_ins->name() != "wait_event")
Paul's avatar
Paul committed
203
        return {};
Paul's avatar
Paul committed
204
    auto wf = *migraphx::any_cast<wait_event>(wait_ins->get_operator()).wait_for;
Paul's avatar
Paul committed
205
206
207
208
    std::sort(wf.begin(), wf.end());
    return wf;
}

Paul's avatar
Paul committed
209
210
211
template <class T>
std::vector<migraphx::instruction_ref>
chain(migraphx::program& p, std::size_t n, T x, migraphx::instruction_ref input)
Paul's avatar
Paul committed
212
213
{
    std::vector<migraphx::instruction_ref> result;
Paul's avatar
Paul committed
214
    for(std::size_t i = 0; i < n; i++)
Paul's avatar
Paul committed
215
216
217
218
219
220
    {
        result.push_back(p.add_instruction(x, input));
        input = result.back();
    }
    return result;
}
Paul's avatar
Paul committed
221
222
TEST_CASE(single_entry)
{
Paul's avatar
Paul committed
223
    schedule_target t{};
Paul's avatar
Paul committed
224
225
    migraphx::program p;
    auto one    = p.add_literal(1);
Paul's avatar
Paul committed
226
227
    auto onep1  = p.add_instruction(unary_op{}, one);
    auto onep2  = p.add_instruction(unary_op{}, one);
Paul's avatar
Paul committed
228
    auto binary = p.add_instruction(nary_op{}, onep1, onep2);
Paul's avatar
Paul committed
229
230
231
232
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(onep1) != t.get_stream(onep2));
    EXPECT(t.get_stream(binary) == 0);
Paul's avatar
Paul committed
233
234
    EXPECT(get_wait_for(binary) ==
           get_wait_for(t.get_stream(binary), {t.get_stream(onep1), t.get_stream(onep2)}));
Paul's avatar
Paul committed
235
236
237
    EXPECT(check_conflicts(p, onep1, onep2));
}

Paul's avatar
Paul committed
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
TEST_CASE(zero_record)
{
    schedule_target t{};
    migraphx::program p;
    auto one    = p.add_literal(1);
    auto onep1  = p.add_instruction(unary_op{}, one);
    auto onep2  = p.add_instruction(unary_op{}, one);
    auto binary = p.add_instruction(nary_op{},
                                    p.add_instruction(migraphx::op::identity{}, onep1),
                                    p.add_instruction(migraphx::op::identity{}, onep2));
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(onep1) != t.get_stream(onep2));
    EXPECT(t.has_stream(binary));
    EXPECT(get_wait_for(binary) ==
           get_wait_for(t.get_stream(binary), {t.get_stream(onep1), t.get_stream(onep2)}));
    EXPECT(check_conflicts(p, onep1, onep2));
}

Paul's avatar
Paul committed
257
258
TEST_CASE(zero_merge1)
{
Paul's avatar
Paul committed
259
    schedule_target t{};
Paul's avatar
Paul committed
260
261
262
263
264
    migraphx::program p;
    auto one    = p.add_literal(1);
    auto onep1  = p.add_instruction(unary_op{}, one);
    auto onep2  = p.add_instruction(unary_op{}, one);
    auto binary = p.add_instruction(migraphx::op::identity{}, onep1, onep2);
Paul's avatar
Paul committed
265
266
267
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(onep1) != t.get_stream(onep2));
Paul's avatar
Paul committed
268
    // No stream assignment
Paul's avatar
Paul committed
269
    EXPECT(not t.has_stream(binary));
Paul's avatar
Paul committed
270
271
272
273
274
275
276
    // There is no wait
    EXPECT(get_wait_for(binary).empty());
    EXPECT(check_conflicts(p, onep1, onep2));
}

TEST_CASE(zero_merge2)
{
Paul's avatar
Paul committed
277
    schedule_target t{};
Paul's avatar
Paul committed
278
279
280
281
    migraphx::program p;
    auto one    = p.add_literal(1);
    auto onep1  = p.add_instruction(unary_op{}, one);
    auto onep2  = p.add_instruction(unary_op{}, one);
Paul's avatar
Paul committed
282
283
284
    auto binary = p.add_instruction(migraphx::op::identity{},
                                    p.add_instruction(migraphx::op::identity{}, onep1),
                                    p.add_instruction(migraphx::op::identity{}, onep2));
Paul's avatar
Paul committed
285
286
287
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(onep1) != t.get_stream(onep2));
Paul's avatar
Paul committed
288
    // No stream assignment
Paul's avatar
Paul committed
289
    EXPECT(not t.has_stream(binary));
Paul's avatar
Paul committed
290
291
292
293
294
    // There is no wait
    EXPECT(get_wait_for(binary).empty());
    EXPECT(check_conflicts(p, onep1, onep2));
}

Paul's avatar
Paul committed
295
296
297
298
TEST_CASE(zero_merge3)
{
    schedule_target t{};
    migraphx::program p;
Paul's avatar
Paul committed
299
300
301
302
    auto one   = p.add_literal(1);
    auto onep1 = p.add_instruction(unary_op{}, one);
    auto onep2 = p.add_instruction(unary_op{}, one);
    auto id    = p.add_instruction(migraphx::op::identity{}, onep1, onep2);
Paul's avatar
Paul committed
303
304
305
306
307
308
309
310
311
312
    auto final = p.add_instruction(unary_op{}, id);
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(onep1) != t.get_stream(onep2));
    // No stream assignment
    EXPECT(not t.has_stream(id));
    // There is no wait
    EXPECT(get_wait_for(id).empty());
    // Stream assignment for final op
    EXPECT(t.get_stream(final) == 0);
Paul's avatar
Paul committed
313
314
    EXPECT(get_wait_for(final) ==
           get_wait_for(t.get_stream(final), {t.get_stream(onep1), t.get_stream(onep2)}));
Paul's avatar
Paul committed
315
316
317
318
319
320
321
    EXPECT(check_conflicts(p, onep1, onep2));
}

TEST_CASE(zero_merge4)
{
    schedule_target t{};
    migraphx::program p;
Paul's avatar
Paul committed
322
323
324
325
326
327
    auto one   = p.add_literal(1);
    auto onep1 = p.add_instruction(unary_op{}, one);
    auto onep2 = p.add_instruction(unary_op{}, one);
    auto id    = p.add_instruction(migraphx::op::identity{},
                                p.add_instruction(migraphx::op::identity{}, onep1),
                                p.add_instruction(migraphx::op::identity{}, onep2));
Paul's avatar
Paul committed
328
329
330
331
332
333
334
335
336
337
    auto final = p.add_instruction(unary_op{}, id);
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(onep1) != t.get_stream(onep2));
    // No stream assignment
    EXPECT(not t.has_stream(id));
    // There is no wait
    EXPECT(get_wait_for(id).empty());
    // Stream assignment for final op
    EXPECT(t.get_stream(final) == 0);
Paul's avatar
Paul committed
338
339
    EXPECT(get_wait_for(final) ==
           get_wait_for(t.get_stream(final), {t.get_stream(onep1), t.get_stream(onep2)}));
Paul's avatar
Paul committed
340
341
342
    EXPECT(check_conflicts(p, onep1, onep2));
}

Paul's avatar
Paul committed
343
TEST_CASE(double_entry)
Paul's avatar
Paul committed
344
{
Paul's avatar
Paul committed
345
    schedule_target t{};
Paul's avatar
Paul committed
346
    migraphx::program p;
Paul's avatar
Paul committed
347
348
349
350
    auto one    = p.add_literal(1);
    auto two    = p.add_literal(2);
    auto onep   = p.add_instruction(unary_op{}, one);
    auto twop   = p.add_instruction(unary_op{}, two);
Paul's avatar
Paul committed
351
    auto binary = p.add_instruction(nary_op{}, onep, twop);
Paul's avatar
Paul committed
352
353
354
355
356
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(not t.has_stream(two));
    EXPECT(t.get_stream(onep) != t.get_stream(twop));
    EXPECT(t.get_stream(binary) == 0);
Paul's avatar
Paul committed
357
358
    EXPECT(get_wait_for(binary) ==
           get_wait_for(t.get_stream(binary), {t.get_stream(onep), t.get_stream(twop)}));
Paul's avatar
Paul committed
359
    EXPECT(check_conflicts(p, onep, twop));
Paul's avatar
Paul committed
360
361
}

Paul's avatar
Paul committed
362
TEST_CASE(two_branches)
Paul's avatar
Paul committed
363
{
Paul's avatar
Paul committed
364
    schedule_target t{};
Paul's avatar
Paul committed
365
366
    migraphx::program p;
    auto one    = p.add_literal(1);
Paul's avatar
Paul committed
367
368
    auto c1     = chain(p, 2, unary_op{}, one);
    auto i1     = p.add_instruction(unary_op{}, one);
Paul's avatar
Paul committed
369
    auto binary = p.add_instruction(nary_op{}, i1, c1.back());
Paul's avatar
Paul committed
370
371
372
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(i1) == 1);
Paul's avatar
Paul committed
373
    for(auto ins : c1)
Paul's avatar
Paul committed
374
375
        EXPECT(t.get_stream(ins) == 0);
    EXPECT(t.get_stream(binary) == 0);
Paul's avatar
Paul committed
376
377
    EXPECT(get_wait_for(binary) ==
           get_wait_for(t.get_stream(binary), {t.get_stream(c1.back()), t.get_stream(i1)}));
Paul's avatar
Paul committed
378
379
380
    check_conflicts(p, {c1, {i1}});
}

Paul's avatar
Paul committed
381
TEST_CASE(four_branches)
Paul's avatar
Paul committed
382
{
Paul's avatar
Paul committed
383
    schedule_target t{};
Paul's avatar
Paul committed
384
385
    migraphx::program p;
    auto one    = p.add_literal(1);
Paul's avatar
Paul committed
386
387
388
389
    auto c1     = chain(p, 4, unary_op{}, one);
    auto c2     = chain(p, 3, unary_op{}, one);
    auto c3     = chain(p, 2, unary_op{}, one);
    auto i1     = p.add_instruction(unary_op{}, one);
Paul's avatar
Paul committed
390
    auto binary = p.add_instruction(nary_op{}, i1, c1.back(), c2.back(), c3.back());
Paul's avatar
Paul committed
391
392
393
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(i1) == 3);
Paul's avatar
Paul committed
394
    for(auto ins : c1)
Paul's avatar
Paul committed
395
        EXPECT(t.get_stream(ins) == 0);
Paul's avatar
Paul committed
396
    for(auto ins : c2)
Paul's avatar
Paul committed
397
        EXPECT(t.get_stream(ins) == 1);
Paul's avatar
Paul committed
398
    for(auto ins : c3)
Paul's avatar
Paul committed
399
400
        EXPECT(t.get_stream(ins) == 2);
    EXPECT(t.get_stream(binary) == 0);
Paul's avatar
Paul committed
401
402
403
404
405
    EXPECT(get_wait_for(binary) == get_wait_for(t.get_stream(binary),
                                                {t.get_stream(c1.back()),
                                                 t.get_stream(c2.back()),
                                                 t.get_stream(c3.back()),
                                                 t.get_stream(i1)}));
Paul's avatar
Paul committed
406
    check_conflicts(p, {c1, c2, c3, {i1}});
Paul's avatar
Paul committed
407
408
}

Paul's avatar
Paul committed
409
TEST_CASE(five_branches)
Paul's avatar
Paul committed
410
{
Paul's avatar
Paul committed
411
    schedule_target t{};
Paul's avatar
Paul committed
412
413
414
415
416
417
418
419
    migraphx::program p;
    auto one    = p.add_literal(1);
    auto c1     = chain(p, 5, unary_op{}, one);
    auto c2     = chain(p, 4, unary_op{}, one);
    auto c3     = chain(p, 3, unary_op{}, one);
    auto c4     = chain(p, 2, unary_op{}, one);
    auto i1     = p.add_instruction(unary_op{}, one);
    auto binary = p.add_instruction(nary_op{}, i1, c1.back(), c2.back(), c3.back(), c4.back());
Paul's avatar
Paul committed
420
421
422
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(i1) == 3);
Paul's avatar
Paul committed
423
    for(auto ins : c1)
Paul's avatar
Paul committed
424
        EXPECT(t.get_stream(ins) == 0);
Paul's avatar
Paul committed
425
    for(auto ins : c2)
Paul's avatar
Paul committed
426
        EXPECT(t.get_stream(ins) == 1);
Paul's avatar
Paul committed
427
    for(auto ins : c3)
Paul's avatar
Paul committed
428
        EXPECT(t.get_stream(ins) == 2);
Paul's avatar
Paul committed
429
    for(auto ins : c4)
Paul's avatar
Paul committed
430
431
        EXPECT(t.get_stream(ins) == 3);
    EXPECT(t.get_stream(binary) == 0);
Paul's avatar
Paul committed
432
433
434
435
436
    EXPECT(get_wait_for(binary) == get_wait_for(t.get_stream(binary),
                                                {t.get_stream(c1.back()),
                                                 t.get_stream(c2.back()),
                                                 t.get_stream(c3.back()),
                                                 t.get_stream(i1)}));
Paul's avatar
Paul committed
437
438
439
440
    check_conflicts(p, {c1, c2, c3, c4});
    check_conflicts(p, {c1, c2, c3, {i1}});
}

Paul's avatar
Paul committed
441
442
TEST_CASE(four_branches_eq)
{
Paul's avatar
Paul committed
443
    schedule_target t{};
Paul's avatar
Paul committed
444
445
446
447
448
449
450
    migraphx::program p;
    auto one    = p.add_literal(1);
    auto onep1  = p.add_instruction(unary_op{}, one);
    auto onep2  = p.add_instruction(unary_op{}, one);
    auto onep3  = p.add_instruction(unary_op{}, one);
    auto onep4  = p.add_instruction(unary_op{}, one);
    auto binary = p.add_instruction(nary_op{}, onep1, onep2, onep3, onep4);
Paul's avatar
Paul committed
451
452
    p.compile(t);
    EXPECT(not t.has_stream(one));
Paul's avatar
Paul committed
453
454
455
456
457
    EXPECT(
        sorted<std::size_t>(
            {t.get_stream(onep1), t.get_stream(onep2), t.get_stream(onep3), t.get_stream(onep4)}) ==
        unique<std::size_t>(
            {t.get_stream(onep1), t.get_stream(onep2), t.get_stream(onep3), t.get_stream(onep4)}));
Paul's avatar
Paul committed
458
    EXPECT(t.get_stream(binary) == 0);
Paul's avatar
Paul committed
459
460
    EXPECT(
        get_wait_for(binary) ==
Paul's avatar
Paul committed
461
462
463
        get_wait_for(
            t.get_stream(binary),
            {t.get_stream(onep1), t.get_stream(onep2), t.get_stream(onep3), t.get_stream(onep4)}));
Paul's avatar
Paul committed
464
465
466
    check_conflicts(p, {{onep1}, {onep2}, {onep3}, {onep4}});
}

Paul's avatar
Paul committed
467
468
TEST_CASE(seq_merge)
{
Paul's avatar
Paul committed
469
    schedule_target t{};
Paul's avatar
Paul committed
470
    migraphx::program p;
Paul's avatar
Paul committed
471
472
473
    auto one     = p.add_literal(1);
    auto c1      = chain(p, 2, unary_op{}, one);
    auto i1      = p.add_instruction(unary_op{}, one);
Paul's avatar
Paul committed
474
475
    auto binary1 = p.add_instruction(nary_op{}, i1, c1.back());

Paul's avatar
Paul committed
476
477
    auto c2      = chain(p, 2, unary_op{}, binary1);
    auto i2      = p.add_instruction(unary_op{}, binary1);
Paul's avatar
Paul committed
478
479
    auto binary2 = p.add_instruction(nary_op{}, i2, c2.back());

Paul's avatar
Paul committed
480
481
    p.compile(t);
    EXPECT(not t.has_stream(one));
Paul's avatar
Paul committed
482

Paul's avatar
Paul committed
483
    EXPECT(t.get_stream(i1) == 2);
Paul's avatar
Paul committed
484
    for(auto ins : c1)
Paul's avatar
Paul committed
485
486
        EXPECT(t.get_stream(ins) == 3);
    EXPECT(t.get_stream(binary1) == 3);
Paul's avatar
Paul committed
487
488
    EXPECT(get_wait_for(binary1) ==
           get_wait_for(t.get_stream(binary1), {t.get_stream(c1.back()), t.get_stream(i1)}));
Paul's avatar
Paul committed
489
490
    check_conflicts(p, {c1, {i1}});

Paul's avatar
Paul committed
491
    EXPECT(t.get_stream(i2) == 3);
Paul's avatar
Paul committed
492
    for(auto ins : c2)
Paul's avatar
Paul committed
493
494
        EXPECT(t.get_stream(ins) == 0);
    EXPECT(t.get_stream(binary2) == 0);
Paul's avatar
Paul committed
495
496
    EXPECT(get_wait_for(binary2) ==
           get_wait_for(t.get_stream(binary2), {t.get_stream(c2.back()), t.get_stream(i2)}));
Paul's avatar
Paul committed
497
498
499
500
501
    check_conflicts(p, {c2, {i2}});
}

TEST_CASE(par_merge)
{
Paul's avatar
Paul committed
502
    schedule_target t{};
Paul's avatar
Paul committed
503
    migraphx::program p;
Paul's avatar
Paul committed
504
505
506
507
    auto one     = p.add_literal(1);
    auto start1  = p.add_instruction(unary_op{}, one);
    auto c1      = chain(p, 3, unary_op{}, start1);
    auto i1      = p.add_instruction(unary_op{}, start1);
Paul's avatar
Paul committed
508
509
    auto binary1 = p.add_instruction(nary_op{}, i1, c1.back());

Paul's avatar
Paul committed
510
511
512
    auto start2  = p.add_instruction(unary_op{}, one);
    auto c2      = chain(p, 2, unary_op{}, start2);
    auto i2      = p.add_instruction(unary_op{}, start2);
Paul's avatar
Paul committed
513
514
515
516
    auto binary2 = p.add_instruction(nary_op{}, i2, c2.back());

    auto binary3 = p.add_instruction(nary_op{}, binary1, binary2);

Paul's avatar
Paul committed
517
518
519
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(binary3) == 0);
Paul's avatar
Paul committed
520

Paul's avatar
Paul committed
521
    EXPECT(t.get_stream(i1) == 2);
Paul's avatar
Paul committed
522
    for(auto ins : c1)
Paul's avatar
Paul committed
523
524
        EXPECT(t.get_stream(ins) == 0);
    EXPECT(t.get_stream(binary1) == 0);
Paul's avatar
Paul committed
525
526
    EXPECT(get_wait_for(binary1) ==
           get_wait_for(t.get_stream(binary1), {t.get_stream(c1.back()), t.get_stream(i1)}));
Paul's avatar
Paul committed
527
528
    check_conflicts(p, {c1, {i1}});

Paul's avatar
Paul committed
529
    EXPECT(t.get_stream(i2) == 1);
Paul's avatar
Paul committed
530
    for(auto ins : c2)
Paul's avatar
Paul committed
531
532
        EXPECT(t.get_stream(ins) == 3);
    EXPECT(t.get_stream(binary2) == 3);
Paul's avatar
Paul committed
533
534
    EXPECT(get_wait_for(binary2) ==
           get_wait_for(t.get_stream(binary2), {t.get_stream(c2.back()), t.get_stream(i2)}));
Paul's avatar
Paul committed
535
536
537
    check_conflicts(p, {c2, {i2}});

    EXPECT(check_conflicts(p, binary1, binary2));
Paul's avatar
Paul committed
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
    check_conflicts(p, {c1, {i1}, c2, {i2}});
}

TEST_CASE(par_merge_multi_entry)
{
    schedule_target t{};
    migraphx::program p;
    auto one     = p.add_literal(1);
    auto start1  = p.add_instruction(unary_op{}, one);
    auto c1      = chain(p, 3, unary_op{}, start1);
    auto i1      = p.add_instruction(unary_op{}, start1);
    auto binary1 = p.add_instruction(nary_op{}, i1, c1.back());

    auto two     = p.add_literal(1);
    auto start2  = p.add_instruction(unary_op{}, two);
    auto c2      = chain(p, 2, unary_op{}, start2);
    auto i2      = p.add_instruction(unary_op{}, start2);
    auto binary2 = p.add_instruction(nary_op{}, i2, c2.back());

    auto binary3 = p.add_instruction(nary_op{}, binary1, binary2);

    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(not t.has_stream(two));
    EXPECT(t.get_stream(binary3) == 0);

    EXPECT(t.get_stream(i1) == 2);
    for(auto ins : c1)
        EXPECT(t.get_stream(ins) == 0);
    EXPECT(t.get_stream(binary1) == 0);
    EXPECT(get_wait_for(binary1) ==
           get_wait_for(t.get_stream(binary1), {t.get_stream(c1.back()), t.get_stream(i1)}));
    check_conflicts(p, {c1, {i1}});

    EXPECT(t.get_stream(i2) == 1);
    for(auto ins : c2)
        EXPECT(t.get_stream(ins) == 3);
    EXPECT(t.get_stream(binary2) == 3);
    EXPECT(get_wait_for(binary2) ==
           get_wait_for(t.get_stream(binary2), {t.get_stream(c2.back()), t.get_stream(i2)}));
    check_conflicts(p, {c2, {i2}});

    EXPECT(check_conflicts(p, binary1, binary2));
    check_conflicts(p, {c1, {i1}, c2, {i2}});
Paul's avatar
Paul committed
582
}
Paul's avatar
Paul committed
583
584
585
586
587
588

TEST_CASE(inception1)
{
    schedule_target t{};
    migraphx::program p;

Paul's avatar
Paul committed
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
    auto i1     = p.add_literal(0);
    auto i2     = p.add_literal(1);
    auto i3     = p.add_literal(1);
    auto i4     = p.add_literal(2);
    auto i7     = p.add_instruction(nary_op{"i7"}, i1, i4, i3, i2);
    auto i8     = p.add_literal(2);
    auto i9     = p.add_instruction(migraphx::op::identity{}, i8);
    auto i10    = p.add_literal(1);
    auto i11    = p.add_instruction(nary_op{"i11"}, i7, i9, i10);
    auto i12    = p.add_literal(2);
    auto i13    = p.add_instruction(migraphx::op::identity{}, i12);
    auto i14    = p.add_literal(1);
    auto i15    = p.add_literal(1);
    auto i16    = p.add_literal(2);
    auto i17    = p.add_instruction(nary_op{"i17"}, i11, i16, i15, i13, i14);
    auto i18    = p.add_literal(2);
    auto i19    = p.add_instruction(migraphx::op::identity{}, i18);
    auto i20    = p.add_literal(1);
    auto i21    = p.add_literal(1);
    auto i22    = p.add_literal(2);
    auto i23    = p.add_instruction(nary_op{"i23"}, i17, i22, i21, i19, i20);
    auto i24    = p.add_literal(1);
    auto i25    = p.add_instruction(nary_op{"i25"}, i23, i24);
    auto i26    = p.add_literal(2);
    auto i27    = p.add_instruction(migraphx::op::identity{}, i26);
    auto i28    = p.add_literal(1);
    auto i29    = p.add_literal(1);
    auto i30    = p.add_literal(2);
    auto i31    = p.add_instruction(nary_op{"i31"}, i25, i30, i29, i27, i28);
    auto i32    = p.add_literal(2);
    auto i33    = p.add_instruction(migraphx::op::identity{}, i32);
    auto i34    = p.add_literal(1);
    auto i35    = p.add_literal(1);
    auto i36    = p.add_literal(2);
    auto i37    = p.add_instruction(nary_op{"i37"}, i31, i36, i35, i33, i34);
    auto i38    = p.add_literal(1);
    auto i39    = p.add_instruction(nary_op{"i39"}, i37, i38);
    auto i41    = p.add_literal(2);
    auto i42    = p.add_instruction(migraphx::op::identity{}, i41);
    auto i43    = p.add_literal(1);
    auto i44    = p.add_literal(1);
    auto i45    = p.add_literal(2);
    auto i48    = p.add_instruction(nary_op{"i48"}, i39, i45, i44, i42, i43);
    auto i49    = p.add_literal(2);
    auto i50    = p.add_instruction(migraphx::op::identity{}, i49);
    auto i51    = p.add_literal(1);
    auto i52    = p.add_literal(1);
    auto i53    = p.add_literal(2);
    auto i54    = p.add_instruction(nary_op{"i54"}, i48, i53, i52, i50, i51);
    auto i55    = p.add_literal(1);
    auto i56    = p.add_instruction(migraphx::op::identity{}, i55);
    auto i57    = p.add_literal(2);
    auto i58    = p.add_instruction(migraphx::op::identity{}, i57);
    auto i59    = p.add_literal(1);
    auto i60    = p.add_literal(2);
    auto i61    = p.add_instruction(nary_op{"i61"}, i54, i60, i59, i58, i56);
    auto i62    = p.add_literal(2);
    auto i63    = p.add_instruction(migraphx::op::identity{}, i62);
    auto i64    = p.add_literal(1);
    auto i65    = p.add_literal(1);
    auto i66    = p.add_literal(2);
    auto i69    = p.add_instruction(nary_op{"i69"}, i39, i66, i65, i63, i64);
    auto i70    = p.add_instruction(migraphx::op::identity{}, i55);
    auto i71    = p.add_literal(2);
    auto i72    = p.add_instruction(migraphx::op::identity{}, i71);
    auto i73    = p.add_literal(1);
    auto i74    = p.add_literal(2);
    auto i75    = p.add_instruction(nary_op{"i75"}, i69, i74, i73, i72, i70);
    auto i77    = p.add_literal(1);
    auto i80    = p.add_instruction(nary_op{"i80"}, i39, i77);
    auto i81    = p.add_instruction(migraphx::op::identity{}, i55);
    auto i82    = p.add_literal(2);
    auto i83    = p.add_instruction(migraphx::op::identity{}, i82);
    auto i84    = p.add_literal(1);
    auto i85    = p.add_literal(2);
    auto i86    = p.add_instruction(nary_op{"i86"}, i80, i85, i84, i83, i81);
    auto i88    = p.add_instruction(migraphx::op::identity{}, i55);
    auto i89    = p.add_literal(2);
    auto i90    = p.add_instruction(migraphx::op::identity{}, i89);
    auto i91    = p.add_literal(1);
    auto i92    = p.add_literal(2);
    auto i94    = p.add_instruction(nary_op{"i94"}, i39, i92, i91, i90, i88);
    auto i96    = p.add_instruction(migraphx::op::identity{}, i55, i94, i75, i61, i86);
    auto i97    = p.add_literal(2);
    auto i98    = p.add_instruction(migraphx::op::identity{}, i97);
    auto i99    = p.add_literal(3);
    auto i100   = p.add_literal(1);
    auto i101   = p.add_literal(2);
    auto output = p.add_instruction(nary_op{"output"}, i96, i101, i100, i98, i99);
Paul's avatar
Paul committed
678
679

    p.compile(t);
Paul's avatar
Paul committed
680
681
682
683
684

    EXPECT(t.get_streams({i7, i11, i17, i23, i25, i31, i37, i39, i94}) ==
           t.get_streams({i7, i7, i7, i7, i7, i7, i7, i7, i7}));
    EXPECT(t.get_streams({i48, i54, i61, output}) ==
           t.get_streams({output, output, output, output}));
Paul's avatar
Paul committed
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
    EXPECT(t.get_streams({i80, i86}) == t.get_streams({i80, i80}));
    EXPECT(t.get_streams({i69, i75}) == t.get_streams({i69, i69}));

    EXPECT(t.get_stream(i7) != t.get_stream(i80));
    EXPECT(t.get_stream(i69) != t.get_stream(i80));
    EXPECT(t.get_stream(i69) != t.get_stream(i7));
    EXPECT(t.get_stream(output) != t.get_stream(i7));
    EXPECT(t.get_stream(output) != t.get_stream(i69));
    EXPECT(t.get_stream(output) != t.get_stream(i80));

    EXPECT(get_wait_for(i48) == get_wait_for({t.get_stream(i39)}));
    EXPECT(get_wait_for(i80) == get_wait_for({t.get_stream(i39)}));
    EXPECT(get_wait_for(i69) == get_wait_for({t.get_stream(i39)}));
    // We dont wait twice
    EXPECT(get_wait_for(i94).empty());
Paul's avatar
Paul committed
700
701
702
703
    EXPECT(
        get_wait_for(output) ==
        get_wait_for(t.get_stream(output),
                     {t.get_stream(i94), t.get_stream(i75), t.get_stream(i61), t.get_stream(i86)}));
Paul's avatar
Paul committed
704

Paul's avatar
Paul committed
705
    check_conflicts(p, {{i80, i86}, {i69, i75}, {i48, i54, i61}, {i94}});
Paul's avatar
Paul committed
706
707
}

Paul's avatar
Paul committed
708
int main(int argc, const char* argv[]) { test::run(argc, argv); }