lowering.cpp 15.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
Shucai Xiao's avatar
Shucai Xiao committed
24
#include <iterator>
Paul's avatar
Paul committed
25
26
27
#include <migraphx/gpu/lowering.hpp>
#include <migraphx/manage_ptr.hpp>
#include <migraphx/instruction.hpp>
28
29
30
31
32
#include <migraphx/make_op.hpp>

#include <migraphx/op/convolution.hpp>
#include <migraphx/op/deconvolution.hpp>
#include <migraphx/op/dot.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
33
#include <migraphx/op/if_op.hpp>
34
35
36
37
38
#include <migraphx/op/reshape.hpp>
#include <migraphx/op/quant_convolution.hpp>
#include <migraphx/op/quant_dot.hpp>

#include <migraphx/gpu/batch_norm_inference.hpp>
Paul's avatar
Paul committed
39
40
#include <migraphx/gpu/context.hpp>
#include <migraphx/gpu/convolution.hpp>
kahmed10's avatar
kahmed10 committed
41
#include <migraphx/gpu/deconvolution.hpp>
42
#include <migraphx/gpu/device_name.hpp>
Paul's avatar
Paul committed
43
#include <migraphx/gpu/gemm.hpp>
44
#include <migraphx/gpu/int8_conv_pack.hpp>
45
46
47
#include <migraphx/gpu/miopen.hpp>
#include <migraphx/gpu/quant_convolution.hpp>
#include <migraphx/gpu/rocblas.hpp>
48
#include <migraphx/gpu/compiler.hpp>
49
#include <migraphx/iterator_for.hpp>
50
#include <migraphx/program.hpp>
Paul's avatar
Paul committed
51
#include <utility>
52
#include <functional>
Khalique's avatar
Khalique committed
53
#include <algorithm>
Shucai Xiao's avatar
Shucai Xiao committed
54
#include <map>
Paul's avatar
Paul committed
55

Paul's avatar
Paul committed
56
namespace migraphx {
Paul's avatar
Paul committed
57
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
58
namespace gpu {
Paul's avatar
Paul committed
59
60
61

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

70
    context& get_context() const
71
72
73
74
75
76
    {
        assert(pass != nullptr);
        assert(pass->ctx != nullptr);
        return *pass->ctx;
    }

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

84
85
    void init()
    {
Shucai Xiao's avatar
Shucai Xiao committed
86
        assert(mod != nullptr);
87
        assert(pass != nullptr);
88

89
90
91
        auto& ctx      = get_context();
        int8_x4_format = get_int8_x4_format(ctx);
        compute_fp32   = get_compute_fp32_flag();
Shucai Xiao's avatar
Shucai Xiao committed
92

Shucai Xiao's avatar
Shucai Xiao committed
93
        offload_copy = (mod->name() == "main") ? pass->offload_copy : false;
Paul's avatar
Paul committed
94

95
96
97
98
        add_generic_op("contiguous");

        add_extend_op("argmax");
        add_extend_op("argmin");
Shucai Xiao's avatar
Shucai Xiao committed
99
        add_extend_op("elu");
100
        add_extend_op("gather");
Shucai Xiao's avatar
Shucai Xiao committed
101
        add_extend_op("leaky_relu");
102
        add_extend_op("logsoftmax");
Shucai Xiao's avatar
Shucai Xiao committed
103
        add_extend_op("lrn");
turneram's avatar
turneram committed
104
        add_extend_op("multinomial");
Shucai Xiao's avatar
Shucai Xiao committed
105
        add_extend_op("nonzero");
106
        add_extend_op("pad");
107
        add_extend_op("pooling");
108
        add_extend_op("prefix_scan_sum");
Cagri Eryilmaz's avatar
Cagri Eryilmaz committed
109
        add_extend_op("reverse");
110
111
112
        add_extend_op("rnn_var_sl_last_output");
        add_extend_op("rnn_var_sl_shift_output");
        add_extend_op("rnn_var_sl_shift_sequence");
113
        add_extend_op("scatter_none");
Shucai Xiao's avatar
Shucai Xiao committed
114
        add_extend_op("topk");
115

Shucai Xiao's avatar
Shucai Xiao committed
116
        add_batch_norm_inference_op();
117
        add_convolution_op();
kahmed10's avatar
kahmed10 committed
118
        add_deconvolution_op();
Shucai Xiao's avatar
Shucai Xiao committed
119
120
        add_gemm_op<op::dot>("dot");
        add_gemm_op<op::quant_dot>("quant_dot");
Shucai Xiao's avatar
Shucai Xiao committed
121
        add_if_op();
Shucai Xiao's avatar
Shucai Xiao committed
122
        add_loop_op();
Shucai Xiao's avatar
Shucai Xiao committed
123
        add_neg_op();
124
        add_nms_op();
Shucai Xiao's avatar
Shucai Xiao committed
125
        add_quant_convolution_op();
126
127
    }

128
    void copy_params() const
129
    {
Shucai Xiao's avatar
Shucai Xiao committed
130
        if(not offload_copy)
131
            return;
132

Shucai Xiao's avatar
Shucai Xiao committed
133
        for(auto ins : iterator_for(*mod))
134
135
136
        {
            if(ins->name() != "@param")
                continue;
137

Shucai Xiao's avatar
Shucai Xiao committed
138
139
140
141
            // parameter no outputs, no need to insert copy to gpu
            if(ins->outputs().empty())
                continue;

142
143
            auto pos = std::next(ins);
            auto a   = insert_allocation(pos, ins->get_shape());
144
            auto c   = mod->insert_instruction(pos, make_op("hip::copy_to_gpu"), ins, a);
Shucai Xiao's avatar
Shucai Xiao committed
145
            mod->replace_instruction(ins, c);
146
        }
147
148

        // return instruction
Shucai Xiao's avatar
Shucai Xiao committed
149
        auto ret = std::prev(mod->end());
150
151
        if(ret->name() == "@return")
        {
152
            const auto& inputs = ret->inputs();
153
154
155

            // each input of ret need to be copied from gpu to host, and replace
            // output with copy output
156
            for(const auto& in : inputs)
157
            {
158
                auto p_output = mod->insert_instruction(ret, make_op("hip::copy_from_gpu"), in);
159
160
161
162
163
164
                instruction::replace_argument(ret, in, p_output);
            }
        }
        // else branch to handle legacy program without the return instruction
        else
        {
165
            mod->add_instruction(make_op("hip::copy_from_gpu"), ret);
166
        }
167
168
    }

Paul's avatar
Paul committed
169
170
    void apply()
    {
171
        init();
Shucai Xiao's avatar
Shucai Xiao committed
172
        for(auto it = mod->begin(); it != mod->end(); it++)
Paul's avatar
Paul committed
173
        {
Paul's avatar
Paul committed
174
            auto s = it->get_shape();
175
            if(apply_map.count(it->name()) > 0)
176
            {
177
                check_shape(s, apply_map.at(it->name())(it));
Paul's avatar
Paul committed
178
            }
179
180
181
182
            else if(has_compiler_for(it->name()))
            {
                check_shape(s, insert_precompile_op(it));
            }
Paul's avatar
Paul committed
183
        }
184

185
        copy_params();
Paul's avatar
Paul committed
186
187
    }

188
    instruction_ref insert_precompile_op(instruction_ref ins) const
189
190
191
192
193
194
195
196
197
198
199
200
    {
        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());
    }

201
    instruction_ref insert_allocation(instruction_ref ins, const shape& s) const
Paul's avatar
Paul committed
202
    {
203
        return mod->insert_instruction(ins, make_op("allocate", {{"shape", to_value(s)}}));
Paul's avatar
Paul committed
204
205
    }

Shucai Xiao's avatar
Shucai Xiao committed
206
    void add_convolution_op()
Paul's avatar
Paul committed
207
    {
208
209
        apply_map.emplace("convolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
210

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

214
            auto workspace = insert_allocation(ins, ws);
215
            auto output    = insert_allocation(ins, ins->get_shape());
kahmed10's avatar
kahmed10 committed
216

Shucai Xiao's avatar
Shucai Xiao committed
217
            return mod->replace_instruction(
kahmed10's avatar
kahmed10 committed
218
219
220
221
222
223
224
225
226
227
                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)};
Paul Fultz II's avatar
Paul Fultz II committed
228
            auto ws   = conv.find(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
kahmed10's avatar
kahmed10 committed
229

230
            auto workspace = insert_allocation(ins, ws);
kahmed10's avatar
kahmed10 committed
231
            auto output    = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
232

Shucai Xiao's avatar
Shucai Xiao committed
233
            return mod->replace_instruction(
234
235
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
Paul's avatar
Paul committed
236
237
    }

238
239
    template <typename Op>
    void add_gemm_op(const std::string& name)
240
241
    {
        apply_map.emplace(name, [=](instruction_ref ins) {
242
            std::vector<instruction_ref> refs = ins->inputs();
243
244
245
            assert(refs.size() == 2);
            auto output = insert_allocation(ins, ins->get_shape());
            refs.push_back(output);
Shucai Xiao's avatar
Shucai Xiao committed
246
            return mod->replace_instruction(
247
                ins, rocblas_gemm<Op>{Op{}, 1, 0, int8_x4_format, compute_fp32}, refs);
248
249
250
        });
    }

251
252
253
254
    void add_quant_convolution_op()
    {
        apply_map.emplace("quant_convolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::quant_convolution>(ins->get_operator());
255
256
257
258
            shape ws;
            miopen_quant_convolution conv;
            auto compile_quant_conv_with_format = [&](bool format) {
                conv = miopen_quant_convolution{op, format, make_conv(op)};
Paul Fultz II's avatar
Paul Fultz II committed
259
                ws   = conv.find(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
260
261
262
263
264
265
266
267
268
            };

            try
            {
                compile_quant_conv_with_format(int8_x4_format);
            }
            catch(migraphx::exception&)
            {
                // In case no solver supports the default format, retry using the other format.
269
                compile_quant_conv_with_format(not int8_x4_format);
270
            }
271

Shucai Xiao's avatar
Shucai Xiao committed
272
            auto args      = ins->inputs();
273
            auto workspace = insert_allocation(ins, ws);
Shucai Xiao's avatar
Shucai Xiao committed
274
275
            auto output    = insert_allocation(ins, ins->get_shape());

Shucai Xiao's avatar
Shucai Xiao committed
276
            return mod->replace_instruction(ins, conv, args[0], args[1], workspace, output);
Shucai Xiao's avatar
Shucai Xiao committed
277
278
279
        });
    }

280
281
282
    // 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

283
284
285
    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
286
    {
287
        apply_map.emplace(op_name, [=](instruction_ref ins) {
288
289
290
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
291

Shucai Xiao's avatar
Shucai Xiao committed
292
            return mod->replace_instruction(ins, make_op(gpu_name), refs);
293
        });
Paul's avatar
Paul committed
294
    }
Paul's avatar
Paul committed
295

296
297
298
    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
299
    {
300
301
        apply_map.emplace(op_name, [=](instruction_ref ins) {
            auto&& op                         = ins->get_operator();
302
303
304
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
305

Shucai Xiao's avatar
Shucai Xiao committed
306
            return mod->replace_instruction(ins, make_op(gpu_name, op.to_value()), refs);
307
        });
Khalique's avatar
Khalique committed
308
309
    }

Shucai Xiao's avatar
Shucai Xiao committed
310
    void add_batch_norm_inference_op()
311
    {
312
313
314
315
        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
316
317
318
319
320
321
322
323
324
325
326
327
328
329
            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};
330
331
            std::vector<instruction_ref> reshapes;
            std::transform(ins->inputs().begin() + 1,
Shucai Xiao's avatar
Shucai Xiao committed
332
333
                           ins->inputs().end(),
                           std::back_inserter(reshapes),
Shucai Xiao's avatar
Shucai Xiao committed
334
                           [&](auto i) { return mod->insert_instruction(ins, reshape_op, i); });
Shucai Xiao's avatar
Shucai Xiao committed
335

Shucai Xiao's avatar
Shucai Xiao committed
336
337
338
339
340
341
342
343
            return mod->replace_instruction(ins,
                                            miopen_batch_norm_inference{op},
                                            input,
                                            reshapes[0],
                                            reshapes[1],
                                            reshapes[2],
                                            reshapes[3],
                                            output);
344
        });
345
    }
Shucai Xiao's avatar
Shucai Xiao committed
346
347
348
349
350
351
352

    // 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
353
            auto l0     = mod->add_literal(literal(s, zeros));
Shucai Xiao's avatar
Shucai Xiao committed
354
            auto output = insert_allocation(ins, s);
Shucai Xiao's avatar
Shucai Xiao committed
355
            return mod->replace_instruction(
356
                ins, make_op("gpu::sub"), l0, ins->inputs().front(), output);
Shucai Xiao's avatar
Shucai Xiao committed
357
358
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
359

Shucai Xiao's avatar
Shucai Xiao committed
360
    // add input and output argument for the if operator
Shucai Xiao's avatar
Shucai Xiao committed
361
362
363
364
    void add_if_op()
    {
        apply_map.emplace("if", [=](instruction_ref ins) {
            std::vector<instruction_ref> inputs = ins->inputs();
365
366
367
            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
368
369
            inputs.front() = sync_cond;

370
            return mod->replace_instruction(ins, ins->get_operator(), inputs, ins->module_inputs());
Shucai Xiao's avatar
Shucai Xiao committed
371
372
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388

    // 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;
389
390
391
392
            std::transform(copy_inputs.begin(),
                           copy_inputs.end(),
                           std::back_inserter(inputs),
                           [&](auto in) { return insert_allocation(ins, in->get_shape()); });
Shucai Xiao's avatar
Shucai Xiao committed
393
394
395
396
397

            auto mod_args = ins->module_inputs();
            auto output   = insert_allocation(ins, ins->get_shape());

            const auto* sub_mod = mod_args.front();
398
399
            auto cond_out       = insert_allocation(ins, sub_mod->get_output_shapes().front());

Shucai Xiao's avatar
Shucai Xiao committed
400
401
402
403
404
405
406
407
            // 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);
        });
    }
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427

    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
428
429
};

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

Paul's avatar
Paul committed
432
} // namespace gpu
Paul's avatar
Paul committed
433
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
434
} // namespace migraphx