miopen_target.cpp 10.7 KB
Newer Older
wsttiger's avatar
wsttiger committed
1
#include <rocblas.h>
Paul's avatar
Paul committed
2
3
4
5
6
7
8
9
10
11
#include <migraph/miopen/miopen_target.hpp>
#include <migraph/manage_ptr.hpp>
#include <migraph/instruction.hpp>
#include <migraph/operators.hpp>
#include <migraph/shape_for_each.hpp>
#include <migraph/miopen/miopen.hpp>
#include <migraph/miopen/hip.hpp>
#include <migraph/dfor.hpp>

namespace migraph {
Paul's avatar
Paul committed
12
namespace miopen {
Paul's avatar
Paul committed
13

Paul's avatar
Paul committed
14
15
16
17
18
struct miopen_context
{
    shared<miopen_handle> handle;
};

Paul's avatar
Paul committed
19
20
21
struct miopen_convolution
{
    convolution op;
Paul's avatar
Paul committed
22
    shared<convolution_descriptor> cd;
Paul's avatar
Paul committed
23
24

    std::string name() const { return "miopen::convolution"; }
Paul's avatar
Paul committed
25
26
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
27
28
        check_shapes{inputs, *this}.has(3);
        return op.compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Paul committed
29
    }
Paul's avatar
Paul committed
30
    argument compute(context& gctx, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
31
    {
Paul's avatar
Paul committed
32
        auto& ctx   = any_cast<miopen_context>(gctx);
Paul's avatar
Paul committed
33
34
        auto x_desc = make_tensor(args[0].get_shape());
        auto w_desc = make_tensor(args[1].get_shape());
Paul's avatar
Paul committed
35
36
        auto y_desc = make_tensor(output_shape);

Paul's avatar
Paul committed
37
        float alpha = 1, beta = 0;
Paul's avatar
Paul committed
38
39
        int algo_count;
        miopenConvAlgoPerf_t perf;
Paul's avatar
Paul committed
40
        miopenFindConvolutionForwardAlgorithm(ctx.handle.get(),
Paul's avatar
Paul committed
41
                                              x_desc.get(),
Paul's avatar
Paul committed
42
                                              args[0].implicit(),
Paul's avatar
Paul committed
43
                                              w_desc.get(),
Paul's avatar
Paul committed
44
                                              args[1].implicit(),
Paul's avatar
Paul committed
45
                                              cd.get(),
Paul's avatar
Paul committed
46
                                              y_desc.get(),
Paul's avatar
Paul committed
47
                                              args[2].implicit(),
Paul's avatar
Paul committed
48
49
50
                                              1,
                                              &algo_count,
                                              &perf,
Paul's avatar
Paul committed
51
52
                                              nullptr,
                                              0,
Paul's avatar
Paul committed
53
                                              false);
Paul's avatar
Paul committed
54
        miopenConvolutionForward(ctx.handle.get(),
Paul's avatar
Paul committed
55
                                 &alpha,
Paul's avatar
Paul committed
56
                                 x_desc.get(),
Paul's avatar
Paul committed
57
                                 args[0].implicit(),
Paul's avatar
Paul committed
58
                                 w_desc.get(),
Paul's avatar
Paul committed
59
                                 args[1].implicit(),
Paul's avatar
Paul committed
60
61
62
                                 cd.get(),
                                 perf.fwd_algo,
                                 &beta,
Paul's avatar
Paul committed
63
                                 y_desc.get(),
Paul's avatar
Paul committed
64
                                 args[2].implicit(),
Paul's avatar
Paul committed
65
66
                                 nullptr,
                                 0);
Paul's avatar
Paul committed
67
        return args[2];
Paul's avatar
Paul committed
68
69
70
    }
};

Paul's avatar
Paul committed
71
72
73
74
75
76
77
78
struct miopen_pooling
{
    pooling op;
    shared<pooling_descriptor> pd;

    std::string name() const { return "miopen::pooling"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
79
        check_shapes{inputs, *this}.has(2);
Paul's avatar
Paul committed
80
81
        return op.compute_shape({inputs.at(1)});
    }
Paul's avatar
Paul committed
82
    argument compute(context& gctx, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
83
    {
Paul's avatar
Paul committed
84
        auto& ctx   = any_cast<miopen_context>(gctx);
Paul's avatar
Paul committed
85
        auto x_desc = make_tensor(args[0].get_shape());
Paul's avatar
Paul committed
86
87
88
89
        auto y_desc = make_tensor(output_shape);

        float alpha = 1, beta = 0;

Paul's avatar
Paul committed
90
        miopenPoolingForward(ctx.handle.get(),
Paul's avatar
Paul committed
91
92
93
                             pd.get(),
                             &alpha,
                             x_desc.get(),
Paul's avatar
Paul committed
94
                             args[0].implicit(),
Paul's avatar
Paul committed
95
96
                             &beta,
                             y_desc.get(),
Paul's avatar
Paul committed
97
                             args[1].implicit(),
Paul's avatar
Paul committed
98
99
100
                             false,
                             nullptr,
                             0);
Paul's avatar
Paul committed
101

Paul's avatar
Paul committed
102
        return args[1];
Paul's avatar
Paul committed
103
104
105
    }
};

Paul's avatar
Paul committed
106
107
108
109
110
struct miopen_add
{
    std::string name() const { return "miopen::add"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
111
112
        check_shapes{inputs, *this}.has(3);
        return inputs.at(0);
Paul's avatar
Paul committed
113
114
    }

Paul's avatar
Paul committed
115
    argument compute(context& gctx, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
116
    {
Paul's avatar
Paul committed
117
        if(args[1].get_shape().broadcasted())
Paul's avatar
Paul committed
118
        {
Paul's avatar
Paul committed
119
120
            argument result{output_shape};

Paul's avatar
Paul committed
121
            visit_all(result, from_gpu(args[0]), from_gpu(args[1]))(
Paul's avatar
Paul committed
122
123
                [&](auto output, auto input1, auto input2) {
                    shape_for_each(output.get_shape(), [&](const auto& idx) {
Paul's avatar
Paul committed
124
125
126
                        output(idx.begin(), idx.end()) =
                            input1(idx.begin(), idx.end()) + input2(idx.begin(), idx.end());
                    });
Paul's avatar
Paul committed
127
                });
Paul's avatar
Paul committed
128
            return to_gpu(result);
Paul's avatar
Paul committed
129
130
131
        }
        else
        {
Paul's avatar
Paul committed
132
            auto& ctx   = any_cast<miopen_context>(gctx);
Paul's avatar
Paul committed
133
            float alpha = 1, beta = 0;
Paul's avatar
Paul committed
134
135
            auto a_desc = make_tensor(args[0].get_shape());
            auto b_desc = make_tensor(args[1].get_shape());
Paul's avatar
Paul committed
136
            auto c_desc = make_tensor(output_shape);
Paul's avatar
Paul committed
137
            miopenOpTensor(ctx.handle.get(),
Paul's avatar
Paul committed
138
139
140
                           miopenTensorOpAdd,
                           &alpha,
                           a_desc.get(),
Paul's avatar
Paul committed
141
                           args[0].implicit(),
Paul's avatar
Paul committed
142
143
                           &alpha,
                           b_desc.get(),
Paul's avatar
Paul committed
144
                           args[1].implicit(),
Paul's avatar
Paul committed
145
146
                           &beta,
                           c_desc.get(),
Paul's avatar
Paul committed
147
148
                           args[2].implicit());
            return args[2];
Paul's avatar
Paul committed
149
150
151
152
        }
    }
};

Paul's avatar
Paul committed
153
154
155
156
157
158
struct miopen_gemm
{
    gemm op;
    std::string name() const { return "miopen::convolution"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
159
160
        check_shapes{inputs, *this}.has(3);
        return op.compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Paul committed
161
    }
Paul's avatar
Paul committed
162
    argument compute(context&, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
163
    {
wsttiger's avatar
wsttiger committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
        rocblas_handle rochandle;
        rocblas_create_handle(&rochandle);
        float alpha     = 1.0f;
        float beta      = 0.0f;
        rocblas_int lda = args[0].get_shape().lens()[1];
        rocblas_int ldb = args[1].get_shape().lens()[1];
        rocblas_int ldc = args[2].get_shape().lens()[1];
        rocblas_int m   = output_shape.lens()[0];
        rocblas_int n   = output_shape.lens()[1];
        rocblas_int k   = args[0].get_shape().lens()[1];
        rocblas_sgemm(rochandle,
                      rocblas_operation_none,
                      rocblas_operation_none,
                      n,
                      m,
                      k,
                      &alpha,
                      args[1].implicit(),
                      ldb,
                      args[0].implicit(),
                      lda,
                      &beta,
                      args[2].implicit(),
                      ldc);
188
        return args[2];
Paul's avatar
Paul committed
189
190
191
    }
};

Paul's avatar
Paul committed
192
193
194
195
struct miopen_relu
{
    shared<activation_descriptor> ad;
    std::string name() const { return "miopen::relu"; }
Paul's avatar
Paul committed
196
    shape compute_shape(std::vector<shape> inputs) const
Paul's avatar
Paul committed
197
    {
Paul's avatar
Paul committed
198
        check_shapes{inputs, *this}.has(2);
Paul's avatar
Paul committed
199
        return inputs.at(1);
Paul's avatar
Paul committed
200
201
    }

Paul's avatar
Paul committed
202
    argument compute(context& gctx, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
203
    {
Paul's avatar
Paul committed
204
        auto& ctx   = any_cast<miopen_context>(gctx);
Paul's avatar
Paul committed
205
        float alpha = 1, beta = 0;
Paul's avatar
Paul committed
206
        auto x_desc = make_tensor(args[0].get_shape());
Paul's avatar
Paul committed
207
        auto y_desc = make_tensor(output_shape);
Paul's avatar
Paul committed
208
        miopenActivationForward(ctx.handle.get(),
Paul's avatar
Paul committed
209
210
211
                                ad.get(),
                                &alpha,
                                x_desc.get(),
Paul's avatar
Paul committed
212
                                args[0].implicit(),
Paul's avatar
Paul committed
213
214
                                &beta,
                                y_desc.get(),
Paul's avatar
Paul committed
215
                                args[1].implicit());
Paul's avatar
Paul committed
216

Paul's avatar
Paul committed
217
        return args[1];
Paul's avatar
Paul committed
218
219
220
    }
};

Paul's avatar
Paul committed
221
222
struct miopen_apply
{
Paul's avatar
Paul committed
223
    program* prog = nullptr;
Paul's avatar
Paul committed
224
225
226

    void apply()
    {
Paul's avatar
Paul committed
227
        prog->insert_instruction(prog->begin(), check_context<miopen_context>{});
Paul's avatar
Paul committed
228
229
230
231
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
            if(it->op.name() == "convolution")
            {
Paul's avatar
Paul committed
232
                apply_convolution(it);
Paul's avatar
Paul committed
233
234
235
            }
            else if(it->op.name() == "activation")
            {
Paul's avatar
Paul committed
236
237
                apply_activation(it);
            }
Paul's avatar
Paul committed
238
239
240
241
            else if(it->op.name() == "pooling")
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
242
243
244
245
            else if(it->op.name() == "add")
            {
                apply_add(it);
            }
Paul's avatar
Paul committed
246
247
248
249
            else if(it->op.name() == "gemm")
            {
                apply_gemm(it);
            }
Paul's avatar
Paul committed
250
251
252
        }
    }

Paul's avatar
Paul committed
253
254
    instruction_ref insert_allocation(instruction_ref ins, const shape& s)
    {
Paul's avatar
Paul committed
255
        if(ins == --prog->end())
Paul's avatar
Paul committed
256
257
258
259
260
        {
            return prog->add_parameter("output", s);
        }
        else
        {
Paul's avatar
Paul committed
261
            auto is     = prog->add_outline(s);
Paul's avatar
Paul committed
262
263
264
265
266
            auto result = prog->insert_instruction(ins, hip_allocate{}, is);
            return result;
        }
    }

Paul's avatar
Paul committed
267
268
    void apply_convolution(instruction_ref ins)
    {
Paul's avatar
Paul committed
269
270
        auto&& op   = any_cast<convolution>(ins->op);
        auto cd     = make_conv(op);
Paul's avatar
Paul committed
271
272
        auto output = insert_allocation(ins, ins->result);

Paul's avatar
Paul committed
273
274
275
276
277
        prog->replace_instruction(ins,
                                  miopen_convolution{op, std::move(cd)},
                                  ins->arguments.at(0),
                                  ins->arguments.at(1),
                                  output);
Paul's avatar
Paul committed
278
279
    }

Paul's avatar
Paul committed
280
281
282
283
284
285
    void apply_pooling(instruction_ref ins)
    {
        auto&& op   = any_cast<pooling>(ins->op);
        auto pd     = make_pooling(op);
        auto output = insert_allocation(ins, ins->result);

Paul's avatar
Paul committed
286
        prog->replace_instruction(
Paul's avatar
Paul committed
287
            ins, miopen_pooling{op, std::move(pd)}, ins->arguments.at(0), output);
Paul's avatar
Paul committed
288
289
    }

Paul's avatar
Paul committed
290
    void apply_activation(instruction_ref ins)
Paul's avatar
Paul committed
291
292
    {
        auto&& op = any_cast<activation>(ins->op);
Paul's avatar
Paul committed
293
294
        auto ad   = make_relu();
        if(op.mode == "relu")
Paul's avatar
Paul committed
295
296
        {
            auto output = insert_allocation(ins, ins->result);
Paul's avatar
Paul committed
297
            prog->replace_instruction(
Paul's avatar
Paul committed
298
                ins, miopen_relu{std::move(ad)}, ins->arguments.at(0), output);
Paul's avatar
Paul committed
299
300
        }
    }
Paul's avatar
Paul committed
301
302
303
304
305

    void apply_add(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->result);
        prog->replace_instruction(
Paul's avatar
Paul committed
306
            ins, miopen_add{}, ins->arguments.at(0), ins->arguments.at(1), output);
Paul's avatar
Paul committed
307
    }
Paul's avatar
Paul committed
308
309
310

    void apply_gemm(instruction_ref ins)
    {
Paul's avatar
Paul committed
311
        auto&& op   = any_cast<gemm>(ins->op);
Paul's avatar
Paul committed
312
313
        auto output = insert_allocation(ins, ins->result);
        prog->replace_instruction(
Paul's avatar
Paul committed
314
            ins, miopen_gemm{op}, ins->arguments.at(0), ins->arguments.at(1), output);
Paul's avatar
Paul committed
315
    }
Paul's avatar
Paul committed
316
317
};

Paul's avatar
Paul committed
318
319
320
321
322
323
324
325
struct miopen_pass
{
    std::string name() const { return "miopen::pass"; }

    void apply(program& p) const { miopen_apply{&p}.apply(); }
};

std::vector<pass> miopen_target::get_passes(context&) const { return {miopen_pass{}}; }
Paul's avatar
Paul committed
326

Paul's avatar
Paul committed
327
std::string miopen_target::name() const { return "miopen"; }
Paul's avatar
Paul committed
328

Paul's avatar
Paul committed
329
330
331
332
333
context miopen_target::get_context() const
{
    return miopen_context{share(make_obj<miopen_handle>(&miopenCreate))};
}

Paul's avatar
Paul committed
334
335
} // namespace miopen

Paul's avatar
Paul committed
336
} // namespace migraph