lowering.cpp 21.2 KB
Newer Older
Shucai Xiao's avatar
Shucai Xiao committed
1
#include <iterator>
Paul's avatar
Paul committed
2
3
4
#include <migraphx/gpu/lowering.hpp>
#include <migraphx/manage_ptr.hpp>
#include <migraphx/instruction.hpp>
5
6
7
8
9
10
11
12
#include <migraphx/make_op.hpp>

#include <migraphx/op/abs.hpp>
#include <migraphx/op/batch_norm_inference.hpp>
#include <migraphx/op/convolution.hpp>
#include <migraphx/op/deconvolution.hpp>
#include <migraphx/op/dot.hpp>
#include <migraphx/op/elu.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
13
#include <migraphx/op/if_op.hpp>
14
15
16
17
18
19
20
21
22
#include <migraphx/op/leaky_relu.hpp>
#include <migraphx/op/lrn.hpp>
#include <migraphx/op/pooling.hpp>
#include <migraphx/op/reshape.hpp>
#include <migraphx/op/quant_convolution.hpp>
#include <migraphx/op/quant_dot.hpp>

#include <migraphx/gpu/abs.hpp>
#include <migraphx/gpu/batch_norm_inference.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
23
#include <migraphx/gpu/compile_roialign.hpp>
Paul's avatar
Paul committed
24
25
#include <migraphx/gpu/context.hpp>
#include <migraphx/gpu/convolution.hpp>
kahmed10's avatar
kahmed10 committed
26
#include <migraphx/gpu/deconvolution.hpp>
Khalique's avatar
Khalique committed
27
#include <migraphx/gpu/elu.hpp>
28
#include <migraphx/gpu/equal.hpp>
Paul's avatar
Paul committed
29
#include <migraphx/gpu/gemm.hpp>
30
#include <migraphx/gpu/greater.hpp>
31
#include <migraphx/gpu/int8_conv_pack.hpp>
32
#include <migraphx/gpu/leaky_relu.hpp>
33
#include <migraphx/gpu/less.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
34
35
36
#include <migraphx/gpu/logical_and.hpp>
#include <migraphx/gpu/logical_or.hpp>
#include <migraphx/gpu/logical_xor.hpp>
37
38
39
40
#include <migraphx/gpu/lrn.hpp>
#include <migraphx/gpu/miopen.hpp>
#include <migraphx/gpu/quant_convolution.hpp>
#include <migraphx/gpu/rocblas.hpp>
41
#include <migraphx/gpu/unary_not.hpp>
turneram's avatar
turneram committed
42
#include <migraphx/gpu/where.hpp>
43
#include <migraphx/iterator_for.hpp>
44
#include <migraphx/program.hpp>
Paul's avatar
Paul committed
45
#include <utility>
46
#include <functional>
Khalique's avatar
Khalique committed
47
#include <algorithm>
Shucai Xiao's avatar
Shucai Xiao committed
48
#include <map>
Paul's avatar
Paul committed
49

Paul's avatar
Paul committed
50
namespace migraphx {
Paul's avatar
Paul committed
51
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
52
namespace gpu {
Paul's avatar
Paul committed
53
54
55

struct miopen_apply
{
Shucai Xiao's avatar
Shucai Xiao committed
56
    module* mod          = nullptr;
57
    const lowering* pass = nullptr;
Shucai Xiao's avatar
Shucai Xiao committed
58
    std::unordered_map<std::string, std::function<instruction_ref(instruction_ref)>> apply_map{};
Shucai Xiao's avatar
Shucai Xiao committed
59
    instruction_ref last{};
60
    std::unordered_map<instruction_ref, std::string> prog_output_names{};
Shucai Xiao's avatar
Shucai Xiao committed
61
62
    bool offload_copy   = false;
    bool int8_x4_format = true;
Paul's avatar
Paul committed
63

64
    context& get_context() const
65
66
67
68
69
70
    {
        assert(pass != nullptr);
        assert(pass->ctx != nullptr);
        return *pass->ctx;
    }

Paul's avatar
Paul committed
71
72
73
74
75
76
77
    void check_shape(shape x, instruction_ref i)
    {
        assert(x == i->get_shape());
        (void)x;
        (void)i;
    }

78
79
    void create_output_names()
    {
Shucai Xiao's avatar
Shucai Xiao committed
80
        this->last = instruction::get_output_alias(std::prev(mod->end()));
81
82
        if(this->last->name() == "@return")
        {
83
            const auto& prog_outputs = last->inputs();
84
            std::vector<instruction_ref> outputs(prog_outputs.size());
85

Shucai Xiao's avatar
Shucai Xiao committed
86
87
88
89
90
91
92
93
94
            std::transform(
                prog_outputs.begin(), prog_outputs.end(), outputs.begin(), [](const auto& i) {
                    auto alias_ins = instruction::get_output_alias(i);
                    auto alias_s   = alias_ins->get_shape();
                    return (alias_s.type() == shape::tuple_type or
                            alias_s.elements() != i->get_shape().elements())
                               ? i
                               : alias_ins;
                });
95
96

            std::size_t index = 0;
97
            for(auto ins : outputs)
98
            {
Shucai Xiao's avatar
Shucai Xiao committed
99
                prog_output_names[ins] = mod->name() + ":#output_" + std::to_string(index++);
100
101
102
103
            }
        }
    }

104
105
    void init()
    {
Shucai Xiao's avatar
Shucai Xiao committed
106
        assert(mod != nullptr);
107
        assert(pass != nullptr);
108

Shucai Xiao's avatar
Shucai Xiao committed
109
110
111
112
113
114
115
#if ROCBLAS_VERSION_MAJOR >= 2 && ROCBLAS_VERSION_MINOR >= 38
        auto& ctx = get_context();
        rocblas_gemm_flags flag;
        rocblas_query_int8_layout_flag(ctx.get_stream().get_rocblas(), &flag);
        int8_x4_format = (flag == rocblas_gemm_flags_pack_int8x4);
#endif

Shucai Xiao's avatar
Shucai Xiao committed
116
        offload_copy = (mod->name() == "main") ? pass->offload_copy : false;
117
        create_output_names();
Paul's avatar
Paul committed
118

119
120
121
122
123
124
125
126
127
128
129
130
        add_generic_op("acos");
        add_generic_op("acosh");
        add_generic_op("add");
        add_generic_op("asin");
        add_generic_op("asinh");
        add_generic_op("atan");
        add_generic_op("atanh");
        add_generic_op("ceil");
        add_generic_op("contiguous");
        add_generic_op("cos");
        add_generic_op("cosh");
        add_generic_op("div");
131
        add_generic_op("equal");
132
133
134
        add_generic_op("erf");
        add_generic_op("exp");
        add_generic_op("floor");
135
136
        add_generic_op("greater");
        add_generic_op("less");
137
        add_generic_op("log");
Shucai Xiao's avatar
Shucai Xiao committed
138
139
140
        add_generic_op("logical_and");
        add_generic_op("logical_or");
        add_generic_op("logical_xor");
141
142
143
        add_generic_op("max");
        add_generic_op("min");
        add_generic_op("mul");
144
        add_generic_op("not");
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
        add_generic_op("pow");
        add_generic_op("prelu");
        add_generic_op("recip");
        add_generic_op("relu");
        add_generic_op("round");
        add_generic_op("rsqrt");
        add_generic_op("sigmoid");
        add_generic_op("sign");
        add_generic_op("sin");
        add_generic_op("sinh");
        add_generic_op("sqdiff");
        add_generic_op("sqrt");
        add_generic_op("sub");
        add_generic_op("tan");
        add_generic_op("tanh");
turneram's avatar
turneram committed
160
        add_generic_op("where");
161

Shucai Xiao's avatar
Shucai Xiao committed
162
        add_extend_op("abs");
163
164
165
166
167
        add_extend_op("argmax");
        add_extend_op("argmin");
        add_extend_op("clip");
        add_extend_op("concat");
        add_extend_op("convert");
Shucai Xiao's avatar
Shucai Xiao committed
168
        add_extend_op("elu");
169
        add_extend_op("gather");
Shucai Xiao's avatar
Shucai Xiao committed
170
        add_extend_op("leaky_relu");
171
        add_extend_op("logsoftmax");
Shucai Xiao's avatar
Shucai Xiao committed
172
        add_extend_op("lrn");
turneram's avatar
turneram committed
173
        add_extend_op("multinomial");
Shucai Xiao's avatar
Shucai Xiao committed
174
        add_extend_op("nonzero");
175
        add_extend_op("pad");
176
        add_extend_op("pooling");
177
        add_extend_op("prefix_scan_sum");
178
179
180
181
182
        add_extend_op("reduce_max");
        add_extend_op("reduce_mean");
        add_extend_op("reduce_min");
        add_extend_op("reduce_prod");
        add_extend_op("reduce_sum");
Cagri Eryilmaz's avatar
Cagri Eryilmaz committed
183
        add_extend_op("reverse");
184
185
186
        add_extend_op("rnn_var_sl_last_output");
        add_extend_op("rnn_var_sl_shift_output");
        add_extend_op("rnn_var_sl_shift_sequence");
187
        add_extend_op("scatter");
188
        add_extend_op("softmax");
Shucai Xiao's avatar
Shucai Xiao committed
189
        add_extend_op("topk");
190

191
192
        add_precompile_op("pointwise");

Shucai Xiao's avatar
Shucai Xiao committed
193
        add_batch_norm_inference_op();
194
        add_convolution_op();
kahmed10's avatar
kahmed10 committed
195
        add_deconvolution_op();
Shucai Xiao's avatar
Shucai Xiao committed
196
197
        add_gemm_op<op::dot>("dot");
        add_gemm_op<op::quant_dot>("quant_dot");
Shucai Xiao's avatar
Shucai Xiao committed
198
        add_if_op();
Shucai Xiao's avatar
Shucai Xiao committed
199
        add_loop_op();
Shucai Xiao's avatar
Shucai Xiao committed
200
        add_neg_op();
201
        add_nms_op();
Shucai Xiao's avatar
Shucai Xiao committed
202
        add_quant_convolution_op();
Shucai Xiao's avatar
Shucai Xiao committed
203
        add_roialign();
204
205
    }

206
207
    void copy_params()
    {
Shucai Xiao's avatar
Shucai Xiao committed
208
        if(not offload_copy)
209
            return;
210

Shucai Xiao's avatar
Shucai Xiao committed
211
        for(auto ins : iterator_for(*mod))
212
213
214
        {
            if(ins->name() != "@param")
                continue;
215

Shucai Xiao's avatar
Shucai Xiao committed
216
217
218
219
            // parameter no outputs, no need to insert copy to gpu
            if(ins->outputs().empty())
                continue;

220
221
            auto pos = std::next(ins);
            auto a   = insert_allocation(pos, ins->get_shape());
222
            auto c   = mod->insert_instruction(pos, make_op("hip::copy_to_gpu"), ins, a);
Shucai Xiao's avatar
Shucai Xiao committed
223
            mod->replace_instruction(ins, c);
224
        }
225
226

        // return instruction
Shucai Xiao's avatar
Shucai Xiao committed
227
        auto ret = std::prev(mod->end());
228
229
        if(ret->name() == "@return")
        {
230
            const auto& inputs = ret->inputs();
231
232
233

            // each input of ret need to be copied from gpu to host, and replace
            // output with copy output
234
            for(const auto& in : inputs)
235
            {
236
                auto p_output = mod->insert_instruction(ret, make_op("hip::copy_from_gpu"), in);
237
238
239
240
241
242
                instruction::replace_argument(ret, in, p_output);
            }
        }
        // else branch to handle legacy program without the return instruction
        else
        {
243
            mod->add_instruction(make_op("hip::copy_from_gpu"), ret);
244
        }
245
246
    }

Paul's avatar
Paul committed
247
248
    void apply()
    {
249
        init();
Shucai Xiao's avatar
Shucai Xiao committed
250
        for(auto it = mod->begin(); it != mod->end(); it++)
Paul's avatar
Paul committed
251
        {
Paul's avatar
Paul committed
252
            auto s = it->get_shape();
253
            if(apply_map.count(it->name()) > 0)
254
            {
255
                check_shape(s, apply_map.at(it->name())(it));
Paul's avatar
Paul committed
256
            }
Paul's avatar
Paul committed
257
        }
258

259
        copy_params();
Paul's avatar
Paul committed
260
261
    }

Paul's avatar
Paul committed
262
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
263
    {
264
        // Instruction's output is an input of the ret instruction
Shucai Xiao's avatar
Shucai Xiao committed
265
        if(offload_copy)
Paul's avatar
Paul committed
266
        {
267
268
            auto result = mod->insert_instruction(
                ins, make_op("hip::allocate", {{"shape", to_value(s)}, {"tag", std::move(tag)}}));
Paul's avatar
Paul committed
269
270
            return result;
        }
271
272

        auto ins_alias = instruction::get_output_alias(ins);
273
        if(last->name() == "@return" and tag.empty())
274
        {
Shucai Xiao's avatar
Shucai Xiao committed
275

276
            auto alias_s = ins_alias->get_shape();
Shucai Xiao's avatar
Shucai Xiao committed
277
278
            if(alias_s.type() == shape::tuple_type or
               alias_s.elements() != ins->get_shape().elements())
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
            {
                if(prog_output_names.count(ins) > 0)
                {
                    auto out_ins = mod->add_parameter(prog_output_names[ins], s);
                    mod->insert_instruction(std::next(ins), make_op("contiguous"), ins, out_ins);
                    mod->replace_instruction(ins, out_ins);
                    return out_ins;
                }
            }
            else
            {
                if(prog_output_names.count(ins_alias) > 0)
                {
                    return mod->add_parameter(prog_output_names[ins_alias], s);
                }
            }
295
296
297
        }
        else if(ins == last and tag.empty())
        {
Shucai Xiao's avatar
Shucai Xiao committed
298
            return mod->add_parameter("output", s);
299
300
        }

301
302
        return mod->insert_instruction(
            ins, make_op("hip::allocate", {{"shape", to_value(s)}, {"tag", std::move(tag)}}));
Paul's avatar
Paul committed
303
304
    }

Shucai Xiao's avatar
Shucai Xiao committed
305
    void add_convolution_op()
Paul's avatar
Paul committed
306
    {
307
308
        apply_map.emplace("convolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
309

310
            auto conv = miopen_convolution{op, make_conv(op)};
311
            auto ws   = conv.find(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
Paul's avatar
Paul committed
312

313
314
            auto workspace = insert_allocation(ins, ws, "workspace");
            auto output    = insert_allocation(ins, ins->get_shape());
kahmed10's avatar
kahmed10 committed
315

Shucai Xiao's avatar
Shucai Xiao committed
316
            return mod->replace_instruction(
kahmed10's avatar
kahmed10 committed
317
318
319
320
321
322
323
324
325
326
327
328
329
330
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
    }

    void add_deconvolution_op()
    {
        apply_map.emplace("deconvolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::deconvolution>(ins->get_operator());

            auto conv = miopen_deconvolution{op, make_deconv(op)};
            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));

            auto workspace = insert_allocation(ins, ws, "workspace");
            auto output    = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
331

Shucai Xiao's avatar
Shucai Xiao committed
332
            return mod->replace_instruction(
333
334
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
Paul's avatar
Paul committed
335
336
    }

337
338
    template <typename Op>
    void add_gemm_op(const std::string& name)
339
340
    {
        apply_map.emplace(name, [=](instruction_ref ins) {
341
            std::vector<instruction_ref> refs = ins->inputs();
Shucai Xiao's avatar
Shucai Xiao committed
342
            if(refs.size() == 2)
343
344
            {
                auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
345
346
347
348
349
350
                refs.push_back(output);
            }
            else
            {
                auto c_alias = instruction::get_output_alias(refs.back());
                if(ins == last or refs.back()->outputs().size() > 1 or c_alias->inputs().empty())
351
                {
352
353
354
355
                    auto output = insert_allocation(ins, ins->get_shape());
                    auto copy_out =
                        mod->insert_instruction(ins, make_op("hip::copy"), refs.back(), output);
                    refs.back() = copy_out;
356
357
                    refs.push_back(copy_out);
                }
Shucai Xiao's avatar
Shucai Xiao committed
358
359
360
361
                else
                {
                    refs.push_back(refs.back());
                }
362
            }
Shucai Xiao's avatar
Shucai Xiao committed
363
            return mod->replace_instruction(
364
                ins, rocblas_gemm<Op>{Op{}, 1, 0, int8_x4_format}, refs);
365
366
367
        });
    }

368
369
370
371
372
    void add_quant_convolution_op()
    {
        apply_map.emplace("quant_convolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::quant_convolution>(ins->get_operator());
            auto conv = miopen_quant_convolution{op, make_conv(op)};
373
            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
374

Shucai Xiao's avatar
Shucai Xiao committed
375
            auto args      = ins->inputs();
376
            auto workspace = insert_allocation(ins, ws, "workspace");
Shucai Xiao's avatar
Shucai Xiao committed
377
378
            auto output    = insert_allocation(ins, ins->get_shape());

Shucai Xiao's avatar
Shucai Xiao committed
379
            return mod->replace_instruction(ins, conv, args[0], args[1], workspace, output);
Shucai Xiao's avatar
Shucai Xiao committed
380
381
382
        });
    }

383
384
385
    void add_generic_op(const std::string& name) { add_generic_op(name, "gpu::" + name); }

    void add_generic_op(const std::string& op_name, const std::string& gpu_name)
Paul's avatar
Paul committed
386
    {
387
        apply_map.emplace(op_name, [=](instruction_ref ins) {
388
389
390
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
391

Shucai Xiao's avatar
Shucai Xiao committed
392
            return mod->replace_instruction(ins, make_op(gpu_name), refs);
393
        });
Paul's avatar
Paul committed
394
    }
Paul's avatar
Paul committed
395

396
397
398
    void add_extend_op(const std::string& name) { add_extend_op(name, "gpu::" + name); }

    void add_extend_op(const std::string& op_name, const std::string& gpu_name)
Khalique's avatar
Khalique committed
399
    {
400
401
        apply_map.emplace(op_name, [=](instruction_ref ins) {
            auto&& op                         = ins->get_operator();
402
403
404
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
405

Shucai Xiao's avatar
Shucai Xiao committed
406
            return mod->replace_instruction(ins, make_op(gpu_name, op.to_value()), refs);
407
        });
Khalique's avatar
Khalique committed
408
409
    }

410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
    void add_precompile_op(const std::string& name)
    {
        apply_map.emplace(name, [=](instruction_ref ins) {
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);

            return mod->replace_instruction(
                ins,
                make_op("gpu::precompile_op", {{"op", to_value(ins->get_operator())}}),
                refs,
                ins->module_inputs());
        });
    }

Shucai Xiao's avatar
Shucai Xiao committed
425
    void add_batch_norm_inference_op()
426
    {
427
428
429
430
        apply_map.emplace("batch_norm_inference", [=](instruction_ref ins) {
            auto&& op       = any_cast<op::batch_norm_inference>(ins->get_operator());
            auto output     = insert_allocation(ins, ins->get_shape());
            shape old_shape = ins->inputs().at(1)->get_shape();
Shucai Xiao's avatar
Shucai Xiao committed
431
432
433
434
435
436
437
438
439
440
441
442
443
444
            auto input      = ins->inputs()[0];
            auto input_lens = input->get_shape().lens();
            std::vector<int64_t> rsp_lens(input_lens.size(), 1);
            // for per_activation case, also need to reshape input
            if(op.bn_mode == op::batch_norm_inference::per_activation)
            {
                std::copy(input_lens.begin() + 1, input_lens.end(), rsp_lens.begin() + 1);
            }
            else
            {
                rsp_lens[1] = static_cast<int64_t>(old_shape.elements());
            }

            auto reshape_op = op::reshape{rsp_lens};
445
446
            std::vector<instruction_ref> reshapes;
            std::transform(ins->inputs().begin() + 1,
Shucai Xiao's avatar
Shucai Xiao committed
447
448
                           ins->inputs().end(),
                           std::back_inserter(reshapes),
Shucai Xiao's avatar
Shucai Xiao committed
449
                           [&](auto i) { return mod->insert_instruction(ins, reshape_op, i); });
Shucai Xiao's avatar
Shucai Xiao committed
450

Shucai Xiao's avatar
Shucai Xiao committed
451
452
453
454
455
456
457
458
            return mod->replace_instruction(ins,
                                            miopen_batch_norm_inference{op},
                                            input,
                                            reshapes[0],
                                            reshapes[1],
                                            reshapes[2],
                                            reshapes[3],
                                            output);
Shucai Xiao's avatar
Shucai Xiao committed
459

460
        });
461
    }
Shucai Xiao's avatar
Shucai Xiao committed
462
463
464
465
466
467
468

    // use 0 - input to represent neg
    void add_neg_op()
    {
        apply_map.emplace("neg", [=](instruction_ref ins) {
            auto s = ins->get_shape();
            std::vector<float> zeros(s.elements(), 0.0f);
Shucai Xiao's avatar
Shucai Xiao committed
469
            auto l0     = mod->add_literal(literal(s, zeros));
Shucai Xiao's avatar
Shucai Xiao committed
470
            auto output = insert_allocation(ins, s);
Shucai Xiao's avatar
Shucai Xiao committed
471
            return mod->replace_instruction(
472
                ins, make_op("gpu::sub"), l0, ins->inputs().front(), output);
Shucai Xiao's avatar
Shucai Xiao committed
473
474
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
475

Shucai Xiao's avatar
Shucai Xiao committed
476
    // add input and output argument for the if operator
Shucai Xiao's avatar
Shucai Xiao committed
477
478
479
480
    void add_if_op()
    {
        apply_map.emplace("if", [=](instruction_ref ins) {
            std::vector<instruction_ref> inputs = ins->inputs();
481
482
483
            auto cpu_cond =
                mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), inputs.front());
            auto sync_cond = mod->insert_instruction(ins, make_op("hip::sync_stream"), cpu_cond);
Shucai Xiao's avatar
Shucai Xiao committed
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
            inputs.front() = sync_cond;

            std::vector<module_ref> mod_args = ins->module_inputs();
            std::map<std::string, shape> name_shapes;
            for(const auto& smod : mod_args)
            {
                auto ps = smod->get_parameter_shapes();
                name_shapes.insert(ps.begin(), ps.end());
            }

            bool ins_output_allocated = false;
            for(auto& pn : name_shapes)
            {
                const auto& s = pn.second;
                instruction_ref output{};
                if(s == ins->get_shape() and not ins_output_allocated)
                {
                    output               = insert_allocation(ins, s);
                    ins_output_allocated = true;
                }
                else
                {
506
507
                    output = mod->insert_instruction(
                        ins, make_op("hip::allocate", {{"shape", to_value(s)}}));
Shucai Xiao's avatar
Shucai Xiao committed
508
509
510
511
512
513
514
                }
                inputs.push_back(output);
            }

            return mod->replace_instruction(ins, ins->get_operator(), inputs, mod_args);
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
515

Shucai Xiao's avatar
Shucai Xiao committed
516
517
518
    void add_roialign()
    {
        apply_map.emplace("roialign", [=](instruction_ref ins) {
Shucai Xiao's avatar
Shucai Xiao committed
519

Shucai Xiao's avatar
Shucai Xiao committed
520
            auto s      = ins->get_shape();
Shucai Xiao's avatar
Shucai Xiao committed
521
            auto op_val = ins->get_operator().to_value();
Shucai Xiao's avatar
Shucai Xiao committed
522
            auto output = insert_allocation(ins, s);
Shucai Xiao's avatar
Shucai Xiao committed
523
524
525
526
527
528
            auto args   = ins->inputs();
            args.push_back(output);

            auto io_shapes = to_shapes(args);
            auto co        = compile_roialign(get_context(), io_shapes, op_val);
            return mod->replace_instruction(ins, co, args);
Shucai Xiao's avatar
Shucai Xiao committed
529
530
531
        });
    }

Shucai Xiao's avatar
Shucai Xiao committed
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
    // replace the loop operator with gpu_loop operator
    void add_loop_op()
    {
        apply_map.emplace("loop", [=](instruction_ref ins) {
            std::vector<instruction_ref> inputs = ins->inputs();
            // copy max_iter from gpu to cpu
            auto cpu_max_iter =
                mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), inputs.at(0));
            auto cpu_cond =
                mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), inputs.at(1));
            auto synced_max_iter =
                mod->insert_instruction(ins, make_op("hip::sync_stream"), cpu_max_iter, cpu_cond);
            inputs.at(0)     = synced_max_iter;
            inputs.at(1)     = cpu_cond;
            auto copy_inputs = inputs;
            std::transform(
                copy_inputs.begin(), copy_inputs.end(), std::back_inserter(inputs), [&](auto in) {
                    return mod->insert_instruction(
                        ins, make_op("hip::allocate", {{"shape", to_value(in->get_shape())}}));
                });

            auto mod_args = ins->module_inputs();
            auto output   = insert_allocation(ins, ins->get_shape());

            const auto* sub_mod = mod_args.front();
            auto cond_out       = mod->insert_instruction(
                ins,
                make_op("hip::allocate",
                        {{"shape", to_value(sub_mod->get_output_shapes().front())}}));
            // add cond and mod outputs to the argument list
            inputs.push_back(cond_out);
            inputs.push_back(output);

            return mod->replace_instruction(
                ins, make_op("gpu::loop", ins->get_operator().to_value()), inputs, mod_args);
        });
    }
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588

    void add_nms_op()
    {
        apply_map.emplace("nonmaxsuppression", [=](instruction_ref ins) {
            auto s      = ins->get_shape();
            auto output = insert_allocation(ins, s);
            std::vector<instruction_ref> cpu_inputs;
            auto inputs = ins->inputs();
            std::transform(
                inputs.begin(), inputs.end(), std::back_inserter(cpu_inputs), [&](auto in) {
                    return mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), in);
                });
            cpu_inputs.front() =
                mod->insert_instruction(ins, make_op("hip::sync_stream"), cpu_inputs);
            auto cpu_out = mod->insert_instruction(ins, ins->get_operator(), cpu_inputs);
            auto gpu_out =
                mod->insert_instruction(ins, make_op("hip::copy_to_gpu"), cpu_out, output);
            return mod->replace_instruction(ins, gpu_out);
        });
    }
Paul's avatar
Paul committed
589
590
};

Shucai Xiao's avatar
Shucai Xiao committed
591
void lowering::apply(module& m) const { miopen_apply{&m, this}.apply(); }
Shucai Xiao's avatar
Shucai Xiao committed
592

Paul's avatar
Paul committed
593
} // namespace gpu
Paul's avatar
Paul committed
594
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
595
} // namespace migraphx