lowering.cpp 11.5 KB
Newer Older
1
#include <rocblas.h>
Paul's avatar
Paul committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#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>
#include <migraphx/gpu/rocblas.hpp>
#include <migraphx/gpu/context.hpp>
#include <migraphx/gpu/convolution.hpp>
17
#include <migraphx/gpu/quant_convolution.hpp>
Paul's avatar
Paul committed
18
19
#include <migraphx/gpu/contiguous.hpp>
#include <migraphx/gpu/relu.hpp>
Khalique's avatar
Khalique committed
20
21
#include <migraphx/gpu/sigmoid.hpp>
#include <migraphx/gpu/abs.hpp>
Paul's avatar
Paul committed
22
#include <migraphx/gpu/leaky_relu.hpp>
Khalique's avatar
Khalique committed
23
#include <migraphx/gpu/elu.hpp>
Paul's avatar
Paul committed
24
#include <migraphx/gpu/softmax.hpp>
25
#include <migraphx/gpu/logsoftmax.hpp>
Paul's avatar
Paul committed
26
#include <migraphx/gpu/add.hpp>
27
#include <migraphx/gpu/sub.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
28
29
#include <migraphx/gpu/exp.hpp>
#include <migraphx/gpu/log.hpp>
30
#include <migraphx/gpu/sin.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
31
32
#include <migraphx/gpu/cos.hpp>
#include <migraphx/gpu/tan.hpp>
33
34
#include <migraphx/gpu/sinh.hpp>
#include <migraphx/gpu/cosh.hpp>
35
#include <migraphx/gpu/tanh.hpp>
36
37
38
#include <migraphx/gpu/asin.hpp>
#include <migraphx/gpu/acos.hpp>
#include <migraphx/gpu/atan.hpp>
Paul's avatar
Paul committed
39
#include <migraphx/gpu/mul.hpp>
Khalique's avatar
Khalique committed
40
41
#include <migraphx/gpu/max.hpp>
#include <migraphx/gpu/min.hpp>
Paul's avatar
Paul committed
42
43
44
#include <migraphx/gpu/batchnorm.hpp>
#include <migraphx/gpu/pooling.hpp>
#include <migraphx/gpu/gemm.hpp>
45
#include <migraphx/gpu/quant_gemm.hpp>
Paul's avatar
Paul committed
46
#include <migraphx/gpu/concat.hpp>
47
#include <migraphx/gpu/pad.hpp>
48
#include <migraphx/gpu/gather.hpp>
Khalique's avatar
Khalique committed
49
#include <migraphx/gpu/lrn.hpp>
Khalique's avatar
Khalique committed
50
#include <migraphx/gpu/clip.hpp>
Paul's avatar
Paul committed
51
#include <utility>
52
#include <functional>
Khalique's avatar
Khalique committed
53
#include <algorithm>
Paul's avatar
Paul committed
54

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

struct miopen_apply
{
Paul's avatar
Paul committed
61
    program* prog = nullptr;
Paul's avatar
Paul committed
62
    context ctx{};
Shucai Xiao's avatar
Shucai Xiao committed
63
    std::unordered_map<std::string, std::function<instruction_ref(instruction_ref)>> apply_map{};
Shucai Xiao's avatar
Shucai Xiao committed
64
    instruction_ref last{};
Paul's avatar
Paul committed
65

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

73
74
    void init()
    {
Shucai Xiao's avatar
Shucai Xiao committed
75
        this->last = instruction::get_output_alias(std::prev(prog->end()));
76
77
78
79
80
81
82
83
84
        add_miopen_simple_op<miopen_relu>("relu", make_relu);
        add_miopen_simple_op<miopen_sigmoid>("sigmoid", make_sigmoid);
        add_miopen_simple_op<miopen_abs>("abs", make_abs);
        add_miopen_simple_op<miopen_tanh>("tanh", make_tanh);

        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");
85
        add_generic_op<hip_sub>("sub");
86
87
        add_generic_op<hip_exp>("exp");
        add_generic_op<hip_log>("log");
88
89
90
91
92
93
94
95
96
        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");
        add_generic_op<hip_asin>("asin");
        add_generic_op<hip_acos>("acos");
        add_generic_op<hip_atan>("atan");
        add_generic_op<hip_mul>("mul");
Khalique's avatar
Khalique committed
97
98
        add_generic_op<hip_max>("max");
        add_generic_op<hip_min>("min");
99
100
101
102
103

        add_extend_op<miopen_gemm, op::dot>("dot");
        add_extend_op<miopen_contiguous, op::contiguous>("contiguous");
        add_extend_op<hip_concat, op::concat>("concat");
        add_extend_op<miopen_softmax, op::softmax>("softmax");
104
        add_extend_op<hip_logsoftmax, op::logsoftmax>("logsoftmax");
Khalique's avatar
Khalique committed
105
        add_extend_op<hip_gather, op::gather>("gather");
106
        add_extend_op<hip_pad, op::pad>("pad");
Khalique's avatar
Khalique committed
107
        add_extend_op<hip_clip, op::clip>("clip");
108

Khalique's avatar
Khalique committed
109
        add_lrn_op();
110
        add_convolution_op();
111
        add_quant_convolution_op();
Shucai Xiao's avatar
Shucai Xiao committed
112
        add_quant_gemm_op();
113
114
        add_pooling_op();
        add_batch_norm_inference_op();
115
116
    }

Paul's avatar
Paul committed
117
118
    void apply()
    {
119
        init();
Paul's avatar
Paul committed
120
121
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
Paul's avatar
Paul committed
122
            auto s = it->get_shape();
123
            if(apply_map.count(it->name()) > 0)
124
            {
125
                check_shape(s, apply_map.at(it->name())(it));
Paul's avatar
Paul committed
126
            }
Paul's avatar
Paul committed
127
128
129
        }
    }

Paul's avatar
Paul committed
130
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
131
    {
Shucai Xiao's avatar
Shucai Xiao committed
132
        if(ins == last and tag.empty())
Paul's avatar
Paul committed
133
134
135
136
137
        {
            return prog->add_parameter("output", s);
        }
        else
        {
138
            auto result = prog->insert_instruction(ins, hip_allocate{s, std::move(tag)});
Paul's avatar
Paul committed
139
140
141
142
            return result;
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
143
    void add_convolution_op()
Paul's avatar
Paul committed
144
    {
145
146
        apply_map.emplace("convolution", [=](instruction_ref ins) {
            auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
147

148
            auto conv = miopen_convolution{op, make_conv(op)};
Paul's avatar
Paul committed
149
            auto ws   = conv.compile(ctx, ins->get_shape(), to_shapes(ins->inputs()));
Paul's avatar
Paul committed
150

151
152
            auto workspace = insert_allocation(ins, ws, "workspace");
            auto output    = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
153

154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
            return prog->replace_instruction(
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
    }

    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)};
            auto ws   = conv.compile(ctx, ins->get_shape(), to_shapes(ins->inputs()));

            auto workspace = insert_allocation(ins, ws, "workspace");
            auto output    = insert_allocation(ins, ins->get_shape());

170
171
172
            return prog->replace_instruction(
                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
        });
Paul's avatar
Paul committed
173
174
    }

Shucai Xiao's avatar
Shucai Xiao committed
175
176
177
    void add_quant_gemm_op()
    {
        apply_map.emplace("quant_dot", [=](instruction_ref ins) {
Shucai Xiao's avatar
Shucai Xiao committed
178
            auto&& op                         = any_cast<op::quant_dot>(ins->get_operator());
Shucai Xiao's avatar
Shucai Xiao committed
179
180
            std::vector<instruction_ref> refs = ins->inputs();

181
182
183
184
185
186
187
188
            // add additional arguments if need packing. Since lowering is added
            // after auto_contiguous and before eliminate contiguous, the shapes 
            // of all inputs are standard, so the input shape cannot be transposed.
            // To avoid that, we need to check whether this argument is an output
            // of contiguous. If true, we should check the shape of the input
            // of the contiguous operator.
            auto prev_ins = refs.at(0);
            if (prev_ins->name() == "gpu::contiguous")
Shucai Xiao's avatar
Shucai Xiao committed
189
            {
190
191
192
193
194
195
196
197
198
199
200
                auto input = prev_ins->inputs().front();
                if (input->get_shape().transposed())
                {
                    auto pack_a = insert_allocation(input, input->get_shape());
                    // replace one of the inputs of quant_gemm from the output to the
                    // input of contiguous. Then the contiguous could become dead code
                    // of prev_ins is its only output
                    refs.at(0) = input;
                    instruction::replace_argument(ins, prev_ins, input);
                    refs.push_back(pack_a);
                }
Shucai Xiao's avatar
Shucai Xiao committed
201
202
            }

Shucai Xiao's avatar
Shucai Xiao committed
203
            if(!refs.at(1)->get_shape().transposed())
Shucai Xiao's avatar
Shucai Xiao committed
204
205
206
207
            {
                auto pack_b = insert_allocation(refs.at(1), refs.at(1)->get_shape());
                refs.push_back(pack_b);
            }
Shucai Xiao's avatar
Shucai Xiao committed
208
            auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
209
210
211
212
213
214
            refs.push_back(output);

            return prog->replace_instruction(ins, miopen_quant_gemm{op}, refs);
        });
    }

Shucai Xiao's avatar
Shucai Xiao committed
215
    void add_pooling_op()
Paul's avatar
Paul committed
216
    {
217
218
219
220
        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());
221

222
223
224
            return prog->replace_instruction(
                ins, miopen_pooling{op, std::move(pd)}, ins->inputs().at(0), output);
        });
Paul's avatar
Paul committed
225
    }
226

Khalique's avatar
Khalique committed
227
    void add_lrn_op()
Khalique's avatar
Khalique committed
228
    {
Khalique's avatar
Khalique committed
229
        apply_map.emplace("lrn", [=](instruction_ref ins) {
Khalique's avatar
Khalique committed
230
231
232
233
234
235
            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
236
    }
Paul's avatar
Paul committed
237

Shucai Xiao's avatar
Shucai Xiao committed
238
    template <class T>
239
    void add_generic_op(std::string name)
Paul's avatar
Paul committed
240
    {
241
242
243
244
        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
245

246
247
            return prog->replace_instruction(ins, T{}, refs);
        });
Paul's avatar
Paul committed
248
    }
Paul's avatar
Paul committed
249

Shucai Xiao's avatar
Shucai Xiao committed
250
    template <class T, class Op>
251
    void add_extend_op(std::string name)
Khalique's avatar
Khalique committed
252
    {
253
254
255
256
257
        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
258

259
260
            return prog->replace_instruction(ins, T{op}, refs);
        });
Khalique's avatar
Khalique committed
261
262
    }

Shucai Xiao's avatar
Shucai Xiao committed
263
    template <class T, class Op, class F>
264
    void add_miopen_extend_op(std::string name, F f)
Paul's avatar
Paul committed
265
    {
Shucai Xiao's avatar
Shucai Xiao committed
266
        apply_map.emplace(name, [=](instruction_ref ins) {
267
268
            auto&& op = any_cast<Op>(ins->get_operator());
            auto ad   = f(op.alpha);
269

270
            auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
271
            return prog->replace_instruction(ins, T{std::move(ad)}, ins->inputs().at(0), output);
272
        });
273
    }
274

Shucai Xiao's avatar
Shucai Xiao committed
275
    template <class T, class F>
276
    void add_miopen_simple_op(std::string name, F f)
277
    {
Shucai Xiao's avatar
Shucai Xiao committed
278
279
        apply_map.emplace(name, [=](instruction_ref ins) {
            auto ad     = f();
280
            auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
281
            return prog->replace_instruction(ins, T{std::move(ad)}, ins->inputs().at(0), output);
282
        });
283
284
    }

Shucai Xiao's avatar
Shucai Xiao committed
285
    void add_batch_norm_inference_op()
286
    {
287
288
289
290
291
292
293
294
        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
295
296
297
                           ins->inputs().end(),
                           std::back_inserter(reshapes),
                           [&](auto i) { return prog->insert_instruction(ins, reshape_op, i); });
298
            return prog->replace_instruction(ins,
Shucai Xiao's avatar
Shucai Xiao committed
299
300
301
302
303
304
305
                                             miopen_batch_norm_inference{op},
                                             ins->inputs().at(0),
                                             reshapes[0],
                                             reshapes[1],
                                             reshapes[2],
                                             reshapes[3],
                                             output);
306
        });
307
    }
Paul's avatar
Paul committed
308
309
};

Paul's avatar
Paul committed
310
void lowering::apply(program& p) const { miopen_apply{&p, ctx}.apply(); }
Paul's avatar
Paul committed
311
} // namespace gpu
Paul's avatar
Paul committed
312
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
313
} // namespace migraphx