lowering.cpp 20.1 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>
Paul's avatar
Paul committed
23
24
#include <migraphx/gpu/context.hpp>
#include <migraphx/gpu/convolution.hpp>
kahmed10's avatar
kahmed10 committed
25
#include <migraphx/gpu/deconvolution.hpp>
26
#include <migraphx/gpu/device_name.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/gpu/compiler.hpp>
44
#include <migraphx/iterator_for.hpp>
45
#include <migraphx/program.hpp>
Paul's avatar
Paul committed
46
#include <utility>
47
#include <functional>
Khalique's avatar
Khalique committed
48
#include <algorithm>
Shucai Xiao's avatar
Shucai Xiao committed
49
#include <map>
Paul's avatar
Paul committed
50

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

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

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

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

80
81
    void create_output_names()
    {
Shucai Xiao's avatar
Shucai Xiao committed
82
        this->last = instruction::get_output_alias(std::prev(mod->end()));
83
84
        if(this->last->name() == "@return")
        {
85
            const auto& prog_outputs = last->inputs();
86
87
88
89
90
91
92
93
94
95
            std::vector<instruction_ref> outputs_alias(prog_outputs.size());

            std::transform(prog_outputs.begin(),
                           prog_outputs.end(),
                           outputs_alias.begin(),
                           [](const auto& i) { return instruction::get_output_alias(i); });

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

101
102
103
104
105
106
    const std::unordered_set<std::string>& get_rocblas_fp32_archs()
    {
        static std::unordered_set<std::string> supported_archs{"gfx908", "gfx90a"};
        return supported_archs;
    }

107
108
    void init()
    {
Shucai Xiao's avatar
Shucai Xiao committed
109
        assert(mod != nullptr);
110
        assert(pass != nullptr);
111

Shucai Xiao's avatar
Shucai Xiao committed
112
#if ROCBLAS_VERSION_MAJOR >= 2 && ROCBLAS_VERSION_MINOR >= 38
113
114
115
116
        auto& ctx              = get_context();
        const auto device_name = trim(split_string(get_device_name(), ':').front());
        if(contains(get_rocblas_fp32_archs(), device_name))
            compute_fp32 = true;
Shucai Xiao's avatar
Shucai Xiao committed
117
118
119
120
121
        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
122
        offload_copy = (mod->name() == "main") ? pass->offload_copy : false;
123
        create_output_names();
Paul's avatar
Paul committed
124

125
126
127
128
129
130
131
132
133
134
135
136
        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");
137
        add_generic_op("equal");
138
139
140
        add_generic_op("erf");
        add_generic_op("exp");
        add_generic_op("floor");
141
142
        add_generic_op("greater");
        add_generic_op("less");
143
        add_generic_op("log");
Shucai Xiao's avatar
Shucai Xiao committed
144
145
146
        add_generic_op("logical_and");
        add_generic_op("logical_or");
        add_generic_op("logical_xor");
147
148
149
        add_generic_op("max");
        add_generic_op("min");
        add_generic_op("mul");
150
        add_generic_op("not");
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
        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
166
        add_generic_op("where");
167

Shucai Xiao's avatar
Shucai Xiao committed
168
        add_extend_op("abs");
169
170
171
172
173
        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
174
        add_extend_op("elu");
175
        add_extend_op("gather");
Shucai Xiao's avatar
Shucai Xiao committed
176
        add_extend_op("leaky_relu");
177
        add_extend_op("logsoftmax");
Shucai Xiao's avatar
Shucai Xiao committed
178
        add_extend_op("lrn");
turneram's avatar
turneram committed
179
        add_extend_op("multinomial");
Shucai Xiao's avatar
Shucai Xiao committed
180
        add_extend_op("nonzero");
181
        add_extend_op("pad");
182
        add_extend_op("pooling");
183
        add_extend_op("prefix_scan_sum");
Cagri Eryilmaz's avatar
Cagri Eryilmaz committed
184
        add_extend_op("reverse");
185
186
187
        add_extend_op("rnn_var_sl_last_output");
        add_extend_op("rnn_var_sl_shift_output");
        add_extend_op("rnn_var_sl_shift_sequence");
188
        add_extend_op("scatter_none");
Shucai Xiao's avatar
Shucai Xiao committed
189
        add_extend_op("topk");
190

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

203
204
    void copy_params()
    {
Shucai Xiao's avatar
Shucai Xiao committed
205
        if(not offload_copy)
206
            return;
207

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

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

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

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

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

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

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

263
264
265
266
267
268
269
270
271
272
273
274
275
    instruction_ref insert_precompile_op(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());
    }

Paul's avatar
Paul committed
276
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
277
    {
278
        // Instruction's output is an input of the ret instruction
Shucai Xiao's avatar
Shucai Xiao committed
279
        if(offload_copy)
Paul's avatar
Paul committed
280
        {
281
282
            auto result = mod->insert_instruction(
                ins, make_op("hip::allocate", {{"shape", to_value(s)}, {"tag", std::move(tag)}}));
Paul's avatar
Paul committed
283
284
            return result;
        }
285
286
287
288

        auto ins_alias = instruction::get_output_alias(ins);
        if(last->name() == "@return" and tag.empty() and prog_output_names.count(ins_alias) > 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
289
            return mod->add_parameter(prog_output_names[ins_alias], s);
290
291
292
        }
        else if(ins == last and tag.empty())
        {
Shucai Xiao's avatar
Shucai Xiao committed
293
            return mod->add_parameter("output", s);
294
295
        }

296
297
        return mod->insert_instruction(
            ins, make_op("hip::allocate", {{"shape", to_value(s)}, {"tag", std::move(tag)}}));
Paul's avatar
Paul committed
298
299
    }

Shucai Xiao's avatar
Shucai Xiao committed
300
    void add_convolution_op()
Paul's avatar
Paul committed
301
    {
302
303
        apply_map.emplace("convolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
304

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

308
309
            auto workspace = insert_allocation(ins, ws, "workspace");
            auto output    = insert_allocation(ins, ins->get_shape());
kahmed10's avatar
kahmed10 committed
310

Shucai Xiao's avatar
Shucai Xiao committed
311
            return mod->replace_instruction(
kahmed10's avatar
kahmed10 committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
                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
326

Shucai Xiao's avatar
Shucai Xiao committed
327
            return mod->replace_instruction(
328
329
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
Paul's avatar
Paul committed
330
331
    }

332
333
    template <typename Op>
    void add_gemm_op(const std::string& name)
334
335
    {
        apply_map.emplace(name, [=](instruction_ref ins) {
336
            std::vector<instruction_ref> refs = ins->inputs();
Shucai Xiao's avatar
Shucai Xiao committed
337
            if(refs.size() == 2)
338
339
            {
                auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
340
341
342
343
344
345
                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())
346
                {
347
348
349
350
                    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;
351
352
                    refs.push_back(copy_out);
                }
Shucai Xiao's avatar
Shucai Xiao committed
353
354
355
356
                else
                {
                    refs.push_back(refs.back());
                }
357
            }
Shucai Xiao's avatar
Shucai Xiao committed
358
            return mod->replace_instruction(
359
                ins, rocblas_gemm<Op>{Op{}, 1, 0, int8_x4_format, compute_fp32}, refs);
360
361
362
        });
    }

363
364
365
366
367
    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)};
368
            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
369

Shucai Xiao's avatar
Shucai Xiao committed
370
            auto args      = ins->inputs();
371
            auto workspace = insert_allocation(ins, ws, "workspace");
Shucai Xiao's avatar
Shucai Xiao committed
372
373
            auto output    = insert_allocation(ins, ins->get_shape());

Shucai Xiao's avatar
Shucai Xiao committed
374
            return mod->replace_instruction(ins, conv, args[0], args[1], workspace, output);
Shucai Xiao's avatar
Shucai Xiao committed
375
376
377
        });
    }

378
379
380
    // add_generic_op just constructs the operator with no fields whereas add_extend_op copies over
    // the fields Since it doesn't have fields its default constructed

381
382
383
    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
384
    {
385
        apply_map.emplace(op_name, [=](instruction_ref ins) {
386
387
388
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
389

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

394
395
396
    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
397
    {
398
399
        apply_map.emplace(op_name, [=](instruction_ref ins) {
            auto&& op                         = ins->get_operator();
400
401
402
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
403

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

Shucai Xiao's avatar
Shucai Xiao committed
408
    void add_batch_norm_inference_op()
409
    {
410
411
412
413
        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
414
415
416
417
418
419
420
421
422
423
424
425
426
427
            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};
428
429
            std::vector<instruction_ref> reshapes;
            std::transform(ins->inputs().begin() + 1,
Shucai Xiao's avatar
Shucai Xiao committed
430
431
                           ins->inputs().end(),
                           std::back_inserter(reshapes),
Shucai Xiao's avatar
Shucai Xiao committed
432
                           [&](auto i) { return mod->insert_instruction(ins, reshape_op, i); });
Shucai Xiao's avatar
Shucai Xiao committed
433

Shucai Xiao's avatar
Shucai Xiao committed
434
435
436
437
438
439
440
441
            return mod->replace_instruction(ins,
                                            miopen_batch_norm_inference{op},
                                            input,
                                            reshapes[0],
                                            reshapes[1],
                                            reshapes[2],
                                            reshapes[3],
                                            output);
442
        });
443
    }
Shucai Xiao's avatar
Shucai Xiao committed
444
445
446
447
448
449
450

    // 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
451
            auto l0     = mod->add_literal(literal(s, zeros));
Shucai Xiao's avatar
Shucai Xiao committed
452
            auto output = insert_allocation(ins, s);
Shucai Xiao's avatar
Shucai Xiao committed
453
            return mod->replace_instruction(
454
                ins, make_op("gpu::sub"), l0, ins->inputs().front(), output);
Shucai Xiao's avatar
Shucai Xiao committed
455
456
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
457

Shucai Xiao's avatar
Shucai Xiao committed
458
    // add input and output argument for the if operator
Shucai Xiao's avatar
Shucai Xiao committed
459
460
461
462
    void add_if_op()
    {
        apply_map.emplace("if", [=](instruction_ref ins) {
            std::vector<instruction_ref> inputs = ins->inputs();
463
464
465
            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
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
            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
                {
488
489
                    output = mod->insert_instruction(
                        ins, make_op("hip::allocate", {{"shape", to_value(s)}}));
Shucai Xiao's avatar
Shucai Xiao committed
490
491
492
493
494
495
496
                }
                inputs.push_back(output);
            }

            return mod->replace_instruction(ins, ins->get_operator(), inputs, mod_args);
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534

    // 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);
        });
    }
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554

    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
555
556
};

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

Paul's avatar
Paul committed
559
} // namespace gpu
Paul's avatar
Paul committed
560
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
561
} // namespace migraphx