lowering.cpp 9.43 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
21
#include <migraphx/gpu/sigmoid.hpp>
#include <migraphx/gpu/tanh.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
25
#include <migraphx/gpu/softmax.hpp>
#include <migraphx/gpu/add.hpp>
26
#include <migraphx/gpu/sin.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
27
28
#include <migraphx/gpu/cos.hpp>
#include <migraphx/gpu/tan.hpp>
Paul's avatar
Paul committed
29
30
31
32
33
#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
34
#include <utility>
35
#include <functional>
Paul's avatar
Paul committed
36

Paul's avatar
Paul committed
37
namespace migraphx {
38
inline namespace MIGRAPH_INLINE_NS {
Paul's avatar
Paul committed
39
namespace gpu {
Paul's avatar
Paul committed
40
41
42

struct miopen_apply
{
Paul's avatar
Paul committed
43
    program* prog = nullptr;
Paul's avatar
Paul committed
44
    context ctx{};
Khalique's avatar
Khalique committed
45
46
    std::unordered_map<std::string, std::function<instruction_ref(miopen_apply&, instruction_ref)>>
        apply_map{};
Paul's avatar
Paul committed
47

Paul's avatar
Paul committed
48
49
50
51
52
53
54
    void check_shape(shape x, instruction_ref i)
    {
        assert(x == i->get_shape());
        (void)x;
        (void)i;
    }

55
56
    void init()
    {
Khalique's avatar
Khalique committed
57
58
59
60
61
62
63
64
        apply_map["convolution"]          = &miopen_apply::apply_convolution;
        apply_map["relu"]                 = &miopen_apply::apply_relu;
        apply_map["sigmoid"]              = &miopen_apply::apply_sigmoid;
        apply_map["tanh"]                 = &miopen_apply::apply_tanh;
        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;
65
66
        apply_map["add"]                  = &miopen_apply::apply_generic_op<hip_add>;
        apply_map["sin"]                  = &miopen_apply::apply_generic_op<hip_sin>;
67
68
        apply_map["cos"]                  = &miopen_apply::apply_generic_op<hip_cos>;
        apply_map["tan"]                  = &miopen_apply::apply_generic_op<hip_tan>;
69
70
        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
71
72
73
74
        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;
75
76
    }

Paul's avatar
Paul committed
77
78
    void apply()
    {
79
        init();
Paul's avatar
Paul committed
80
81
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
Paul's avatar
Paul committed
82
            auto s = it->get_shape();
83
            if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
84
            {
85
                check_shape(s, apply_map.at(it->name())(*this, it));
Paul's avatar
Paul committed
86
            }
Paul's avatar
Paul committed
87
88
89
        }
    }

Paul's avatar
Paul committed
90
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
91
    {
Paul's avatar
Paul committed
92
        if(ins == --prog->end() and tag.empty())
Paul's avatar
Paul committed
93
94
95
96
97
        {
            return prog->add_parameter("output", s);
        }
        else
        {
Paul's avatar
Paul committed
98
            auto is     = prog->add_outline(s);
Paul's avatar
Paul committed
99
            auto result = prog->insert_instruction(ins, hip_allocate{std::move(tag)}, is);
Paul's avatar
Paul committed
100
101
102
103
            return result;
        }
    }

Paul's avatar
Paul committed
104
    instruction_ref apply_convolution(instruction_ref ins)
Paul's avatar
Paul committed
105
    {
wsttiger's avatar
wsttiger committed
106
        auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
107

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

111
        auto workspace = insert_allocation(ins, ws, "workspace");
Paul's avatar
Paul committed
112
        auto output    = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
113

Paul's avatar
Paul committed
114
        return prog->replace_instruction(
Paul's avatar
Paul committed
115
            ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
Paul's avatar
Paul committed
116
117
    }

Paul's avatar
Paul committed
118
    instruction_ref apply_pooling(instruction_ref ins)
Paul's avatar
Paul committed
119
    {
wsttiger's avatar
wsttiger committed
120
        auto&& op   = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
121
        auto pd     = make_pooling(op);
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, miopen_pooling{op, std::move(pd)}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
126
127
    }

Khalique's avatar
Khalique committed
128
    instruction_ref apply_relu(instruction_ref ins)
Paul's avatar
Paul committed
129
    {
Khalique's avatar
Khalique committed
130
        auto ad = make_relu();
Khalique's avatar
Khalique committed
131
132
133
134

        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
135
    }
136

Khalique's avatar
Khalique committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
    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
164
165
166
167
    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);
168

Khalique's avatar
Khalique committed
169
170
171
172
        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
173

Khalique's avatar
Khalique committed
174
175
176
177
178
179
180
181
182
183
    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
184
185
    instruction_ref apply_softmax(instruction_ref ins)
    {
wsttiger's avatar
wsttiger committed
186
        auto&& op   = any_cast<op::softmax>(ins->get_operator());
Paul's avatar
Paul committed
187
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
188
        return prog->replace_instruction(ins, miopen_softmax{op}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
189
190
    }

191
    /*
Paul's avatar
Paul committed
192
    instruction_ref apply_add(instruction_ref ins)
Paul's avatar
Paul committed
193
    {
Paul's avatar
Paul committed
194
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
195
        return prog->replace_instruction(
Paul's avatar
Paul committed
196
            ins, hip_add{}, ins->inputs().at(0), ins->inputs().at(1), output);
Paul's avatar
Paul committed
197
    }
198
    */
Paul's avatar
Paul committed
199

200
    /*
201
202
203
    instruction_ref apply_sin(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->get_shape());
Shucai Xiao's avatar
Shucai Xiao committed
204
        return prog->replace_instruction(ins, hip_sin{}, ins->inputs().at(0), output);
205
    }
206
    */
207

Shucai Xiao's avatar
Shucai Xiao committed
208
209
210
211
    template <class T>
    instruction_ref apply_generic_op(instruction_ref ins)
    {
        auto output                       = insert_allocation(ins, ins->get_shape());
212
213
214
215
216
217
        std::vector<instruction_ref> refs = ins->inputs();
        refs.push_back(output);

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

Paul's avatar
Paul committed
218
    instruction_ref apply_contiguous(instruction_ref ins)
219
    {
wsttiger's avatar
wsttiger committed
220
        auto&& op   = any_cast<op::contiguous>(ins->get_operator());
Paul's avatar
Paul committed
221
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
222
        return prog->replace_instruction(ins, miopen_contiguous{op}, ins->inputs().at(0), output);
223
    }
224

225
226
    instruction_ref apply_concat(instruction_ref ins)
    {
wsttiger's avatar
wsttiger committed
227
228
        auto&& op                         = any_cast<op::concat>(ins->get_operator());
        auto output                       = insert_allocation(ins, ins->get_shape());
229
230
231
232
233
        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
234
    instruction_ref apply_batch_norm_inference(instruction_ref ins)
235
    {
wsttiger's avatar
wsttiger committed
236
        auto&& op       = any_cast<op::batch_norm_inference>(ins->get_operator());
Paul's avatar
Paul committed
237
        auto output     = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
238
        shape old_shape = ins->inputs().at(1)->get_shape();
wsttiger's avatar
wsttiger committed
239
        std::vector<int64_t> new_shape{1, static_cast<int64_t>(old_shape.elements()), 1, 1};
wsttiger's avatar
wsttiger committed
240
        auto reshape_op = op::reshape{new_shape};
Paul's avatar
Paul committed
241
        std::vector<instruction_ref> reshapes;
Paul's avatar
Paul committed
242
243
        std::transform(ins->inputs().begin() + 1,
                       ins->inputs().end(),
Paul's avatar
Paul committed
244
245
                       std::back_inserter(reshapes),
                       [&](auto i) { return prog->insert_instruction(ins, reshape_op, i); });
Paul's avatar
Paul committed
246
        return prog->replace_instruction(ins,
Paul's avatar
Paul committed
247
                                         miopen_batch_norm_inference{op},
Paul's avatar
Paul committed
248
                                         ins->inputs().at(0),
Paul's avatar
Paul committed
249
250
251
252
253
                                         reshapes[0],
                                         reshapes[1],
                                         reshapes[2],
                                         reshapes[3],
                                         output);
254
    }
Paul's avatar
Paul committed
255
256
};

Paul's avatar
Paul committed
257
void lowering::apply(program& p) const { miopen_apply{&p, ctx}.apply(); }
Paul's avatar
Paul committed
258
} // namespace gpu
259
} // namespace MIGRAPH_INLINE_NS
Paul's avatar
Paul committed
260
} // namespace migraphx