lowering.cpp 15.9 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
#include <migraphx/gpu/lowering.hpp>
#include <migraphx/manage_ptr.hpp>
#include <migraphx/instruction.hpp>
4
5
6
7
8
9
10
11
#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
12
#include <migraphx/op/if_op.hpp>
13
14
15
16
17
18
19
20
21
#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
22
23
#include <migraphx/gpu/context.hpp>
#include <migraphx/gpu/convolution.hpp>
kahmed10's avatar
kahmed10 committed
24
#include <migraphx/gpu/deconvolution.hpp>
Khalique's avatar
Khalique committed
25
#include <migraphx/gpu/elu.hpp>
26
#include <migraphx/gpu/equal.hpp>
Paul's avatar
Paul committed
27
#include <migraphx/gpu/gemm.hpp>
28
#include <migraphx/gpu/greater.hpp>
29
#include <migraphx/gpu/int8_conv_pack.hpp>
30
#include <migraphx/gpu/leaky_relu.hpp>
31
#include <migraphx/gpu/less.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
32
33
34
#include <migraphx/gpu/logical_and.hpp>
#include <migraphx/gpu/logical_or.hpp>
#include <migraphx/gpu/logical_xor.hpp>
35
36
37
38
#include <migraphx/gpu/lrn.hpp>
#include <migraphx/gpu/miopen.hpp>
#include <migraphx/gpu/quant_convolution.hpp>
#include <migraphx/gpu/rocblas.hpp>
39
#include <migraphx/gpu/unary_not.hpp>
40
#include <migraphx/iterator_for.hpp>
41
#include <migraphx/program.hpp>
Paul's avatar
Paul committed
42
#include <utility>
43
#include <functional>
Khalique's avatar
Khalique committed
44
#include <algorithm>
Shucai Xiao's avatar
Shucai Xiao committed
45
#include <map>
Paul's avatar
Paul committed
46

Paul's avatar
Paul committed
47
namespace migraphx {
Paul's avatar
Paul committed
48
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
49
namespace gpu {
Paul's avatar
Paul committed
50
51
52

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

60
    context& get_context() const
61
62
63
64
65
66
    {
        assert(pass != nullptr);
        assert(pass->ctx != nullptr);
        return *pass->ctx;
    }

Paul's avatar
Paul committed
67
68
69
70
71
72
73
    void check_shape(shape x, instruction_ref i)
    {
        assert(x == i->get_shape());
        (void)x;
        (void)i;
    }

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

95
96
    void init()
    {
Shucai Xiao's avatar
Shucai Xiao committed
97
        assert(mod != nullptr);
98
        assert(pass != nullptr);
99

Shucai Xiao's avatar
Shucai Xiao committed
100
        offload_copy = (mod->name() == "main") ? pass->offload_copy : false;
101
        create_output_names();
Paul's avatar
Paul committed
102

103
104
105
106
107
108
109
110
111
112
113
114
        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");
115
        add_generic_op("equal");
116
117
118
        add_generic_op("erf");
        add_generic_op("exp");
        add_generic_op("floor");
119
120
        add_generic_op("greater");
        add_generic_op("less");
121
        add_generic_op("log");
Shucai Xiao's avatar
Shucai Xiao committed
122
123
124
        add_generic_op("logical_and");
        add_generic_op("logical_or");
        add_generic_op("logical_xor");
125
126
127
        add_generic_op("max");
        add_generic_op("min");
        add_generic_op("mul");
128
        add_generic_op("not");
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
        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");

Shucai Xiao's avatar
Shucai Xiao committed
145
        add_extend_op("abs");
146
147
148
149
150
        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
151
        add_extend_op("elu");
152
        add_extend_op("gather");
Shucai Xiao's avatar
Shucai Xiao committed
153
        add_extend_op("leaky_relu");
154
        add_extend_op("logsoftmax");
Shucai Xiao's avatar
Shucai Xiao committed
155
        add_extend_op("lrn");
156
        add_extend_op("pad");
157
        add_extend_op("pooling");
158
159
160
161
162
        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
163
        add_extend_op("reverse");
164
165
166
167
168
        add_extend_op("rnn_var_sl_last_output");
        add_extend_op("rnn_var_sl_shift_output");
        add_extend_op("rnn_var_sl_shift_sequence");
        add_extend_op("softmax");

169
170
        add_gemm_op<op::dot>("dot");
        add_gemm_op<op::quant_dot>("quant_dot");
171
        add_convolution_op();
kahmed10's avatar
kahmed10 committed
172
        add_deconvolution_op();
173
        add_quant_convolution_op();
174
        add_batch_norm_inference_op();
Shucai Xiao's avatar
Shucai Xiao committed
175
        add_neg_op();
Shucai Xiao's avatar
Shucai Xiao committed
176
        add_if_op();
177
178
    }

179
180
    void copy_params()
    {
Shucai Xiao's avatar
Shucai Xiao committed
181
        if(not offload_copy)
182
            return;
183

Shucai Xiao's avatar
Shucai Xiao committed
184
        for(auto ins : iterator_for(*mod))
185
186
187
        {
            if(ins->name() != "@param")
                continue;
188

189
190
            auto pos = std::next(ins);
            auto a   = insert_allocation(pos, ins->get_shape());
191
            auto c   = mod->insert_instruction(pos, make_op("hip::copy_to_gpu"), ins, a);
Shucai Xiao's avatar
Shucai Xiao committed
192
            mod->replace_instruction(ins, c);
193
        }
194
195

        // return instruction
Shucai Xiao's avatar
Shucai Xiao committed
196
        auto ret = std::prev(mod->end());
197
198
        if(ret->name() == "@return")
        {
199
            const auto& inputs = ret->inputs();
200
201
202

            // each input of ret need to be copied from gpu to host, and replace
            // output with copy output
203
            for(const auto& in : inputs)
204
            {
205
                auto p_output = mod->insert_instruction(ret, make_op("hip::copy_from_gpu"), in);
206
207
208
209
210
211
                instruction::replace_argument(ret, in, p_output);
            }
        }
        // else branch to handle legacy program without the return instruction
        else
        {
212
            mod->add_instruction(make_op("hip::copy_from_gpu"), ret);
213
        }
214
215
    }

Paul's avatar
Paul committed
216
217
    void apply()
    {
218
        init();
Shucai Xiao's avatar
Shucai Xiao committed
219
        for(auto it = mod->begin(); it != mod->end(); it++)
Paul's avatar
Paul committed
220
        {
Paul's avatar
Paul committed
221
            auto s = it->get_shape();
222
            if(apply_map.count(it->name()) > 0)
223
            {
224
                check_shape(s, apply_map.at(it->name())(it));
Paul's avatar
Paul committed
225
            }
Paul's avatar
Paul committed
226
        }
227

228
        copy_params();
Paul's avatar
Paul committed
229
230
    }

Paul's avatar
Paul committed
231
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
232
    {
233
        // Instruction's output is an input of the ret instruction
Shucai Xiao's avatar
Shucai Xiao committed
234
        if(offload_copy)
Paul's avatar
Paul committed
235
        {
236
237
            auto result = mod->insert_instruction(
                ins, make_op("hip::allocate", {{"shape", to_value(s)}, {"tag", std::move(tag)}}));
Paul's avatar
Paul committed
238
239
            return result;
        }
240
241
242
243

        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
244
            return mod->add_parameter(prog_output_names[ins_alias], s);
245
246
247
        }
        else if(ins == last and tag.empty())
        {
Shucai Xiao's avatar
Shucai Xiao committed
248
            return mod->add_parameter("output", s);
249
250
        }

251
252
        return mod->insert_instruction(
            ins, make_op("hip::allocate", {{"shape", to_value(s)}, {"tag", std::move(tag)}}));
Paul's avatar
Paul committed
253
254
    }

Shucai Xiao's avatar
Shucai Xiao committed
255
    void add_convolution_op()
Paul's avatar
Paul committed
256
    {
257
258
        apply_map.emplace("convolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
259

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

263
264
            auto workspace = insert_allocation(ins, ws, "workspace");
            auto output    = insert_allocation(ins, ins->get_shape());
kahmed10's avatar
kahmed10 committed
265

Shucai Xiao's avatar
Shucai Xiao committed
266
            return mod->replace_instruction(
kahmed10's avatar
kahmed10 committed
267
268
269
270
271
272
273
274
275
276
277
278
279
280
                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
281

Shucai Xiao's avatar
Shucai Xiao committed
282
            return mod->replace_instruction(
283
284
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
Paul's avatar
Paul committed
285
286
    }

287
288
289
290
291
292
293
    template <class Op>
    void add_gemm_op(std::string name)
    {
        apply_map.emplace(name, [=](instruction_ref ins) {
            auto&& op                         = any_cast<Op>(ins->get_operator());
            auto beta                         = op.beta;
            std::vector<instruction_ref> refs = ins->inputs();
Shucai Xiao's avatar
Shucai Xiao committed
294
            if(refs.size() == 2)
295
296
            {
                auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
297
298
299
300
301
302
303
                beta        = 0;
                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())
304
                {
305
306
307
308
                    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;
309
310
                    refs.push_back(copy_out);
                }
Shucai Xiao's avatar
Shucai Xiao committed
311
312
313
314
                else
                {
                    refs.push_back(refs.back());
                }
315
316
            }

Shucai Xiao's avatar
Shucai Xiao committed
317
            return mod->replace_instruction(ins, rocblas_gemm<Op>{Op{op.alpha, beta}}, refs);
318
319
320
        });
    }

321
322
323
324
325
    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)};
326
            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
327

Shucai Xiao's avatar
Shucai Xiao committed
328
            auto args      = ins->inputs();
329
            auto workspace = insert_allocation(ins, ws, "workspace");
Shucai Xiao's avatar
Shucai Xiao committed
330
331
            auto output    = insert_allocation(ins, ins->get_shape());

Shucai Xiao's avatar
Shucai Xiao committed
332
            return mod->replace_instruction(ins, conv, args[0], args[1], workspace, output);
Shucai Xiao's avatar
Shucai Xiao committed
333
334
335
        });
    }

336
337
338
    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
339
    {
340
        apply_map.emplace(op_name, [=](instruction_ref ins) {
341
342
343
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
344

Shucai Xiao's avatar
Shucai Xiao committed
345
            return mod->replace_instruction(ins, make_op(gpu_name), refs);
346
        });
Paul's avatar
Paul committed
347
    }
Paul's avatar
Paul committed
348

349
350
351
    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
352
    {
353
354
        apply_map.emplace(op_name, [=](instruction_ref ins) {
            auto&& op                         = ins->get_operator();
355
356
357
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
358

Shucai Xiao's avatar
Shucai Xiao committed
359
            return mod->replace_instruction(ins, make_op(gpu_name, op.to_value()), refs);
360
        });
Khalique's avatar
Khalique committed
361
362
    }

Shucai Xiao's avatar
Shucai Xiao committed
363
    void add_batch_norm_inference_op()
364
    {
365
366
367
368
        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
369
370
371
372
373
374
375
376
377
378
379
380
381
382
            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};
383
384
            std::vector<instruction_ref> reshapes;
            std::transform(ins->inputs().begin() + 1,
Shucai Xiao's avatar
Shucai Xiao committed
385
386
                           ins->inputs().end(),
                           std::back_inserter(reshapes),
Shucai Xiao's avatar
Shucai Xiao committed
387
                           [&](auto i) { return mod->insert_instruction(ins, reshape_op, i); });
Shucai Xiao's avatar
Shucai Xiao committed
388

Shucai Xiao's avatar
Shucai Xiao committed
389
390
391
392
393
394
395
396
            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
397

398
        });
399
    }
Shucai Xiao's avatar
Shucai Xiao committed
400
401
402
403
404
405
406

    // 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
407
            auto l0     = mod->add_literal(literal(s, zeros));
Shucai Xiao's avatar
Shucai Xiao committed
408
            auto output = insert_allocation(ins, s);
Shucai Xiao's avatar
Shucai Xiao committed
409
            return mod->replace_instruction(
410
                ins, make_op("gpu::sub"), l0, ins->inputs().front(), output);
Shucai Xiao's avatar
Shucai Xiao committed
411
412
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
413
414
415
416
417
418

    // replace the if operator with gpu_if operator
    void add_if_op()
    {
        apply_map.emplace("if", [=](instruction_ref ins) {
            std::vector<instruction_ref> inputs = ins->inputs();
419
420
421
            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
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
            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
                {
444
445
                    output = mod->insert_instruction(
                        ins, make_op("hip::allocate", {{"shape", to_value(s)}}));
Shucai Xiao's avatar
Shucai Xiao committed
446
447
448
449
450
451
452
                }
                inputs.push_back(output);
            }

            return mod->replace_instruction(ins, ins->get_operator(), inputs, mod_args);
        });
    }
Paul's avatar
Paul committed
453
454
};

Shucai Xiao's avatar
Shucai Xiao committed
455
void lowering::apply(module& m) const { miopen_apply{&m, this}.apply(); }
Paul's avatar
Paul committed
456
} // namespace gpu
Paul's avatar
Paul committed
457
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
458
} // namespace migraphx