lowering.cpp 23.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>
Shucai Xiao's avatar
Shucai Xiao committed
23
#include <migraphx/gpu/compile_roialign.hpp>
turneram's avatar
turneram committed
24
#include <migraphx/gpu/compile_scatternd.hpp>
Paul's avatar
Paul committed
25
26
#include <migraphx/gpu/context.hpp>
#include <migraphx/gpu/convolution.hpp>
kahmed10's avatar
kahmed10 committed
27
#include <migraphx/gpu/deconvolution.hpp>
28
#include <migraphx/gpu/device_name.hpp>
Khalique's avatar
Khalique committed
29
#include <migraphx/gpu/elu.hpp>
30
#include <migraphx/gpu/equal.hpp>
Paul's avatar
Paul committed
31
#include <migraphx/gpu/gemm.hpp>
32
#include <migraphx/gpu/greater.hpp>
33
#include <migraphx/gpu/int8_conv_pack.hpp>
34
#include <migraphx/gpu/leaky_relu.hpp>
35
#include <migraphx/gpu/less.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
36
37
38
#include <migraphx/gpu/logical_and.hpp>
#include <migraphx/gpu/logical_or.hpp>
#include <migraphx/gpu/logical_xor.hpp>
39
40
41
42
#include <migraphx/gpu/lrn.hpp>
#include <migraphx/gpu/miopen.hpp>
#include <migraphx/gpu/quant_convolution.hpp>
#include <migraphx/gpu/rocblas.hpp>
43
#include <migraphx/gpu/unary_not.hpp>
turneram's avatar
turneram committed
44
#include <migraphx/gpu/where.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");
185
186
187
188
189
        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
190
        add_extend_op("reverse");
191
192
193
        add_extend_op("rnn_var_sl_last_output");
        add_extend_op("rnn_var_sl_shift_output");
        add_extend_op("rnn_var_sl_shift_sequence");
194
        add_extend_op("scatter");
195
        add_extend_op("softmax");
Shucai Xiao's avatar
Shucai Xiao committed
196
        add_extend_op("topk");
197

198
199
        add_precompile_op("pointwise");

Shucai Xiao's avatar
Shucai Xiao committed
200
        add_batch_norm_inference_op();
201
        add_convolution_op();
kahmed10's avatar
kahmed10 committed
202
        add_deconvolution_op();
Shucai Xiao's avatar
Shucai Xiao committed
203
204
        add_gemm_op<op::dot>("dot");
        add_gemm_op<op::quant_dot>("quant_dot");
Shucai Xiao's avatar
Shucai Xiao committed
205
        add_if_op();
Shucai Xiao's avatar
Shucai Xiao committed
206
        add_loop_op();
Shucai Xiao's avatar
Shucai Xiao committed
207
        add_neg_op();
208
        add_nms_op();
Shucai Xiao's avatar
Shucai Xiao committed
209
        add_quant_convolution_op();
Shucai Xiao's avatar
Shucai Xiao committed
210
        add_roialign();
turneram's avatar
turneram committed
211
        add_scatternd();
212
213
    }

214
215
    void copy_params()
    {
Shucai Xiao's avatar
Shucai Xiao committed
216
        if(not offload_copy)
217
            return;
218

Shucai Xiao's avatar
Shucai Xiao committed
219
        for(auto ins : iterator_for(*mod))
220
221
222
        {
            if(ins->name() != "@param")
                continue;
223

Shucai Xiao's avatar
Shucai Xiao committed
224
225
226
227
            // parameter no outputs, no need to insert copy to gpu
            if(ins->outputs().empty())
                continue;

228
229
            auto pos = std::next(ins);
            auto a   = insert_allocation(pos, ins->get_shape());
230
            auto c   = mod->insert_instruction(pos, make_op("hip::copy_to_gpu"), ins, a);
Shucai Xiao's avatar
Shucai Xiao committed
231
            mod->replace_instruction(ins, c);
232
        }
233
234

        // return instruction
Shucai Xiao's avatar
Shucai Xiao committed
235
        auto ret = std::prev(mod->end());
236
237
        if(ret->name() == "@return")
        {
238
            const auto& inputs = ret->inputs();
239
240
241

            // each input of ret need to be copied from gpu to host, and replace
            // output with copy output
242
            for(const auto& in : inputs)
243
            {
244
                auto p_output = mod->insert_instruction(ret, make_op("hip::copy_from_gpu"), in);
245
246
247
248
249
250
                instruction::replace_argument(ret, in, p_output);
            }
        }
        // else branch to handle legacy program without the return instruction
        else
        {
251
            mod->add_instruction(make_op("hip::copy_from_gpu"), ret);
252
        }
253
254
    }

Paul's avatar
Paul committed
255
256
    void apply()
    {
257
        init();
Shucai Xiao's avatar
Shucai Xiao committed
258
        for(auto it = mod->begin(); it != mod->end(); it++)
Paul's avatar
Paul committed
259
        {
Paul's avatar
Paul committed
260
            auto s = it->get_shape();
261
            if(apply_map.count(it->name()) > 0)
262
            {
263
                check_shape(s, apply_map.at(it->name())(it));
Paul's avatar
Paul committed
264
            }
Paul's avatar
Paul committed
265
        }
266

267
        copy_params();
Paul's avatar
Paul committed
268
269
    }

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

        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
283
            return mod->add_parameter(prog_output_names[ins_alias], s);
284
285
286
        }
        else if(ins == last and tag.empty())
        {
Shucai Xiao's avatar
Shucai Xiao committed
287
            return mod->add_parameter("output", s);
288
289
        }

290
291
        return mod->insert_instruction(
            ins, make_op("hip::allocate", {{"shape", to_value(s)}, {"tag", std::move(tag)}}));
Paul's avatar
Paul committed
292
293
    }

Shucai Xiao's avatar
Shucai Xiao committed
294
    void add_convolution_op()
Paul's avatar
Paul committed
295
    {
296
297
        apply_map.emplace("convolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
298

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

302
303
            auto workspace = insert_allocation(ins, ws, "workspace");
            auto output    = insert_allocation(ins, ins->get_shape());
kahmed10's avatar
kahmed10 committed
304

Shucai Xiao's avatar
Shucai Xiao committed
305
            return mod->replace_instruction(
kahmed10's avatar
kahmed10 committed
306
307
308
309
310
311
312
313
314
315
316
317
318
319
                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
320

Shucai Xiao's avatar
Shucai Xiao committed
321
            return mod->replace_instruction(
322
323
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
Paul's avatar
Paul committed
324
325
    }

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

357
358
359
360
361
    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)};
362
            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
363

Shucai Xiao's avatar
Shucai Xiao committed
364
            auto args      = ins->inputs();
365
            auto workspace = insert_allocation(ins, ws, "workspace");
Shucai Xiao's avatar
Shucai Xiao committed
366
367
            auto output    = insert_allocation(ins, ins->get_shape());

Shucai Xiao's avatar
Shucai Xiao committed
368
            return mod->replace_instruction(ins, conv, args[0], args[1], workspace, output);
Shucai Xiao's avatar
Shucai Xiao committed
369
370
371
        });
    }

372
373
374
    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
375
    {
376
        apply_map.emplace(op_name, [=](instruction_ref ins) {
377
378
379
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
380

Shucai Xiao's avatar
Shucai Xiao committed
381
            return mod->replace_instruction(ins, make_op(gpu_name), refs);
382
        });
Paul's avatar
Paul committed
383
    }
Paul's avatar
Paul committed
384

385
386
387
    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
388
    {
389
390
        apply_map.emplace(op_name, [=](instruction_ref ins) {
            auto&& op                         = ins->get_operator();
391
392
393
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
394

Shucai Xiao's avatar
Shucai Xiao committed
395
            return mod->replace_instruction(ins, make_op(gpu_name, op.to_value()), refs);
396
        });
Khalique's avatar
Khalique committed
397
398
    }

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

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
504
505
506
507
    void add_roialign()
    {
        apply_map.emplace("roialign", [=](instruction_ref ins) {
            auto s      = ins->get_shape();
Shucai Xiao's avatar
Shucai Xiao committed
508
            auto op_val = ins->get_operator().to_value();
Shucai Xiao's avatar
Shucai Xiao committed
509
            auto output = insert_allocation(ins, s);
Shucai Xiao's avatar
Shucai Xiao committed
510
511
512
513
514
515
            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
516
517
518
        });
    }

turneram's avatar
turneram committed
519
520
521
522
523
524
525
526
527
528
529
530
531
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
569
570
571
572
    void add_scatternd()
    {
        apply_map.emplace("scatternd_none", [=](instruction_ref ins) {
            auto s      = ins->get_shape();
            auto op_val = ins->get_operator().to_value();
            auto output = insert_allocation(ins, s);
            auto args   = ins->inputs();
            args.push_back(output);

            auto io_shapes = to_shapes(args);
            io_shapes.erase(io_shapes.begin());
            const std::string reduction = "none";
            auto co                     = compile_scatternd(get_context(), io_shapes, reduction);
            auto copy   = mod->insert_instruction(ins, make_op("hip::copy"), args.front(), output);
            args.back() = copy;
            args.erase(args.begin());
            return mod->replace_instruction(ins, co, args);
        });

        apply_map.emplace("scatternd_add", [=](instruction_ref ins) {
            auto s      = ins->get_shape();
            auto op_val = ins->get_operator().to_value();
            auto output = insert_allocation(ins, s);
            auto args   = ins->inputs();
            args.push_back(output);

            auto io_shapes = to_shapes(args);
            io_shapes.erase(io_shapes.begin());
            const std::string reduction = "add";
            auto co                     = compile_scatternd(get_context(), io_shapes, reduction);
            auto copy   = mod->insert_instruction(ins, make_op("hip::copy"), args.front(), output);
            args.back() = copy;
            args.erase(args.begin());
            return mod->replace_instruction(ins, co, args);
        });

        apply_map.emplace("scatternd_mul", [=](instruction_ref ins) {
            auto s      = ins->get_shape();
            auto op_val = ins->get_operator().to_value();
            auto output = insert_allocation(ins, s);
            auto args   = ins->inputs();
            args.push_back(output);

            auto io_shapes = to_shapes(args);
            io_shapes.erase(io_shapes.begin());
            const std::string reduction = "mul";
            auto co                     = compile_scatternd(get_context(), io_shapes, reduction);
            auto copy   = mod->insert_instruction(ins, make_op("hip::copy"), args.front(), output);
            args.back() = copy;
            args.erase(args.begin());
            return mod->replace_instruction(ins, co, args);
        });
    }

Shucai Xiao's avatar
Shucai Xiao committed
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
    // 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);
        });
    }
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629

    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
630
631
};

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

Paul's avatar
Paul committed
634
} // namespace gpu
Paul's avatar
Paul committed
635
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
636
} // namespace migraphx