schedule_test.cpp 30 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
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)
{
Paul's avatar
Paul committed
185
    return unique(wait_for);
Paul's avatar
Paul committed
186
187
}

Paul's avatar
Paul committed
188
189
190
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));
Paul's avatar
Paul committed
191
    return unique(wait_for);
Paul's avatar
Paul committed
192
193
194
195
196
}

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

Paul's avatar
Paul committed
207
208
209
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
210
211
{
    std::vector<migraphx::instruction_ref> result;
Paul's avatar
Paul committed
212
    for(std::size_t i = 0; i < n; i++)
Paul's avatar
Paul committed
213
214
215
216
217
218
    {
        result.push_back(p.add_instruction(x, input));
        input = result.back();
    }
    return result;
}
Paul's avatar
Paul committed
219
220
TEST_CASE(single_entry)
{
Paul's avatar
Paul committed
221
    schedule_target t{};
Paul's avatar
Paul committed
222
223
    migraphx::program p;
    auto one    = p.add_literal(1);
Paul's avatar
Paul committed
224
225
    auto onep1  = p.add_instruction(unary_op{}, one);
    auto onep2  = p.add_instruction(unary_op{}, one);
Paul's avatar
Paul committed
226
    auto binary = p.add_instruction(nary_op{}, onep1, onep2);
Paul's avatar
Paul committed
227
228
229
230
    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
231
232
    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
233
234
235
    EXPECT(check_conflicts(p, onep1, onep2));
}

Paul's avatar
Paul committed
236
237
238
239
240
241
242
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);
Paul's avatar
Paul committed
243
244
    auto onei1  = p.add_instruction(migraphx::op::identity{}, onep1);
    auto onei2  = p.add_instruction(migraphx::op::identity{}, onep2);
Paul's avatar
Paul committed
245
    auto binary = p.add_instruction(nary_op{}, onei1, onei2);
Paul's avatar
Paul committed
246
247
248
249
250
251
252
    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
253
    check_conflicts(p, {{onep1, onei1}, {onep2, onei2}});
Paul's avatar
Paul committed
254
255
}

Paul's avatar
Paul committed
256
257
TEST_CASE(zero_merge1)
{
Paul's avatar
Paul committed
258
    schedule_target t{};
Paul's avatar
Paul committed
259
260
261
262
263
    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
264
265
266
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(onep1) != t.get_stream(onep2));
Paul's avatar
Paul committed
267
    // No stream assignment
Paul's avatar
Paul committed
268
    EXPECT(not t.has_stream(binary));
Paul's avatar
Paul committed
269
270
271
272
273
274
275
    // 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
276
    schedule_target t{};
Paul's avatar
Paul committed
277
278
279
280
    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
281
282
283
    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
284
285
286
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(onep1) != t.get_stream(onep2));
Paul's avatar
Paul committed
287
    // No stream assignment
Paul's avatar
Paul committed
288
    EXPECT(not t.has_stream(binary));
Paul's avatar
Paul committed
289
290
291
292
293
    // There is no wait
    EXPECT(get_wait_for(binary).empty());
    EXPECT(check_conflicts(p, onep1, onep2));
}

Paul's avatar
Paul committed
294
295
296
297
TEST_CASE(zero_merge3)
{
    schedule_target t{};
    migraphx::program p;
Paul's avatar
Paul committed
298
299
300
301
    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
302
303
304
305
306
307
308
309
310
311
    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
312
313
    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
314
315
316
317
318
319
320
    EXPECT(check_conflicts(p, onep1, onep2));
}

TEST_CASE(zero_merge4)
{
    schedule_target t{};
    migraphx::program p;
Paul's avatar
Paul committed
321
322
323
324
325
326
    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
327
328
329
330
331
332
333
334
335
336
    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
337
338
    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
339
340
341
    EXPECT(check_conflicts(p, onep1, onep2));
}

Paul's avatar
Paul committed
342
TEST_CASE(double_entry)
Paul's avatar
Paul committed
343
{
Paul's avatar
Paul committed
344
    schedule_target t{};
Paul's avatar
Paul committed
345
    migraphx::program p;
Paul's avatar
Paul committed
346
347
    auto one    = p.add_instruction(migraphx::op::identity{}, p.add_literal(1));
    auto two    = p.add_instruction(migraphx::op::identity{}, p.add_literal(2));
Paul's avatar
Paul committed
348
349
    auto onep   = p.add_instruction(unary_op{}, one);
    auto twop   = p.add_instruction(unary_op{}, two);
Paul's avatar
Paul committed
350
    auto binary = p.add_instruction(nary_op{}, onep, twop);
Paul's avatar
Paul committed
351
352
353
354
355
    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
356
357
    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
358
    check_conflicts(p, {{onep, one}, {twop, two}});
Paul's avatar
Paul committed
359
360
}

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

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

Paul's avatar
Paul committed
408
TEST_CASE(five_branches)
Paul's avatar
Paul committed
409
{
Paul's avatar
Paul committed
410
    schedule_target t{};
Paul's avatar
Paul committed
411
412
413
414
415
416
417
418
    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
419
420
421
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(i1) == 3);
Paul's avatar
Paul committed
422
    for(auto ins : c1)
Paul's avatar
Paul committed
423
        EXPECT(t.get_stream(ins) == 0);
Paul's avatar
Paul committed
424
    for(auto ins : c2)
Paul's avatar
Paul committed
425
        EXPECT(t.get_stream(ins) == 1);
Paul's avatar
Paul committed
426
    for(auto ins : c3)
Paul's avatar
Paul committed
427
        EXPECT(t.get_stream(ins) == 2);
Paul's avatar
Paul committed
428
    for(auto ins : c4)
Paul's avatar
Paul committed
429
430
        EXPECT(t.get_stream(ins) == 3);
    EXPECT(t.get_stream(binary) == 0);
Paul's avatar
Paul committed
431
432
433
434
435
    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
436
437
438
439
    check_conflicts(p, {c1, c2, c3, c4});
    check_conflicts(p, {c1, c2, c3, {i1}});
}

Paul's avatar
Paul committed
440
441
TEST_CASE(four_branches_eq)
{
Paul's avatar
Paul committed
442
    schedule_target t{};
Paul's avatar
Paul committed
443
444
445
446
447
448
449
    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
450
451
    p.compile(t);
    EXPECT(not t.has_stream(one));
Paul's avatar
Paul committed
452
453
454
455
456
    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
457
    EXPECT(t.get_stream(binary) == 0);
Paul's avatar
Paul committed
458
459
    EXPECT(
        get_wait_for(binary) ==
Paul's avatar
Paul committed
460
461
462
        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
463
464
465
    check_conflicts(p, {{onep1}, {onep2}, {onep3}, {onep4}});
}

Paul's avatar
Paul committed
466
467
TEST_CASE(seq_merge)
{
Paul's avatar
Paul committed
468
    schedule_target t{};
Paul's avatar
Paul committed
469
    migraphx::program p;
Paul's avatar
Paul committed
470
471
472
    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
473
474
    auto binary1 = p.add_instruction(nary_op{}, i1, c1.back());

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

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

Paul's avatar
Paul committed
482
    EXPECT(t.get_stream(i1) == 2);
Paul's avatar
Paul committed
483
    for(auto ins : c1)
Paul's avatar
Paul committed
484
485
        EXPECT(t.get_stream(ins) == 3);
    EXPECT(t.get_stream(binary1) == 3);
Paul's avatar
Paul committed
486
487
    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
488
489
    check_conflicts(p, {c1, {i1}});

Paul's avatar
Paul committed
490
    EXPECT(t.get_stream(i2) == 3);
Paul's avatar
Paul committed
491
    for(auto ins : c2)
Paul's avatar
Paul committed
492
493
        EXPECT(t.get_stream(ins) == 0);
    EXPECT(t.get_stream(binary2) == 0);
Paul's avatar
Paul committed
494
495
    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
496
497
498
499
500
    check_conflicts(p, {c2, {i2}});
}

TEST_CASE(par_merge)
{
Paul's avatar
Paul committed
501
    schedule_target t{};
Paul's avatar
Paul committed
502
    migraphx::program p;
Paul's avatar
Paul committed
503
504
505
506
    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
507
508
    auto binary1 = p.add_instruction(nary_op{}, i1, c1.back());

Paul's avatar
Paul committed
509
510
511
    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
512
513
514
515
    auto binary2 = p.add_instruction(nary_op{}, i2, c2.back());

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

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

Paul's avatar
Paul committed
520
    EXPECT(t.get_stream(i1) == 2);
Paul's avatar
Paul committed
521
    for(auto ins : c1)
Paul's avatar
Paul committed
522
523
        EXPECT(t.get_stream(ins) == 0);
    EXPECT(t.get_stream(binary1) == 0);
Paul's avatar
Paul committed
524
525
    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
526
527
    check_conflicts(p, {c1, {i1}});

Paul's avatar
Paul committed
528
    EXPECT(t.get_stream(i2) == 1);
Paul's avatar
Paul committed
529
    for(auto ins : c2)
Paul's avatar
Paul committed
530
531
        EXPECT(t.get_stream(ins) == 3);
    EXPECT(t.get_stream(binary2) == 3);
Paul's avatar
Paul committed
532
533
    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
534
535
536
    check_conflicts(p, {c2, {i2}});

    EXPECT(check_conflicts(p, binary1, binary2));
Paul's avatar
Paul committed
537
538
539
    check_conflicts(p, {c1, {i1}, c2, {i2}});
}

Paul's avatar
Paul committed
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
TEST_CASE(inner_par_merge)
{
    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 start2  = p.add_instruction(unary_op{}, one);
    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 outer1 = p.add_instruction(unary_op{}, one);
    auto outer2 = p.add_instruction(unary_op{}, one);

    auto output = p.add_instruction(nary_op{}, binary1, binary2, outer1, outer2);

    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(output) == 0);
Paul's avatar
Paul committed
563
564
565
566
567
    EXPECT(get_wait_for(output) == get_wait_for(t.get_stream(output),
                                                {t.get_stream(binary1),
                                                 t.get_stream(binary2),
                                                 t.get_stream(outer1),
                                                 t.get_stream(outer2)}));
Paul's avatar
Paul committed
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590

    EXPECT(t.get_stream(outer1) == 1);
    EXPECT(t.get_stream(outer2) == 2);

    EXPECT(t.get_stream(i1) != t.get_stream(i2));
    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}});

    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}, {outer1}, {outer2}});
}

Paul's avatar
Paul committed
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
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);

Paul's avatar
Paul committed
614
    EXPECT(t.get_stream(i1) != t.get_stream(i2));
Paul's avatar
Paul committed
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
    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}});

    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
631
}
Paul's avatar
Paul committed
632

Paul's avatar
Paul committed
633
634
635
636
637
638
639
TEST_CASE(inner_split1)
{
    schedule_target t{};
    migraphx::program p;
    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
640
641
    auto s1     = p.add_instruction(unary_op{}, c1);
    auto s2     = p.add_instruction(unary_op{}, c1);
Paul's avatar
Paul committed
642
643
644
645
646
647
648
649
650
651
    auto output = p.add_instruction(nary_op{}, i1, s1, s2);
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(i1) == 3);
    for(auto ins : c1)
        EXPECT(t.get_stream(ins) != t.get_stream(i1));
    EXPECT(t.get_stream(s1) != t.get_stream(s2));

    EXPECT(t.get_stream(output) == 0);
    EXPECT(get_wait_for(output) ==
Paul's avatar
Paul committed
652
653
654
           get_wait_for(
               t.get_stream(output),
               {t.get_stream(c1.back()), t.get_stream(i1), t.get_stream(s1), t.get_stream(s2)}));
Paul's avatar
Paul committed
655
656
657
658
659
660
661
662
663
664
    check_conflicts(p, {c1, {i1}, {s1}, {s2}});
}

TEST_CASE(inner_split2)
{
    schedule_target t{};
    migraphx::program p;
    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
665
666
    auto s1     = chain(p, 3, unary_op{}, c1.back());
    auto s2     = chain(p, 4, unary_op{}, c1.back());
Paul's avatar
Paul committed
667
668
669
670
671
672
673
674
675
    auto output = p.add_instruction(nary_op{}, i1, s1.back(), s2.back());
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(i1) == 2);
    for(auto ins : c1)
        EXPECT(t.get_stream(ins) != t.get_stream(i1));
    EXPECT(t.get_stream(s1.back()) != t.get_stream(s2.back()));

    EXPECT(t.get_stream(output) == 0);
Paul's avatar
Paul committed
676
677
678
679
680
    EXPECT(get_wait_for(output) == get_wait_for(t.get_stream(output),
                                                {t.get_stream(c1.back()),
                                                 t.get_stream(i1),
                                                 t.get_stream(s1.back()),
                                                 t.get_stream(s2.back())}));
Paul's avatar
Paul committed
681
682
683
    check_conflicts(p, {c1, {i1}, s1, s2});
}

Paul's avatar
Paul committed
684
685
686
687
688
TEST_CASE(inception_resnet)
{
    schedule_target t{};
    migraphx::program p;
    auto one    = p.add_literal(1);
Paul's avatar
Paul committed
689
    auto input  = p.add_instruction(unary_op{}, one);
Paul's avatar
Paul committed
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
    auto c1     = chain(p, 2, unary_op{}, input);
    auto i1     = p.add_instruction(unary_op{}, input);
    auto binary = p.add_instruction(nary_op{}, i1, c1.back());
    auto output = p.add_instruction(nary_op{}, binary, input);
    p.compile(t);
    EXPECT(not t.has_stream(one));
    EXPECT(t.get_stream(i1) == 2);
    for(auto ins : c1)
        EXPECT(t.get_stream(ins) == 0);
    EXPECT(t.get_stream(binary) == 0);
    EXPECT(get_wait_for(binary) ==
           get_wait_for(t.get_stream(binary), {t.get_stream(c1.back()), t.get_stream(i1)}));
    EXPECT(t.get_stream(output) == 0);
    EXPECT(get_wait_for(output).empty());
    check_conflicts(p, {c1, {i1}});
}

Paul's avatar
Paul committed
707
708
709
710
711
TEST_CASE(inception1)
{
    schedule_target t{};
    migraphx::program p;

Paul's avatar
Paul committed
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
    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
801
802

    p.compile(t);
Paul's avatar
Paul committed
803
804
805
806
807

    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
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
    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
823
824
825
826
    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
827

Paul's avatar
Paul committed
828
    check_conflicts(p, {{i80, i86}, {i69, i75}, {i48, i54, i61}, {i94}});
Paul's avatar
Paul committed
829
830
}

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