lowering.cpp 15.6 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/hip.hpp>
30
#include <migraphx/gpu/int8_conv_pack.hpp>
31
#include <migraphx/gpu/leaky_relu.hpp>
32
#include <migraphx/gpu/less.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
33
34
35
#include <migraphx/gpu/logical_and.hpp>
#include <migraphx/gpu/logical_or.hpp>
#include <migraphx/gpu/logical_xor.hpp>
36
37
38
39
#include <migraphx/gpu/lrn.hpp>
#include <migraphx/gpu/miopen.hpp>
#include <migraphx/gpu/quant_convolution.hpp>
#include <migraphx/gpu/rocblas.hpp>
40
#include <migraphx/gpu/unary_not.hpp>
41
#include <migraphx/iterator_for.hpp>
42
#include <migraphx/program.hpp>
Paul's avatar
Paul committed
43
#include <utility>
44
#include <functional>
Khalique's avatar
Khalique committed
45
#include <algorithm>
Shucai Xiao's avatar
Shucai Xiao committed
46
#include <map>
Paul's avatar
Paul committed
47

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

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

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

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

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

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

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

104
105
106
107
108
109
110
111
112
113
114
115
        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");
116
        add_generic_op("equal");
117
118
119
        add_generic_op("erf");
        add_generic_op("exp");
        add_generic_op("floor");
120
121
        add_generic_op("greater");
        add_generic_op("less");
122
        add_generic_op("log");
Shucai Xiao's avatar
Shucai Xiao committed
123
124
125
        add_generic_op("logical_and");
        add_generic_op("logical_or");
        add_generic_op("logical_xor");
126
127
128
        add_generic_op("max");
        add_generic_op("min");
        add_generic_op("mul");
129
        add_generic_op("not");
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
        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
146
        add_extend_op("abs");
147
148
149
150
151
        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
152
        add_extend_op("elu");
153
        add_extend_op("gather");
Shucai Xiao's avatar
Shucai Xiao committed
154
        add_extend_op("leaky_relu");
155
        add_extend_op("logsoftmax");
Shucai Xiao's avatar
Shucai Xiao committed
156
        add_extend_op("lrn");
157
        add_extend_op("pad");
158
        add_extend_op("pooling");
159
160
161
162
163
164
165
166
167
168
        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");
        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());
Shucai Xiao's avatar
Shucai Xiao committed
191
192
            auto c   = mod->insert_instruction(pos, hip_copy_to_gpu{}, ins, a);
            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
            {
Shucai Xiao's avatar
Shucai Xiao committed
205
                auto p_output = mod->insert_instruction(ret, 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
        {
Shucai Xiao's avatar
Shucai Xiao committed
212
            mod->add_instruction(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
        {
Shucai Xiao's avatar
Shucai Xiao committed
236
            auto result = mod->insert_instruction(ins, hip_allocate{s, std::move(tag)});
Paul's avatar
Paul committed
237
238
            return result;
        }
239
240
241
242

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

Shucai Xiao's avatar
Shucai Xiao committed
250
        return mod->insert_instruction(ins, hip_allocate{s, std::move(tag)});
Paul's avatar
Paul committed
251
252
    }

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

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

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

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

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

285
286
287
288
289
290
291
    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
292
            if(refs.size() == 2)
293
294
            {
                auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
295
296
297
298
299
300
301
                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())
302
                {
Shucai Xiao's avatar
Shucai Xiao committed
303
                    auto output   = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
304
                    auto copy_out = mod->insert_instruction(ins, hip_copy{}, refs.back(), output);
305
306
307
                    refs.back()   = copy_out;
                    refs.push_back(copy_out);
                }
Shucai Xiao's avatar
Shucai Xiao committed
308
309
310
311
                else
                {
                    refs.push_back(refs.back());
                }
312
313
            }

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

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
342
            return mod->replace_instruction(ins, make_op(gpu_name), refs);
343
        });
Paul's avatar
Paul committed
344
    }
Paul's avatar
Paul committed
345

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
386
387
388
389
390
391
392
393
            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
394

395
        });
396
    }
Shucai Xiao's avatar
Shucai Xiao committed
397
398
399
400
401
402
403

    // 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
404
            auto l0     = mod->add_literal(literal(s, zeros));
Shucai Xiao's avatar
Shucai Xiao committed
405
            auto output = insert_allocation(ins, s);
Shucai Xiao's avatar
Shucai Xiao committed
406
            return mod->replace_instruction(
407
                ins, make_op("gpu::sub"), l0, ins->inputs().front(), output);
Shucai Xiao's avatar
Shucai Xiao committed
408
409
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447

    // 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();
            auto cpu_cond  = mod->insert_instruction(ins, hip_copy_from_gpu{}, inputs.front());
            auto sync_cond = mod->insert_instruction(ins, hip_sync_device{}, cpu_cond);
            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
                {
                    output = mod->insert_instruction(ins, hip_allocate{s});
                }
                inputs.push_back(output);
            }

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

Shucai Xiao's avatar
Shucai Xiao committed
450
void lowering::apply(module& m) const { miopen_apply{&m, this}.apply(); }
Paul's avatar
Paul committed
451
} // namespace gpu
Paul's avatar
Paul committed
452
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
453
} // namespace migraphx