lowering.cpp 9.82 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
19
20
21
#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>
#include <migraphx/gpu/leaky_relu.hpp>
#include <migraphx/gpu/softmax.hpp>
#include <migraphx/gpu/add.hpp>
22
#include <migraphx/gpu/sin.hpp>
23
24
#include <migraphx/gpu/sinh.hpp>
#include <migraphx/gpu/cosh.hpp>
25
26
27
#include <migraphx/gpu/asin.hpp>
#include <migraphx/gpu/acos.hpp>
#include <migraphx/gpu/atan.hpp>
Paul's avatar
Paul committed
28
29
30
31
32
#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
33
#include <utility>
Paul's avatar
Paul committed
34

Paul's avatar
Paul committed
35
namespace migraphx {
36
inline namespace MIGRAPH_INLINE_NS {
Paul's avatar
Paul committed
37
namespace gpu {
Paul's avatar
Paul committed
38
39
40

struct miopen_apply
{
Paul's avatar
Paul committed
41
    program* prog = nullptr;
Paul's avatar
Paul committed
42
    context ctx{};
Paul's avatar
Paul committed
43

Paul's avatar
Paul committed
44
45
46
47
48
49
50
    void check_shape(shape x, instruction_ref i)
    {
        assert(x == i->get_shape());
        (void)x;
        (void)i;
    }

Paul's avatar
Paul committed
51
52
    void apply()
    {
Paul's avatar
Paul committed
53
54
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
Paul's avatar
Paul committed
55
            auto s = it->get_shape();
Paul's avatar
Paul committed
56
            if(it->name() == "convolution")
Paul's avatar
Paul committed
57
            {
Paul's avatar
Paul committed
58
                check_shape(s, apply_convolution(it));
Paul's avatar
Paul committed
59
            }
Khalique's avatar
Khalique committed
60
            else if(it->name() == "relu")
Paul's avatar
Paul committed
61
            {
Khalique's avatar
Khalique committed
62
                check_shape(s, apply_relu(it));
Paul's avatar
Paul committed
63
            }
Khalique's avatar
Khalique committed
64
65
66
67
            else if(it->name() == "leaky_relu")
            {
                check_shape(s, apply_leaky_relu(it));
            }
Paul's avatar
Paul committed
68
            else if(it->name() == "pooling")
Paul's avatar
Paul committed
69
            {
Paul's avatar
Paul committed
70
                check_shape(s, apply_pooling(it));
Paul's avatar
Paul committed
71
            }
Paul's avatar
Paul committed
72
            else if(it->name() == "add")
Paul's avatar
Paul committed
73
            {
Paul's avatar
Paul committed
74
                check_shape(s, apply_add(it));
Paul's avatar
Paul committed
75
            }
76
77
78
79
            else if(it->name() == "sin")
            {
                check_shape(s, apply_sin(it));
            }
80
81
82
83
84
85
86
87
            else if(it->name() == "sinh")
            {
                check_shape(s, apply_sinh(it));
            }
            else if(it->name() == "cosh")
            {
                check_shape(s, apply_cosh(it));
            }
88
89
90
91
92
93
94
95
96
97
98
99
            else if(it->name() == "asin")
            {
                check_shape(s, apply_asin(it));
            }
            else if(it->name() == "acos")
            {
                check_shape(s, apply_acos(it));
            }
            else if(it->name() == "atan")
            {
                check_shape(s, apply_atan(it));
            }
Khalique's avatar
Khalique committed
100
101
102
103
            else if(it->name() == "mul")
            {
                check_shape(s, apply_mul(it));
            }
Shucai Xiao's avatar
Shucai Xiao committed
104
            else if(it->name() == "dot")
Paul's avatar
Paul committed
105
            {
Paul's avatar
Paul committed
106
                check_shape(s, apply_gemm(it));
Paul's avatar
Paul committed
107
            }
Paul's avatar
Paul committed
108
            else if(it->name() == "contiguous")
109
            {
Paul's avatar
Paul committed
110
                check_shape(s, apply_contiguous(it));
111
            }
112
113
114
115
            else if(it->name() == "concat")
            {
                check_shape(s, apply_concat(it));
            }
Paul's avatar
Paul committed
116
            else if(it->name() == "batch_norm_inference")
117
            {
Paul's avatar
Paul committed
118
                check_shape(s, apply_batch_norm_inference(it));
119
            }
Paul's avatar
Paul committed
120
121
122
123
            else if(it->name() == "softmax")
            {
                check_shape(s, apply_softmax(it));
            }
Paul's avatar
Paul committed
124
125
126
        }
    }

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

Paul's avatar
Paul committed
141
    instruction_ref apply_convolution(instruction_ref ins)
Paul's avatar
Paul committed
142
    {
wsttiger's avatar
wsttiger committed
143
        auto&& op = any_cast<op::convolution>(ins->get_operator());
Paul's avatar
Paul committed
144

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

148
        auto workspace = insert_allocation(ins, ws, "workspace");
Paul's avatar
Paul committed
149
        auto output    = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
150

Paul's avatar
Paul committed
151
        return prog->replace_instruction(
Paul's avatar
Paul committed
152
            ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
Paul's avatar
Paul committed
153
154
    }

Paul's avatar
Paul committed
155
    instruction_ref apply_pooling(instruction_ref ins)
Paul's avatar
Paul committed
156
    {
wsttiger's avatar
wsttiger committed
157
        auto&& op   = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
158
        auto pd     = make_pooling(op);
Paul's avatar
Paul committed
159
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
160

Paul's avatar
Paul committed
161
        return prog->replace_instruction(
Paul's avatar
Paul committed
162
            ins, miopen_pooling{op, std::move(pd)}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
163
164
    }

Khalique's avatar
Khalique committed
165
    instruction_ref apply_relu(instruction_ref ins)
Paul's avatar
Paul committed
166
    {
Khalique's avatar
Khalique committed
167
        auto ad = make_relu();
Khalique's avatar
Khalique committed
168
169
170
171

        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
172
    }
173

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

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
    }

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

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

204
205
206
207
208
209
210
211
212
213
214
215
    instruction_ref apply_sinh(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(ins, hip_sinh{}, ins->inputs().at(0), output);
    }

    instruction_ref apply_cosh(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(ins, hip_cosh{}, ins->inputs().at(0), output);
    }

216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
    instruction_ref apply_asin(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(ins, hip_asin{}, ins->inputs().at(0), output);
    }

    instruction_ref apply_acos(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(ins, hip_acos{}, ins->inputs().at(0), output);
    }

    instruction_ref apply_atan(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(ins, hip_atan{}, ins->inputs().at(0), output);
    }

Khalique's avatar
Khalique committed
234
235
236
237
238
239
240
    instruction_ref apply_mul(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->get_shape());
        return prog->replace_instruction(
            ins, hip_mul{}, ins->inputs().at(0), ins->inputs().at(1), output);
    }

Paul's avatar
Paul committed
241
    instruction_ref apply_gemm(instruction_ref ins)
Paul's avatar
Paul committed
242
    {
Shucai Xiao's avatar
Shucai Xiao committed
243
        auto&& op   = any_cast<op::dot>(ins->get_operator());
Paul's avatar
Paul committed
244
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
245
        return prog->replace_instruction(
Paul's avatar
Paul committed
246
            ins, miopen_gemm{op}, ins->inputs().at(0), ins->inputs().at(1), output);
Paul's avatar
Paul committed
247
    }
248

Paul's avatar
Paul committed
249
    instruction_ref apply_contiguous(instruction_ref ins)
250
    {
wsttiger's avatar
wsttiger committed
251
        auto&& op   = any_cast<op::contiguous>(ins->get_operator());
Paul's avatar
Paul committed
252
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
253
        return prog->replace_instruction(ins, miopen_contiguous{op}, ins->inputs().at(0), output);
254
    }
255

256
257
    instruction_ref apply_concat(instruction_ref ins)
    {
wsttiger's avatar
wsttiger committed
258
259
        auto&& op                         = any_cast<op::concat>(ins->get_operator());
        auto output                       = insert_allocation(ins, ins->get_shape());
260
261
262
263
264
        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
265
    instruction_ref apply_batch_norm_inference(instruction_ref ins)
266
    {
wsttiger's avatar
wsttiger committed
267
        auto&& op       = any_cast<op::batch_norm_inference>(ins->get_operator());
Paul's avatar
Paul committed
268
        auto output     = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
269
        shape old_shape = ins->inputs().at(1)->get_shape();
wsttiger's avatar
wsttiger committed
270
        std::vector<int64_t> new_shape{1, static_cast<int64_t>(old_shape.elements()), 1, 1};
wsttiger's avatar
wsttiger committed
271
        auto reshape_op = op::reshape{new_shape};
Paul's avatar
Paul committed
272
        std::vector<instruction_ref> reshapes;
Paul's avatar
Paul committed
273
274
        std::transform(ins->inputs().begin() + 1,
                       ins->inputs().end(),
Paul's avatar
Paul committed
275
276
                       std::back_inserter(reshapes),
                       [&](auto i) { return prog->insert_instruction(ins, reshape_op, i); });
Paul's avatar
Paul committed
277
        return prog->replace_instruction(ins,
Paul's avatar
Paul committed
278
                                         miopen_batch_norm_inference{op},
Paul's avatar
Paul committed
279
                                         ins->inputs().at(0),
Paul's avatar
Paul committed
280
281
282
283
284
                                         reshapes[0],
                                         reshapes[1],
                                         reshapes[2],
                                         reshapes[3],
                                         output);
285
    }
Paul's avatar
Paul committed
286
287
};

Paul's avatar
Paul committed
288
void lowering::apply(program& p) const { miopen_apply{&p, ctx}.apply(); }
Paul's avatar
Paul committed
289
} // namespace gpu
290
} // namespace MIGRAPH_INLINE_NS
Paul's avatar
Paul committed
291
} // namespace migraphx