schedule.cpp 20.3 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
#include <migraphx/schedule.hpp>
#include <migraphx/program.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/iterator_for.hpp>
5
#include <migraphx/iterator.hpp>
Paul's avatar
Paul committed
6
#include <migraphx/dfor.hpp>
7
#include <migraphx/par_for.hpp>
Paul's avatar
Paul committed
8
9
#include <migraphx/functional.hpp>
#include <migraphx/ranges.hpp>
Paul Fultz II's avatar
Paul Fultz II committed
10
#include <migraphx/dom_info.hpp>
Paul's avatar
Paul committed
11
12
#include <unordered_map>
#include <unordered_set>
13
#include <queue>
14
15
#include <thread>
#include <mutex>
16
17
#include <migraphx/make_op.hpp>

Paul's avatar
Paul committed
18
19
#include <set>
#include <deque>
20
#include <chrono>
Paul Fultz II's avatar
Paul Fultz II committed
21
#include <iomanip>
Paul's avatar
Paul committed
22
23
24
25

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {

26
27
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_SCHEDULE)

Paul's avatar
Paul committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
auto get_inputs()
{
    return [](auto i) { return i->inputs(); };
}

auto get_outputs()
{
    return [](auto i) { return i->outputs(); };
}

struct stream_info
{
    std::unordered_map<instruction_ref, std::size_t> ins2stream;
    std::unordered_map<instruction_ref, std::size_t> weights;
    std::unordered_map<instruction_ref, std::size_t> iweights;
Shucai Xiao's avatar
Shucai Xiao committed
43
44
45
    ins_dep_map mod_implicit_deps;

    void calc_implicit_deps(const module& p) { mod_implicit_deps = p.calc_implicit_deps(); }
Paul's avatar
Paul committed
46
47
48
49
50
51
52
53
54
55

    void accumulate_weights(instruction_ref last, const schedule_model& model)
    {
        fix<std::size_t>([&](auto self, auto ins) -> std::size_t {
            if(not contains(weights, ins))
            {
                std::size_t weight = 0;
                auto&& op          = ins->get_operator();
                if(not is_context_free(op) and op.name()[0] != '@')
                    weight = model.weight(op);
56
57
58
                // This will ensure a stream will be assigned to return
                if(op.name() == "@return")
                    weight = 1;
Paul's avatar
Paul committed
59
                iweights[ins] = weight;
Shucai Xiao's avatar
Shucai Xiao committed
60
61
62
63
64
65
66
67
68
69
70
                auto inputs   = ins->inputs();
                if(contains(mod_implicit_deps, ins))
                {
                    const auto& impl_deps = mod_implicit_deps.at(ins);
                    inputs.insert(inputs.end(), impl_deps.begin(), impl_deps.end());
                }

                weights[ins] = std::accumulate(
                    inputs.begin(), inputs.end(), weight, [&](std::size_t w, instruction_ref i) {
                        return w + self(i);
                    });
Paul's avatar
Paul committed
71
72
73
74
75
            }
            return weights[ins];
        })(last);
    }

76
77
78
79
80
81
82
83
84
85
86
    template <class Compare>
    void sort_args_by_weight(std::vector<instruction_ref>& args, Compare compare) const
    {
        if(args.size() < 2)
            return;
        std::sort(args.begin(), args.end(), by(compare, [this](auto x) {
                      return std::make_tuple(
                          this->weights.at(x), x->inputs().size(), std::addressof(*x));
                  }));
    }

Paul's avatar
Paul committed
87
88
89
90
91
92
93
    std::vector<instruction_ref>::iterator sort_args(std::vector<instruction_ref>& args)
    {
        if(args.size() < 2)
        {
            return args.end();
        }

Paul Fultz II's avatar
Paul Fultz II committed
94
        const std::size_t min_partition_threshold = 2;
95
        sort_args_by_weight(args, std::greater<>{});
Paul's avatar
Paul committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

        auto it = std::lower_bound(std::next(args.begin()),
                                   args.end(),
                                   min_partition_threshold,
                                   [&](auto i, std::size_t w) { return this->weights[i] > w; });
        assert(it == args.end() or this->weights[*it] <= min_partition_threshold);
        assert(it == args.end() or std::prev(it) == args.begin() or
               this->weights[*std::prev(it)] > min_partition_threshold);
        return it;
    }

    struct partition
    {
        std::size_t weight = 0;
        std::vector<instruction_ref> instructions{};

        void add(instruction_ref ins, std::size_t w)
        {
            weight += w;
            instructions.push_back(ins);
        }
    };

119
    std::size_t assign_streams(module& p, std::size_t n)
Paul's avatar
Paul committed
120
    {
121
        assert(n > 0);
Paul's avatar
Paul committed
122
123
124
125
        partition critical;
        std::unordered_map<instruction_ref, std::deque<partition>> partitions;
        partitions.reserve(weights.size());
        fix([&](auto self, auto ins, auto& part) {
126
            assert(not is_end(ins, p.end()));
Shucai Xiao's avatar
Shucai Xiao committed
127
128
            if(not p.has_instruction(ins))
                return;
129
130
            if(contains(partitions, ins))
                return;
Shucai Xiao's avatar
Shucai Xiao committed
131

Paul's avatar
Paul committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
            // Add an entry so we know the instruction was visited
            partitions[ins];
            part.add(ins, this->iweights[ins]);

            auto args         = ins->inputs();
            auto threshold_it = this->sort_args(args);

            if(not args.empty())
            {
                assert(threshold_it != args.begin());
                self(args.front(), part);
                for(auto i : range(std::next(args.begin()), threshold_it))
                {
                    partitions[ins].emplace_back();
                    self(i, partitions[ins].back());
                }
                for(auto i : range(threshold_it, args.end()))
                {
                    self(i, part);
                }
            }
            // Sort instructions
            p.move_instruction(ins, p.end());
        })(std::prev(p.end()), critical);

        // Set the critical partition to stream 0
        set_stream(critical, 0);
159
160
161
162
163
164
165
166
167
        if(n == 1)
        {
            // Assign streams for the other partitions
            for(auto&& ins_part : partitions)
                for(auto&& part : ins_part.second)
                    set_stream(part, 0);
            return 1;
        }
        else
Paul's avatar
Paul committed
168
        {
169
170
171
            std::vector<std::size_t> streams(n - 1);
            // Assign streams for the other partitions
            for(auto&& ins_part : partitions)
Paul's avatar
Paul committed
172
            {
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
                std::sort(ins_part.second.begin(),
                          ins_part.second.end(),
                          by(std::greater<>{}, [](auto&& x) {
                              return std::make_tuple(x.weight, x.instructions.size());
                          }));
                for(auto&& part : ins_part.second)
                {
                    auto stream =
                        std::min_element(streams.begin(), streams.end()) - streams.begin();
                    set_stream(part, stream + 1);
                    streams[stream] += part.weight;
                }
            }
            return 1 + std::count_if(streams.begin(), streams.end(), [](auto x) { return x > 0; });
        }
    }

    using weight_ins = std::pair<std::size_t, instruction_ref>;
    struct compare_weight_ins
    {
        bool operator()(const weight_ins& x, const weight_ins& y) const
        {
            return std::make_pair(x.first, std::addressof(*x.second)) <
                   std::make_pair(y.first, std::addressof(*y.second));
        }
    };

200
    void sort(module& p, std::size_t)
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
    {
        std::set<weight_ins, compare_weight_ins> children;
        std::unordered_map<instruction_ref, std::size_t> visited;
        auto last      = std::prev(p.end());
        auto mw        = this->weights.at(last);
        auto nw        = mw / (p.size() + 1);
        auto add_child = [&](auto ins) {
            auto x  = 1 + (mw - this->weights.at(ins)) / (nw + 1);
            auto w  = x * this->iweights.at(ins);
            auto& v = visited[ins];
            auto it = children.find(std::make_pair(v * w, ins));
            if(it == children.end())
            {
                v++;
                children.insert(std::make_pair(v * w, ins));
            }
        };
        add_child(last);

        while(not children.empty())
        {
            // Pop the first element
            auto top = children.begin()->second;
            children.erase(children.begin());
            p.move_instruction(top, p.begin());
            for(auto ins : top->inputs())
            {
Shucai Xiao's avatar
Shucai Xiao committed
228
229
                if(not p.has_instruction(ins))
                    continue;
230
                add_child(ins);
Paul's avatar
Paul committed
231
            }
Shucai Xiao's avatar
Shucai Xiao committed
232
233
234
235
236
237
238
239
240

            if(contains(mod_implicit_deps, top))
            {
                for(auto ins : mod_implicit_deps.at(top))
                {
                    assert(p.has_instruction(ins));
                    add_child(ins);
                }
            }
Paul's avatar
Paul committed
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
        }
    }

    void set_stream(const partition& p, std::size_t n)
    {
        for(auto ins : p.instructions)
            if(iweights[ins] > 0)
                set_stream(ins, n);
    }

    void set_stream(instruction_ref ins, std::size_t n)
    {
        assert(iweights[ins] > 0);
        ins2stream[ins] = n;
    }

    std::size_t get_stream(instruction_ref ins) const { return ins2stream.at(ins); }

    bool has_stream(instruction_ref ins) const { return contains(ins2stream, ins); }

    template <class F>
    bool different(F f, std::size_t stream) const
    {
        bool result = false;
        f([&](auto s) {
            if(s != stream)
            {
                result = true;
                return false;
            }
            // cppcheck-suppress uselessAssignmentArg
            stream = s;
            return true;
        });
        return result;
    }

    template <class F>
    bool different(F f) const
    {
        bool result = false;
        f([&](auto s) {
            result = this->different(f, s);
            return false;
        });
        return result;
    }

    template <class Selector>
    auto get_streams_from(instruction_ref start, Selector select) const
    {
        return [=](auto f) {
            return fix<bool>([&](auto self, auto ins) {
294
                return all_of(select(ins), [&](auto i) {
Paul's avatar
Paul committed
295
                    if(iweights.at(i) == 0)
296
                        return self(i);
Paul's avatar
Paul committed
297
                    else
298
299
                        return f(this->get_stream(i));
                });
Paul's avatar
Paul committed
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
            })(start);
        };
    }

    std::unordered_set<std::size_t> get_streams(instruction_ref ins) const
    {
        if(has_stream(ins))
            return {get_stream(ins)};
        std::unordered_set<std::size_t> result;
        get_streams_from(ins, get_inputs())([&](auto s) {
            result.insert(s);
            return true;
        });
        return result;
    }

    template <class... Ts>
    bool is_merge_point(instruction_ref ins, Ts... xs) const
    {
        return different(get_streams_from(ins, get_inputs()), xs...);
    }

    template <class... Ts>
    bool is_split_point(instruction_ref ins, Ts... xs) const
    {
        return different(get_streams_from(ins, get_outputs()), xs...);
    }

    std::vector<instruction_ref> get_recorded_instructions(instruction_ref start)
    {
        std::vector<instruction_ref> result;
        std::unordered_map<std::size_t, instruction_ref> m;
        fix([&](auto self, auto ins) {
            for(auto i : ins->inputs())
            {
                if(iweights.at(i) == 0)
                {
                    self(i);
                    continue;
                }
                auto stream = this->get_stream(i);
                if(not contains(m, stream))
                    m[stream] = i;
                else
                    m[stream] = std::min(m[stream], i, by(std::less<>{}, [&](auto x) {
                                             return std::distance(x, start);
                                         }));
            }
        })(start);
        std::transform(
            m.begin(), m.end(), std::back_inserter(result), [](auto&& p) { return p.second; });
        return result;
    }

    std::unordered_map<instruction_ref, std::vector<std::vector<instruction_ref>>>
355
    find_concurrent_instructions(module& p) const
Paul's avatar
Paul committed
356
357
358
    {
        std::unordered_map<instruction_ref, std::vector<std::vector<instruction_ref>>> result;
        std::unordered_map<instruction_ref, std::unordered_set<instruction_ref>> merge_from;
Paul Fultz II's avatar
Paul Fultz II committed
359
        dominator_info di = compute_dominator(p);
Paul's avatar
Paul committed
360
361
362
363
364
365
        result.reserve(p.size());
        merge_from.reserve(p.size());
        for(auto ins : reverse_iterator_for(p))
        {
            for(auto&& arg : ins->outputs())
            {
Shucai Xiao's avatar
Shucai Xiao committed
366
367
                if(not p.has_instruction(arg))
                    continue;
Paul's avatar
Paul committed
368
369
370
371
372
                if(is_merge_point(arg))
                    merge_from[ins].insert(arg);
                merge_from[ins].insert(merge_from[arg].begin(), merge_from[arg].end());
            }

Paul Fultz II's avatar
Paul Fultz II committed
373
374
375
376
377
            if(is_split_point(ins))
            {
                erase_if(merge_from[ins],
                         [&](auto merge) { return di.strictly_dominate(ins, merge); });
            }
Paul's avatar
Paul committed
378

Paul Fultz II's avatar
Paul Fultz II committed
379
            auto streams = this->get_streams(ins);
Paul's avatar
Paul committed
380
            // Collect concur instructions for each merge point.
381
            for(const auto& merge : merge_from[ins])
Paul's avatar
Paul committed
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
            {
                for(auto stream : streams)
                {
                    if(result[merge].size() <= stream)
                        result[merge].resize(stream + 1);
                    auto&& r = result[merge][stream];
                    r.push_back(ins);
                    // Copy inputs if they dont have a stream(and are not a builtin and context
                    // free). Inputs without a stream can have a implicit dependency
                    std::copy_if(ins->inputs().begin(),
                                 ins->inputs().end(),
                                 std::back_inserter(r),
                                 [&](auto x) {
                                     return not this->has_stream(x) and
                                            not is_context_free(x->get_operator()) and
                                            x->name().front() != '@';
                                 });
                }
            }
        }
        return result;
    }

    std::unordered_map<instruction_ref, std::unordered_set<instruction_ref>>
406
    get_conflicts(module& p)
Paul's avatar
Paul committed
407
    {
Paul Fultz II's avatar
Paul Fultz II committed
408

409
410
411
        using conflict_table_type =
            std::unordered_map<instruction_ref, std::unordered_set<instruction_ref>>;
        conflict_table_type conflict_table;
Paul's avatar
Paul committed
412
        auto concur_ins = this->find_concurrent_instructions(p);
413

Paul Fultz II's avatar
Paul Fultz II committed
414
415
416
417
418
419
        // Compute an index for each instruction
        std::unordered_map<instruction_ref, std::size_t> ins2index;
        std::size_t index_total = 0;
        for(auto ins : iterator_for(p))
            ins2index[ins] = index_total++;

420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
        std::vector<conflict_table_type> thread_conflict_tables(
            std::thread::hardware_concurrency());
        std::vector<instruction_ref> index_to_ins;
        index_to_ins.reserve(concur_ins.size());
        std::transform(concur_ins.begin(),
                       concur_ins.end(),
                       std::back_inserter(index_to_ins),
                       [](auto&& it) { return it.first; });

        par_for(concur_ins.size(), [&](auto ins_index, auto tid) {
            auto merge_first = index_to_ins[ins_index];
            assert(concur_ins.count(merge_first) > 0);
            auto& merge_second = concur_ins.at(merge_first);

            // ensure there are enough elements for different threads
            assert(tid < thread_conflict_tables.size());
            auto& thrd_table = thread_conflict_tables.at(tid);

            std::unordered_set<instruction_ref> checked_ins_set;
            auto range_i = range(merge_second.begin(), std::prev(merge_second.end()));
            for(auto it_i : iterator_for(range_i))
            {
                std::unordered_set<instruction_ref> ins1_set;
                std::copy_if(it_i->begin(),
                             it_i->end(),
                             std::inserter(ins1_set, ins1_set.end()),
                             [&](auto i) { return not contains(checked_ins_set, i); });
                checked_ins_set.insert(ins1_set.begin(), ins1_set.end());

                auto range_j = range(std::next(it_i), merge_second.end());
                std::unordered_set<instruction_ref> ins2_set;
                for(auto it_j : iterator_for(range_j))
                {
                    std::copy_if(it_j->begin(),
                                 it_j->end(),
                                 std::inserter(ins2_set, ins2_set.end()),
                                 [&](auto i) { return not contains(checked_ins_set, i); });
                }

                for(auto ins1 : ins1_set)
Paul's avatar
Paul committed
460
                {
Paul Fultz II's avatar
Paul Fultz II committed
461
                    auto p1 = ins2index.at(ins1);
462
                    for(auto ins2 : ins2_set)
Paul's avatar
Paul committed
463
464
465
                    {
                        if(ins1 == ins2)
                            continue;
Paul Fultz II's avatar
Paul Fultz II committed
466
467
                        auto p2 = ins2index.at(ins2);
                        if(p2 > p1)
468
                            thrd_table[ins2].insert(ins1);
Paul's avatar
Paul committed
469
                        else
470
                            thrd_table[ins1].insert(ins2);
Paul's avatar
Paul committed
471
472
                    }
                }
473
474
475
476
477
478
479
480
481
482
            }
        });

        // merge thread_conflict_tables together
        for(auto& tbl : thread_conflict_tables)
        {
            for(auto& it : tbl)
            {
                conflict_table[it.first].insert(it.second.begin(), it.second.end());
            }
Paul's avatar
Paul committed
483
        }
484
485

        // Remove instructions from the conflict table of an ealier instruction
Paul's avatar
Paul committed
486
487
488
489
490
491
492
        for(auto&& ip : conflict_table)
        {
            auto ins1 = ip.first;
            for(auto ins2 : ip.second)
                if(contains(conflict_table[ins2], ins1))
                    conflict_table[ins2].erase(ins1);
        }
493

Paul's avatar
Paul committed
494
495
496
497
        return conflict_table;
    }
};

498
void schedule::apply(module& p) const
Paul's avatar
Paul committed
499
{
Paul's avatar
Paul committed
500
    if(not enable)
Paul's avatar
Paul committed
501
        return;
Shucai Xiao's avatar
Shucai Xiao committed
502

Paul's avatar
Paul committed
503
    stream_info si;
Shucai Xiao's avatar
Shucai Xiao committed
504
    si.calc_implicit_deps(p);
Paul's avatar
Paul committed
505
506
    auto last = std::prev(p.end());
    si.accumulate_weights(last, model);
507
508
    auto nstreams = si.assign_streams(p, model.concurrency());
    si.sort(p, model.concurrency());
Paul's avatar
Paul committed
509

510
    if(enabled(MIGRAPHX_TRACE_COMPILE{}) or enabled(MIGRAPHX_TRACE_SCHEDULE{}))
Paul's avatar
Paul committed
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
    {
        p.annotate(std::cout, [&](auto ins) {
            std::cout << ":";
            std::cout << " weight=" << si.weights.at(ins);
            std::cout << " input={";
            si.get_streams_from(ins, get_inputs())([&](auto s) {
                std::cout << s << ",";
                return true;
            });
            std::cout << "}";
            if(si.has_stream(ins))
                std::cout << " stream=" << si.get_stream(ins);
        });
        std::cout << std::endl;
    }

527
528
529
530
    // No concurrency
    if(nstreams < 2)
        return;

Paul's avatar
Paul committed
531
532
533
534
535
536
537
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
582
583
584
585
586
587
588
589
590
591
    // Schedule instructions
    std::size_t wait_id = 0;
    std::unordered_map<instruction_ref, std::size_t> ins2wait;
    std::unordered_map<std::size_t, std::unordered_set<std::size_t>> waited_for;
    std::unordered_map<instruction_ref, std::unordered_set<std::size_t>> ins2waited;
    ins2wait.reserve(p.size());
    ins2waited.reserve(p.size());
    for(auto ins : iterator_for(p))
    {
        // Only schedule instructions that have a stream
        if(not si.has_stream(ins))
            continue;
        assert(si.weights[ins] > 0);
        // Schedule instruction on the stream
        auto stream = si.get_stream(ins);
        assert(stream < model.concurrency());
        model.sched(p, ins, stream);
        // Insert wait instructions
        if(si.is_merge_point(ins, stream))
        {
            for(auto i : si.get_recorded_instructions(ins))
            {
                if(not si.has_stream(i))
                    continue;
                auto istream = si.get_stream(i);
                if(stream == istream)
                    continue;
                // Create a new event if it hasn't been recorded
                if(not contains(ins2wait, i))
                {
                    ins2wait[i] = wait_id;
                    model.record(p, i, wait_id);
                    wait_id++;
                }
                auto w = ins2wait.at(i);
                // If we already waited for the event on this stream then dont
                // insert another wait event
                if(not contains(waited_for[stream], w))
                    model.wait(p, ins, w);
                // Store the event as waited
                waited_for[stream].insert(w);
                // Store all wait events that have been waited on prior to the recorded instruction
                waited_for[stream].insert(ins2waited[i].begin(), ins2waited[i].end());
            }
        }
        // Store wait events that have already been waited on
        if(si.is_split_point(ins, stream))
        {
            ins2waited[ins] = waited_for[stream];
        }
    }

    // Add memory conflicts
    auto conflict_table = si.get_conflicts(p);
    for(auto&& ip : conflict_table)
    {
        if(ip.second.empty())
            continue;
        std::vector<instruction_ref> args;
        args.push_back(ip.first);
        args.insert(args.end(), ip.second.begin(), ip.second.end());
592
        p.insert_instruction(std::next(ip.first), make_op("identity"), args);
Paul's avatar
Paul committed
593
594
595
596
597
    }
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx