miopen_target.cpp 9.99 KB
Newer Older
Paul's avatar
Paul committed
1
2
#include <rtg/miopen/miopen_target.hpp>
#include <rtg/manage_ptr.hpp>
Paul's avatar
Paul committed
3
4
#include <rtg/instruction.hpp>
#include <rtg/operators.hpp>
Paul's avatar
Paul committed
5
#include <rtg/shape_for_each.hpp>
Paul's avatar
Paul committed
6
7
#include <rtg/miopen/miopen.hpp>
#include <rtg/miopen/hip.hpp>
Paul's avatar
Paul committed
8
#include <rtg/dfor.hpp>
Paul's avatar
Paul committed
9

Paul's avatar
Paul committed
10
11
namespace rtg {
namespace miopen {
Paul's avatar
Paul committed
12

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

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

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

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

Paul's avatar
Paul committed
70
71
72
73
74
75
76
77
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
78
        check_shapes{inputs, *this}.has(2);
Paul's avatar
Paul committed
79
80
        return op.compute_shape({inputs.at(1)});
    }
Paul's avatar
Paul committed
81
    argument compute(context& gctx, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
82
    {
Paul's avatar
Paul committed
83
        auto& ctx   = any_cast<miopen_context>(gctx);
Paul's avatar
Paul committed
84
        auto x_desc = make_tensor(args[0].get_shape());
Paul's avatar
Paul committed
85
86
87
88
        auto y_desc = make_tensor(output_shape);

        float alpha = 1, beta = 0;

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

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

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

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

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

Paul's avatar
Paul committed
152
153
154
155
156
157
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
158
159
        check_shapes{inputs, *this}.has(3);
        return op.compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Paul committed
160
    }
Paul's avatar
Paul committed
161
    argument compute(context&, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
162
163
164
    {
        argument result{output_shape};

Paul's avatar
Paul committed
165
        visit_all(result, from_gpu(args[0]), from_gpu(args[1]))(
Paul's avatar
Paul committed
166
            [&](auto output, auto input1, auto input2) {
Paul's avatar
Paul committed
167
168
169
170
                dfor(input1.get_shape().lens()[0],
                     input2.get_shape().lens()[1],
                     input2.get_shape().lens()[0])(
                    [&](auto i, auto j, auto k) { output(i, j) += input1(i, k) * input2(k, j); });
Paul's avatar
Paul committed
171
172
173
174
175
            });
        return to_gpu(result);
    }
};

Paul's avatar
Paul committed
176
177
178
179
struct miopen_relu
{
    shared<activation_descriptor> ad;
    std::string name() const { return "miopen::relu"; }
Paul's avatar
Paul committed
180
    shape compute_shape(std::vector<shape> inputs) const
Paul's avatar
Paul committed
181
    {
Paul's avatar
Paul committed
182
        check_shapes{inputs, *this}.has(2);
Paul's avatar
Paul committed
183
        return inputs.at(1);
Paul's avatar
Paul committed
184
185
    }

Paul's avatar
Paul committed
186
    argument compute(context& gctx, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
187
    {
Paul's avatar
Paul committed
188
        auto& ctx   = any_cast<miopen_context>(gctx);
Paul's avatar
Paul committed
189
        float alpha = 1, beta = 0;
Paul's avatar
Paul committed
190
        auto x_desc = make_tensor(args[0].get_shape());
Paul's avatar
Paul committed
191
        auto y_desc = make_tensor(output_shape);
Paul's avatar
Paul committed
192
        miopenActivationForward(ctx.handle.get(),
Paul's avatar
Paul committed
193
194
195
                                ad.get(),
                                &alpha,
                                x_desc.get(),
Paul's avatar
Paul committed
196
                                args[0].implicit(),
Paul's avatar
Paul committed
197
198
                                &beta,
                                y_desc.get(),
Paul's avatar
Paul committed
199
                                args[1].implicit());
Paul's avatar
Paul committed
200

Paul's avatar
Paul committed
201
        return args[1];
Paul's avatar
Paul committed
202
203
204
    }
};

Paul's avatar
Paul committed
205
206
struct miopen_apply
{
Paul's avatar
Paul committed
207
    program* prog = nullptr;
Paul's avatar
Paul committed
208
209
210

    void apply()
    {
Paul's avatar
Paul committed
211
        prog->insert_instruction(prog->begin(), check_context<miopen_context>{});
Paul's avatar
Paul committed
212
213
214
215
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
            if(it->op.name() == "convolution")
            {
Paul's avatar
Paul committed
216
                apply_convolution(it);
Paul's avatar
Paul committed
217
218
219
            }
            else if(it->op.name() == "activation")
            {
Paul's avatar
Paul committed
220
221
                apply_activation(it);
            }
Paul's avatar
Paul committed
222
223
224
225
            else if(it->op.name() == "pooling")
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
226
227
228
229
            else if(it->op.name() == "add")
            {
                apply_add(it);
            }
Paul's avatar
Paul committed
230
231
232
233
            else if(it->op.name() == "gemm")
            {
                apply_gemm(it);
            }
Paul's avatar
Paul committed
234
235
236
        }
    }

Paul's avatar
Paul committed
237
238
    instruction_ref insert_allocation(instruction_ref ins, const shape& s)
    {
Paul's avatar
Paul committed
239
        if(ins == --prog->end())
Paul's avatar
Paul committed
240
241
242
243
244
        {
            return prog->add_parameter("output", s);
        }
        else
        {
Paul's avatar
Paul committed
245
            auto is     = prog->add_outline(s);
Paul's avatar
Paul committed
246
247
248
249
250
            auto result = prog->insert_instruction(ins, hip_allocate{}, is);
            return result;
        }
    }

Paul's avatar
Paul committed
251
252
    void apply_convolution(instruction_ref ins)
    {
Paul's avatar
Paul committed
253
254
        auto&& op   = any_cast<convolution>(ins->op);
        auto cd     = make_conv(op);
Paul's avatar
Paul committed
255
256
        auto output = insert_allocation(ins, ins->result);

Paul's avatar
Paul committed
257
258
259
260
261
        prog->replace_instruction(ins,
                                  miopen_convolution{op, std::move(cd)},
                                  ins->arguments.at(0),
                                  ins->arguments.at(1),
                                  output);
Paul's avatar
Paul committed
262
263
    }

Paul's avatar
Paul committed
264
265
266
267
268
269
    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
270
        prog->replace_instruction(
Paul's avatar
Paul committed
271
            ins, miopen_pooling{op, std::move(pd)}, ins->arguments.at(0), output);
Paul's avatar
Paul committed
272
273
    }

Paul's avatar
Paul committed
274
    void apply_activation(instruction_ref ins)
Paul's avatar
Paul committed
275
276
    {
        auto&& op = any_cast<activation>(ins->op);
Paul's avatar
Paul committed
277
278
        auto ad   = make_relu();
        if(op.mode == "relu")
Paul's avatar
Paul committed
279
280
        {
            auto output = insert_allocation(ins, ins->result);
Paul's avatar
Paul committed
281
            prog->replace_instruction(
Paul's avatar
Paul committed
282
                ins, miopen_relu{std::move(ad)}, ins->arguments.at(0), output);
Paul's avatar
Paul committed
283
284
        }
    }
Paul's avatar
Paul committed
285
286
287
288
289

    void apply_add(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->result);
        prog->replace_instruction(
Paul's avatar
Paul committed
290
            ins, miopen_add{}, ins->arguments.at(0), ins->arguments.at(1), output);
Paul's avatar
Paul committed
291
    }
Paul's avatar
Paul committed
292
293
294

    void apply_gemm(instruction_ref ins)
    {
Paul's avatar
Paul committed
295
        auto&& op   = any_cast<gemm>(ins->op);
Paul's avatar
Paul committed
296
297
        auto output = insert_allocation(ins, ins->result);
        prog->replace_instruction(
Paul's avatar
Paul committed
298
            ins, miopen_gemm{op}, ins->arguments.at(0), ins->arguments.at(1), output);
Paul's avatar
Paul committed
299
    }
Paul's avatar
Paul committed
300
301
};

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

Paul's avatar
Paul committed
304
void miopen_target::apply(program& p) const { miopen_apply{&p}.apply(); }
Paul's avatar
Paul committed
305

Paul's avatar
Paul committed
306
307
308
309
310
context miopen_target::get_context() const
{
    return miopen_context{share(make_obj<miopen_handle>(&miopenCreate))};
}

Paul's avatar
Paul committed
311
312
313
} // namespace miopen

} // namespace rtg