instruction.cpp 12.6 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
#include <migraphx/instruction.hpp>
#include <migraphx/builtin.hpp>
#include <migraphx/erase.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
4
5
#include <migraphx/module.hpp>
#include <migraphx/ranges.hpp>
Paul's avatar
Paul committed
6

Paul's avatar
Paul committed
7
namespace migraphx {
Paul's avatar
Paul committed
8
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
9

10
11
12
13
14
15
template <class T>
auto equal_to(const T& x)
{
    return [&](const T& y) { return std::equal_to<T>{}(x, y); };
}

Paul's avatar
Paul committed
16
17
18
19
instruction::instruction(operation o, shape r, std::vector<instruction_ref> args)
    : op(std::move(o)), result(std::move(r)), arguments(std::move(args))
{
}
Paul's avatar
Paul committed
20

Shucai Xiao's avatar
Shucai Xiao committed
21
22
23
24
25
26
27
28
29
30
31
instruction::instruction(operation o,
                         shape r,
                         std::vector<instruction_ref> args,
                         std::vector<module_ref> modules)
    : op(std::move(o)),
      result(std::move(r)),
      arguments(std::move(args)),
      module_args(std::move(modules))
{
}

Paul's avatar
Paul committed
32
33
34
35
instruction::instruction(literal l)
    : op(builtin::literal{}), result(l.get_shape()), lit(std::move(l))
{
}
Paul's avatar
Paul committed
36

Paul's avatar
Paul committed
37
38
39
void instruction::replace(const shape& r)
{
    if(r != result)
Paul's avatar
Paul committed
40
    {
Paul's avatar
Paul committed
41
42
        result = r;
        for(auto&& ins : output)
Paul's avatar
Paul committed
43
        {
44
45
46
            if(ins->name() == "@return")
                continue;

Paul's avatar
Paul committed
47
48
            assert(ins->name().front() != '@');
            ins->recompute_shape();
Paul's avatar
Paul committed
49
50
        }
    }
Paul's avatar
Paul committed
51
}
Paul's avatar
Paul committed
52

Paul's avatar
Paul committed
53
void instruction::replace(operation o)
Paul's avatar
Paul committed
54
{
55
56
    normalized = false;
    op         = std::move(o);
Paul's avatar
Paul committed
57
58
59
    recompute_shape();
}

Shucai Xiao's avatar
Shucai Xiao committed
60
void instruction::recompute_shape() { replace(compute_shape(op, arguments, module_args)); }
Paul's avatar
Paul committed
61

Paul's avatar
Paul committed
62
63
64
void instruction::clear_arguments()
{
    for(auto&& arg : arguments)
Paul's avatar
Paul committed
65
    {
Paul's avatar
Paul committed
66
        arg->remove_output(*this);
Paul's avatar
Paul committed
67
    }
Paul's avatar
Paul committed
68
    arguments.clear();
Shucai Xiao's avatar
Shucai Xiao committed
69
    module_args.clear();
Paul's avatar
Paul committed
70
71
72
73
74
75
}

bool operator==(const instruction& i, instruction_ref ref)
{
    return std::addressof(i) == std::addressof(*ref);
}
Paul's avatar
Paul committed
76

Shucai Xiao's avatar
Shucai Xiao committed
77
bool instruction::valid(instruction_ref start, bool check_order) const
Paul's avatar
Paul committed
78
79
80
{
    return valid() && std::all_of(arguments.begin(), arguments.end(), [&](instruction_ref i) {
               auto self = std::find(i->outputs().begin(), i->outputs().end(), *this);
Shucai Xiao's avatar
Shucai Xiao committed
81
82
83
84
85
86
               bool ret  = self != i->outputs().end();
               if(check_order)
               {
                   ret = ret and (std::distance(start, i) < std::distance(start, *self));
               }
               return ret;
Paul's avatar
Paul committed
87
88
89
90
91
92
93
           });
}

bool instruction::valid() const
{
    shape computed;
    if(op.name() == "@literal")
Paul's avatar
Paul committed
94
    {
Paul's avatar
Paul committed
95
        computed = lit.get_shape();
Paul's avatar
Paul committed
96
    }
Paul's avatar
Paul committed
97
    else if(op.name() == "@param")
Paul's avatar
Paul committed
98
    {
Paul's avatar
Paul committed
99
        computed = result;
Paul's avatar
Paul committed
100
    }
101
102
103
104
    else if(op.name() == "@return")
    {
        computed = {};
    }
Paul's avatar
Paul committed
105
    else
Paul's avatar
Paul committed
106
    {
Paul's avatar
Paul committed
107
        try
Paul's avatar
Paul committed
108
        {
Shucai Xiao's avatar
Shucai Xiao committed
109
            computed = compute_shape(op, arguments, module_args);
Paul's avatar
Paul committed
110
        }
Paul's avatar
Paul committed
111
        catch(migraphx::exception&)
Paul's avatar
Paul committed
112
        {
Paul's avatar
Paul committed
113
            return false;
Paul's avatar
Paul committed
114
115
        }
    }
116

Shucai Xiao's avatar
Shucai Xiao committed
117
118
    return (result == computed) &&
           std::all_of(output.begin(), output.end(), [&](instruction_ref i) {
Paul's avatar
Paul committed
119
120
121
               return std::find(i->inputs().begin(), i->inputs().end(), *this) != i->inputs().end();
           });
}
Paul's avatar
Paul committed
122

Paul's avatar
Paul committed
123
124
125
126
127
128
shape instruction::get_shape() const { return result; }
const literal& instruction::get_literal() const
{
    assert(op.name() == "@literal");
    return lit;
}
Paul's avatar
Paul committed
129

Paul's avatar
Paul committed
130
const operation& instruction::get_operator() const { return op; }
Paul's avatar
Paul committed
131

Paul's avatar
Paul committed
132
std::string instruction::name() const { return op.name(); }
Paul's avatar
Paul committed
133

Paul's avatar
Paul committed
134
const std::vector<instruction_ref>& instruction::inputs() const { return arguments; }
Paul's avatar
Paul committed
135

Shucai Xiao's avatar
Shucai Xiao committed
136
137
const std::vector<module_ref>& instruction::module_inputs() const { return module_args; }

Paul's avatar
Paul committed
138
const std::vector<instruction_ref>& instruction::outputs() const { return output; }
Paul's avatar
Paul committed
139

Paul's avatar
Paul committed
140
141
bool operator==(const instruction& x, const instruction& y)
{
142
143
144
145
146
147
148
    if(not std::equal(x.arguments.begin(),
                      x.arguments.end(),
                      y.arguments.begin(),
                      y.arguments.end(),
                      std::equal_to<instruction_ref>{}))
        return false;
    if(std::tie(x.result, x.op, x.module_args) != std::tie(y.result, y.op, y.module_args))
Paul's avatar
Paul committed
149
150
151
152
153
        return false;
    if(x.name() == "@literal")
        return x.lit == y.lit;
    return true;
}
Paul's avatar
Paul committed
154

Paul's avatar
Paul committed
155
156
bool operator!=(const instruction& x, const instruction& y) { return !(x == y); }

Paul's avatar
Paul committed
157
bool operator==(instruction_ref ref, const instruction& i) { return i == ref; }
Paul's avatar
Paul committed
158

Paul's avatar
Paul committed
159
bool operator!=(const instruction& i, instruction_ref ref) { return !(i == ref); }
Paul's avatar
Paul committed
160

Paul's avatar
Paul committed
161
bool operator!=(instruction_ref ref, const instruction& i) { return !(i == ref); }
Paul's avatar
Paul committed
162

Paul's avatar
Paul committed
163
164
void instruction::add_output(instruction_ref ins)
{
165
    if(std::find_if(output.begin(), output.end(), equal_to(ins)) == output.end())
Paul's avatar
Paul committed
166
167
        output.push_back(ins);
}
Paul's avatar
Paul committed
168

Paul's avatar
Paul committed
169
170
171
172
173
void instruction::backreference(instruction_ref ref)
{
    for(auto&& arg : ref->inputs())
        arg->add_output(ref);
}
Paul's avatar
Paul committed
174

Paul's avatar
Paul committed
175
176
177
178
179
180
181
182
void instruction::replace_argument(instruction_ref ins,
                                   instruction_ref old,
                                   instruction_ref new_ins)
{
    ins->replace_argument(old, new_ins);
    backreference(ins);
    ins->recompute_shape();
}
Paul's avatar
Paul committed
183

Shucai Xiao's avatar
Shucai Xiao committed
184
185
186
187
188
189
190
void instruction::replace_mod_argument(instruction_ref ins, module_ref old, module_ref new_mod)
{
    ins->replace_mod_argument(old, new_mod);
    backreference(ins);
    ins->recompute_shape();
}

Paul's avatar
Paul committed
191
192
193
194
195
196
197
198
void instruction::replace(instruction_ref ins,
                          operation o,
                          const shape& r,
                          std::vector<instruction_ref> args)
{
    ins->replace(std::move(o), r, std::move(args));
    backreference(ins);
}
Paul's avatar
Paul committed
199

Shucai Xiao's avatar
Shucai Xiao committed
200
201
202
203
204
205
206
207
208
209
void instruction::replace(instruction_ref ins,
                          operation o,
                          const shape& r,
                          std::vector<instruction_ref> args,
                          std::vector<module_ref> module_args)
{
    ins->replace(std::move(o), r, std::move(args), std::move(module_args));
    backreference(ins);
}

Paul's avatar
Paul committed
210
211
void instruction::replace(operation o, const shape& r, std::vector<instruction_ref> args)
{
212
213
    normalized = false;
    op         = std::move(o);
Paul's avatar
Paul committed
214
215
216
    replace(r);
    replace(std::move(args));
}
Paul's avatar
Paul committed
217

Shucai Xiao's avatar
Shucai Xiao committed
218
219
220
221
222
223
224
225
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
void instruction::replace(operation o,
                          const shape& r,
                          std::vector<instruction_ref> args,
                          std::vector<module_ref> mdl_args)
{
    op = std::move(o);
    replace(r);
    replace(std::move(args), std::move(mdl_args));
}

void instruction::replace_refs(
    instruction_ref ins,
    const std::unordered_map<instruction_ref, instruction_ref>& map_insts,
    const std::unordered_map<module_ref, module_ref>& map_mods)
{
    const auto& args = ins->inputs();
    for(const auto& arg : args)
    {
        if(contains(map_insts, arg))
        {
            instruction::replace_argument(ins, arg, map_insts.at(arg));
        }
    }

    const auto& module_args = ins->module_inputs();
    if(module_args.empty())
        return;

    for(const auto& mod : module_args)
    {
        if(contains(map_mods, mod))
        {
            instruction::replace_mod_argument(ins, mod, map_mods.at(mod));
        }
    }
}

Paul's avatar
Paul committed
255
256
257
258
259
void instruction::replace(std::vector<instruction_ref> args)
{
    clear_arguments();
    arguments = std::move(args);
}
Paul's avatar
Paul committed
260

Shucai Xiao's avatar
Shucai Xiao committed
261
262
263
264
265
266
267
void instruction::replace(std::vector<instruction_ref> args, std::vector<module_ref> mdl_args)
{
    clear_arguments();
    arguments   = std::move(args);
    module_args = std::move(mdl_args);
}

Paul's avatar
Paul committed
268
269
void instruction::replace_argument(instruction_ref old, instruction_ref new_ins)
{
270
271
    assert(std::any_of(arguments.begin(), arguments.end(), equal_to(old)));
    std::replace_if(arguments.begin(), arguments.end(), equal_to(old), new_ins);
Paul's avatar
Paul committed
272
273
    old->remove_output(*this);
}
Paul's avatar
Paul committed
274

Shucai Xiao's avatar
Shucai Xiao committed
275
276
277
278
279
280
void instruction::replace_mod_argument(module_ref old, module_ref new_mod)
{
    assert(std::any_of(module_args.begin(), module_args.end(), [&](auto i) { return i == old; }));
    std::replace(module_args.begin(), module_args.end(), old, new_mod);
}

Paul's avatar
Paul committed
281
282
bool instruction::can_eval() const
{
283
    if(op.name() == "@literal" or op.name() == "shape")
Paul's avatar
Paul committed
284
285
286
    {
        return true;
    }
Paul's avatar
Paul committed
287
    else if(is_context_free(op))
Paul's avatar
Paul committed
288
    {
Paul's avatar
Paul committed
289
290
        return std::all_of(
            this->inputs().begin(), this->inputs().end(), [](auto arg) { return arg->can_eval(); });
Paul's avatar
Paul committed
291
292
293
294
295
296
297
    }
    else
    {
        return false;
    }
}

Paul's avatar
Paul committed
298
argument instruction::eval(bool check_eval) const
Paul's avatar
Paul committed
299
300
301
302
303
{
    if(op.name() == "@literal")
    {
        return this->get_literal().get_argument();
    }
Shucai Xiao's avatar
Shucai Xiao committed
304
    else if(op.name() == "shape")
305
306
307
308
309
    {
        argument arg{this->inputs().front()->get_shape()};
        return normalized_operator().compute(result, {arg});
    }

Paul's avatar
Paul committed
310
    if(is_context_free(op))
Paul's avatar
Paul committed
311
    {
Paul's avatar
Paul committed
312
        if(check_eval and not this->can_eval())
313
        {
Paul's avatar
Paul committed
314
            return {};
315
        }
Paul's avatar
Paul committed
316
        std::vector<argument> args;
Paul's avatar
Paul committed
317
318
319
        std::transform(this->inputs().begin(),
                       this->inputs().end(),
                       std::back_inserter(args),
Paul's avatar
Paul committed
320
                       [](auto arg) { return arg->eval(false); });
321
        return normalized_operator().compute(result, args);
Paul's avatar
Paul committed
322
323
324
325
    }
    return {};
}

Paul's avatar
Paul committed
326
327
void instruction::finalize(context& ctx)
{
Paul's avatar
Paul committed
328
    if(has_finalize(this->op))
Paul's avatar
Paul committed
329
330
331
        this->op.finalize(ctx, this->get_shape(), to_shapes(this->inputs()));
}

332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
void instruction::print(std::ostream& os,
                        instruction_ref ins,
                        const std::unordered_map<instruction_ref, std::string>& names)
{
    os << names.at(ins) << " = ";

    os << ins->get_operator();

    if(ins->name() == "@literal")
    {
        if(ins->get_literal().get_shape().elements() > 10)
            os << "{ ... }";
        else
            os << "{" << ins->get_literal() << "}";
    }

    if(!ins->inputs().empty())
    {
        char delim = '(';
        for(auto&& arg : ins->inputs())
        {
Shucai Xiao's avatar
Shucai Xiao committed
353
354
            std::string arg_name = contains(names, arg) ? names.at(arg) : "?";
            os << delim << arg_name;
355
356
357
358
359
            delim = ',';
        }
        os << ")";
    }

Shucai Xiao's avatar
Shucai Xiao committed
360
361
362
363
364
365
366
367
368
369
370
371
    // print module inputs
    if(!ins->module_inputs().empty())
    {
        std::string delim = ", [";
        for(auto&& mod_arg : ins->module_inputs())
        {
            os << delim << mod_arg->name();
            delim = ", ";
        }
        os << "]";
    }

372
373
374
375
376
    // skip return instruction shape
    if(ins->name() != "@return")
        os << " -> " << ins->get_shape();
}

Paul Fultz II's avatar
Paul Fultz II committed
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
static void debug_name(std::ostream& os, const instruction& ins)
{
    if(ins.name() == "@literal")
    {
        os << "@literal";
        if(ins.get_literal().get_shape().elements() > 10)
            os << "{ ... }";
        else
            os << "{" << ins.get_literal() << "}";
    }
    else
    {
        os << ins.get_operator();
    }
}

void instruction::debug_print() const
{
    debug_name(std::cout, *this);
    std::string delim = "(";
    for(auto arg : this->inputs())
    {
        std::cout << delim;
        debug_name(std::cout, *arg);
        delim = ", ";
    }
    if(not this->inputs().empty())
        std::cout << ")";
    std::cout << " -> " << this->get_shape() << std::endl;
}

Paul's avatar
Paul committed
408
instruction_ref instruction::get_output_alias(instruction_ref ins, bool shallow)
Paul's avatar
Paul committed
409
{
Paul's avatar
Paul committed
410
    auto i = ins->get_operator().output_alias(to_shapes(ins->inputs()));
Paul's avatar
Paul committed
411
412
    if(i < 0)
        return ins;
Paul's avatar
Paul committed
413
    if(shallow)
Paul's avatar
Paul committed
414
        return ins->inputs().at(i);
Paul's avatar
Paul committed
415
416
417
    return get_output_alias(ins->inputs().at(i));
}

418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
void instruction::set_normalized(bool value) { normalized = value; }

bool instruction::is_normalized() const { return normalized; }

bool instruction::need_normalization() const
{
    return this->get_operator().need_normalization() and not normalized;
}

operation instruction::normalized_operator() const
{
    operation o = this->get_operator();
    if(this->need_normalization())
    {
        auto lens = this->inputs().front()->get_shape().lens();
        if(!normalize_attributes(o, lens))
            return this->get_operator();
    }
    return o;
}

Paul's avatar
Paul committed
439
440
441
442
443
444
445
446
std::vector<shape> to_shapes(const std::vector<instruction_ref>& args)
{
    std::vector<shape> shapes(args.size());
    std::transform(
        args.begin(), args.end(), shapes.begin(), [](instruction_ref i) { return i->get_shape(); });
    return shapes;
}

Paul's avatar
Paul committed
447
448
shape compute_shape(const operation& op, const std::vector<instruction_ref>& args)
{
Paul's avatar
Paul committed
449
    return op.compute_shape(to_shapes(args));
Paul's avatar
Paul committed
450
451
}

Shucai Xiao's avatar
Shucai Xiao committed
452
453
454
455
456
457
458
459
460
461
462
463
464
shape compute_shape(const operation& op,
                    const std::vector<instruction_ref>& args,
                    const std::vector<module_ref>& mods)
{
    if(mods.empty())
    {
        return op.compute_shape(to_shapes(args));
    }
    else
    {
        return op.compute_shape(to_shapes(args), mods);
    }
}
465
466
467
468
469
470
471
472
473
474
475
476
477
478

std::vector<shape> try_compute_shape(const operation& op, const std::vector<shape>& inputs)
{
    shape new_shape;
    try
    {
        new_shape = op.compute_shape(inputs);
    }
    catch(...)
    {
        return {};
    }
    return {new_shape};
}
479
480
481
482
483
484

migraphx::instruction* as_address(const instruction_ref& ins) noexcept
{
    return std::addressof(*ins);
}

Paul's avatar
Paul committed
485
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
486
} // namespace migraphx