lowering.cpp 10.4 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
17
18
#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>
#include <migraphx/gpu/contiguous.hpp>
#include <migraphx/gpu/relu.hpp>
Khalique's avatar
Khalique committed
19
20
#include <migraphx/gpu/sigmoid.hpp>
#include <migraphx/gpu/abs.hpp>
Paul's avatar
Paul committed
21
#include <migraphx/gpu/leaky_relu.hpp>
Khalique's avatar
Khalique committed
22
#include <migraphx/gpu/elu.hpp>
Paul's avatar
Paul committed
23
24
#include <migraphx/gpu/softmax.hpp>
#include <migraphx/gpu/add.hpp>
25
#include <migraphx/gpu/sin.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
26
27
#include <migraphx/gpu/cos.hpp>
#include <migraphx/gpu/tan.hpp>
28
29
#include <migraphx/gpu/sinh.hpp>
#include <migraphx/gpu/cosh.hpp>
30
#include <migraphx/gpu/tanh.hpp>
31
32
33
#include <migraphx/gpu/asin.hpp>
#include <migraphx/gpu/acos.hpp>
#include <migraphx/gpu/atan.hpp>
Paul's avatar
Paul committed
34
35
36
37
38
#include <migraphx/gpu/mul.hpp>
#include <migraphx/gpu/batchnorm.hpp>
#include <migraphx/gpu/pooling.hpp>
#include <migraphx/gpu/gemm.hpp>
#include <migraphx/gpu/concat.hpp>
Paul's avatar
Paul committed
39
#include <utility>
40
#include <functional>
Paul's avatar
Paul committed
41

Paul's avatar
Paul committed
42
namespace migraphx {
43
inline namespace MIGRAPH_INLINE_NS {
Paul's avatar
Paul committed
44
namespace gpu {
Paul's avatar
Paul committed
45
46
47

struct miopen_apply
{
Paul's avatar
Paul committed
48
    program* prog = nullptr;
Paul's avatar
Paul committed
49
    context ctx{};
Khalique's avatar
Khalique committed
50
51
    std::unordered_map<std::string, std::function<instruction_ref(miopen_apply&, instruction_ref)>>
        apply_map{};
Paul's avatar
Paul committed
52

Paul's avatar
Paul committed
53
54
55
56
57
58
59
    void check_shape(shape x, instruction_ref i)
    {
        assert(x == i->get_shape());
        (void)x;
        (void)i;
    }

60
61
    void init()
    {
Khalique's avatar
Khalique committed
62
63
64
65
66
67
68
        apply_map["convolution"]          = &miopen_apply::apply_convolution;
        apply_map["relu"]                 = &miopen_apply::apply_relu;
        apply_map["sigmoid"]              = &miopen_apply::apply_sigmoid;
        apply_map["abs"]                  = &miopen_apply::apply_abs;
        apply_map["leaky_relu"]           = &miopen_apply::apply_leaky_relu;
        apply_map["elu"]                  = &miopen_apply::apply_elu;
        apply_map["pooling"]              = &miopen_apply::apply_pooling;
69
70
        apply_map["add"]                  = &miopen_apply::apply_generic_op<hip_add>;
        apply_map["sin"]                  = &miopen_apply::apply_generic_op<hip_sin>;
71
72
        apply_map["cos"]                  = &miopen_apply::apply_generic_op<hip_cos>;
        apply_map["tan"]                  = &miopen_apply::apply_generic_op<hip_tan>;
73
74
75
76
77
78
        apply_map["sinh"]                 = &miopen_apply::apply_generic_op<hip_sinh>;
        apply_map["cosh"]                 = &miopen_apply::apply_generic_op<hip_cosh>;
        apply_map["tanh"]                 = &miopen_apply::apply_tanh;
        apply_map["asin"]                 = &miopen_apply::apply_generic_op<hip_asin>;
        apply_map["acos"]                 = &miopen_apply::apply_generic_op<hip_acos>;
        apply_map["atan"]                 = &miopen_apply::apply_generic_op<hip_atan>;
79
80
        apply_map["mul"]                  = &miopen_apply::apply_generic_op<hip_mul>;
        apply_map["dot"]                  = &miopen_apply::apply_generic_op<miopen_gemm>;
Khalique's avatar
Khalique committed
81
82
83
84
        apply_map["contiguous"]           = &miopen_apply::apply_contiguous;
        apply_map["concat"]               = &miopen_apply::apply_concat;
        apply_map["batch_norm_inference"] = &miopen_apply::apply_batch_norm_inference;
        apply_map["softmax"]              = &miopen_apply::apply_softmax;
85
86
    }

Paul's avatar
Paul committed
87
88
    void apply()
    {
89
        init();
Paul's avatar
Paul committed
90
91
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
Paul's avatar
Paul committed
92
            auto s = it->get_shape();
93
            if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
94
            {
95
                check_shape(s, apply_map.at(it->name())(*this, it));
Paul's avatar
Paul committed
96
            }
Paul's avatar
Paul committed
97
98
99
        }
    }

Paul's avatar
Paul committed
100
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
101
    {
Paul's avatar
Paul committed
102
        if(ins == --prog->end() and tag.empty())
Paul's avatar
Paul committed
103
104
105
106
107
        {
            return prog->add_parameter("output", s);
        }
        else
        {
Paul's avatar
Paul committed
108
            auto is     = prog->add_outline(s);
Paul's avatar
Paul committed
109
            auto result = prog->insert_instruction(ins, hip_allocate{std::move(tag)}, is);
Paul's avatar
Paul committed
110
111
112
113
            return result;
        }
    }

Paul's avatar
Paul committed
114
    instruction_ref apply_convolution(instruction_ref ins)
Paul's avatar
Paul committed
115
    {
wsttiger's avatar
wsttiger committed
116
        auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
117

Paul's avatar
Paul committed
118
        auto conv = miopen_convolution{op, make_conv(op)};
Paul's avatar
Paul committed
119
        auto ws   = conv.compile(ctx, ins->get_shape(), ins->inputs());
Paul's avatar
Paul committed
120

121
        auto workspace = insert_allocation(ins, ws, "workspace");
Paul's avatar
Paul committed
122
        auto output    = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
123

Paul's avatar
Paul committed
124
        return prog->replace_instruction(
Paul's avatar
Paul committed
125
            ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
Paul's avatar
Paul committed
126
127
    }

Paul's avatar
Paul committed
128
    instruction_ref apply_pooling(instruction_ref ins)
Paul's avatar
Paul committed
129
    {
wsttiger's avatar
wsttiger committed
130
        auto&& op   = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
131
        auto pd     = make_pooling(op);
Paul's avatar
Paul committed
132
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
133

Paul's avatar
Paul committed
134
        return prog->replace_instruction(
Paul's avatar
Paul committed
135
            ins, miopen_pooling{op, std::move(pd)}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
136
137
    }

Khalique's avatar
Khalique committed
138
    instruction_ref apply_relu(instruction_ref ins)
Paul's avatar
Paul committed
139
    {
Khalique's avatar
Khalique committed
140
        auto ad = make_relu();
Khalique's avatar
Khalique committed
141
142
143
144

        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(
            ins, miopen_relu{std::move(ad)}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
145
    }
146

Khalique's avatar
Khalique committed
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
    instruction_ref apply_sigmoid(instruction_ref ins)
    {
        auto ad = make_sigmoid();

        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(
            ins, miopen_sigmoid{std::move(ad)}, ins->inputs().at(0), output);
    }

    instruction_ref apply_tanh(instruction_ref ins)
    {
        auto ad = make_tanh();

        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(
            ins, miopen_tanh{std::move(ad)}, ins->inputs().at(0), output);
    }

    instruction_ref apply_abs(instruction_ref ins)
    {
        auto ad = make_abs();

        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(
            ins, miopen_abs{std::move(ad)}, ins->inputs().at(0), output);
    }

Khalique's avatar
Khalique committed
174
175
176
177
    instruction_ref apply_leaky_relu(instruction_ref ins)
    {
        auto&& op = any_cast<op::leaky_relu>(ins->get_operator());
        auto ad   = make_leaky_relu(op.alpha);
178

Khalique's avatar
Khalique committed
179
180
181
182
        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(
            ins, miopen_leaky_relu{std::move(ad)}, ins->inputs().at(0), output);
    }
Paul's avatar
Paul committed
183

Khalique's avatar
Khalique committed
184
185
186
187
188
189
190
191
192
193
    instruction_ref apply_elu(instruction_ref ins)
    {
        auto&& op = any_cast<op::leaky_relu>(ins->get_operator());
        auto ad   = make_elu(op.alpha);

        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(
            ins, miopen_elu{std::move(ad)}, ins->inputs().at(0), output);
    }

Paul's avatar
Paul committed
194
195
    instruction_ref apply_softmax(instruction_ref ins)
    {
wsttiger's avatar
wsttiger committed
196
        auto&& op   = any_cast<op::softmax>(ins->get_operator());
Paul's avatar
Paul committed
197
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
198
        return prog->replace_instruction(ins, miopen_softmax{op}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
199
200
    }

Shucai Xiao's avatar
Shucai Xiao committed
201
202
203
204
    template <class T>
    instruction_ref apply_generic_op(instruction_ref ins)
    {
        auto output                       = insert_allocation(ins, ins->get_shape());
205
206
207
208
209
210
        std::vector<instruction_ref> refs = ins->inputs();
        refs.push_back(output);

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

211
212
213
    template <class T, class Op>
    instruction_ref apply_extend_op(instruction_ref ins)
    {
Shucai Xiao's avatar
Shucai Xiao committed
214
215
        auto&& op                         = any_cast<Op>(ins->get_operator());
        auto output                       = insert_allocation(ins, ins->get_shape());
216
217
218
219
220
221
        std::vector<instruction_ref> refs = ins->inputs();
        refs.push_back(output);

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

Shucai Xiao's avatar
Shucai Xiao committed
222
223
224
225
226
227
228
229
    /*
        template<class T>
        void apply_generic_op_test(std::string name, instruction_ref ins)
        {
            apply_map.emplace(name, [&]() {
                auto output                       = insert_allocation(ins, ins->get_shape());
                std::vector<instruction_ref> refs = ins->inputs();
                refs.push_back(output);
230

Shucai Xiao's avatar
Shucai Xiao committed
231
232
233
234
                return prog->replace_instruction(ins, T{}, refs);
            });
        }
    */
235

Paul's avatar
Paul committed
236
    instruction_ref apply_contiguous(instruction_ref ins)
237
    {
wsttiger's avatar
wsttiger committed
238
        auto&& op   = any_cast<op::contiguous>(ins->get_operator());
Paul's avatar
Paul committed
239
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
240
        return prog->replace_instruction(ins, miopen_contiguous{op}, ins->inputs().at(0), output);
241
    }
242

243
244
    instruction_ref apply_concat(instruction_ref ins)
    {
wsttiger's avatar
wsttiger committed
245
246
        auto&& op                         = any_cast<op::concat>(ins->get_operator());
        auto output                       = insert_allocation(ins, ins->get_shape());
247
248
249
250
251
        std::vector<instruction_ref> refs = ins->inputs();
        refs.push_back(output);
        return prog->replace_instruction(ins, hip_concat{op}, refs);
    }

Paul's avatar
Paul committed
252
    instruction_ref apply_batch_norm_inference(instruction_ref ins)
253
    {
wsttiger's avatar
wsttiger committed
254
        auto&& op       = any_cast<op::batch_norm_inference>(ins->get_operator());
Paul's avatar
Paul committed
255
        auto output     = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
256
        shape old_shape = ins->inputs().at(1)->get_shape();
wsttiger's avatar
wsttiger committed
257
        std::vector<int64_t> new_shape{1, static_cast<int64_t>(old_shape.elements()), 1, 1};
wsttiger's avatar
wsttiger committed
258
        auto reshape_op = op::reshape{new_shape};
Paul's avatar
Paul committed
259
        std::vector<instruction_ref> reshapes;
Paul's avatar
Paul committed
260
261
        std::transform(ins->inputs().begin() + 1,
                       ins->inputs().end(),
Paul's avatar
Paul committed
262
263
                       std::back_inserter(reshapes),
                       [&](auto i) { return prog->insert_instruction(ins, reshape_op, i); });
Paul's avatar
Paul committed
264
        return prog->replace_instruction(ins,
Paul's avatar
Paul committed
265
                                         miopen_batch_norm_inference{op},
Paul's avatar
Paul committed
266
                                         ins->inputs().at(0),
Paul's avatar
Paul committed
267
268
269
270
271
                                         reshapes[0],
                                         reshapes[1],
                                         reshapes[2],
                                         reshapes[3],
                                         output);
272
    }
Paul's avatar
Paul committed
273
274
};

Paul's avatar
Paul committed
275
void lowering::apply(program& p) const { miopen_apply{&p, ctx}.apply(); }
Paul's avatar
Paul committed
276
} // namespace gpu
277
} // namespace MIGRAPH_INLINE_NS
Paul's avatar
Paul committed
278
} // namespace migraphx