onnx.cpp 46.3 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
5
6
7
8
#include <google/protobuf/text_format.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <onnx.pb.h>
#include <iostream>
#include <fstream>
#include <unordered_map>
#include <functional>
#include <array>
Paul's avatar
Paul committed
9
#include <utility>
10
#include <vector>
Paul's avatar
Paul committed
11

Paul's avatar
Paul committed
12
13
14
15
16
17
#include <migraphx/fallthrough.hpp>
#include <migraphx/program.hpp>
#include <migraphx/operators.hpp>
#include <migraphx/ranges.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/config.hpp>
18
#include <migraphx/onnx.hpp>
Paul's avatar
Paul committed
19
20

namespace migraphx {
Paul's avatar
Paul committed
21
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
22
23
24
25
26

struct onnx_parser
{
    using attribute_map = std::unordered_map<std::string, onnx::AttributeProto>;
    using node_map      = std::unordered_map<std::string, onnx::NodeProto>;
Paul's avatar
Paul committed
27
28
    using op_func =
        std::function<std::vector<instruction_ref>(attribute_map, std::vector<instruction_ref>)>;
Paul's avatar
Paul committed
29
30
    node_map nodes;
    std::unordered_map<std::string, instruction_ref> instructions;
Scott Thornton's avatar
Scott Thornton committed
31
    program prog    = program();
32
    bool is_pytorch = false;
Paul's avatar
Paul committed
33
34

    std::unordered_map<std::string, op_func> ops;
35
    std::unordered_map<std::string, operation> map_actv_funcs;
Paul's avatar
Paul committed
36
37
38

    onnx_parser()
    {
Shucai Xiao's avatar
Shucai Xiao committed
39
        add_generic_op("MatMul", op::dot{});
Khalique's avatar
Khalique committed
40
        add_generic_op("Relu", op::relu{});
Khalique's avatar
Khalique committed
41
42
        add_generic_op("Sigmoid", op::sigmoid{});
        add_generic_op("Abs", op::abs{});
Shucai Xiao's avatar
Shucai Xiao committed
43
44
        add_generic_op("Exp", op::exp{});
        add_generic_op("Log", op::log{});
Khalique's avatar
Khalique committed
45
46
        // disable dropout for inference
        add_generic_op("Dropout", op::identity{});
Khalique's avatar
Khalique committed
47
        add_generic_op("Identity", op::identity{});
Shucai Xiao's avatar
Shucai Xiao committed
48
49
50
        add_generic_op("Sin", op::sin{});
        add_generic_op("Cos", op::cos{});
        add_generic_op("Tan", op::tan{});
51
52
        add_generic_op("Sinh", op::sinh{});
        add_generic_op("Cosh", op::cosh{});
53
        add_generic_op("Tanh", op::tanh{});
54
55
56
        add_generic_op("Asin", op::asin{});
        add_generic_op("Acos", op::acos{});
        add_generic_op("Atan", op::atan{});
Paul's avatar
Paul committed
57

Khalique's avatar
Khalique committed
58
59
60
61
62
        add_binary_op("Add", op::add{});
        add_binary_op("Div", op::div{});
        add_binary_op("Mul", op::mul{});
        add_binary_op("Sub", op::sub{});

Khalique's avatar
Khalique committed
63
64
65
        add_variadic_op("Sum", op::add{});
        add_variadic_op("Max", op::max{});
        add_variadic_op("Min", op::min{});
Paul's avatar
Paul committed
66

Khalique's avatar
Khalique committed
67
        add_mem_op("LRN", &onnx_parser::parse_lrn);
Khalique's avatar
Khalique committed
68
        add_mem_op("ImageScaler", &onnx_parser::parse_imagescaler);
69
        add_mem_op("LeakyRelu", &onnx_parser::parse_leaky_relu);
Khalique's avatar
Khalique committed
70
        add_mem_op("Elu", &onnx_parser::parse_elu);
Paul's avatar
Paul committed
71
72
        add_mem_op("Constant", &onnx_parser::parse_constant);
        add_mem_op("Conv", &onnx_parser::parse_conv);
Paul's avatar
Paul committed
73
74
        add_mem_op("MaxPool", &onnx_parser::parse_pooling);
        add_mem_op("AveragePool", &onnx_parser::parse_pooling);
75
76
        add_mem_op("GlobalMaxPool", &onnx_parser::parse_pooling);
        add_mem_op("GlobalAveragePool", &onnx_parser::parse_pooling);
Paul's avatar
Paul committed
77
        add_mem_op("Reshape", &onnx_parser::parse_reshape);
Paul's avatar
Paul committed
78
79
        add_mem_op("Flatten", &onnx_parser::parse_flatten);
        add_mem_op("Gemm", &onnx_parser::parse_gemm);
80
        add_mem_op("BatchNormalization", &onnx_parser::parse_batchnorm);
Paul's avatar
Paul committed
81
        add_mem_op("Softmax", &onnx_parser::parse_softmax);
82
83
84
        add_mem_op("Squeeze", &onnx_parser::parse_squeeze);
        add_mem_op("Unsqueeze", &onnx_parser::parse_unsqueeze);
        add_mem_op("Slice", &onnx_parser::parse_slice);
Scott Thornton's avatar
Scott Thornton committed
85
        add_mem_op("Concat", &onnx_parser::parse_concat);
86
87
88
        add_mem_op("Gather", &onnx_parser::parse_gather);
        add_mem_op("Shape", &onnx_parser::parse_shape);
        add_mem_op("ConstantFill", &onnx_parser::parse_constant_fill);
Khalique's avatar
Khalique committed
89
        add_mem_op("Transpose", &onnx_parser::parse_transpose);
Shucai Xiao's avatar
Shucai Xiao committed
90
        add_mem_op("RNN", &onnx_parser::parse_rnn);
91
        add_mem_op("GRU", &onnx_parser::parse_gru);
Khalique's avatar
Khalique committed
92
        add_mem_op("Pad", &onnx_parser::parse_pad);
93
94
95
96
97
98
99

        // init the activation function map
        init_actv_func();
    }

    void init_actv_func()
    {
100
101
102
103
104
        map_actv_funcs.insert(std::make_pair("tanh", op::tanh{}));
        map_actv_funcs.insert(std::make_pair("relu", op::relu{}));
        map_actv_funcs.insert(std::make_pair("sigmoid", op::sigmoid{}));
        map_actv_funcs.insert(std::make_pair("leakyrelu", op::leaky_relu{}));
        map_actv_funcs.insert(std::make_pair("elu", op::elu{}));
Paul's avatar
Paul committed
105
106
107
108
    }

    template <class F>
    void add_op(std::string name, F f)
Paul's avatar
Paul committed
109
110
111
112
113
114
115
116
117
    {
        ops.emplace(name, [=](auto&&... xs) {
            return std::vector<instruction_ref>{f(std::forward<decltype(xs)>(xs)...)};
        });
    }

    // Multi output op
    template <class F>
    void add_multi_op(std::string name, F f)
Paul's avatar
Paul committed
118
119
120
121
122
123
124
    {
        ops.emplace(name, f);
    }

    template <class F>
    void add_mem_op(std::string name, F f)
    {
Paul's avatar
Paul committed
125
        add_op(name, [=](auto&&... xs) {
Paul's avatar
Paul committed
126
127
128
            return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
        });
    }
Khalique's avatar
Khalique committed
129

130
    template <class T>
Khalique's avatar
Khalique committed
131
    void add_binary_op(std::string name, T x)
132
    {
Paul's avatar
Paul committed
133
        add_op(name, [this, x](attribute_map attributes, std::vector<instruction_ref> args) {
Scott Thornton's avatar
Scott Thornton committed
134
            if(args.size() != 2)
Paul's avatar
Paul committed
135
                MIGRAPHX_THROW("binary operators should have 2 operands");
136
            if(contains(attributes, "broadcast") and contains(attributes, "axis"))
137
138
139
140
            {
                uint64_t broadcasted = parse_value(attributes.at("broadcast")).at<uint64_t>();
                if(broadcasted != 0)
                {
141
                    uint64_t axis = parse_value(attributes.at("axis")).at<uint64_t>();
142
143
144
145
                    auto l =
                        prog.add_instruction(op::broadcast{axis, args[0]->get_shape()}, args[1]);
                    return prog.add_instruction(x, args[0], l);
                }
146
                return prog.add_instruction(x, args);
147
            }
Paul's avatar
Paul committed
148
            else
149
            {
Khalique's avatar
Khalique committed
150
                return add_broadcastable_binary_op(args[0], args[1], x);
151
152
153
154
            }
        });
    }

Khalique's avatar
Khalique committed
155
156
157
158
159
    template <class T>
    instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
    {
        if(arg0->get_shape() != arg1->get_shape())
        {
Khalique's avatar
Khalique committed
160
161
162
163
164
165
166
167
168
169
170
171
172
            // Example:
            // s0 = (3,2,4,5) and s1 = (2,1,1)
            //
            // In this case we need to broadcast (:,1,1) portion of
            // s1 plus broadcast the 1st dimension of s1
            // giving output_lens = (3,2,4,5)
            //
            // Another example:
            // s0 = (3,2,1,5) and s1 = (2,7,5)
            // In this case we need to broadcast the (:,:,1:,:) axis
            // of s0 plus the 1st dimension of s1 giving
            // output_lens = (3,2,7,5)
            //
Khalique's avatar
Khalique committed
173
174
175
176
177
178
179
180
            // Get lengths for both arguments
            const std::vector<std::size_t>* s0 = &arg0->get_shape().lens();
            const std::vector<std::size_t>* s1 = &arg1->get_shape().lens();

            // Make sure s0 is the smaller size
            if(s0->size() > s1->size())
                std::swap(s0, s1);

Khalique's avatar
Khalique committed
181
            std::vector<std::size_t> output_lens(*s1);
Khalique's avatar
Khalique committed
182
183
            auto offset = s1->size() - s0->size();
            std::transform(s0->begin(),
Khalique's avatar
Khalique committed
184
185
186
187
                           s0->end(),
                           s1->begin() + offset,
                           output_lens.begin() + offset,
                           [](auto a, auto b) { return std::max(a, b); });
Khalique's avatar
Khalique committed
188
189
190
191
192
193
194
195
196

            auto l0 = prog.add_instruction(op::multibroadcast{output_lens}, arg0);
            auto l1 = prog.add_instruction(op::multibroadcast{output_lens}, arg1);
            return prog.add_instruction(x, l0, l1);
        }
        else
        {
            return prog.add_instruction(x, {arg0, arg1});
        }
197
198
    }

Paul's avatar
Paul committed
199
    template <class T>
Paul's avatar
Paul committed
200
201
    void add_generic_op(std::string name, T x)
    {
Paul's avatar
Paul committed
202
        add_op(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
Paul's avatar
Paul committed
203
204
205
206
            return prog.add_instruction(x, args);
        });
    }

Khalique's avatar
Khalique committed
207
    template <class T>
Khalique's avatar
Khalique committed
208
    void add_variadic_op(std::string name, T x)
Khalique's avatar
Khalique committed
209
    {
Paul's avatar
Paul committed
210
        add_op(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
Khalique's avatar
Khalique committed
211
            return std::accumulate(std::next(args.begin()),
Khalique's avatar
Khalique committed
212
213
214
215
216
                                   args.end(),
                                   args.front(),
                                   [this, x](instruction_ref a, instruction_ref b) {
                                       return add_broadcastable_binary_op(a, b, x);
                                   });
Khalique's avatar
Khalique committed
217
        });
Khalique's avatar
Khalique committed
218
219
    }

Paul's avatar
Paul committed
220
    instruction_ref
Paul's avatar
Paul committed
221
    parse_softmax(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
222
223
    {
        auto dims = args.front()->get_shape().lens();
Scott Thornton's avatar
Scott Thornton committed
224
225
        auto r =
            prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1]), 1, 1}}, args.front());
226
227
        auto s = prog.add_instruction(op::softmax{}, r);
        return prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1])}}, s);
Paul's avatar
Paul committed
228
229
    }

Paul's avatar
Paul committed
230
    instruction_ref
Paul's avatar
Paul committed
231
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
232
    {
233
        op::convolution op;
234
        auto l0 = args[0];
Paul's avatar
Paul committed
235
236
        if(contains(attributes, "pads"))
        {
Scott Thornton's avatar
Scott Thornton committed
237
            if(contains(attributes, "auto_pad"))
238
            {
Paul's avatar
Paul committed
239
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
240
            }
241
242
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
243
            if(padding.size() != 4)
244
            {
Paul's avatar
Paul committed
245
                MIGRAPHX_THROW("padding should have 4 values");
246
            }
Scott Thornton's avatar
Scott Thornton committed
247
            if(padding[0] != padding[2] || padding[1] != padding[3])
248
            {
249
250
                // insert zeros for pad op (args[0] has 4 dims)
                padding = {0, 0, padding[0], padding[1], 0, 0, padding[2], padding[3]};
Khalique's avatar
Khalique committed
251
                l0      = prog.add_instruction(op::pad{padding}, l0);
252
            }
253
254
255
256
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
257
            }
Paul's avatar
Paul committed
258
        }
Paul's avatar
Paul committed
259
260
261
262
263
264
265
266
        if(contains(attributes, "strides"))
        {
            copy(attributes["strides"].ints(), op.stride.begin());
        }
        if(contains(attributes, "dilations"))
        {
            copy(attributes["dilations"].ints(), op.dilation.begin());
        }
Scott Thornton's avatar
Scott Thornton committed
267
        if(contains(attributes, "auto_pad"))
268
269
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
270
            if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
271
            {
Paul's avatar
Paul committed
272
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
273
274
            }

wsttiger's avatar
fixes  
wsttiger committed
275
            if(s.find("SAME") != std::string::npos)
276
            {
277
                op.padding_mode = op::padding_mode_t::same;
278
279
            }
        }
Khalique's avatar
Khalique committed
280
281
282
283
        if(contains(attributes, "group"))
        {
            op.group = parse_value(attributes.at("group")).at<int>();
        }
Paul's avatar
Paul committed
284
285
286
287
        if(args.size() == 3)
        {
            uint64_t axis = 1;
            auto l1       = prog.add_instruction(op, args[0], args[1]);
Scott Thornton's avatar
Scott Thornton committed
288
            auto l2       = prog.add_instruction(op::broadcast{axis, l1->get_shape()}, args[2]);
289
            return prog.add_instruction(op::add{}, l1, l2);
Paul's avatar
Paul committed
290
        }
291
        return prog.add_instruction(op, l0, args[1]);
Paul's avatar
Paul committed
292
    }
Paul's avatar
Paul committed
293

Paul's avatar
Paul committed
294
295
296
    instruction_ref parse_pooling(const std::string& name,
                                  attribute_map attributes,
                                  std::vector<instruction_ref> args)
Paul's avatar
Paul committed
297
    {
Khalique's avatar
Khalique committed
298
        op::pooling op{ends_with(name, "MaxPool") ? "max" : "average"};
299
        auto l0 = args[0];
Khalique's avatar
Khalique committed
300
        if(starts_with(name, "Global"))
301
        {
Khalique's avatar
Khalique committed
302
303
            auto lens  = args.front()->get_shape().lens();
            op.lengths = {lens[2], lens[3]};
304
        }
Paul's avatar
Paul committed
305
306
        if(contains(attributes, "pads"))
        {
307
308
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
309
            if(padding.size() != 4)
310
            {
Paul's avatar
Paul committed
311
                MIGRAPHX_THROW("padding should have 4 values");
312
            }
Scott Thornton's avatar
Scott Thornton committed
313
            if(padding[0] != padding[2] || padding[1] != padding[3])
314
            {
315
316
                // insert zeros for pad op (args[0] has 4 dims)
                padding = {0, 0, padding[0], padding[1], 0, 0, padding[2], padding[3]};
Khalique's avatar
Khalique committed
317
                l0      = prog.add_instruction(op::pad{padding}, l0);
318
319
320
321
322
            }
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
323
            }
Paul's avatar
Paul committed
324
325
326
327
328
329
330
331
332
        }
        if(contains(attributes, "strides"))
        {
            copy(attributes["strides"].ints(), op.stride.begin());
        }
        if(contains(attributes, "kernel_shape"))
        {
            copy(attributes["kernel_shape"].ints(), op.lengths.begin());
        }
Scott Thornton's avatar
Scott Thornton committed
333
        if(contains(attributes, "auto_pad"))
334
335
        {
            auto s = attributes["auto_pad"].s();
336
            if(s.find("SAME_UPPER") == std::string::npos)
337
            {
338
                MIGRAPHX_THROW("auto_pad only supports SAME_UPPER for pooling");
339
            }
340
            op.padding_mode = op::padding_mode_t::same;
341
342
        }

343
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
344
345
    }

Paul's avatar
Paul committed
346
    instruction_ref
Paul's avatar
Paul committed
347
    parse_reshape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
348
    {
349
        op::reshape op;
Paul's avatar
Paul committed
350
351
352
353
354
355
356
        if(args.size() == 1)
        {
            literal s = parse_value(attributes.at("shape"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
        }
        if(args.size() == 2)
        {
Paul's avatar
Paul committed
357
            auto s = args[1]->eval();
Paul's avatar
Paul committed
358
            if(s.empty())
Paul's avatar
Paul committed
359
                MIGRAPHX_THROW("Dynamic shape is not supported.");
Paul's avatar
Paul committed
360
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
361
        }
Paul's avatar
Paul committed
362
363
364
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
365
    instruction_ref
Paul's avatar
Paul committed
366
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
367
    {
368
        uint64_t axis = 1;
Paul's avatar
Paul committed
369
370
371
372
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
373
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
374
375
    }

376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
    instruction_ref
    parse_squeeze(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::squeeze op;
        literal s = parse_value(attributes.at("axes"));
        s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        return prog.add_instruction(op, args[0]);
    }

    instruction_ref
    parse_unsqueeze(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::unsqueeze op;
        literal s = parse_value(attributes.at("axes"));
        s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        return prog.add_instruction(op, args[0]);
    }

Scott Thornton's avatar
Scott Thornton committed
394
395
396
397
398
399
400
    instruction_ref
    parse_concat(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        std::size_t axis = parse_value(attributes.at("axis")).at<int>();
        op::concat op{axis};
        return prog.add_instruction(op, std::move(args));
    }
401

402
403
404
    instruction_ref
    parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
405
        int axis = 0;
406
407
408
409
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
410
        op::gather op{axis};
411
412
413
        return prog.add_instruction(op, std::move(args));
    }

414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
    instruction_ref
    parse_slice(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::slice op;
        if(contains(attributes, "axes"))
        {
            literal s = parse_value(attributes.at("axes"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        }
        {
            literal s = parse_value(attributes.at("ends"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.ends)); });
        }
        {
            literal s = parse_value(attributes.at("starts"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.starts)); });
        }
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
434
435
436
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
437
438
439
440
    {
        literal v = parse_value(attributes.at("value"));
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
441

Paul's avatar
Paul committed
442
    instruction_ref
Paul's avatar
Paul committed
443
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
444
445
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
446
        float beta  = 1.0f;
Paul's avatar
Paul committed
447
448
449
450
451
452
453
454
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
455
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
456
457
458
459
460
461
462
463
464
465
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
        std::vector<int64_t> perm = {1, 0};
466
467
        auto l1 = (transa) ? prog.add_instruction(op::transpose{perm}, args[0]) : args[0];
        auto l2 = (transb) ? prog.add_instruction(op::transpose{perm}, args[1]) : args[1];
Paul's avatar
Paul committed
468
469
        if(args.size() == 3)
        {
Khalique's avatar
Khalique committed
470
            if(beta != 0.f)
471
            {
Khalique's avatar
Khalique committed
472
                auto l3 = prog.add_instruction(op::dot{alpha}, l1, l2);
Khalique's avatar
Khalique committed
473
                auto l4 = args[2];
Khalique's avatar
Khalique committed
474
                if(l4->get_shape().scalar()) // ignore args[2] (no C value added to alpha*A*B)
Khalique's avatar
Khalique committed
475
                    return l3;
Khalique's avatar
Khalique committed
476
                if(beta != 1.f)
Khalique's avatar
Khalique committed
477
478
                {
                    auto beta_val = prog.add_literal(beta);
Khalique's avatar
Khalique committed
479
480
                    auto l5 = prog.add_instruction(op::scalar{args[2]->get_shape()}, beta_val);
                    l4      = prog.add_instruction(op::mul{}, args[2], l5);
Khalique's avatar
Khalique committed
481
482
                }
                return add_broadcastable_binary_op(l3, l4, op::add{});
483
            }
Paul's avatar
Paul committed
484
        }
Shucai Xiao's avatar
Shucai Xiao committed
485
        return prog.add_instruction(op::dot{alpha, beta}, l1, l2);
Paul's avatar
Paul committed
486
487
    }

488
    instruction_ref
Paul's avatar
Paul committed
489
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
490
    {
Scott Thornton's avatar
Scott Thornton committed
491
492
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
493
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Scott Thornton's avatar
Scott Thornton committed
494
        bool is_test                                      = false;
495
496
497
498
499
500
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
501
            momentum = parse_value(attributes.at("momentum")).at<float>();
502
503
504
        }
        if(contains(attributes, "is_test"))
        {
wsttiger's avatar
wsttiger committed
505
            is_test = parse_value(attributes.at("is_test")).at<uint64_t>() > 0;
506
507
508
        }
        if(contains(attributes, "spatial"))
        {
509
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
510
511
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
512
        }
Paul's avatar
Paul committed
513
        (void)is_test;
Paul's avatar
Paul committed
514
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
515
        return prog.add_instruction(op, std::move(args));
516
517
    }

518
519
520
521
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
522
        float alpha = 0.01; // default alpha val for leaky relu
523
524
525
526
527
528
529
530
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        op::leaky_relu op{alpha};
        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
531
532
    instruction_ref
    parse_elu(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
533
534
535
536
537
538
539
540
541
542
    {
        float alpha = 1.0; // default alpha val for elu
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        op::elu op{alpha};
        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
543
544
    instruction_ref
    parse_lrn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
545
546
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
547
548
549
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
Khalique's avatar
Khalique committed
550
551
552
553
554
555
556
557
558
559
560
561
        if(contains(attributes, "alpha"))
            alpha = parse_value(attributes.at("alpha")).at<float>();
        if(contains(attributes, "beta"))
            beta = parse_value(attributes.at("beta")).at<float>();
        if(contains(attributes, "bias"))
            bias = parse_value(attributes.at("bias")).at<float>();
        if(contains(attributes, "size"))
            size = parse_value(attributes.at("size")).at<int>();
        op::lrn op{alpha, beta, bias, size};
        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
    instruction_ref parse_imagescaler(const std::string&,
                                      attribute_map attributes,
                                      std::vector<instruction_ref> args)
    {
        float scale = 1.0;
        std::vector<float> bias{};
        if(contains(attributes, "scale"))
        {
            scale = parse_value(attributes.at("scale")).at<float>();
        }

        if(contains(attributes, "bias"))
        {
            auto&& bias_floats = attributes["bias"].floats();
            bias               = std::vector<float>(bias_floats.begin(), bias_floats.end());
        }
        auto input_shape = args.front()->get_shape();
Khalique's avatar
Khalique committed
579

Khalique's avatar
Khalique committed
580
581
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
582
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
583

Paul's avatar
Paul committed
584
585
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_shape}, scale_val);
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Paul's avatar
Paul committed
586
        auto bias_bcast = prog.add_instruction(migraphx::op::broadcast{1, input_shape}, bias_vals);
Paul's avatar
Paul committed
587
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
588
    }
Khalique's avatar
Khalique committed
589

Khalique's avatar
Khalique committed
590
591
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
592
593
594
595
596
597
598
    {
        std::vector<int64_t> perm{};
        if(contains(attributes, "perm"))
        {
            auto&& perm_vals = attributes["perm"].ints();
            perm             = std::vector<int64_t>(perm_vals.begin(), perm_vals.end());
        }
Paul's avatar
Paul committed
599
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
600
601
    }

Khalique's avatar
Khalique committed
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
    instruction_ref
    parse_pad(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        std::vector<int64_t> pads{};
        float value = 0.0f;
        if(contains(attributes, "pads"))
        {
            auto&& pad_vals = attributes["pads"].ints();
            pads            = std::vector<int64_t>(pad_vals.begin(), pad_vals.end());
        }
        if(contains(attributes, "value"))
        {
            value = parse_value(attributes.at("value")).at<float>();
        }
        if(contains(attributes, "mode"))
        {
            auto mode = attributes.at("mode").s();
            if(mode != "constant")
                MIGRAPHX_THROW("migraphx currently only supports constant padding");
        }
        return prog.add_instruction(migraphx::op::pad{pads, value}, args.front());
    }
624
625
626
    // Use a literal instruction to replace the shape since, output of
    // shape operator are literals in migraphx
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
627
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
628
629
    {
        if(args.size() != 1)
630
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
        std::vector<std::size_t> arg_shape = args[0]->get_shape().lens();
        std::vector<int64_t> vec_shape(arg_shape.size());
        migraphx::shape s(migraphx::shape::int64_type, {arg_shape.size()});
        std::transform(arg_shape.begin(), arg_shape.end(), vec_shape.begin(), [](auto i) {
            return int64_t(i);
        });
        return prog.add_literal(migraphx::literal{s, vec_shape});
    }

    // Use a literal instruction to replace the constantFill operator. In RNN, input shape
    // and value are fixed, so no need to do the actual computation for the constantFill
    // operator
    instruction_ref parse_constant_fill(const std::string&,
                                        attribute_map attributes,
                                        std::vector<instruction_ref> args)
    {
        int input_as_shape = 0;
        int dtype          = 1;
        float value        = 0.0f;

        if(contains(attributes, "dtype"))
        {
            dtype = parse_value(attributes.at("dtype")).at<int>();
        }
        migraphx::shape::type_t type = get_type(dtype);

        if(contains(attributes, "input_as_shape"))
        {
            input_as_shape = parse_value(attributes.at("input_as_shape")).at<int>();
        }

        if(contains(attributes, "value"))
        {
            value = parse_value(attributes.at("value")).at<float>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
667
668
        if(contains(attributes, "extra_shape"))
        {
669
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
670
671
        }

672
673
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
674
            if(args.size() != 1)
675
            {
676
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
677
678
            }

Shucai Xiao's avatar
Shucai Xiao committed
679
680
            if(contains(attributes, "shape"))
            {
681
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
682
                               "at the same time");
683
684
            }

685
686
687
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
688
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
689
            }
690

691
692
693
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
694
695
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
696
697
698
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
699
700
            if(!contains(attributes, "shape"))
            {
701
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
702
703
704
            }

            literal ls = parse_value(attributes.at("shape"));
705
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
706
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
707
            migraphx::shape s{type, dims};
708
709
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
710
711
712
        }
        else
        {
713
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
714
715
716
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
717
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
718
719
720
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
721
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
722
723
724

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
725
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
726
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
727
728
729
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
730
731
732
733
734
735
736
737
738
        }

        // Handling of direction to be added later
        std::string direction{"forward"};
        if(contains(attributes, "direction"))
        {
            direction = attributes.at("direction").s();
        }

739
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
740
741
        if(direction == "bidirectional")
        {
742
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
743
744
745
        }
        else if(direction == "reverse")
        {
746
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
747
748
        }

749
750
751
752
753
        std::vector<std::string> vec_names{"tanh"};
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
754
            for_each(names.begin(), names.end(), [&](auto& fn) { vec_names.push_back(fn); });
755
756
        }

Shucai Xiao's avatar
Shucai Xiao committed
757
        for_each(vec_names.begin(), vec_names.end(), [&](auto& fn) {
758
759
            if(map_actv_funcs.count(fn) == 0)
            {
Shucai Xiao's avatar
Shucai Xiao committed
760
                MIGRAPHX_THROW("RNN: activation function " + std::string(fn) + " not supported");
761
762
763
            }
        });

Shucai Xiao's avatar
Shucai Xiao committed
764
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
765
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
766
        // if only one actv function is provided, we use it in both
767
        // forward and reverse direction
768
        if(dirct == op::rnn_direction::bidirectional)
769
        {
Shucai Xiao's avatar
Shucai Xiao committed
770
            if(vec_names.size() == 1)
771
772
773
774
775
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
776
777
778
        std::vector<operation> vec_actv_funcs(vec_names.size());
        std::transform(vec_names.begin(), vec_names.end(), vec_actv_funcs.begin(), [&](auto& fn) {
            return map_actv_funcs[fn];
779
        });
Shucai Xiao's avatar
Shucai Xiao committed
780

Shucai Xiao's avatar
Shucai Xiao committed
781
782
783
784
785
786
787
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

788
789
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
790
        if(args.size() < 6)
791
792
793
794
795
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
796
797
        // first output for the concatenation of hidden states
        auto hidden_states = prog.add_instruction(op::rnn{hidden_size, vec_actv_funcs, dirct, clip},
Shucai Xiao's avatar
Shucai Xiao committed
798
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
799

800
        // second output for the last hidden state
Shucai Xiao's avatar
Shucai Xiao committed
801
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
802

Shucai Xiao's avatar
Shucai Xiao committed
803
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
804
805
    }

806
    std::vector<instruction_ref>
807
808
809
810
811
812
813
    parse_gru(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
        std::size_t hidden_size     = args[2]->get_shape().lens()[2];

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
814
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
815
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
816
817
818
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
819
820
821
822
823
824
825
826
827
        }

        // Handling of direction to be added later
        std::string direction{"forward"};
        if(contains(attributes, "direction"))
        {
            direction = attributes.at("direction").s();
        }

828
        op::rnn_direction dirct = op::rnn_direction::forward;
829
830
        if(direction == "bidirectional")
        {
831
            dirct = op::rnn_direction::bidirectional;
832
833
834
        }
        else if(direction == "reverse")
        {
835
            dirct = op::rnn_direction::reverse;
836
837
        }

838
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
839
840
        if(contains(attributes, "activations"))
        {
841
            auto names = attributes.at("activations").strings();
842
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
843
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
844
845
            std::transform(
                names.begin(), names.end(), vec_names.begin(), [](auto& str) { return str; });
846
847
        }

848
        // need 4 activation functions
849
        if(dirct == op::rnn_direction::bidirectional)
850
        {
Shucai Xiao's avatar
Shucai Xiao committed
851
            // 4 activation functions are used in the bidirectional
852
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
853
854
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
855
856
            // assume forward and reverse use the same pair of actv
            // functions. For the case of 3 actv functions provided,
Shucai Xiao's avatar
Shucai Xiao committed
857
858
859
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
860
            if(vec_names.size() == 1)
861
            {
862
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
863
            }
864
            else if(vec_names.size() == 2)
865
            {
866
867
868
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
869
            }
870
            else if(vec_names.size() == 3)
871
            {
872
                vec_names.push_back(vec_names.at(2));
873
874
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
875
        else
876
        {
877
            if(vec_names.size() == 1)
878
            {
879
                vec_names.push_back(vec_names.at(0));
880
881
882
            }
        }

883
        for_each(vec_names.begin(), vec_names.end(), [&](auto& name) {
Shucai Xiao's avatar
Shucai Xiao committed
884
885
            if(map_actv_funcs.count(name) == 0)
            {
Shucai Xiao's avatar
Shucai Xiao committed
886
                MIGRAPHX_THROW("GRU: activation function " + std::string(name) + " not supported");
Shucai Xiao's avatar
Shucai Xiao committed
887
888
            }
        });
889

Shucai Xiao's avatar
Shucai Xiao committed
890
891
892
        std::vector<operation> vec_actv_funcs(vec_names.size());
        std::transform(vec_names.begin(), vec_names.end(), vec_actv_funcs.begin(), [&](auto& name) {
            return map_actv_funcs[name];
Shucai Xiao's avatar
Shucai Xiao committed
893
        });
894
895
896
897
898
899
900
901

        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

        int linear_before_reset = 0;
Shucai Xiao's avatar
Shucai Xiao committed
902
        if(contains(attributes, "linear_before_reset"))
903
904
905
906
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
907
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
908
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
909
910
911
912
913
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

914
915
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
916
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
917
            std::move(args));
918
919

        // second output for last gru output
920
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
921

Shucai Xiao's avatar
Shucai Xiao committed
922
        return {hidden_states, last_output};
923
924
    }

Paul's avatar
Paul committed
925
926
927
928
929
930
931
932
933
934
935
936
    void parse_from(std::istream& is)
    {
        onnx::ModelProto model;
        if(model.ParseFromIstream(&is))
        {
            if(model.has_graph())
            {
                this->parse_graph(model.graph());
            }
        }
        else
        {
Paul's avatar
Paul committed
937
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
938
939
940
941
942
943
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
944
945
946
947
948
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
949
950
951
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
952
953
954
955
956
957
958
959
960
961
962
963
            // Does the input have an initializer?
            if(contains(initializer_data, name))
            {
                auto t             = initializer_data[name];
                instructions[name] = prog.add_literal(parse_tensor(t));
            }
            else
            {
                // TODO: Get shape of input parameter
                shape s            = parse_type(input.type());
                instructions[name] = prog.add_parameter(name, s);
            }
Paul's avatar
Paul committed
964
965
966
        }
        for(auto&& p : nodes)
        {
Paul's avatar
Paul committed
967
            this->parse_node(p.first);
Paul's avatar
Paul committed
968
969
970
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
971
    void parse_undefined(const std::string& name)
972
    {
Shucai Xiao's avatar
Shucai Xiao committed
973
        auto ins           = prog.add_instruction(op::undefined{});
974
975
976
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
977
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
978
    {
Paul's avatar
Paul committed
979
        if(name.empty())
Paul's avatar
Paul committed
980
            MIGRAPHX_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
981
982
983
984
985
986
987
988
        if(instructions.count(name) == 0)
        {
            auto&& node = nodes.at(name);
            std::vector<instruction_ref> args;
            for(auto&& input : node.input())
            {
                if(nodes.count(input) > 0)
                {
Paul's avatar
Paul committed
989
990
                    assert(name != input);
                    this->parse_node(input);
Paul's avatar
Paul committed
991
                }
Shucai Xiao's avatar
Shucai Xiao committed
992
                else if(input.empty())
Paul's avatar
Paul committed
993
                {
994
                    this->parse_undefined(input);
Paul's avatar
Paul committed
995
                }
996
                args.push_back(instructions.at(input));
Paul's avatar
Paul committed
997
            }
Paul's avatar
Paul committed
998
            std::vector<instruction_ref> result;
Paul's avatar
Paul committed
999
1000
            if(ops.count(node.op_type()) == 0)
            {
Paul's avatar
Paul committed
1001
                result.push_back(prog.add_instruction(unknown{node.op_type()}, args));
Paul's avatar
Paul committed
1002
1003
1004
            }
            else
            {
Paul's avatar
Paul committed
1005
                result = ops[node.op_type()](get_attributes(node), args);
Paul's avatar
Paul committed
1006
            }
Paul's avatar
Paul committed
1007
            // Even no output nodes produce output in migraphx
Paul's avatar
Paul committed
1008
            if(node.output().empty() and result.size() == 1)
Paul's avatar
Paul committed
1009
1010
            {
                instructions[name] = result.front();
Paul's avatar
Paul committed
1011
1012
1013
            }
            else
            {
Paul's avatar
Paul committed
1014
1015
1016
1017
1018
1019
                assert(node.output().size() >= result.size());
                std::transform(result.begin(),
                               result.end(),
                               node.output().begin(),
                               std::inserter(instructions, instructions.end()),
                               [](auto&& x, auto&& y) { return std::make_pair(y, x); });
Paul's avatar
Paul committed
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
            }
        }
    }

    static attribute_map get_attributes(const onnx::NodeProto& node)
    {
        std::unordered_map<std::string, onnx::AttributeProto> result;
        for(auto&& attr : node.attribute())
        {
            result[attr.name()] = attr;
        }
        return result;
    }

    static node_map get_nodes(const onnx::GraphProto& graph)
    {
        std::unordered_map<std::string, onnx::NodeProto> result;
Paul's avatar
Paul committed
1037
        std::size_t n = 0;
Paul's avatar
Paul committed
1038
1039
        for(auto&& node : graph.node())
        {
Paul's avatar
Paul committed
1040
            if(node.output().empty())
Paul's avatar
Paul committed
1041
            {
Paul's avatar
Paul committed
1042
                if(node.name().empty())
Paul's avatar
Paul committed
1043
1044
1045
1046
1047
1048
1049
1050
1051
                {
                    result["migraphx_unamed_node_" + std::to_string(n)] = node;
                    n++;
                }
                else
                {
                    result[node.name()] = node;
                }
            }
Paul's avatar
Paul committed
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
            for(auto&& output : node.output())
            {
                result[output] = node;
            }
        }
        return result;
    }

    template <class T>
    static literal from_repeated(shape::type_t t, const T& r)
    {
        std::size_t size = r.size();
        return literal{{t, {size}}, r.begin(), r.end()};
    }

    static literal parse_value(const onnx::AttributeProto& attr)
    {
        switch(attr.type())
        {
        case onnx::AttributeProto::UNDEFINED: return {};
        case onnx::AttributeProto::FLOAT: return literal{attr.f()};
        case onnx::AttributeProto::INT: return literal{attr.i()};
        case onnx::AttributeProto::STRING: return {};
        case onnx::AttributeProto::TENSOR: return parse_tensor(attr.t());
        case onnx::AttributeProto::GRAPH: return {};
Paul's avatar
Paul committed
1077
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
1078
1079
1080
1081
1082
        case onnx::AttributeProto::INTS: return from_repeated(shape::int64_type, attr.ints());
        case onnx::AttributeProto::STRINGS: return {};
        case onnx::AttributeProto::TENSORS: return {};
        case onnx::AttributeProto::GRAPHS: return {};
        }
Paul's avatar
Paul committed
1083
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
1084
1085
1086
1087
1088
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
Khalique's avatar
Khalique committed
1089
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1090
        if(dims.empty())
Khalique's avatar
Khalique committed
1091
1092
1093
        {
            dims = {1};
        }
1094
1095
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
1096
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
            switch(t.data_type())
            {
            case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
            case onnx::TensorProto::FLOAT: return literal{{shape::float_type, dims}, s.data()};
            case onnx::TensorProto::UINT8: throw std::runtime_error("");
            case onnx::TensorProto::INT8: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::UINT16: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT16: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT32: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT64: return literal{{shape::int64_type, dims}, s.data()};
            case onnx::TensorProto::STRING: throw std::runtime_error("");
            case onnx::TensorProto::BOOL: return literal{{shape::int32_type, dims}, s.data()};
Paul's avatar
Paul committed
1109
            case onnx::TensorProto::FLOAT16: return literal{{shape::half_type, dims}, s.data()};
Scott Thornton's avatar
Scott Thornton committed
1110
1111
1112
1113
1114
1115
            case onnx::TensorProto::DOUBLE: return literal{{shape::double_type, dims}, s.data()};
            case onnx::TensorProto::UINT32: throw std::runtime_error("");
            case onnx::TensorProto::UINT64: throw std::runtime_error("");
            case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
            case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
            }
Paul's avatar
Paul committed
1116
            MIGRAPHX_THROW("Invalid tensor type");
1117
        }
Paul's avatar
Paul committed
1118
1119
1120
1121
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Paul's avatar
Paul committed
1122
            return literal{{shape::float_type, dims}, t.float_data().begin(), t.float_data().end()};
Paul's avatar
Paul committed
1123
1124
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Paul's avatar
Paul committed
1125
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1126
        case onnx::TensorProto::UINT16:
Paul's avatar
Paul committed
1127
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1128
        case onnx::TensorProto::INT16:
Paul's avatar
Paul committed
1129
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1130
        case onnx::TensorProto::INT32:
Paul's avatar
Paul committed
1131
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1132
        case onnx::TensorProto::INT64:
Paul's avatar
Paul committed
1133
            return literal{{shape::int64_type, dims}, t.int64_data().begin(), t.int64_data().end()};
Paul's avatar
Paul committed
1134
1135
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Paul's avatar
Paul committed
1136
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1137
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
1138
        {
Khalique's avatar
Khalique committed
1139
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
1140
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
1141
1142
1143
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
1144
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
1145
            return literal{{shape::half_type, dims}, data_half.begin(), data_half.end()};
Khalique's avatar
Khalique committed
1146
        }
Paul's avatar
Paul committed
1147
1148
1149
1150
1151
1152
1153
1154
        case onnx::TensorProto::DOUBLE:
            return literal{
                {shape::double_type, dims}, t.double_data().begin(), t.double_data().end()};
        case onnx::TensorProto::UINT32: throw std::runtime_error("");
        case onnx::TensorProto::UINT64: throw std::runtime_error("");
        case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
        case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
        }
Paul's avatar
Paul committed
1155
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
    }

    static shape parse_type(const onnx::TypeProto& t)
    {
        shape::type_t shape_type{};
        switch(t.tensor_type().elem_type())
        {
        case onnx::TensorProto::UNDEFINED:
            break; // throw std::runtime_error("Unsupported type UNDEFINED");
        case onnx::TensorProto::FLOAT: shape_type = shape::float_type; break;
        case onnx::TensorProto::UINT8:
            break; // throw std::runtime_error("Unsupported type UINT8");
        case onnx::TensorProto::INT8: shape_type = shape::int8_type; break;
        case onnx::TensorProto::UINT16: shape_type = shape::uint16_type; break;
        case onnx::TensorProto::INT16: shape_type = shape::int16_type; break;
        case onnx::TensorProto::INT32: shape_type = shape::int32_type; break;
        case onnx::TensorProto::INT64: shape_type = shape::int64_type; break;
        case onnx::TensorProto::STRING:
            break; // throw std::runtime_error("Unsupported type STRING");
        case onnx::TensorProto::BOOL:
            break; // throw std::runtime_error("Unsupported type BOOL");
Paul's avatar
Paul committed
1177
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1178
1179
1180
1181
1182
1183
1184
1185
1186
        case onnx::TensorProto::DOUBLE: shape_type = shape::double_type; break;
        case onnx::TensorProto::UINT32: shape_type = shape::uint32_type; break;
        case onnx::TensorProto::UINT64: shape_type = shape::uint64_type; break;
        case onnx::TensorProto::COMPLEX64:
            break; // throw std::runtime_error("Unsupported type COMPLEX64");
        case onnx::TensorProto::COMPLEX128:
            break; // throw std::runtime_error("Unsupported type COMPLEX128");
        }
        std::vector<std::size_t> dims;
Paul's avatar
Paul committed
1187
        auto&& tensor_dims = t.tensor_type().shape().dim();
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
        std::transform(tensor_dims.begin(),
                       tensor_dims.end(),
                       std::back_inserter(dims),
                       [](auto&& d) -> std::size_t {
                           if(not d.has_dim_value())
                           {
                               long default_batch_size = 1; // FIXME
                               return default_batch_size;
                           }
                           return d.dim_value();
                       });
Paul's avatar
Paul committed
1199
1200
        return {shape_type, dims};
    }
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222

    shape::type_t get_type(int dtype)
    {
        switch(dtype)
        {
        case 1: return shape::float_type;
        case 2: return shape::uint8_type;
        case 3: return shape::int8_type;
        case 4: return shape::uint16_type;
        case 5: return shape::int16_type;
        case 6: return shape::int32_type;
        case 7: return shape::int64_type;
        case 10: return shape::half_type;
        case 11: return shape::double_type;
        case 12: return shape::uint32_type;
        case 13: return shape::uint64_type;
        default:
        {
            MIGRAPHX_THROW("Prototensor data type " + std::to_string(dtype) + " not supported");
        }
        }
    }
Paul's avatar
Paul committed
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
};

program parse_onnx(const std::string& name)
{
    std::fstream input(name.c_str(), std::ios::in | std::ios::binary);
    onnx_parser parser;
#ifndef NDEBUG
    // Log the program when it can't be parsed
    try
    {
        parser.parse_from(input);
    }
    catch(...)
    {
        std::cerr << parser.prog << std::endl;
        throw;
    }
#else
    parser.parse_from(input);
#endif
    return std::move(parser.prog);
}

Paul's avatar
Paul committed
1246
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
1247
} // namespace migraphx