miopen_target.cpp 7.46 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
7

#include <miopen/miopen.h>

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
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

struct hip_allocate
{
    std::string name() const { return "hip::allocate"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs}.has(1);
        return inputs.front();
    }
    argument compute(shape output_shape, std::vector<argument>) const
    {
        char * data = nullptr;
        // TODO: Check return status
        hipMalloc(&data, output_shape.bytes());
        return {output_shape, data};
    }
};

struct hip_free
{
    std::string name() const { return "hip::free"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs}.has(1);
        return {};
    }
    argument compute(shape output_shape, std::vector<argument> args) const
    {
        // TODO: Check return status
        hipFree(args.front().data());
        return {};
    }
};


Paul's avatar
Paul committed
46
using miopen_handle     = RTG_MANAGE_PTR(miopenHandle_t, miopenDestroy);
Paul's avatar
Paul committed
47
using tensor_descriptor = RTG_MANAGE_PTR(miopenTensorDescriptor_t, miopenDestroyTensorDescriptor);
Paul's avatar
Paul committed
48
49
50
51
using convolution_descriptor = RTG_MANAGE_PTR(miopenConvolutionDescriptor_t,
                                              miopenDestroyConvolutionDescriptor);
using activation_descriptor  = RTG_MANAGE_PTR(miopenActivationDescriptor_t,
                                             miopenDestroyActivationDescriptor);
Paul's avatar
Paul committed
52

Paul's avatar
Paul committed
53
template <class Result, class F, class... Ts>
Paul's avatar
Paul committed
54
55
56
Result make_obj(F f, Ts... xs)
{
    typename Result::pointer x = nullptr;
Paul's avatar
Paul committed
57
    auto status                = f(&x, xs...);
Paul's avatar
Paul committed
58
    Result r{x};
Paul's avatar
Paul committed
59
    if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
60
61
62
63
64
65
66
67
68
69
70
        RTG_THROW("MIOpen call failed");
    return r;
}

tensor_descriptor make_tensor(const rtg::shape& s)
{
    auto t = make_obj<tensor_descriptor>(&miopenCreateTensorDescriptor);
    // Convert to ints
    std::vector<int> lens(s.lens().begin(), s.lens().end());
    std::vector<int> strides(s.strides().begin(), s.strides().end());
    miopenDataType_t d;
Paul's avatar
Paul committed
71
72
73
74
    if(s.type() == shape::float_type)
        d = miopenFloat;
    else
        RTG_THROW("Unsupported type");
Paul's avatar
Paul committed
75
76
77
78
79
80
81
    miopenSetTensorDescriptor(t.get(), d, s.lens().size(), lens.data(), strides.data());
    return t;
}

convolution_descriptor make_conv(const rtg::convolution& op)
{
    auto c = make_obj<convolution_descriptor>(&miopenCreateConvolutionDescriptor);
Paul's avatar
Paul committed
82
83
84
85
86
87
88
89
    miopenInitConvolutionDescriptor(c.get(),
                                    miopenConvolution,
                                    op.padding[0],
                                    op.padding[1],
                                    op.stride[0],
                                    op.stride[1],
                                    op.dilation[0],
                                    op.dilation[1]);
Paul's avatar
Paul committed
90
91
92
    return c;
}

Paul's avatar
Paul committed
93
94
95
96
97
98
99
activation_descriptor make_relu()
{
    auto ad = make_obj<activation_descriptor>(&miopenCreateActivationDescriptor);
    miopenSetActivationDescriptor(ad.get(), miopenActivationRELU, 0, 0, 0);
    return ad;
}

Paul's avatar
Paul committed
100
101
102
struct miopen_convolution
{
    convolution op;
Paul's avatar
Paul committed
103
    shared<convolution_descriptor> cd;
Paul's avatar
Paul committed
104
105

    std::string name() const { return "miopen::convolution"; }
Paul's avatar
Paul committed
106
107
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
108
        check_shapes{inputs}.has(4);
Paul's avatar
Paul committed
109
        return op.compute_shape({inputs.at(1), inputs.at(2)});
Paul's avatar
Paul committed
110
111
112
113
114
115
116
    }
    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
117
        float alpha = 1, beta = 0;
Paul's avatar
Paul committed
118
119
        int algo_count;
        miopenConvAlgoPerf_t perf;
Paul's avatar
Paul committed
120
        miopenFindConvolutionForwardAlgorithm(args[0].get(),
Paul's avatar
Paul committed
121
                                              x_desc.get(),
Paul's avatar
Paul committed
122
123
124
                                              args[1].get(),
                                              w_desc.get(),
                                              args[2].get(),
Paul's avatar
Paul committed
125
                                              cd.get(),
Paul's avatar
Paul committed
126
                                              y_desc.get(),
Paul's avatar
Paul committed
127
                                              args[3].get(),
Paul's avatar
Paul committed
128
129
130
                                              1,
                                              &algo_count,
                                              &perf,
Paul's avatar
Paul committed
131
132
                                              nullptr,
                                              0,
Paul's avatar
Paul committed
133
                                              false);
Paul's avatar
Paul committed
134
        miopenConvolutionForward(args[0].get(),
Paul's avatar
Paul committed
135
                                 &alpha,
Paul's avatar
Paul committed
136
137
138
139
                                 x_desc.get(),
                                 args[1].get(),
                                 w_desc.get(),
                                 args[2].get(),
Paul's avatar
Paul committed
140
141
142
                                 cd.get(),
                                 perf.fwd_algo,
                                 &beta,
Paul's avatar
Paul committed
143
144
                                 y_desc.get(),
                                 args[3].get(),
Paul's avatar
Paul committed
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
                                 nullptr,
                                 0);
        return args[3];
    }
};

struct miopen_relu
{
    shared<activation_descriptor> ad;
    std::string name() const { return "miopen::relu"; }
    shape compute_shape(std::vector<shape> inputs) const 
    {
        check_shapes{inputs}.has(3); 
        return inputs.at(1); 
    }

    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);
        miopenActivationForward(args[0].get(), ad.get(), &alpha, x_desc.get(), args[1].get(), &beta, y_desc.get(), args[2].get());

        return args[2];
Paul's avatar
Paul committed
169
170
171
    }
};

Paul's avatar
Paul committed
172
173
struct miopen_apply
{
Paul's avatar
Paul committed
174
    program* prog;
Paul's avatar
Paul committed
175
    instruction_ref handle;
Paul's avatar
Paul committed
176
177
178

    void apply()
    {
Paul's avatar
Paul committed
179
        handle = prog->add_parameter("handle", shape{shape::any_type});
Paul's avatar
Paul committed
180
181
182
183
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
            if(it->op.name() == "convolution")
            {
Paul's avatar
Paul committed
184
                apply_convolution(it);
Paul's avatar
Paul committed
185
186
187
            }
            else if(it->op.name() == "activation")
            {
Paul's avatar
Paul committed
188
189
190
191
192
                apply_activation(it);
            }
        }
    }

Paul's avatar
Paul committed
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
    instruction_ref insert_allocation(instruction_ref ins, const shape& s)
    {
        if (ins == --prog->end())
        {
            return prog->add_parameter("output", s);
        }
        else
        {
            auto is = prog->add_outline(s);
            auto result = prog->insert_instruction(ins, hip_allocate{}, is);
            prog->insert_instruction(++ins, hip_free{}, result);
            return result;
        }
    }

Paul's avatar
Paul committed
208
209
    void apply_convolution(instruction_ref ins)
    {
Paul's avatar
Paul committed
210
211
212
213
214
        auto&& op = any_cast<convolution>(ins->op);
        auto cd = make_conv(op);
        auto output = insert_allocation(ins, ins->result);

        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
215
216
    }

Paul's avatar
Paul committed
217
218
219
220
221
222
223
224
225
226
    void apply_activation(instruction_ref ins) 
    {
        auto&& op = any_cast<activation>(ins->op);
        auto ad = make_relu();
        if(op.mode == "relu") 
        {
            auto output = insert_allocation(ins, ins->result);
            prog->replace_instruction(ins, miopen_relu{std::move(ad)}, handle, ins->arguments.at(0), output);
        }
    }
Paul's avatar
Paul committed
227
228
};

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

Paul's avatar
Paul committed
231
void miopen_target::apply(program& p) const { miopen_apply{&p}.apply(); }
Paul's avatar
Paul committed
232
233
234
235

} // namespace miopen

} // namespace rtg