lowering.cpp 13.3 KB
Newer Older
1
#include <rocblas.h>
Paul's avatar
Paul committed
2
3
4
5
6
7
8
9
10
11
12
13
#include <migraphx/gpu/lowering.hpp>
#include <migraphx/manage_ptr.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/operators.hpp>
#include <migraphx/generate.hpp>
#include <migraphx/shape_for_each.hpp>
#include <migraphx/gpu/miopen.hpp>
#include <migraphx/gpu/hip.hpp>
#include <migraphx/dfor.hpp>
#include <migraphx/gpu/device/contiguous.hpp>
#include <migraphx/gpu/device/add.hpp>
#include <migraphx/iterator_for.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
14
15
#include <migraphx/gpu/argmax.hpp>
#include <migraphx/gpu/argmin.hpp>
Paul's avatar
Paul committed
16
17
18
#include <migraphx/gpu/rocblas.hpp>
#include <migraphx/gpu/context.hpp>
#include <migraphx/gpu/convolution.hpp>
19
#include <migraphx/gpu/quant_convolution.hpp>
Paul's avatar
Paul committed
20
21
#include <migraphx/gpu/contiguous.hpp>
#include <migraphx/gpu/relu.hpp>
Khalique's avatar
Khalique committed
22
23
#include <migraphx/gpu/sigmoid.hpp>
#include <migraphx/gpu/abs.hpp>
Paul's avatar
Paul committed
24
#include <migraphx/gpu/leaky_relu.hpp>
Khalique's avatar
Khalique committed
25
#include <migraphx/gpu/elu.hpp>
Paul's avatar
Paul committed
26
#include <migraphx/gpu/softmax.hpp>
27
#include <migraphx/gpu/logsoftmax.hpp>
Paul's avatar
Paul committed
28
#include <migraphx/gpu/add.hpp>
29
#include <migraphx/gpu/sub.hpp>
30
#include <migraphx/gpu/div.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
31
#include <migraphx/gpu/exp.hpp>
32
#include <migraphx/gpu/erf.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
33
#include <migraphx/gpu/log.hpp>
34
#include <migraphx/gpu/sin.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
35
#include <migraphx/gpu/sign.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
36
37
#include <migraphx/gpu/cos.hpp>
#include <migraphx/gpu/tan.hpp>
38
39
#include <migraphx/gpu/sinh.hpp>
#include <migraphx/gpu/cosh.hpp>
40
#include <migraphx/gpu/tanh.hpp>
41
42
43
#include <migraphx/gpu/asin.hpp>
#include <migraphx/gpu/acos.hpp>
#include <migraphx/gpu/atan.hpp>
Paul's avatar
Paul committed
44
#include <migraphx/gpu/mul.hpp>
Khalique's avatar
Khalique committed
45
46
#include <migraphx/gpu/max.hpp>
#include <migraphx/gpu/min.hpp>
Paul's avatar
Paul committed
47
48
49
50
#include <migraphx/gpu/batchnorm.hpp>
#include <migraphx/gpu/pooling.hpp>
#include <migraphx/gpu/gemm.hpp>
#include <migraphx/gpu/concat.hpp>
51
#include <migraphx/gpu/pad.hpp>
52
#include <migraphx/gpu/gather.hpp>
Khalique's avatar
Khalique committed
53
#include <migraphx/gpu/lrn.hpp>
54
#include <migraphx/gpu/convert.hpp>
Khalique's avatar
Khalique committed
55
#include <migraphx/gpu/clip.hpp>
56
#include <migraphx/gpu/round.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
57
58
#include <migraphx/gpu/ceil.hpp>
#include <migraphx/gpu/floor.hpp>
Khalique's avatar
Khalique committed
59
#include <migraphx/gpu/rsqrt.hpp>
60
#include <migraphx/gpu/sqrt.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
61
#include <migraphx/gpu/reduce_max.hpp>
62
#include <migraphx/gpu/reduce_mean.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
63
#include <migraphx/gpu/reduce_min.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
64
65
#include <migraphx/gpu/reduce_prod.hpp>
#include <migraphx/gpu/reduce_sum.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
66
#include <migraphx/gpu/pow.hpp>
Khalique's avatar
Khalique committed
67
#include <migraphx/gpu/sqdiff.hpp>
68
#include <migraphx/gpu/int8_conv_pack.hpp>
Paul's avatar
Paul committed
69
#include <utility>
70
#include <functional>
Khalique's avatar
Khalique committed
71
#include <algorithm>
Paul's avatar
Paul committed
72

Paul's avatar
Paul committed
73
namespace migraphx {
Paul's avatar
Paul committed
74
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
75
namespace gpu {
Paul's avatar
Paul committed
76
77
78

struct miopen_apply
{
79
80
    program* prog        = nullptr;
    const lowering* pass = nullptr;
Shucai Xiao's avatar
Shucai Xiao committed
81
    std::unordered_map<std::string, std::function<instruction_ref(instruction_ref)>> apply_map{};
Shucai Xiao's avatar
Shucai Xiao committed
82
    instruction_ref last{};
Paul's avatar
Paul committed
83

84
85
86
87
88
89
90
    context& get_context()
    {
        assert(pass != nullptr);
        assert(pass->ctx != nullptr);
        return *pass->ctx;
    }

Paul's avatar
Paul committed
91
92
93
94
95
96
97
    void check_shape(shape x, instruction_ref i)
    {
        assert(x == i->get_shape());
        (void)x;
        (void)i;
    }

98
99
    void init()
    {
100
101
        assert(prog != nullptr);
        assert(pass != nullptr);
Shucai Xiao's avatar
Shucai Xiao committed
102
        this->last = instruction::get_output_alias(std::prev(prog->end()));
Paul's avatar
Paul committed
103

104
105
106
107
108
109
        add_miopen_simple_op<miopen_abs>("abs", make_abs);

        add_miopen_extend_op<miopen_leaky_relu, op::leaky_relu>("leaky_relu", make_leaky_relu);
        add_miopen_extend_op<miopen_elu, op::elu>("elu", make_elu);

        add_generic_op<hip_add>("add");
110
        add_generic_op<hip_sub>("sub");
111
        add_generic_op<hip_exp>("exp");
Shucai Xiao's avatar
Shucai Xiao committed
112
        add_generic_op<hip_erf>("erf");
113
        add_generic_op<hip_log>("log");
114
115
116
117
118
        add_generic_op<hip_sin>("sin");
        add_generic_op<hip_cos>("cos");
        add_generic_op<hip_tan>("tan");
        add_generic_op<hip_sinh>("sinh");
        add_generic_op<hip_cosh>("cosh");
119
        add_generic_op<hip_tanh>("tanh");
120
121
122
        add_generic_op<hip_asin>("asin");
        add_generic_op<hip_acos>("acos");
        add_generic_op<hip_atan>("atan");
123
        add_generic_op<hip_sqrt>("sqrt");
124
        add_generic_op<hip_mul>("mul");
125
        add_generic_op<hip_div>("div");
Khalique's avatar
Khalique committed
126
127
        add_generic_op<hip_max>("max");
        add_generic_op<hip_min>("min");
Khalique's avatar
Khalique committed
128
        add_generic_op<hip_rsqrt>("rsqrt");
129
        add_generic_op<hip_round>("round");
130
        add_generic_op<hip_pow>("pow");
Khalique's avatar
Khalique committed
131
        add_generic_op<hip_sqdiff>("sqdiff");
132
        add_generic_op<hip_relu>("relu");
Shucai Xiao's avatar
Shucai Xiao committed
133
        add_generic_op<hip_sign>("sign");
134
        add_generic_op<hip_sigmoid>("sigmoid");
Shucai Xiao's avatar
Shucai Xiao committed
135
136
        add_generic_op<hip_ceil>("ceil");
        add_generic_op<hip_floor>("floor");
137
138
139

        add_extend_op<miopen_contiguous, op::contiguous>("contiguous");
        add_extend_op<hip_concat, op::concat>("concat");
Khalique's avatar
Khalique committed
140
        add_extend_op<hip_softmax, op::softmax>("softmax");
141
        add_extend_op<hip_logsoftmax, op::logsoftmax>("logsoftmax");
142
143
        add_extend_op<hip_argmax, op::argmax>("argmax");
        add_extend_op<hip_argmin, op::argmin>("argmin");
Khalique's avatar
Khalique committed
144
        add_extend_op<hip_gather, op::gather>("gather");
145
        add_extend_op<hip_pad, op::pad>("pad");
146
        add_extend_op<hip_convert, op::convert>("convert");
Khalique's avatar
Khalique committed
147
        add_extend_op<hip_clip, op::clip>("clip");
Shucai Xiao's avatar
Shucai Xiao committed
148
        add_extend_op<hip_reduce_max, op::reduce_max>("reduce_max");
149
        add_extend_op<hip_reduce_mean, op::reduce_mean>("reduce_mean");
Shucai Xiao's avatar
Shucai Xiao committed
150
        add_extend_op<hip_reduce_min, op::reduce_min>("reduce_min");
Shucai Xiao's avatar
Shucai Xiao committed
151
152
        add_extend_op<hip_reduce_prod, op::reduce_prod>("reduce_prod");
        add_extend_op<hip_reduce_sum, op::reduce_sum>("reduce_sum");
153
154
        add_gemm_op<op::dot>("dot");
        add_gemm_op<op::quant_dot>("quant_dot");
155

Khalique's avatar
Khalique committed
156
        add_lrn_op();
157
        add_convolution_op();
158
        add_quant_convolution_op();
159
160
        add_pooling_op();
        add_batch_norm_inference_op();
161
162
    }

163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
    void copy_params()
    {
        if(not pass->offload_copy)
            return;
        for(auto ins : iterator_for(*prog))
        {
            if(ins->name() != "@param")
                continue;
            auto pos = std::next(ins);
            auto a   = insert_allocation(pos, ins->get_shape());
            auto c   = prog->insert_instruction(pos, hip_copy_to_gpu{}, ins, a);
            prog->replace_instruction(ins, c);
        }
        auto end = std::prev(prog->end());
        prog->add_instruction(hip_copy_from_gpu{}, end);
    }

Paul's avatar
Paul committed
180
181
    void apply()
    {
182
        init();
Paul's avatar
Paul committed
183
184
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
Paul's avatar
Paul committed
185
            auto s = it->get_shape();
186
            if(apply_map.count(it->name()) > 0)
187
            {
188
                check_shape(s, apply_map.at(it->name())(it));
Paul's avatar
Paul committed
189
            }
Paul's avatar
Paul committed
190
        }
191
        copy_params();
Paul's avatar
Paul committed
192
193
    }

Paul's avatar
Paul committed
194
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
195
    {
196
        if(not pass->offload_copy and ins == last and tag.empty())
Paul's avatar
Paul committed
197
198
199
200
201
        {
            return prog->add_parameter("output", s);
        }
        else
        {
202
            auto result = prog->insert_instruction(ins, hip_allocate{s, std::move(tag)});
Paul's avatar
Paul committed
203
204
205
206
            return result;
        }
    }

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

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

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

218
219
220
            return prog->replace_instruction(
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
Paul's avatar
Paul committed
221
222
    }

223
224
225
226
227
228
229
    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
230
            if(refs.size() == 2)
231
232
            {
                auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
233
234
235
236
237
238
239
                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())
240
                {
Shucai Xiao's avatar
Shucai Xiao committed
241
                    auto output   = insert_allocation(ins, ins->get_shape());
242
243
244
245
                    auto copy_out = prog->insert_instruction(ins, hip_copy{}, refs.back(), output);
                    refs.back()   = copy_out;
                    refs.push_back(copy_out);
                }
Shucai Xiao's avatar
Shucai Xiao committed
246
247
248
249
                else
                {
                    refs.push_back(refs.back());
                }
250
251
252
253
254
255
            }

            return prog->replace_instruction(ins, rocblas_gemm<Op>{Op{op.alpha, beta}}, refs);
        });
    }

256
257
258
259
260
    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)};
261
            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
262

Shucai Xiao's avatar
Shucai Xiao committed
263
            auto args      = ins->inputs();
264
            auto workspace = insert_allocation(ins, ws, "workspace");
Shucai Xiao's avatar
Shucai Xiao committed
265
266
            auto output    = insert_allocation(ins, ins->get_shape());

Shucai Xiao's avatar
Shucai Xiao committed
267
            return prog->replace_instruction(ins, conv, args[0], args[1], workspace, output);
Shucai Xiao's avatar
Shucai Xiao committed
268
269
270
        });
    }

Shucai Xiao's avatar
Shucai Xiao committed
271
    void add_pooling_op()
Paul's avatar
Paul committed
272
    {
273
274
275
276
        apply_map.emplace("pooling", [=](instruction_ref ins) {
            auto&& op   = any_cast<op::pooling>(ins->get_operator());
            auto pd     = make_pooling(op);
            auto output = insert_allocation(ins, ins->get_shape());
277

278
279
280
            return prog->replace_instruction(
                ins, miopen_pooling{op, std::move(pd)}, ins->inputs().at(0), output);
        });
Paul's avatar
Paul committed
281
    }
282

Khalique's avatar
Khalique committed
283
    void add_lrn_op()
Khalique's avatar
Khalique committed
284
    {
Khalique's avatar
Khalique committed
285
        apply_map.emplace("lrn", [=](instruction_ref ins) {
Khalique's avatar
Khalique committed
286
287
288
289
290
291
            auto&& op   = any_cast<op::lrn>(ins->get_operator());
            auto ldesc  = make_lrn(op);
            auto output = insert_allocation(ins, ins->get_shape());
            return prog->replace_instruction(
                ins, miopen_lrn{std::move(ldesc)}, ins->inputs().at(0), output);
        });
Khalique's avatar
Khalique committed
292
    }
Paul's avatar
Paul committed
293

Shucai Xiao's avatar
Shucai Xiao committed
294
    template <class T>
295
    void add_generic_op(std::string name)
Paul's avatar
Paul committed
296
    {
297
298
299
300
        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);
Paul's avatar
Paul committed
301

302
303
            return prog->replace_instruction(ins, T{}, refs);
        });
Paul's avatar
Paul committed
304
    }
Paul's avatar
Paul committed
305

Shucai Xiao's avatar
Shucai Xiao committed
306
    template <class T, class Op>
307
    void add_extend_op(std::string name)
Khalique's avatar
Khalique committed
308
    {
309
310
311
312
313
        apply_map.emplace(name, [=](instruction_ref ins) {
            auto&& op                         = any_cast<Op>(ins->get_operator());
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
314

315
316
            return prog->replace_instruction(ins, T{op}, refs);
        });
Khalique's avatar
Khalique committed
317
318
    }

Shucai Xiao's avatar
Shucai Xiao committed
319
    template <class T, class Op, class F>
320
    void add_miopen_extend_op(std::string name, F f)
Paul's avatar
Paul committed
321
    {
Shucai Xiao's avatar
Shucai Xiao committed
322
        apply_map.emplace(name, [=](instruction_ref ins) {
323
324
            auto&& op = any_cast<Op>(ins->get_operator());
            auto ad   = f(op.alpha);
325

326
            auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
327
            return prog->replace_instruction(ins, T{std::move(ad)}, ins->inputs().at(0), output);
328
        });
329
    }
330

Shucai Xiao's avatar
Shucai Xiao committed
331
    template <class T, class F>
332
    void add_miopen_simple_op(std::string name, F f)
333
    {
Shucai Xiao's avatar
Shucai Xiao committed
334
335
        apply_map.emplace(name, [=](instruction_ref ins) {
            auto ad     = f();
336
            auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
337
            return prog->replace_instruction(ins, T{std::move(ad)}, ins->inputs().at(0), output);
338
        });
339
340
    }

Shucai Xiao's avatar
Shucai Xiao committed
341
    void add_batch_norm_inference_op()
342
    {
343
344
345
346
347
348
349
350
        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();
            std::vector<int64_t> new_shape{1, static_cast<int64_t>(old_shape.elements()), 1, 1};
            auto reshape_op = op::reshape{new_shape};
            std::vector<instruction_ref> reshapes;
            std::transform(ins->inputs().begin() + 1,
Shucai Xiao's avatar
Shucai Xiao committed
351
352
353
                           ins->inputs().end(),
                           std::back_inserter(reshapes),
                           [&](auto i) { return prog->insert_instruction(ins, reshape_op, i); });
354
            return prog->replace_instruction(ins,
Shucai Xiao's avatar
Shucai Xiao committed
355
356
357
358
359
360
361
                                             miopen_batch_norm_inference{op},
                                             ins->inputs().at(0),
                                             reshapes[0],
                                             reshapes[1],
                                             reshapes[2],
                                             reshapes[3],
                                             output);
362
        });
363
    }
Paul's avatar
Paul committed
364
365
};

366
void lowering::apply(program& p) const { miopen_apply{&p, this}.apply(); }
Paul's avatar
Paul committed
367
} // namespace gpu
Paul's avatar
Paul committed
368
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
369
} // namespace migraphx