miopen_target.cpp 9.57 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
struct miopen_convolution
{
    convolution op;
Paul's avatar
Paul committed
16
    shared<convolution_descriptor> cd;
Paul's avatar
Paul committed
17
18

    std::string name() const { return "miopen::convolution"; }
Paul's avatar
Paul committed
19
20
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
21
        check_shapes{inputs}.has(4);
Paul's avatar
Paul committed
22
        return op.compute_shape({inputs.at(1), inputs.at(2)});
Paul's avatar
Paul committed
23
24
25
26
27
28
29
    }
    argument compute(shape output_shape, std::vector<argument> args) const
    {
        auto x_desc = make_tensor(args[1].get_shape());
        auto w_desc = make_tensor(args[2].get_shape());
        auto y_desc = make_tensor(output_shape);

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

Paul's avatar
Paul committed
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
struct miopen_pooling
{
    pooling op;
    shared<pooling_descriptor> pd;

    std::string name() const { return "miopen::pooling"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs}.has(3);
        return op.compute_shape({inputs.at(1)});
    }
    argument compute(shape output_shape, std::vector<argument> args) const
    {
        auto x_desc = make_tensor(args[1].get_shape());
        auto y_desc = make_tensor(output_shape);

        float alpha = 1, beta = 0;

        miopenPoolingForward(args[0].implicit(),
Paul's avatar
Paul committed
83
84
85
86
87
88
89
90
91
92
                             pd.get(),
                             &alpha,
                             x_desc.get(),
                             args[1].implicit(),
                             &beta,
                             y_desc.get(),
                             args[2].implicit(),
                             false,
                             nullptr,
                             0);
Paul's avatar
Paul committed
93
94
95
96
97

        return args[2];
    }
};

Paul's avatar
Paul committed
98
99
100
101
102
103
104
105
106
107
108
struct miopen_add
{
    std::string name() const { return "miopen::add"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs}.has(4);
        return inputs.at(1);
    }

    argument compute(shape output_shape, std::vector<argument> args) const
    {
Paul's avatar
Paul committed
109
110
        if(args[2].get_shape().broadcasted())
        {
Paul's avatar
Paul committed
111
112
            argument result{output_shape};

Paul's avatar
Paul committed
113
114
115
            visit_all(result, from_gpu(args[1]), from_gpu(args[2]))(
                [&](auto output, auto input1, auto input2) {
                    shape_for_each(output.get_shape(), [&](const auto& idx) {
Paul's avatar
Paul committed
116
117
118
                        output(idx.begin(), idx.end()) =
                            input1(idx.begin(), idx.end()) + input2(idx.begin(), idx.end());
                    });
Paul's avatar
Paul committed
119
                });
Paul's avatar
Paul committed
120
            return to_gpu(result);
Paul's avatar
Paul committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
        }
        else
        {
            float alpha = 1, beta = 0;
            auto a_desc = make_tensor(args[1].get_shape());
            auto b_desc = make_tensor(args[2].get_shape());
            auto c_desc = make_tensor(output_shape);
            miopenOpTensor(args[0].implicit(),
                           miopenTensorOpAdd,
                           &alpha,
                           a_desc.get(),
                           args[1].implicit(),
                           &alpha,
                           b_desc.get(),
                           args[2].implicit(),
                           &beta,
                           c_desc.get(),
                           args[3].implicit());
            return args[3];
Paul's avatar
Paul committed
140
141
142
143
        }
    }
};

Paul's avatar
Paul committed
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
struct miopen_gemm
{
    gemm op;
    std::string name() const { return "miopen::convolution"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs}.has(4);
        return op.compute_shape({inputs.at(1), inputs.at(2)});
    }
    argument compute(shape output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};

        visit_all(result, from_gpu(args[1]), from_gpu(args[2]))(
            [&](auto output, auto input1, auto input2) {
                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);
                });
            });
        return to_gpu(result);
    }
};

Paul's avatar
Paul committed
167
168
169
170
struct miopen_relu
{
    shared<activation_descriptor> ad;
    std::string name() const { return "miopen::relu"; }
Paul's avatar
Paul committed
171
    shape compute_shape(std::vector<shape> inputs) const
Paul's avatar
Paul committed
172
    {
Paul's avatar
Paul committed
173
174
        check_shapes{inputs}.has(3);
        return inputs.at(1);
Paul's avatar
Paul committed
175
176
177
178
179
180
181
    }

    argument compute(shape output_shape, std::vector<argument> args) const
    {
        float alpha = 1, beta = 0;
        auto x_desc = make_tensor(args[1].get_shape());
        auto y_desc = make_tensor(output_shape);
182
        miopenActivationForward(args[0].implicit(),
Paul's avatar
Paul committed
183
184
185
                                ad.get(),
                                &alpha,
                                x_desc.get(),
186
                                args[1].implicit(),
Paul's avatar
Paul committed
187
188
                                &beta,
                                y_desc.get(),
189
                                args[2].implicit());
Paul's avatar
Paul committed
190
191

        return args[2];
Paul's avatar
Paul committed
192
193
194
    }
};

Paul's avatar
Paul committed
195
196
struct miopen_apply
{
Paul's avatar
Paul committed
197
198
    program* prog = nullptr;
    instruction_ref handle{};
Paul's avatar
Paul committed
199
200
201

    void apply()
    {
Paul's avatar
Paul committed
202
        handle = prog->add_parameter("handle", shape{shape::any_type});
Paul's avatar
Paul committed
203
204
205
206
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
            if(it->op.name() == "convolution")
            {
Paul's avatar
Paul committed
207
                apply_convolution(it);
Paul's avatar
Paul committed
208
209
210
            }
            else if(it->op.name() == "activation")
            {
Paul's avatar
Paul committed
211
212
                apply_activation(it);
            }
Paul's avatar
Paul committed
213
214
215
216
            else if(it->op.name() == "pooling")
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
217
218
219
220
            else if(it->op.name() == "add")
            {
                apply_add(it);
            }
Paul's avatar
Paul committed
221
222
223
224
            else if(it->op.name() == "gemm")
            {
                apply_gemm(it);
            }
Paul's avatar
Paul committed
225
226
227
        }
    }

Paul's avatar
Paul committed
228
229
    instruction_ref insert_allocation(instruction_ref ins, const shape& s)
    {
Paul's avatar
Paul committed
230
        if(ins == --prog->end())
Paul's avatar
Paul committed
231
232
233
234
235
        {
            return prog->add_parameter("output", s);
        }
        else
        {
Paul's avatar
Paul committed
236
            auto is     = prog->add_outline(s);
Paul's avatar
Paul committed
237
238
239
240
241
            auto result = prog->insert_instruction(ins, hip_allocate{}, is);
            return result;
        }
    }

Paul's avatar
Paul committed
242
243
    void apply_convolution(instruction_ref ins)
    {
Paul's avatar
Paul committed
244
245
        auto&& op   = any_cast<convolution>(ins->op);
        auto cd     = make_conv(op);
Paul's avatar
Paul committed
246
247
        auto output = insert_allocation(ins, ins->result);

Paul's avatar
Paul committed
248
249
250
251
252
253
        prog->replace_instruction(ins,
                                  miopen_convolution{op, std::move(cd)},
                                  handle,
                                  ins->arguments.at(0),
                                  ins->arguments.at(1),
                                  output);
Paul's avatar
Paul committed
254
255
    }

Paul's avatar
Paul committed
256
257
258
259
260
261
    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
262
263
        prog->replace_instruction(
            ins, miopen_pooling{op, std::move(pd)}, handle, ins->arguments.at(0), output);
Paul's avatar
Paul committed
264
265
    }

Paul's avatar
Paul committed
266
    void apply_activation(instruction_ref ins)
Paul's avatar
Paul committed
267
268
    {
        auto&& op = any_cast<activation>(ins->op);
Paul's avatar
Paul committed
269
270
        auto ad   = make_relu();
        if(op.mode == "relu")
Paul's avatar
Paul committed
271
272
        {
            auto output = insert_allocation(ins, ins->result);
Paul's avatar
Paul committed
273
274
            prog->replace_instruction(
                ins, miopen_relu{std::move(ad)}, handle, ins->arguments.at(0), output);
Paul's avatar
Paul committed
275
276
        }
    }
Paul's avatar
Paul committed
277
278
279
280
281
282
283

    void apply_add(instruction_ref ins)
    {
        auto output = insert_allocation(ins, ins->result);
        prog->replace_instruction(
            ins, miopen_add{}, handle, ins->arguments.at(0), ins->arguments.at(1), output);
    }
Paul's avatar
Paul committed
284
285
286
287
288
289
290
291

    void apply_gemm(instruction_ref ins)
    {
        auto&& op = any_cast<gemm>(ins->op);
        auto output = insert_allocation(ins, ins->result);
        prog->replace_instruction(
            ins, miopen_gemm{op}, handle, ins->arguments.at(0), ins->arguments.at(1), output);
    }
Paul's avatar
Paul committed
292
293
};

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

Paul's avatar
Paul committed
296
void miopen_target::apply(program& p) const { miopen_apply{&p}.apply(); }
Paul's avatar
Paul committed
297
298
299
300

} // namespace miopen

} // namespace rtg