lowering.cpp 20.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>
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/layernorm.hpp>
33
#include <migraphx/gpu/leaky_relu.hpp>
34
#include <migraphx/gpu/less.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
35
36
37
#include <migraphx/gpu/logical_and.hpp>
#include <migraphx/gpu/logical_or.hpp>
#include <migraphx/gpu/logical_xor.hpp>
38
39
40
41
#include <migraphx/gpu/lrn.hpp>
#include <migraphx/gpu/miopen.hpp>
#include <migraphx/gpu/quant_convolution.hpp>
#include <migraphx/gpu/rocblas.hpp>
42
#include <migraphx/gpu/unary_not.hpp>
turneram's avatar
turneram committed
43
#include <migraphx/gpu/where.hpp>
44
#include <migraphx/gpu/compiler.hpp>
45
#include <migraphx/iterator_for.hpp>
46
#include <migraphx/program.hpp>
Paul's avatar
Paul committed
47
#include <utility>
48
#include <functional>
Khalique's avatar
Khalique committed
49
#include <algorithm>
Shucai Xiao's avatar
Shucai Xiao committed
50
#include <map>
Paul's avatar
Paul committed
51

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

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

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

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

81
82
    void create_output_names()
    {
Shucai Xiao's avatar
Shucai Xiao committed
83
        this->last = instruction::get_output_alias(std::prev(mod->end()));
84
85
        if(this->last->name() == "@return")
        {
86
            const auto& prog_outputs = last->inputs();
87
88
89
90
91
92
93
94
95
96
            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
97
                prog_output_names[ins] = mod->name() + ":#output_" + std::to_string(index++);
98
99
100
101
            }
        }
    }

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

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

Shucai Xiao's avatar
Shucai Xiao committed
113
#if ROCBLAS_VERSION_MAJOR >= 2 && ROCBLAS_VERSION_MINOR >= 38
114
115
116
117
        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
118
119
120
121
122
        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
123
        offload_copy = (mod->name() == "main") ? pass->offload_copy : false;
124
        create_output_names();
Paul's avatar
Paul committed
125

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

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

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();
203
204
    }

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

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

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

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

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

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

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

262
        copy_params();
Paul's avatar
Paul committed
263
264
    }

265
266
267
268
269
270
271
272
273
274
275
276
277
    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
278
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
279
    {
280
        // Instruction's output is an input of the ret instruction
Shucai Xiao's avatar
Shucai Xiao committed
281
        if(offload_copy)
Paul's avatar
Paul committed
282
        {
283
284
            auto result = mod->insert_instruction(
                ins, make_op("hip::allocate", {{"shape", to_value(s)}, {"tag", std::move(tag)}}));
Paul's avatar
Paul committed
285
286
            return result;
        }
287
288
289
290

        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
291
            return mod->add_parameter(prog_output_names[ins_alias], s);
292
293
294
        }
        else if(ins == last and tag.empty())
        {
Shucai Xiao's avatar
Shucai Xiao committed
295
            return mod->add_parameter("output", s);
296
297
        }

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

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

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

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

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

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

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

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

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

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

380
381
382
    // 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

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
    }

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

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

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

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

            return mod->replace_instruction(ins, ins->get_operator(), inputs, mod_args);
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
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
535
536

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

    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
557
558
};

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

Paul's avatar
Paul committed
561
} // namespace gpu
Paul's avatar
Paul committed
562
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
563
} // namespace migraphx