program.cpp 41.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
24
#include <migraphx/version.h>
25
#include <migraphx/compile_options.hpp>
Paul's avatar
Paul committed
26
27
28
#include <migraphx/program.hpp>
#include <migraphx/stringutils.hpp>
#include <migraphx/instruction.hpp>
29
#include <migraphx/op/identity.hpp>
Paul's avatar
Paul committed
30
#include <migraphx/target.hpp>
Paul's avatar
Paul committed
31
32
33
#include <migraphx/env.hpp>
#include <migraphx/ranges.hpp>
#include <migraphx/time.hpp>
34
#include <migraphx/pass_manager.hpp>
35
#include <migraphx/register_target.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
36
#include <migraphx/iterator_for.hpp>
37
#include <migraphx/iterator.hpp>
38
#include <migraphx/algorithm.hpp>
39
#include <migraphx/output_iterator.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
40
#include <migraphx/make_op.hpp>
41
#include <migraphx/marker.hpp>
varunsh's avatar
varunsh committed
42
#include <migraphx/supported_segments.hpp>
Paul's avatar
Paul committed
43
#include <iostream>
umangyadav's avatar
umangyadav committed
44
#include <queue>
Paul's avatar
Paul committed
45
#include <sstream>
Paul's avatar
Paul committed
46
#include <algorithm>
47
#include <set>
48
#include <unordered_map>
Paul's avatar
Paul committed
49
#include <utility>
50
#include <unordered_set>
Shucai Xiao's avatar
Shucai Xiao committed
51
52
#include <map>
#include <cassert>
Paul's avatar
Paul committed
53

Paul's avatar
Paul committed
54
namespace migraphx {
Paul's avatar
Paul committed
55
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
56

57
58
using milliseconds = std::chrono::duration<double, std::milli>;

59
60
61
62
63
64
65
66
67
68
69
struct mark_instruction_target
{
    std::size_t target_id = 0;
    std::string name() const { return "mark_instruction_target"; }
    void apply(module& m) const
    {
        for(auto& ins : m)
            ins.set_target_id(target_id);
    }
};

Paul's avatar
Paul committed
70
71
struct program_impl
{
Shucai Xiao's avatar
Shucai Xiao committed
72
    // A map is used to keep references to modules of the program
73
    std::unordered_map<std::string, module> modules;
74
    std::vector<context> contexts;
75
    std::vector<target> targets;
Paul's avatar
Paul committed
76
77
};

78
program::program() : impl(std::make_unique<program_impl>()) { this->create_module("main"); }
Paul's avatar
Paul committed
79

Paul's avatar
Paul committed
80
program::program(program&&) noexcept = default;
Shucai Xiao's avatar
Shucai Xiao committed
81
program::~program() noexcept         = default;
Paul's avatar
Paul committed
82

83
// copy constructor
Shucai Xiao's avatar
Shucai Xiao committed
84
program::program(const program& p) { assign(p); }
85
86

// copy assignment operator
Shucai Xiao's avatar
Shucai Xiao committed
87
program& program::operator=(program p)
88
{
Shucai Xiao's avatar
Shucai Xiao committed
89
    std::swap(p.impl, this->impl);
90
91
92
    return *this;
}

Shucai Xiao's avatar
Shucai Xiao committed
93
void program::assign(const program& p)
94
{
95
    if(not impl)
96
97
98
    {
        impl = std::make_unique<program_impl>();
    }
Shucai Xiao's avatar
Shucai Xiao committed
99

100
    *impl = *p.impl;
Shucai Xiao's avatar
Shucai Xiao committed
101
102
103
104

    // build a map from old ins to new ins
    // Build a map from old module to new module
    std::unordered_map<module_ref, module_ref> mod_map;
Paul's avatar
Paul committed
105
106
107
108
109
    std::transform(
        impl->modules.begin(),
        impl->modules.end(),
        std::inserter(mod_map, mod_map.begin()),
        [&](auto&& xp) { return std::make_pair(&p.impl->modules.at(xp.first), &xp.second); });
Shucai Xiao's avatar
Shucai Xiao committed
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

    std::unordered_map<instruction_ref, instruction_ref> ins_map;
    for(auto&& pp : mod_map)
    {
        auto old_ins = iterator_for(*pp.first);
        auto new_ins = iterator_for(*pp.second);
        std::transform(old_ins.begin(),
                       old_ins.end(),
                       new_ins.begin(),
                       std::inserter(ins_map, ins_map.begin()),
                       [](auto x, auto y) { return std::make_pair(x, y); });
    }

    // Update all references from all modules
    for(auto&& mp : impl->modules)
    {
126
        for(auto ins : iterator_for(mp.second))
Shucai Xiao's avatar
Shucai Xiao committed
127
128
            instruction::replace_refs(ins, ins_map, mod_map);
    }
129
130
}

Paul's avatar
Paul committed
131
shape program::get_parameter_shape(std::string name) const
Paul's avatar
Paul committed
132
{
Shucai Xiao's avatar
Shucai Xiao committed
133
134
    const auto* mm = this->get_main_module();
    return mm->get_parameter_shape(std::move(name));
Paul's avatar
Paul committed
135
136
}

137
138
std::vector<std::string> program::get_parameter_names() const
{
Shucai Xiao's avatar
Shucai Xiao committed
139
140
    const auto* mm = this->get_main_module();
    return mm->get_parameter_names();
141
142
}

mei-ye's avatar
mei-ye committed
143
144
instruction_ref program::get_parameter(std::string name) const
{
Shucai Xiao's avatar
Shucai Xiao committed
145
146
    const auto* mm = this->get_main_module();
    return mm->get_parameter(std::move(name));
mei-ye's avatar
mei-ye committed
147
148
}

Paul's avatar
Paul committed
149
150
std::unordered_map<std::string, shape> program::get_parameter_shapes() const
{
Shucai Xiao's avatar
Shucai Xiao committed
151
152
    const auto* mm = this->get_main_module();
    return mm->get_parameter_shapes();
Paul's avatar
Paul committed
153
154
}

Shucai Xiao's avatar
Shucai Xiao committed
155
std::size_t program::size() const { return impl->modules.size(); }
156

157
158
std::vector<shape> program::get_output_shapes() const
{
Shucai Xiao's avatar
Shucai Xiao committed
159
160
    const auto* mm = this->get_main_module();
    return mm->get_output_shapes();
161
}
Paul's avatar
Paul committed
162

163
164
165
166
167
context& program::get_context() const
{
    assert(impl->contexts.size() == 1);
    return impl->contexts.front();
}
Paul's avatar
Paul committed
168

Paul's avatar
Paul committed
169
170
instruction_ref program::validate() const
{
Shucai Xiao's avatar
Shucai Xiao committed
171
172
    const auto* mm = this->get_main_module();
    return mm->validate();
Paul's avatar
Paul committed
173
174
}

175
176
177
178
179
180
/*
Assigns each instruction inside program to a target.
It does it by first finding subgraphs supported on a given target based on assignment options.
It is possible that instructions have multiple target assignments and part of multiple subgraphs.
Current logic is simple and assigns entire subgraph containing supported instruction to a particular
target on first seen basis and doesn't find the "best" target assignment.
umangyadav's avatar
umangyadav committed
181
Assumes that all instructions will have target_assignment after this.
182
*/
183
184
185
186
187
188
189
190
target_assignments program::get_target_assignments(const std::vector<target>& targets,
                                                   assignment_options options)
{
    const auto m = options.metric;

    target_assignments p;

    const auto* mod = get_main_module();
umangyadav's avatar
umangyadav committed
191
    std::vector<std::pair<std::size_t, supported_segments>> target_subgraphs;
varunsh's avatar
varunsh committed
192
    target_subgraphs.reserve(targets.size());
umangyadav's avatar
umangyadav committed
193
194
195
196
    for(auto tid : range(targets.size()))
    {
        target_subgraphs.push_back(std::make_pair(tid, targets[tid].find_supported(mod, m)));
    }
varunsh's avatar
varunsh committed
197
    for(const auto ins : iterator_for(*mod))
198
    {
varunsh's avatar
varunsh committed
199
200
201
202
203
        if(contains(p, ins))
        {
            continue;
        }

umangyadav's avatar
umangyadav committed
204
        for(const auto& [tid, subgraph] : target_subgraphs)
varunsh's avatar
varunsh committed
205
206
        {
            // can't pass a structured binding into lambda in C++17 so create a variable for it
umangyadav's avatar
umangyadav committed
207
            const auto& t = tid;
varunsh's avatar
varunsh committed
208
209
210
            for(const auto& segment : subgraph)
            {
                const auto& instructions = segment.instructions;
umangyadav's avatar
umangyadav committed
211
                if(contains(instructions, ins))
varunsh's avatar
varunsh committed
212
                {
umangyadav's avatar
umangyadav committed
213
214
215
216
                    std::transform(instructions.begin(),
                                   instructions.end(),
                                   std::inserter(p, p.end()),
                                   [&](auto instr) { return std::make_pair(instr, t); });
varunsh's avatar
varunsh committed
217
218
219
                }
            }
        }
220
    }
umangyadav's avatar
umangyadav committed
221

222
223
224
    return p;
}

225
bool program::is_compiled() const { return not this->impl->contexts.empty(); }
226

227
228
229
230
231
232
233
234
235
236
237
238
239
240
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
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
void program::compile(const std::vector<target>& targets, std::vector<compile_options> compile_opts)
{
    // Gather all the target roots
    std::unordered_multimap<std::size_t, module_ref> roots;
    auto mods = this->get_modules();
    for(auto* mod : mods)
    {
        for(const auto& ins : *mod)
        {
            if(ins.name() != "run_on_target")
                continue;
            auto v                     = ins.get_operator().to_value();
            module_ref root            = ins.module_inputs().front();
            std::size_t root_target_id = v.at("target_id").to<std::size_t>();
            assert(root_target_id < targets.size());
            roots.insert({root_target_id, root});
        }
    }

    auto trace = tracer{};
    // TODO: Add tracer based on compile options
    if(enabled(MIGRAPHX_TRACE_COMPILE{}))
        trace = tracer{std::cout};

    trace(*this);
    trace();
    // It is assumed that all instructions outside of any root module would run on "ref" target
    // Ref target may or may not be passed as one of the target for the "compile()".
    // If it is not passed, Create one and add context of it into the map.
    auto target_idx = [&](const std::string& t_name) {
        return static_cast<std::size_t>(
            std::find_if(
                targets.begin(), targets.end(), [&](const auto& t) { return t.name() == t_name; }) -
            targets.begin());
    };

    std::size_t ref_target_id = target_idx("ref");
    if(ref_target_id == targets.size())
    {
        this->impl->contexts.resize(targets.size() + 1);
        this->impl->contexts[ref_target_id] = migraphx::make_target("ref").get_context();
        // users could pass lessers compile_ops than targets, in that case use default compile_opts
        compile_opts.resize(targets.size() + 1, migraphx::compile_options{});
    }
    else
    {
        this->impl->contexts.resize(targets.size());
        compile_opts.resize(targets.size(), migraphx::compile_options{});
    }
    // mark all the instruction as ref target first, later change target_id based on root-target
    run_passes(*this, {mark_instruction_target{ref_target_id}});

    // Run passes on each root target
    for(const auto i : range(targets.size()))
    {
        const auto& root_target              = targets.at(i);
        auto root_target_id                  = i;
        auto root_modules_range              = roots.equal_range(root_target_id);
        this->impl->contexts[root_target_id] = root_target.get_context();
        for(const auto& [id, current_mod] : range(root_modules_range))
        {
            auto passes = root_target.get_passes(this->impl->contexts[root_target_id],
                                                 compile_opts[root_target_id]);
            passes.push_back(mark_instruction_target{static_cast<size_t>(root_target_id)});
            run_passes(*this, current_mod, passes, trace);

            auto invalid = current_mod->validate();
            if(invalid != current_mod->end())
            {
                MIGRAPHX_THROW("Invalid module " + current_mod->name() +
                               " from compilation at instruction " +
                               std::to_string(std::distance(current_mod->begin(), invalid)));
            }
            auto dangling = current_mod->find_dangling_reference();
            if(dangling != current_mod->end())
            {
                auto index = std::distance(current_mod->begin(), dangling);
                MIGRAPHX_THROW("Dangling reference in module " + current_mod->name() +
                               " from instruction " + std::to_string(index));
            }
        }
    }
309
    this->finalize();
310
311
}

312
void program::compile(const target& t, compile_options options)
Paul's avatar
Paul committed
313
{
314
    // todo: combine with multi-target compile method
315
    assert(not this->is_compiled());
316
317
    this->impl->targets  = {t};
    this->impl->contexts = {t.get_context()};
318

Paul's avatar
Paul committed
319
    if(enabled(MIGRAPHX_TRACE_COMPILE{}))
320
        options.trace = tracer{std::cout};
Shucai Xiao's avatar
Shucai Xiao committed
321

322
323
    options.trace(*this);
    options.trace();
324
    auto&& passes = t.get_passes(this->impl->contexts.front(), options);
325
326
327
328
    run_passes(*this, passes, options.trace);
    auto mods = this->get_modules();
    // Validate and finalize
    for(const auto& mod : reverse(mods))
Paul's avatar
Paul committed
329
    {
Shucai Xiao's avatar
Shucai Xiao committed
330
331
332
333
334
335
        auto invalid = mod->validate();
        if(invalid != mod->end())
        {
            MIGRAPHX_THROW("Invalid module " + mod->name() + " from compilation at instruction " +
                           std::to_string(std::distance(mod->begin(), invalid)));
        }
336
337
338
339
340
341
342
        auto dangling = mod->find_dangling_reference();
        if(dangling != mod->end())
        {
            auto index = std::distance(mod->begin(), dangling);
            MIGRAPHX_THROW("Dangling reference in module " + mod->name() + " from instruction " +
                           std::to_string(index));
        }
343
        mod->finalize(this->impl->contexts);
Paul's avatar
Paul committed
344
    }
Paul's avatar
Paul committed
345
346
347
348
}

void program::finalize()
{
Shucai Xiao's avatar
Shucai Xiao committed
349
    auto* mm = this->get_main_module();
350
    mm->finalize(this->impl->contexts);
Paul's avatar
Paul committed
351
352
}

353
354
355
356
357
358
359
360
361
362
363
364
365
366
template <class T>
std::string classify(T x)
{
    switch(std::fpclassify(x))
    {
    case FP_INFINITE: return "inf";
    case FP_NAN: return "nan";
    case FP_NORMAL: return "normal";
    case FP_SUBNORMAL: return "subnormal";
    case FP_ZERO: return "zero";
    default: return "unknown";
    }
}

367
368
369
370
371
372
373
void print_statistics(std::ostream& os, const argument& a)
{
    a.visit(
        [&](auto t) {
            os << "Min value: " << *std::min_element(t.begin(), t.end()) << ", ";
            os << "Max value: " << *std::max_element(t.begin(), t.end()) << ", ";
            double num_elements = t.size();
Umang Yadav's avatar
Umang Yadav committed
374
            auto mean           = std::accumulate(t.begin(), t.end(), 0.0) / num_elements;
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
            auto stddev         = std::sqrt(
                std::accumulate(t.begin(),
                                t.end(),
                                0.0,
                                [&](auto r, auto v) { return r + std::pow((v - mean), 2.0); }) /
                num_elements);
            os << "Mean: " << mean << ", ";
            os << "StdDev: " << stddev << "\n";
        },
        [&](const auto& xs) {
            for(const auto& x : xs)
            {
                print_statistics(os, x);
            }
        });
}

392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
std::unordered_set<std::string> classify_argument(const argument& a)
{
    std::unordered_set<std::string> result;
    a.visit(
        [&](auto t) {
            for(const auto& x : t)
                result.insert(classify(x));
        },
        [&](const auto& xs) {
            for(const auto& x : xs)
            {
                auto r = classify_argument(x);
                result.insert(r.begin(), r.end());
            }
        });
    return result;
}

void preview_argument(std::ostream& os, const argument& a)
{
    a.visit(
        [&](auto t) {
            if(t.size() <= 10)
            {
                os << t;
            }
            else
            {
                os << to_string_range(t.begin(), t.begin() + 5);
                os << ", ..., ";
                os << to_string_range(t.end() - 5, t.end());
            }
        },
        [&](const auto& xs) {
            for(const auto& x : xs)
            {
                os << '{';
                preview_argument(os, x);
                os << '}';
            }
        });
}

Paul's avatar
Paul committed
435
template <class F>
Shucai Xiao's avatar
Shucai Xiao committed
436
std::vector<argument> generic_eval(const module* mod,
437
                                   std::vector<context>& ctx,
438
                                   std::unordered_map<std::string, argument> params,
Shucai Xiao's avatar
Shucai Xiao committed
439
                                   std::unordered_map<instruction_ref, argument> results,
440
                                   F trace)
Paul's avatar
Paul committed
441
{
Shucai Xiao's avatar
Shucai Xiao committed
442
443
    assert(mod->validate() == mod->end());
    results.reserve(mod->size() * 2);
Paul's avatar
Paul committed
444
445
    std::vector<argument> values;
    values.reserve(16);
Shucai Xiao's avatar
Shucai Xiao committed
446
    for(auto ins : iterator_for(*mod))
Paul's avatar
Paul committed
447
    {
448
        assert(results.find(ins) == results.end());
449
450
        const auto& name = ins->name();
        if(name == "@literal")
Paul's avatar
Paul committed
451
        {
Paul's avatar
Paul committed
452
            results.emplace(ins, trace(ins, [&] { return ins->get_literal().get_argument(); }));
Paul's avatar
Paul committed
453
        }
454
        else if(name == "@param")
Paul's avatar
Paul committed
455
        {
Paul's avatar
Paul committed
456
457
458
459
460
            results.emplace(
                ins, trace(ins, [&] {
                    auto param_name = any_cast<builtin::param>(ins->get_operator()).parameter;
                    if(not contains(params, param_name))
                        MIGRAPHX_THROW("Parameter not found: " + param_name);
461
                    auto param = params[param_name];
462
                    // TODO: may want to check correct number of dimensions and/or was within bounds
463
464
                    if(not ins->get_shape().any_of_dynamic() and
                       param.get_shape() != ins->get_shape())
465
                    {
Paul's avatar
Paul committed
466
                        MIGRAPHX_THROW("Incorrect shape {" + to_string(param.get_shape()) +
Brian Pickrell's avatar
Brian Pickrell committed
467
468
                                       "} for parameter: " + param_name +
                                       " should be: " + to_string(ins->get_shape()));
469
                    }
Paul's avatar
Paul committed
470
471
                    return param;
                }));
Paul's avatar
Paul committed
472
        }
473
        else if(name == "@outline")
Paul's avatar
Paul committed
474
        {
Paul's avatar
Paul committed
475
            results.emplace(ins, trace(ins, [&] { return argument{ins->get_shape(), nullptr}; }));
Paul's avatar
Paul committed
476
        }
477
478
479
480
481
482
483
484
485
486
487
488
489
        else if(name == "@return")
        {
            std::vector<argument> prog_outputs;
            std::transform(ins->inputs().begin(),
                           ins->inputs().end(),
                           std::back_inserter(prog_outputs),
                           [&](instruction_ref i) {
                               assert(results.find(i) != results.end());
                               return results[i];
                           });

            return prog_outputs;
        }
Paul's avatar
Paul committed
490
491
        else
        {
Paul's avatar
Paul committed
492
            values.resize(ins->inputs().size());
Paul's avatar
Paul committed
493
494
495
496
497
            std::transform(
                ins->inputs().begin(), ins->inputs().end(), values.begin(), [&](instruction_ref i) {
                    assert(results.find(i) != results.end());
                    return results[i];
                });
Shucai Xiao's avatar
Shucai Xiao committed
498
499
500
            const auto& mod_args = ins->module_inputs();
            auto module_eval     = [&](module_ref smod,
                                   const std::unordered_map<std::string, argument>& inputs) {
501
                return generic_eval(smod, ctx, inputs, results, trace);
Shucai Xiao's avatar
Shucai Xiao committed
502
503
            };

504
505
506
507
508
509
510
511
512
513
            results.emplace(
                ins, trace(ins, [&] {
                    auto op = ins->normalized_operator();
                    if(op.is_context_free())
                        return op.compute(ins->get_shape(), values, mod_args, module_eval);
                    if(ins->get_target_id() >= ctx.size())
                        MIGRAPHX_THROW("No context available for " + op.name());
                    return op.compute(
                        ctx[ins->get_target_id()], ins->get_shape(), values, mod_args, module_eval);
                }));
Paul's avatar
Paul committed
514
        }
515
        assert(results.find(ins) != results.end());
Charlie Lin's avatar
Charlie Lin committed
516
        if(not ins->get_shape().any_of_dynamic())
517
518
519
        {
            assert(results.at(ins).get_shape() == ins->get_shape());
        }
Paul's avatar
Paul committed
520
    }
Shucai Xiao's avatar
Shucai Xiao committed
521
    return {results.at(std::prev(mod->end()))};
Paul's avatar
Paul committed
522
523
}

Shucai Xiao's avatar
Shucai Xiao committed
524
525
template <class F>
std::vector<argument> generic_eval(const program& p,
526
                                   std::vector<context>& ctx,
Shucai Xiao's avatar
Shucai Xiao committed
527
                                   std::unordered_map<std::string, argument> params,
528
                                   F trace)
Shucai Xiao's avatar
Shucai Xiao committed
529
{
Shucai Xiao's avatar
Shucai Xiao committed
530
    const module* mm = p.get_main_module();
531
    return generic_eval(mm, ctx, params, {}, trace);
Shucai Xiao's avatar
Shucai Xiao committed
532
533
}

534
std::vector<argument> program::eval(parameter_map params, execution_environment exec_env) const
Paul's avatar
Paul committed
535
{
536
    auto& contexts = this->impl->contexts;
Paul's avatar
Paul committed
537
538

    auto trace_level = value_of(MIGRAPHX_TRACE_EVAL{});
539
540
541
542
    std::vector<argument> ret;

    if(exec_env.async)
    {
543
544
        assert(contexts.size() == 1);
        contexts.front().wait_for(exec_env.queue);
545
    }
Paul's avatar
Paul committed
546
547

    if(trace_level > 0)
Paul's avatar
Paul committed
548
    {
Shucai Xiao's avatar
Shucai Xiao committed
549
550
551
552
553
554
555
        std::unordered_map<instruction_ref, std::string> ins_out;
        // get instruction names
        this->print([&](auto x, auto ins_names) {
            std::stringstream ss;
            instruction::print(ss, x, ins_names);
            ins_out[x] = ss.str();
        });
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
        ret = generic_eval(*this, contexts, std::move(params), [&](instruction_ref ins, auto f) {
            auto& ctx = contexts[ins->get_target_id()];
            ctx.finish();
            std::cout << "Run instruction: " << ins_out.at(ins) << std::endl;
            timer t{};
            auto result = f();
            double t1   = t.record<milliseconds>();
            ctx.finish();
            double t2 = t.record<milliseconds>();
            std::cout << "Time: " << t1 << "ms, " << t2 << "ms" << std::endl;
            if(trace_level > 1 and ins->name().front() != '@' and ins->name() != "load" and
               not result.empty())
            {
                migraphx::argument buffer;
                try
Umang Yadav's avatar
Umang Yadav committed
571
                {
572
573
                    const target& tgt = this->impl->targets.at(ins->get_target_id());
                    buffer            = tgt.copy_from(result);
Umang Yadav's avatar
Umang Yadav committed
574
                }
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
                catch(const migraphx::exception&)
                {
                    // instruction was run on host then no need to copy buffer from target
                    buffer = result;
                }
                catch(...)
                {
                    MIGRAPHX_THROW("MIGraphX program execution with MIGRAPHX_TRACE_EVAL failed.\n");
                }
                if(trace_level == 2)
                {
                    std::cout << "Output has " << to_string_range(classify_argument(buffer))
                              << std::endl;
                    std::cout << "Output: ";
                    preview_argument(std::cout, buffer);
                    std::cout << std::endl;
                    print_statistics(std::cout, buffer);
                }
                else
                {
                    std::cout << "Output: " << buffer << std::endl;
                }
            }
            return result;
        });
Paul's avatar
Paul committed
600
601
602
    }
    else
    {
603
        ret = generic_eval(*this, contexts, std::move(params), [&](auto&&, auto f) { return f(); });
Paul's avatar
Paul committed
604
    }
605
606
607

    if(exec_env.async)
    {
608
609
        assert(contexts.size() == 1);
        contexts.front().finish_on(exec_env.queue);
610
611
612
    }

    return ret;
Paul's avatar
Paul committed
613
614
}

615
616
617
618
619
620
void program::finish() const
{
    for(const auto& ctx : this->impl->contexts)
        ctx.finish();
}

621
622
623
624
625
626
627
628
629
630
631
632
std::string get_migraphx_version()
{
    std::stringstream ss;
    ss << std::to_string(MIGRAPHX_VERSION_MAJOR) << "." << std::to_string(MIGRAPHX_VERSION_MINOR)
       << "." << std::to_string(MIGRAPHX_VERSION_PATCH);
    return ss.str();
}

/*
program file version is for the data structure or format of the MXR file. Version should be bumped
if any changes occur to the format of the MXR file.
*/
633
const int program_file_version = 6;
634
635
636
637

value program::to_value() const
{
    value result;
638
639
640
641
642
    result["version"]          = program_file_version;
    result["migraphx_version"] = get_migraphx_version();
    result["targets"]          = migraphx::to_value(this->impl->targets);
    result["contexts"]         = migraphx::to_value(this->impl->contexts);
    value module_vals          = value::object{};
Shucai Xiao's avatar
Shucai Xiao committed
643
    std::unordered_map<instruction_ref, std::string> names;
644
    for(auto& mod : this->get_modules())
Shucai Xiao's avatar
Shucai Xiao committed
645
    {
Shucai Xiao's avatar
Shucai Xiao committed
646
647
        value mod_val;
        value nodes;
648
649
        mod_val["name"] = mod->name();
        names           = mod->print(
Shucai Xiao's avatar
Shucai Xiao committed
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
678
679
680
681
682
683
            [&](auto ins, auto ins_names) {
                value node;
                node["output"]     = ins_names.at(ins);
                node["name"]       = ins->name();
                node["shape"]      = migraphx::to_value(ins->get_shape());
                node["normalized"] = ins->is_normalized();
                if(ins->name() == "@literal")
                    node["literal"] = migraphx::to_value(ins->get_literal());
                node["operator"] = ins->get_operator().to_value();
                std::vector<std::string> inputs;
                std::transform(ins->inputs().begin(),
                               ins->inputs().end(),
                               std::back_inserter(inputs),
                               [&](auto i) {
                                   assert(contains(ins_names, i));
                                   return ins_names.at(i);
                               });
                node["inputs"]   = inputs;
                auto module_args = ins->module_inputs();
                if(not module_args.empty())
                {
                    std::vector<std::string> module_inputs;
                    std::transform(module_args.begin(),
                                   module_args.end(),
                                   std::back_inserter(module_inputs),
                                   [&](auto mod_ref) { return mod_ref->name(); });
                    node["module_inputs"] = module_inputs;
                }

                nodes.push_back(node);
            },
            names);
        mod_val["nodes"] = nodes;

684
        module_vals[mod->name()] = mod_val;
Shucai Xiao's avatar
Shucai Xiao committed
685
    }
Shucai Xiao's avatar
Shucai Xiao committed
686
687
688

    result["modules"] = module_vals;

689
690
    return result;
}
Shucai Xiao's avatar
Shucai Xiao committed
691

Shucai Xiao's avatar
Shucai Xiao committed
692
693
694
695
696
static void mod_from_val(module_ref mod,
                         const value& v,
                         std::unordered_map<std::string, instruction_ref>& instructions,
                         const std::unordered_map<std::string, module_ref>& map_mods)
{
697
    const auto& module_val = v.at(mod->name());
Shucai Xiao's avatar
Shucai Xiao committed
698
699
700
701
702
703
704
705
706
    for(const value& node : module_val.at("nodes"))
    {
        instruction_ref output;
        auto name       = node.at("name").to<std::string>();
        auto fields     = node.at("operator");
        auto normalized = node.at("normalized").to<bool>();

        if(name == "@param")
        {
Paul Fultz II's avatar
Paul Fultz II committed
707
708
709
            output = mod->insert_parameter(mod->end(),
                                           fields["parameter"].to<std::string>(),
                                           migraphx::from_value<shape>(node.at("shape")));
Shucai Xiao's avatar
Shucai Xiao committed
710
711
712
        }
        else if(name == "@literal")
        {
Paul Fultz II's avatar
Paul Fultz II committed
713
714
            output =
                mod->insert_literal(mod->end(), migraphx::from_value<literal>(node.at("literal")));
Shucai Xiao's avatar
Shucai Xiao committed
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
        }
        else
        {
            auto op = make_op(name, fields);
            std::vector<instruction_ref> inputs;
            std::transform(node.at("inputs").begin(),
                           node.at("inputs").end(),
                           std::back_inserter(inputs),
                           [&](const value& i) {
                               auto i_name = i.to<std::string>();
                               assert(contains(instructions, i_name));
                               return instructions.at(i_name);
                           });

            std::vector<module_ref> module_inputs;
            if(node.contains("module_inputs"))
            {
                std::transform(node.at("module_inputs").begin(),
                               node.at("module_inputs").end(),
                               std::back_inserter(module_inputs),
                               [&](const value& i) { return map_mods.at(i.to<std::string>()); });

                for(auto& smod : module_inputs)
                {
                    mod_from_val(smod, v, instructions, map_mods);
                }
            }

            if(name == "@return")
            {
                output = mod->add_return(inputs);
            }
            else if(module_inputs.empty())
            {
Paul Fultz II's avatar
Paul Fultz II committed
749
                output = mod->insert_instruction(mod->end(), op, inputs);
Shucai Xiao's avatar
Shucai Xiao committed
750
751
752
            }
            else
            {
Paul Fultz II's avatar
Paul Fultz II committed
753
                output = mod->insert_instruction(mod->end(), op, inputs, module_inputs);
Shucai Xiao's avatar
Shucai Xiao committed
754
755
756
757
758
759
760
            }
        }
        output->set_normalized(normalized);
        instructions[node.at("output").to<std::string>()] = output;
    }
}

761
762
763
764
void program::from_value(const value& v)
{
    auto version = v.at("version").to<int>();
    if(version != program_file_version)
Shucai Xiao's avatar
Shucai Xiao committed
765
    {
766
767
768
769
770
771
772
773
774
775
776
777
778
        MIGRAPHX_THROW(
            "Error: Program version mismatch. MXR file was created using program file version: " +
            std::to_string(version) + ", while installed MIGraphX is using program file version: " +
            std::to_string(program_file_version) +
            ", Try regenerating MXR file using installed MIGraphX and running again.");
    }

    auto migx_version = v.at("migraphx_version").to<std::string>();
    if(migx_version != get_migraphx_version())
    {
        std::cout << "WARNING: MXR File was created using MIGraphX version: " << migx_version
                  << ", while installed MIGraphX is at version: " << get_migraphx_version()
                  << ", operators implementation could be mismatched.";
Shucai Xiao's avatar
Shucai Xiao committed
779
780
    }

781
782
783
    migraphx::from_value(v.at("targets"), this->impl->targets);

    for(auto i : range(this->impl->targets.size()))
784
    {
785
786
        this->impl->contexts.push_back(this->impl->targets[i].get_context());
        this->impl->contexts.back().from_value(v.at("contexts")[i]);
787
788
    }

Shucai Xiao's avatar
Shucai Xiao committed
789
790
    auto module_vals = v.at("modules");
    for(const auto& vv : module_vals)
791
    {
792
        const auto& name = vv.get_key();
Shucai Xiao's avatar
Shucai Xiao committed
793
794
        if(name == "main")
            continue;
795
        impl->modules.emplace(name, name);
796
    }
797
    std::unordered_map<std::string, module_ref> map_mods;
Paul's avatar
Paul committed
798
799
800
801
    std::transform(impl->modules.begin(),
                   impl->modules.end(),
                   std::inserter(map_mods, map_mods.end()),
                   [&](auto&& pp) { return std::make_pair(pp.first, &pp.second); });
Shucai Xiao's avatar
Shucai Xiao committed
802
803
804
805
806

    std::unordered_map<std::string, instruction_ref> map_insts;
    auto* mm = get_main_module();
    mod_from_val(mm, module_vals, map_insts, map_mods);

807
808
809
    // Finalize a compiled model
    if(not this->impl->contexts.empty())
        this->finalize();
810
811
}

Paul's avatar
Paul committed
812
813
814
double common_average(const std::vector<double>& v)
{
    std::size_t n = v.size() / 4;
Paul's avatar
Paul committed
815
816
    double total  = std::accumulate(v.begin() + n, v.end() - n, 0.0);
    return total / std::distance(v.begin() + n, v.end() - n);
Paul's avatar
Paul committed
817
818
}

Paul Fultz II's avatar
Paul Fultz II committed
819
820
821
822
823
824
825
826
std::string perf_group(const operation& op)
{
    auto attr = op.attributes();
    if(attr.contains("group"))
        return attr.at("group").to<std::string>();
    return op.name();
}

827
828
void program::mark(const parameter_map& params, marker&& m)
{
829
    auto& ctx = this->impl->contexts;
830
831
    // Run once by itself
    eval(params);
832
    this->finish();
833
834
    // Start marking
    m.mark_start(*this);
835
    generic_eval(*this, ctx, params, [&](auto ins, auto f) {
836
837
838
839
840
        argument result;
        m.mark_start(ins);
        result = f();
        m.mark_stop(ins);
        return result;
841
    });
842
843
844
    m.mark_stop(*this);
}

845
846
847
848
void program::perf_report(std::ostream& os,
                          std::size_t n,
                          parameter_map params,
                          std::size_t batch) const
Paul's avatar
Paul committed
849
{
850
    auto& ctx = this->impl->contexts;
Paul's avatar
Paul committed
851
852
    // Run once by itself
    eval(params);
853
    this->finish();
Paul's avatar
Paul committed
854
    // Run and time entire program
Paul's avatar
Paul committed
855
856
    std::vector<double> total_vec;
    total_vec.reserve(n);
Paul's avatar
Paul committed
857
    for(std::size_t i = 0; i < n; i++)
Paul's avatar
Paul committed
858
    {
Paul's avatar
Paul committed
859
860
        total_vec.push_back(time<milliseconds>([&] {
            eval(params);
861
            this->finish();
Paul's avatar
Paul committed
862
        }));
Paul's avatar
Paul committed
863
    }
Paul's avatar
Paul committed
864
865
    std::sort(total_vec.begin(), total_vec.end());
    std::unordered_map<instruction_ref, std::vector<double>> ins_vec;
Paul's avatar
Paul committed
866
    // Fill the map
867
    generic_eval(*this, ctx, params, [&](auto ins, auto) {
Paul's avatar
Paul committed
868
        ins_vec[ins].reserve(n);
869
        return argument{ins->get_shape(), nullptr};
870
    });
871

Paul's avatar
Paul committed
872
    // Run and time each instruction
Paul's avatar
Paul committed
873
    for(std::size_t i = 0; i < n; i++)
Paul's avatar
Paul committed
874
    {
875
        generic_eval(*this, ctx, params, [&](auto ins, auto f) {
876
            argument result;
Paul's avatar
Paul committed
877
878
            ins_vec[ins].push_back(time<milliseconds>([&] {
                result = f();
879
                this->impl->contexts[ins->get_target_id()].finish();
Paul's avatar
Paul committed
880
            }));
881
            return result;
882
        });
Paul's avatar
Paul committed
883
    }
Paul's avatar
Paul committed
884
885
    for(auto&& p : ins_vec)
        std::sort(p.second.begin(), p.second.end());
Paul's avatar
Paul committed
886
    // Run and time implicit overhead
Paul's avatar
Paul committed
887
888
    std::vector<double> overhead_vec;
    overhead_vec.reserve(n);
Paul's avatar
Paul committed
889
    for(std::size_t i = 0; i < n; i++)
Paul's avatar
Paul committed
890
    {
Paul's avatar
Paul committed
891
        overhead_vec.push_back(time<milliseconds>([&] { dry_run(params); }));
Paul's avatar
Paul committed
892
893
    }

Paul's avatar
Paul committed
894
    double total_time             = common_average(total_vec);
Paul's avatar
Paul committed
895
    double rate                   = 1000.0 / total_time;
Paul's avatar
Paul committed
896
    double overhead_time          = common_average(overhead_vec);
Paul's avatar
Paul committed
897
    double overhead_percent       = overhead_time * 100.0 / total_time;
Paul's avatar
Paul committed
898
    double total_instruction_time = 0.0;
Paul's avatar
Paul committed
899
    std::unordered_map<std::string, double> op_times;
900
    std::unordered_map<std::string, std::size_t> op_n;
Paul's avatar
Paul committed
901
    for(auto&& p : ins_vec)
Paul's avatar
Paul committed
902
903
    {
        double avg = common_average(p.second);
Paul Fultz II's avatar
Paul Fultz II committed
904
        op_times[perf_group(p.first->get_operator())] += avg;
Paul's avatar
Paul committed
905
        total_instruction_time += avg;
906
        op_n[perf_group(p.first->get_operator())]++;
Paul's avatar
Paul committed
907
    }
Paul's avatar
Paul committed
908
909
    double calculate_overhead_time    = total_time - total_instruction_time;
    double calculate_overhead_percent = calculate_overhead_time * 100.0 / total_time;
Paul's avatar
Paul committed
910

Shucai Xiao's avatar
Shucai Xiao committed
911
912
913
    std::unordered_map<instruction_ref, std::string> names;
    this->print(names, [&](auto ins, auto ins_names) {
        instruction::print(std::cout, ins, ins_names);
914
915
916
917
918

        // skip return instruction
        if(ins->name() == "@return")
            return;

Paul's avatar
Paul committed
919
920
921
        double avg     = common_average(ins_vec[ins]);
        double percent = std::ceil(100.0 * avg / total_instruction_time);
        os << ": " << avg << "ms, " << percent << "%";
922
        os << std::endl;
Paul's avatar
Paul committed
923
    });
Paul's avatar
Paul committed
924
925
926

    os << std::endl;
    os << "Summary:" << std::endl;
927
928
929
930
931
932
    std::vector<std::tuple<double, std::size_t, std::string>> op_times_sorted;
    std::transform(
        op_times.begin(), op_times.end(), std::back_inserter(op_times_sorted), [&](auto p) {
            auto&& name = p.first;
            return std::make_tuple(p.second, op_n.at(name), name);
        });
933
    std::sort(op_times_sorted.begin(), op_times_sorted.end(), std::greater<>{});
934
    for(auto&& [avg, nn, name] : op_times_sorted)
Paul's avatar
Paul committed
935
936
    {
        double percent = std::ceil(100.0 * avg / total_instruction_time);
937
938
939
        double per_ins = avg / nn;
        os << name << ": " << avg << "ms / " << nn << " = " << per_ins << "ms, " << percent << "%"
           << std::endl;
Paul's avatar
Paul committed
940
941
942
    }

    os << std::endl;
Paul's avatar
Paul committed
943

944
945
    os << "Batch size: " << batch << std::endl;
    os << "Rate: " << rate * batch << "/sec" << std::endl;
Paul's avatar
Paul committed
946
947
    os << "Total time: " << total_time << "ms" << std::endl;
    os << "Total instructions time: " << total_instruction_time << "ms" << std::endl;
Paul's avatar
Paul committed
948
949
950
951
    os << "Overhead time: " << overhead_time << "ms"
       << ", " << calculate_overhead_time << "ms" << std::endl;
    os << "Overhead: " << std::round(overhead_percent) << "%"
       << ", " << std::round(calculate_overhead_percent) << "%" << std::endl;
Paul's avatar
Paul committed
952
953
}

Paul's avatar
Paul committed
954
955
void program::debug_print() const { std::cout << *this << std::endl; }
void program::debug_print(instruction_ref ins) const
Paul's avatar
Paul committed
956
{
Shucai Xiao's avatar
Shucai Xiao committed
957
    std::unordered_map<instruction_ref, std::string> names;
958
    if(std::any_of(this->impl->modules.begin(), this->impl->modules.end(), [&](const auto& pp) {
959
           return is_end(pp.second.end(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
960
       }))
Paul's avatar
Paul committed
961
962
963
964
    {
        std::cout << "End instruction" << std::endl;
        return;
    }
965
966
    else if(std::none_of(this->impl->modules.begin(),
                         this->impl->modules.end(),
967
                         [&](const auto& pp) { return pp.second.has_instruction(ins); }))
Paul's avatar
Paul committed
968
969
970
971
    {
        std::cout << "Instruction not part of program" << std::endl;
        return;
    }
Shucai Xiao's avatar
Shucai Xiao committed
972

Paul's avatar
Paul committed
973
    std::stringstream ss;
Shucai Xiao's avatar
Shucai Xiao committed
974
    this->print(names, [&](auto x, auto ins_names) {
Paul's avatar
Paul committed
975
        if(x == ins)
Paul's avatar
Paul committed
976
        {
Shucai Xiao's avatar
Shucai Xiao committed
977
            instruction::print(std::cout, x, ins_names);
Paul's avatar
Paul committed
978
979
980
981
982
            std::cout << std::endl;
        }
    });
}

Shucai Xiao's avatar
Shucai Xiao committed
983
984
985
986
void program::print(
    std::unordered_map<instruction_ref, std::string>& names,
    const std::function<void(instruction_ref, std::unordered_map<instruction_ref, std::string>)>&
        print_func) const
987
{
988
    for(const auto& pp : this->impl->modules)
989
    {
990
        names = pp.second.print(print_func, names);
991
992
993
    }
}

Shucai Xiao's avatar
Shucai Xiao committed
994
995
996
997
998
999
1000
1001
void program::print(
    const std::function<void(instruction_ref ins,
                             std::unordered_map<instruction_ref, std::string>)>& print_func) const
{
    std::unordered_map<instruction_ref, std::string> names;
    this->print(names, print_func);
}

Shucai Xiao's avatar
Shucai Xiao committed
1002
void program::print_graph(std::ostream& os, bool brief) const
1003
{
Shucai Xiao's avatar
Shucai Xiao committed
1004
1005
    const auto* mm = this->get_main_module();
    mm->print_graph(os, brief);
1006
1007
}

1008
1009
1010
1011
1012
1013
1014
void program::print_py(std::ostream& os) const
{
    auto vec_modules = this->get_modules();
    std::unordered_map<instruction_ref, std::string> names;
    os << "p = migraphx.program()\n";
    for(auto& mod : vec_modules)
    {
1015
1016
1017
        std::string var_name = "m";
        if(mod->name() != "main")
            var_name += mod->name();
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
        os << var_name << " = ";
        if(mod->name() == "main")
            os << "p.get_main_module()";
        else
            os << "p.create_module(\"" << mod->name() << "\");";
        os << std::endl;
        names = mod->print_py(os, var_name, names);
        os << std::endl;
    }
}

1029
1030
void program::print_cpp(std::ostream& os) const
{
Shucai Xiao's avatar
Shucai Xiao committed
1031
1032
    auto vec_modules = this->get_modules();
    std::unordered_map<instruction_ref, std::string> names;
1033
    os << "migraphx::program p;\n";
Shucai Xiao's avatar
Shucai Xiao committed
1034
1035
    for(auto& mod : vec_modules)
    {
1036
1037
1038
1039
1040
1041
1042
1043
        std::string var_name = "m" + mod->name();
        os << "migraphx::module_ref " << var_name << " = ";
        if(mod->name() == "main")
            os << "p.get_main_module();";
        else
            os << "p.create_module(\"" << mod->name() << "\");";
        os << std::endl;
        names = mod->print_cpp(os, var_name, names);
Shucai Xiao's avatar
Shucai Xiao committed
1044
1045
        os << std::endl;
    }
1046
1047
}

Paul's avatar
Paul committed
1048
1049
void program::dry_run(std::unordered_map<std::string, argument> params) const
{
1050
1051
    auto& ctx = this->impl->contexts;
    generic_eval(*this, ctx, std::move(params), [](auto ins, auto&&...) {
1052
        return argument{ins->get_shape(), nullptr};
1053
    });
Paul's avatar
Paul committed
1054
1055
}

Shucai Xiao's avatar
Shucai Xiao committed
1056
void program::annotate(std::ostream& os, const std::function<void(instruction_ref)>& a) const
Paul's avatar
Paul committed
1057
{
1058
    for(auto& pp : this->impl->modules)
Shucai Xiao's avatar
Shucai Xiao committed
1059
    {
1060
1061
        std::cout << pp.first << ":" << std::endl;
        pp.second.annotate(os, a);
Shucai Xiao's avatar
Shucai Xiao committed
1062
    }
Paul's avatar
Paul committed
1063
1064
}

Paul's avatar
Paul committed
1065
const module* program::get_module(const std::string& name) const { return &impl->modules.at(name); }
Shucai Xiao's avatar
Shucai Xiao committed
1066
1067
1068

module* program::create_module(const std::string& name)
{
1069
    assert(not contains(impl->modules, name));
1070
1071
    auto r = impl->modules.emplace(name, name);
    return &(r.first->second);
Shucai Xiao's avatar
Shucai Xiao committed
1072
1073
}

Paul's avatar
Paul committed
1074
module* program::get_module(const std::string& name) { return &impl->modules.at(name); }
Shucai Xiao's avatar
Shucai Xiao committed
1075
1076
1077
1078
1079

module* program::get_main_module() { return get_module("main"); }

const module* program::get_main_module() const { return get_module("main"); }

Paul's avatar
Paul committed
1080
template <class T>
1081
std::vector<T*> generic_get_modules(T* mm)
Shucai Xiao's avatar
Shucai Xiao committed
1082
{
1083
    std::vector<T*> vec_modules;
Shucai Xiao's avatar
Shucai Xiao committed
1084
1085
1086
1087
1088
    vec_modules.push_back(mm);
    auto sub_modules = mm->get_sub_modules();
    vec_modules.insert(vec_modules.end(), sub_modules.begin(), sub_modules.end());
    return vec_modules;
}
Shucai Xiao's avatar
Shucai Xiao committed
1089

Paul's avatar
Paul committed
1090
template <class Map, class T, class OutputIterator>
1091
void generic_get_unused_modules(Map& m, const std::vector<T*>& mods, OutputIterator out)
Shucai Xiao's avatar
Shucai Xiao committed
1092
{
1093
1094
1095
1096
    std::unordered_set<std::string> used;
    std::transform(mods.begin(), mods.end(), std::inserter(used, used.end()), [](auto&& mod) {
        return mod->name();
    });
bpickrel's avatar
bpickrel committed
1097
1098
1099
1100
1101
1102
    transform_if(
        m.begin(),
        m.end(),
        out,
        [&](auto&& pp) { return not contains(used, pp.first); },
        [](auto&& pp) { return &pp.second; });
1103
}
Shucai Xiao's avatar
Shucai Xiao committed
1104

1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
std::vector<const module*> program::get_modules() const
{
    auto result = generic_get_modules(this->get_main_module());
    generic_get_unused_modules(impl->modules, result, std::back_inserter(result));
    return result;
}

std::vector<module*> program::get_modules()
{
    auto result = generic_get_modules(this->get_main_module());
    generic_get_unused_modules(impl->modules, result, std::back_inserter(result));
    return result;
Shucai Xiao's avatar
Shucai Xiao committed
1117
1118
}

1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
template <class Module, class Map>
void generic_insert_module_tree(Module* pm, Map& m)
{
    for(auto* sm : pm->get_sub_modules(true))
    {
        m.insert(std::make_pair(sm, pm));
        generic_insert_module_tree(sm, m);
    }
}

std::unordered_multimap<module_ref, module_ref> program::get_module_tree()
{
    std::unordered_multimap<module_ref, module_ref> result;
    generic_insert_module_tree(this->get_main_module(), result);
    return result;
}

1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
template <class Map, class T>
bool is_unused_module(Map& m, const std::vector<T*>& mods, const std::string& name)
{
    bool is_unused = false;
    generic_get_unused_modules(m, mods, make_function_output_iterator([&](auto* mod) {
                                   if(mod->name() == name)
                                       is_unused = true;
                               }));
    return is_unused;
}

template <class Map>
bool references_instruction(Map& m, const instruction& ins, const std::string& name)
{
    return std::any_of(m.begin(), m.end(), [&](auto&& p) {
        if(p.first == name)
            return false;
        return std::any_of(p.second.begin(), p.second.end(), [&](auto&& i) {
            return std::any_of(i.inputs().begin(), i.inputs().end(), [&](auto&& j) {
                return std::addressof(*j) == std::addressof(ins);
            });
        });
    });
}

void program::remove_module(const std::string& name)
{
    // cppcheck-suppress assertWithSideEffect
    assert(is_unused_module(impl->modules, generic_get_modules(this->get_main_module()), name) &&
           "Module used in program");
    assert(std::none_of(
               impl->modules.at(name).begin(),
               impl->modules.at(name).end(),
               [&](auto&& ins) { return references_instruction(impl->modules, ins, name); }) &&
           "Instruction referenced in another module");
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186

    // if an instruction has an input out side of the current module, need to remove
    // the instruction from its input's outputs
    auto& mod = impl->modules.at(name);
    for(auto ins : iterator_for(mod))
    {
        auto inputs = ins->inputs();
        for(auto in : inputs)
        {
            if(not mod.has_instruction(in))
            {
                in->remove_output(ins);
            }
        }
    }

1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
    impl->modules.erase(name);
}

void program::remove_unused_modules()
{
    std::vector<module*> unused;
    generic_get_unused_modules(
        impl->modules, generic_get_modules(this->get_main_module()), std::back_inserter(unused));
    for(auto* m : unused)
        this->remove_module(m->name());
}

1199
1200
program& program::sort()
{
umangyadav's avatar
umangyadav committed
1201
1202
1203
    std::queue<migraphx::module_ref> mqueue;
    mqueue.push(get_main_module());
    while(!mqueue.empty())
Shucai Xiao's avatar
Shucai Xiao committed
1204
    {
umangyadav's avatar
umangyadav committed
1205
1206
1207
1208
1209
1210
1211
1212
        module_ref current_mod = mqueue.front();
        current_mod->sort();
        mqueue.pop();
        auto child_mods = current_mod->get_sub_modules(true);
        for(auto& sub_mod : child_mods)
        {
            mqueue.push(sub_mod);
        }
Shucai Xiao's avatar
Shucai Xiao committed
1213
    }
1214
1215
1216
    return *this;
}

Paul's avatar
Paul committed
1217
bool operator==(const program& x, const program& y) { return to_string(x) == to_string(y); }
Paul's avatar
Paul committed
1218

Paul's avatar
Paul committed
1219
std::ostream& operator<<(std::ostream& os, const program& p)
Paul's avatar
Paul committed
1220
{
Shucai Xiao's avatar
Shucai Xiao committed
1221
1222
1223
    auto vec_modules = p.get_modules();
    std::unordered_map<instruction_ref, std::string> names;
    for(auto& mod : vec_modules)
Shucai Xiao's avatar
Shucai Xiao committed
1224
    {
Shucai Xiao's avatar
Shucai Xiao committed
1225
1226
1227
1228
1229
1230
1231
        os << "module: \"" << mod->name() << "\"" << std::endl;
        names = mod->print(
            [&](auto ins, auto ins_names) {
                instruction::print(os, ins, ins_names);
                os << std::endl;
            },
            names);
1232
        os << std::endl;
Shucai Xiao's avatar
Shucai Xiao committed
1233
1234
    }

Paul's avatar
Paul committed
1235
    return os;
Paul's avatar
Paul committed
1236
}
Paul's avatar
Paul committed
1237

Paul's avatar
Paul committed
1238
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
1239
} // namespace migraphx