lowering.cpp 9.54 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);
    }

Paul's avatar
Paul committed
211
    instruction_ref apply_contiguous(instruction_ref ins)
212
    {
wsttiger's avatar
wsttiger committed
213
        auto&& op   = any_cast<op::contiguous>(ins->get_operator());
Paul's avatar
Paul committed
214
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
215
        return prog->replace_instruction(ins, miopen_contiguous{op}, ins->inputs().at(0), output);
216
    }
217

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

Paul's avatar
Paul committed
250
void lowering::apply(program& p) const { miopen_apply{&p, ctx}.apply(); }
Paul's avatar
Paul committed
251
} // namespace gpu
252
} // namespace MIGRAPH_INLINE_NS
Paul's avatar
Paul committed
253
} // namespace migraphx