miopen_target.cpp 6.47 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
6
#include <rtg/miopen/miopen.hpp>
#include <rtg/miopen/hip.hpp>
Paul's avatar
Paul committed
7

Paul's avatar
Paul committed
8
9
namespace rtg {
namespace miopen {
Paul's avatar
Paul committed
10

Paul's avatar
Paul committed
11
12
13
struct miopen_convolution
{
    convolution op;
Paul's avatar
Paul committed
14
    shared<convolution_descriptor> cd;
Paul's avatar
Paul committed
15
16

    std::string name() const { return "miopen::convolution"; }
Paul's avatar
Paul committed
17
18
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
19
        check_shapes{inputs}.has(4);
Paul's avatar
Paul committed
20
        return op.compute_shape({inputs.at(1), inputs.at(2)});
Paul's avatar
Paul committed
21
22
23
24
25
26
27
    }
    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
28
        float alpha = 1, beta = 0;
Paul's avatar
Paul committed
29
30
        int algo_count;
        miopenConvAlgoPerf_t perf;
31
        miopenFindConvolutionForwardAlgorithm(args[0].implicit(),
Paul's avatar
Paul committed
32
                                              x_desc.get(),
33
                                              args[1].implicit(),
Paul's avatar
Paul committed
34
                                              w_desc.get(),
35
                                              args[2].implicit(),
Paul's avatar
Paul committed
36
                                              cd.get(),
Paul's avatar
Paul committed
37
                                              y_desc.get(),
38
                                              args[3].implicit(),
Paul's avatar
Paul committed
39
40
41
                                              1,
                                              &algo_count,
                                              &perf,
Paul's avatar
Paul committed
42
43
                                              nullptr,
                                              0,
Paul's avatar
Paul committed
44
                                              false);
45
        miopenConvolutionForward(args[0].implicit(),
Paul's avatar
Paul committed
46
                                 &alpha,
Paul's avatar
Paul committed
47
                                 x_desc.get(),
48
                                 args[1].implicit(),
Paul's avatar
Paul committed
49
                                 w_desc.get(),
50
                                 args[2].implicit(),
Paul's avatar
Paul committed
51
52
53
                                 cd.get(),
                                 perf.fwd_algo,
                                 &beta,
Paul's avatar
Paul committed
54
                                 y_desc.get(),
55
                                 args[3].implicit(),
Paul's avatar
Paul committed
56
57
58
59
60
61
                                 nullptr,
                                 0);
        return args[3];
    }
};

Paul's avatar
Paul committed
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
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
81
82
83
84
85
86
87
88
89
90
                             pd.get(),
                             &alpha,
                             x_desc.get(),
                             args[1].implicit(),
                             &beta,
                             y_desc.get(),
                             args[2].implicit(),
                             false,
                             nullptr,
                             0);
Paul's avatar
Paul committed
91
92
93
94
95

        return args[2];
    }
};

Paul's avatar
Paul committed
96
97
98
99
struct miopen_relu
{
    shared<activation_descriptor> ad;
    std::string name() const { return "miopen::relu"; }
Paul's avatar
Paul committed
100
    shape compute_shape(std::vector<shape> inputs) const
Paul's avatar
Paul committed
101
    {
Paul's avatar
Paul committed
102
103
        check_shapes{inputs}.has(3);
        return inputs.at(1);
Paul's avatar
Paul committed
104
105
106
107
108
109
110
    }

    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);
111
        miopenActivationForward(args[0].implicit(),
Paul's avatar
Paul committed
112
113
114
                                ad.get(),
                                &alpha,
                                x_desc.get(),
115
                                args[1].implicit(),
Paul's avatar
Paul committed
116
117
                                &beta,
                                y_desc.get(),
118
                                args[2].implicit());
Paul's avatar
Paul committed
119
120

        return args[2];
Paul's avatar
Paul committed
121
122
123
    }
};

Paul's avatar
Paul committed
124
125
struct miopen_apply
{
Paul's avatar
Paul committed
126
127
    program* prog = nullptr;
    instruction_ref handle{};
Paul's avatar
Paul committed
128
129
130

    void apply()
    {
Paul's avatar
Paul committed
131
        handle = prog->add_parameter("handle", shape{shape::any_type});
Paul's avatar
Paul committed
132
133
134
135
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
            if(it->op.name() == "convolution")
            {
Paul's avatar
Paul committed
136
                apply_convolution(it);
Paul's avatar
Paul committed
137
138
139
            }
            else if(it->op.name() == "activation")
            {
Paul's avatar
Paul committed
140
141
                apply_activation(it);
            }
Paul's avatar
Paul committed
142
143
144
145
            else if(it->op.name() == "pooling")
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
146
147
148
        }
    }

Paul's avatar
Paul committed
149
150
    instruction_ref insert_allocation(instruction_ref ins, const shape& s)
    {
Paul's avatar
Paul committed
151
        if(ins == --prog->end())
Paul's avatar
Paul committed
152
153
154
155
156
        {
            return prog->add_parameter("output", s);
        }
        else
        {
Paul's avatar
Paul committed
157
            auto is     = prog->add_outline(s);
Paul's avatar
Paul committed
158
159
160
161
162
            auto result = prog->insert_instruction(ins, hip_allocate{}, is);
            return result;
        }
    }

Paul's avatar
Paul committed
163
164
    void apply_convolution(instruction_ref ins)
    {
Paul's avatar
Paul committed
165
166
        auto&& op   = any_cast<convolution>(ins->op);
        auto cd     = make_conv(op);
Paul's avatar
Paul committed
167
168
        auto output = insert_allocation(ins, ins->result);

Paul's avatar
Paul committed
169
170
171
172
173
174
        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
175
176
    }

Paul's avatar
Paul committed
177
178
179
180
181
182
    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
183
184
        prog->replace_instruction(
            ins, miopen_pooling{op, std::move(pd)}, handle, ins->arguments.at(0), output);
Paul's avatar
Paul committed
185
186
    }

Paul's avatar
Paul committed
187
    void apply_activation(instruction_ref ins)
Paul's avatar
Paul committed
188
189
    {
        auto&& op = any_cast<activation>(ins->op);
Paul's avatar
Paul committed
190
191
        auto ad   = make_relu();
        if(op.mode == "relu")
Paul's avatar
Paul committed
192
193
        {
            auto output = insert_allocation(ins, ins->result);
Paul's avatar
Paul committed
194
195
            prog->replace_instruction(
                ins, miopen_relu{std::move(ad)}, handle, ins->arguments.at(0), output);
Paul's avatar
Paul committed
196
197
        }
    }
Paul's avatar
Paul committed
198
199
};

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

Paul's avatar
Paul committed
202
void miopen_target::apply(program& p) const { miopen_apply{&p}.apply(); }
Paul's avatar
Paul committed
203
204
205
206

} // namespace miopen

} // namespace rtg