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

Khalique's avatar
Khalique committed
57
58
59
60
61
        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
62
63
64
        add_variadic_op("Sum", op::add{});
        add_variadic_op("Max", op::max{});
        add_variadic_op("Min", op::min{});
Paul's avatar
Paul committed
65

66
67
        add_mem_op("ArgMax", &onnx_parser::parse_argmax);
        add_mem_op("ArgMin", &onnx_parser::parse_argmin);
68
        add_mem_op("Cast", &onnx_parser::parse_cast);
Khalique's avatar
Khalique committed
69
        add_mem_op("Clip", &onnx_parser::parse_clip);
Khalique's avatar
Khalique committed
70
        add_mem_op("LRN", &onnx_parser::parse_lrn);
Khalique's avatar
Khalique committed
71
        add_mem_op("ImageScaler", &onnx_parser::parse_imagescaler);
72
        add_mem_op("LeakyRelu", &onnx_parser::parse_leaky_relu);
Khalique's avatar
Khalique committed
73
        add_mem_op("Elu", &onnx_parser::parse_elu);
74
        add_mem_op("Expand", &onnx_parser::parse_expand);
Paul's avatar
Paul committed
75
76
        add_mem_op("Constant", &onnx_parser::parse_constant);
        add_mem_op("Conv", &onnx_parser::parse_conv);
Paul's avatar
Paul committed
77
78
        add_mem_op("MaxPool", &onnx_parser::parse_pooling);
        add_mem_op("AveragePool", &onnx_parser::parse_pooling);
79
80
        add_mem_op("GlobalMaxPool", &onnx_parser::parse_pooling);
        add_mem_op("GlobalAveragePool", &onnx_parser::parse_pooling);
Paul's avatar
Paul committed
81
        add_mem_op("Reshape", &onnx_parser::parse_reshape);
Paul's avatar
Paul committed
82
83
        add_mem_op("Flatten", &onnx_parser::parse_flatten);
        add_mem_op("Gemm", &onnx_parser::parse_gemm);
84
        add_mem_op("MatMul", &onnx_parser::parse_matmul);
85
        add_mem_op("BatchNormalization", &onnx_parser::parse_batchnorm);
Paul's avatar
Paul committed
86
        add_mem_op("Softmax", &onnx_parser::parse_softmax);
Shucai Xiao's avatar
Shucai Xiao committed
87
        add_mem_op("LogSoftmax", &onnx_parser::parse_logsoftmax);
88
89
90
        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
91
        add_mem_op("Concat", &onnx_parser::parse_concat);
92
93
94
        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);
95
        add_mem_op("ConstantOfShape", &onnx_parser::parse_constant_of_shape);
Khalique's avatar
Khalique committed
96
        add_mem_op("Transpose", &onnx_parser::parse_transpose);
Shucai Xiao's avatar
Shucai Xiao committed
97
        add_mem_op("RNN", &onnx_parser::parse_rnn);
98
        add_mem_op("GRU", &onnx_parser::parse_gru);
Shucai Xiao's avatar
Shucai Xiao committed
99
        add_mem_op("LSTM", &onnx_parser::parse_lstm);
Khalique's avatar
Khalique committed
100
        add_mem_op("Pad", &onnx_parser::parse_pad);
101
        add_mem_op("ReduceSum", &onnx_parser::parse_reduce_sum);
102
103
104
105
106
107
108

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

    void init_actv_func()
    {
109
110
111
112
113
114
        // Support name format of all lower case or the first letter capital
        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
115
116
117
118
    }

    template <class F>
    void add_op(std::string name, F f)
Paul's avatar
Paul committed
119
120
121
122
123
124
125
126
127
    {
        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
128
129
130
131
132
133
134
    {
        ops.emplace(name, f);
    }

    template <class F>
    void add_mem_op(std::string name, F f)
    {
Paul's avatar
Paul committed
135
        add_op(name, [=](auto&&... xs) {
Paul's avatar
Paul committed
136
137
138
            return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
        });
    }
Khalique's avatar
Khalique committed
139

140
    template <class T>
Khalique's avatar
Khalique committed
141
    void add_binary_op(std::string name, T x)
142
    {
Paul's avatar
Paul committed
143
        add_op(name, [this, x](attribute_map attributes, std::vector<instruction_ref> args) {
Scott Thornton's avatar
Scott Thornton committed
144
            if(args.size() != 2)
Paul's avatar
Paul committed
145
                MIGRAPHX_THROW("binary operators should have 2 operands");
146
            if(contains(attributes, "broadcast") and contains(attributes, "axis"))
147
148
149
150
            {
                uint64_t broadcasted = parse_value(attributes.at("broadcast")).at<uint64_t>();
                if(broadcasted != 0)
                {
151
                    uint64_t axis = parse_value(attributes.at("axis")).at<uint64_t>();
Shucai Xiao's avatar
Shucai Xiao committed
152
153
                    auto l = prog.add_instruction(op::broadcast{axis, args[0]->get_shape().lens()},
                                                  args[1]);
154
155
                    return prog.add_instruction(x, args[0], l);
                }
156
                return prog.add_instruction(x, args);
157
            }
Paul's avatar
Paul committed
158
            else
159
            {
Khalique's avatar
Khalique committed
160
                return add_broadcastable_binary_op(args[0], args[1], x);
161
162
163
164
            }
        });
    }

Shucai Xiao's avatar
Shucai Xiao committed
165
166
    std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
                                                      std::vector<std::size_t> s1)
167
168
169
170
171
172
173
174
175
176
177
178
179
    {
        // 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)
Shucai Xiao's avatar
Shucai Xiao committed
180
        if(s0.size() > s1.size())
181
182
183
184
185
186
        {
            s0.swap(s1);
        }

        std::vector<std::size_t> out_lens(s1);
        auto offset = s1.size() - s0.size();
Shucai Xiao's avatar
Shucai Xiao committed
187
188
189
190
        std::transform(s0.begin(),
                       s0.end(),
                       s1.begin() + offset,
                       out_lens.begin() + offset,
191
                       [](auto a, auto b) {
Shucai Xiao's avatar
Shucai Xiao committed
192
                           if(a != b and a != 1 and b != 1)
193
194
195
                           {
                               MIGRAPHX_THROW("COMPUTE_BROADCASTED_LEN: input shapes mismatch!");
                           }
Shucai Xiao's avatar
Shucai Xiao committed
196
197
                           return std::max(a, b);
                       });
198
199
200
201

        return out_lens;
    }

Khalique's avatar
Khalique committed
202
203
204
    template <class T>
    instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
    {
Khalique's avatar
Khalique committed
205
        if(arg0->get_shape().lens() != arg1->get_shape().lens())
Khalique's avatar
Khalique committed
206
207
        {
            // Get lengths for both arguments
Shucai Xiao's avatar
Shucai Xiao committed
208
209
            auto s0       = arg0->get_shape().lens();
            auto s1       = arg1->get_shape().lens();
210
            auto out_lens = compute_broadcasted_lens(s0, s1);
Shucai Xiao's avatar
Shucai Xiao committed
211
212
            auto l0       = prog.add_instruction(op::multibroadcast{out_lens}, arg0);
            auto l1       = prog.add_instruction(op::multibroadcast{out_lens}, arg1);
Khalique's avatar
Khalique committed
213
214
215
216
217
218
            return prog.add_instruction(x, l0, l1);
        }
        else
        {
            return prog.add_instruction(x, {arg0, arg1});
        }
219
220
    }

Paul's avatar
Paul committed
221
    template <class T>
Paul's avatar
Paul committed
222
223
    void add_generic_op(std::string name, T x)
    {
Paul's avatar
Paul committed
224
        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
Paul's avatar
Paul committed
225
226
227
228
            return prog.add_instruction(x, args);
        });
    }

Khalique's avatar
Khalique committed
229
    template <class T>
Khalique's avatar
Khalique committed
230
    void add_variadic_op(std::string name, T x)
Khalique's avatar
Khalique committed
231
    {
Paul's avatar
Paul committed
232
        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
Khalique's avatar
Khalique committed
233
            return std::accumulate(std::next(args.begin()),
Khalique's avatar
Khalique committed
234
235
236
237
238
                                   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
239
        });
Khalique's avatar
Khalique committed
240
241
    }

Khalique's avatar
Khalique committed
242
243
244
    instruction_ref parse_clip(const std::string&,
                               const attribute_map& attributes,
                               std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
245
246
247
248
249
250
251
252
253
254
255
256
257
    {
        op::clip op;
        if(contains(attributes, "max"))
        {
            op.max_val = parse_value(attributes.at("max")).at<float>();
        }
        if(contains(attributes, "min"))
        {
            op.min_val = parse_value(attributes.at("min")).at<float>();
        }
        return prog.add_instruction(op, std::move(args));
    }

Paul's avatar
Paul committed
258
    instruction_ref
Paul's avatar
Paul committed
259
    parse_softmax(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
260
261
    {
        auto dims = args.front()->get_shape().lens();
Scott Thornton's avatar
Scott Thornton committed
262
263
        auto r =
            prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1]), 1, 1}}, args.front());
264
265
        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
266
267
    }

Shucai Xiao's avatar
Shucai Xiao committed
268
269
270
    instruction_ref parse_logsoftmax(const std::string&,
                                     const attribute_map& attributes,
                                     std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
271
272
273
274
275
276
277
278
279
280
    {
        int axis = 1;
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }

        return prog.add_instruction(op::logsoftmax{axis}, std::move(args));
    }

281
    instruction_ref parse_argmax(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
282
283
                                 const attribute_map& attributes,
                                 std::vector<instruction_ref> args)
284
285
286
287
288
289
290
    {
        int axis = 0;
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
291
        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
292
        if(contains(attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
293
294
295
296
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
297
        if(keep_dims == 0)
298
299
300
301
302
303
304
305
        {
            auto ins = prog.add_instruction(op::argmax{axis}, std::move(args));
            return prog.add_instruction(op::squeeze{{static_cast<int64_t>(axis)}}, ins);
        }
        else
        {
            return prog.add_instruction(op::argmax{axis}, std::move(args));
        }
306
307
308
    }

    instruction_ref parse_argmin(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
309
310
                                 const attribute_map& attributes,
                                 std::vector<instruction_ref> args)
311
312
313
314
315
316
317
    {
        int axis = 0;
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
318
        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
319
        if(contains(attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
320
321
322
323
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
324
        if(keep_dims == 0)
325
326
327
328
329
330
331
332
        {
            auto ins = prog.add_instruction(op::argmin{axis}, std::move(args));
            return prog.add_instruction(op::squeeze{{static_cast<int64_t>(axis)}}, ins);
        }
        else
        {
            return prog.add_instruction(op::argmin{axis}, std::move(args));
        }
333
334
    }

Paul's avatar
Paul committed
335
    instruction_ref
Paul's avatar
Paul committed
336
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
337
    {
338
        op::convolution op;
339
        auto l0 = args[0];
Paul's avatar
Paul committed
340
341
        if(contains(attributes, "pads"))
        {
Scott Thornton's avatar
Scott Thornton committed
342
            if(contains(attributes, "auto_pad"))
343
            {
Paul's avatar
Paul committed
344
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
345
            }
346
347
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
348
            if(padding.size() != 4)
349
            {
Paul's avatar
Paul committed
350
                MIGRAPHX_THROW("padding should have 4 values");
351
            }
Scott Thornton's avatar
Scott Thornton committed
352
            if(padding[0] != padding[2] || padding[1] != padding[3])
353
            {
354
355
                // 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
356
                l0      = prog.add_instruction(op::pad{padding}, l0);
357
            }
358
359
360
361
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
362
            }
Paul's avatar
Paul committed
363
        }
Paul's avatar
Paul committed
364
365
366
367
368
369
370
371
        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
372
        if(contains(attributes, "auto_pad"))
373
374
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
375
            if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
376
            {
Paul's avatar
Paul committed
377
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
378
379
            }

wsttiger's avatar
fixes  
wsttiger committed
380
            if(s.find("SAME") != std::string::npos)
381
            {
382
                op.padding_mode = op::padding_mode_t::same;
383
384
            }
        }
Khalique's avatar
Khalique committed
385
386
387
388
        if(contains(attributes, "group"))
        {
            op.group = parse_value(attributes.at("group")).at<int>();
        }
Paul's avatar
Paul committed
389
390
391
392
        if(args.size() == 3)
        {
            uint64_t axis = 1;
            auto l1       = prog.add_instruction(op, args[0], args[1]);
Shucai Xiao's avatar
Shucai Xiao committed
393
            auto l2 = prog.add_instruction(op::broadcast{axis, l1->get_shape().lens()}, args[2]);
394
            return prog.add_instruction(op::add{}, l1, l2);
Paul's avatar
Paul committed
395
        }
396
        return prog.add_instruction(op, l0, args[1]);
Paul's avatar
Paul committed
397
    }
Paul's avatar
Paul committed
398

Paul's avatar
Paul committed
399
400
401
    instruction_ref parse_pooling(const std::string& name,
                                  attribute_map attributes,
                                  std::vector<instruction_ref> args)
Paul's avatar
Paul committed
402
    {
Khalique's avatar
Khalique committed
403
        op::pooling op{ends_with(name, "MaxPool") ? "max" : "average"};
404
        auto l0 = args[0];
Khalique's avatar
Khalique committed
405
        if(starts_with(name, "Global"))
406
        {
Khalique's avatar
Khalique committed
407
408
            auto lens  = args.front()->get_shape().lens();
            op.lengths = {lens[2], lens[3]};
409
        }
Paul's avatar
Paul committed
410
411
        if(contains(attributes, "pads"))
        {
412
413
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
414
            if(padding.size() != 4)
415
            {
Paul's avatar
Paul committed
416
                MIGRAPHX_THROW("padding should have 4 values");
417
            }
Scott Thornton's avatar
Scott Thornton committed
418
            if(padding[0] != padding[2] || padding[1] != padding[3])
419
            {
420
421
                // 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
422
423
                l0 = prog.add_instruction(op::pad{padding, std::numeric_limits<float>::lowest()},
                                          l0);
424
425
426
427
428
            }
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
429
            }
Paul's avatar
Paul committed
430
431
432
433
434
435
436
437
438
        }
        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
439
        if(contains(attributes, "auto_pad"))
440
441
        {
            auto s = attributes["auto_pad"].s();
442
            if(s.find("SAME_UPPER") == std::string::npos)
443
            {
444
                MIGRAPHX_THROW("auto_pad only supports SAME_UPPER for pooling");
445
            }
446
            op.padding_mode = op::padding_mode_t::same;
447
448
        }

449
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
450
451
    }

Paul's avatar
Paul committed
452
    instruction_ref
Paul's avatar
Paul committed
453
    parse_reshape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
454
    {
455
        op::reshape op;
Paul's avatar
Paul committed
456
457
458
459
460
461
462
        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
463
            auto s = args[1]->eval();
Paul's avatar
Paul committed
464
            if(s.empty())
Paul's avatar
Paul committed
465
                MIGRAPHX_THROW("Dynamic shape is not supported.");
Paul's avatar
Paul committed
466
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
467
        }
Paul's avatar
Paul committed
468
469
470
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
471
    instruction_ref
Paul's avatar
Paul committed
472
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
473
    {
474
        uint64_t axis = 1;
Paul's avatar
Paul committed
475
476
477
478
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
479
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
480
481
    }

482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
    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
500
501
502
503
504
505
506
    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));
    }
507

508
509
510
    instruction_ref
    parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
511
        int axis = 0;
512
513
514
515
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
516
        op::gather op{axis};
517
518
519
        return prog.add_instruction(op, std::move(args));
    }

520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
    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
540
541
542
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
543
    {
Shucai Xiao's avatar
Shucai Xiao committed
544
        literal v     = parse_value(attributes.at("value"));
545
546
547
        auto dim_size = attributes.at("value").t().dims_size();
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
548
        {
549
            migraphx::shape scalar_shape{v.get_shape().type()};
550
551
552
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

Paul's avatar
Paul committed
553
554
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
555

Paul's avatar
Paul committed
556
    instruction_ref
Paul's avatar
Paul committed
557
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
558
559
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
560
        float beta  = 1.0f;
Paul's avatar
Paul committed
561
562
563
564
565
566
567
568
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
569
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
570
571
572
573
574
575
576
577
578
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
579
580
581
582
583
584

        std::vector<int64_t> perm(args[0]->get_shape().lens().size());
        std::iota(perm.begin(), perm.end(), int64_t{0});
        // swap the last two elements
        std::swap(*perm.rbegin(), *(perm.rbegin() + 1));

585
586
        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
587
588
        if(args.size() == 3)
        {
589
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
590
            {
Shucai Xiao's avatar
Shucai Xiao committed
591
                auto out_lens   = l1->get_shape().lens();
592
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
593
                auto l3         = args[2];
Shucai Xiao's avatar
Shucai Xiao committed
594
595
                auto l3_lens    = l3->get_shape().lens();
                if(!std::equal(out_lens.begin(), out_lens.end(), l3_lens.begin(), l3_lens.end()))
Khalique's avatar
Khalique committed
596
                {
597
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
598
                }
599
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
600
            }
Paul's avatar
Paul committed
601
        }
602
603

        return prog.add_instruction(op::dot{alpha, beta}, l1, l2);
Paul's avatar
Paul committed
604
605
    }

606
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
607
    parse_matmul(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
608
    {
Shucai Xiao's avatar
Shucai Xiao committed
609
610
        auto l0      = args[0];
        auto l1      = args[1];
611
612
613
614
615
        auto l0_lens = l0->get_shape().lens();
        auto l1_lens = l1->get_shape().lens();

        // args[0] is a vector, prepend 1 to the shape
        bool is_a_prepended = false;
Shucai Xiao's avatar
Shucai Xiao committed
616
        if(l0_lens.size() == 1)
617
618
619
620
621
622
623
        {
            is_a_prepended = true;
            l0_lens.insert(l0_lens.begin(), 1);
            l0 = prog.add_instruction(op::unsqueeze{{0}}, args[0]);
        }

        bool is_b_appended = false;
Shucai Xiao's avatar
Shucai Xiao committed
624
        if(l1_lens.size() == 1)
625
626
627
628
629
630
631
632
        {
            is_b_appended = true;
            l1_lens.push_back(1);
            l1 = prog.add_instruction(op::unsqueeze{{1}}, args[1]);
        }

        instruction_ref bl0 = l0;
        instruction_ref bl1 = l1;
Shucai Xiao's avatar
Shucai Xiao committed
633
        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
634
635
636
637
638
639
        {
            auto l0_it = l0_lens.begin() + l0_lens.size() - 2;
            std::vector<std::size_t> l0_broadcasted_lens(l0_lens.begin(), l0_it);
            auto l1_it = l1_lens.begin() + l1_lens.size() - 2;
            std::vector<std::size_t> l1_broadcasted_lens(l1_lens.begin(), l1_it);
            auto output_lens = compute_broadcasted_lens(l0_broadcasted_lens, l1_broadcasted_lens);
640
            l0_broadcasted_lens = output_lens;
641
            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
642
            l1_broadcasted_lens = output_lens;
643
            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
Shucai Xiao's avatar
Shucai Xiao committed
644
            if(l0_lens != l0_broadcasted_lens)
645
646
647
            {
                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
            }
Shucai Xiao's avatar
Shucai Xiao committed
648
            if(l1_lens != l1_broadcasted_lens)
649
650
651
652
653
            {
                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
654
        auto dot_res     = prog.add_instruction(op::dot{1.0f, 0.0f}, bl0, bl1);
655
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
656
        if(is_a_prepended)
657
658
659
660
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
661
        if(is_b_appended)
662
663
664
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
665

666
667
668
        return dot_res;
    }

669
    instruction_ref
Paul's avatar
Paul committed
670
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
671
    {
Scott Thornton's avatar
Scott Thornton committed
672
673
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
674
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Scott Thornton's avatar
Scott Thornton committed
675
        bool is_test                                      = false;
676
677
678
679
680
681
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
682
            momentum = parse_value(attributes.at("momentum")).at<float>();
683
684
685
        }
        if(contains(attributes, "is_test"))
        {
wsttiger's avatar
wsttiger committed
686
            is_test = parse_value(attributes.at("is_test")).at<uint64_t>() > 0;
687
688
689
        }
        if(contains(attributes, "spatial"))
        {
690
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
691
692
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
693
        }
Paul's avatar
Paul committed
694
        (void)is_test;
Paul's avatar
Paul committed
695
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
696
        return prog.add_instruction(op, std::move(args));
697
698
    }

699
700
701
702
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
703
        float alpha = 0.01; // default alpha val for leaky relu
704
705
706
707
708
709
710
711
        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
712
713
    instruction_ref
    parse_elu(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
714
715
716
717
718
719
720
721
722
723
    {
        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
724
725
    instruction_ref
    parse_lrn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
726
727
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
728
729
730
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
Khalique's avatar
Khalique committed
731
732
733
734
735
736
737
738
739
740
741
742
        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
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
    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());
        }
759
        auto input_lens = args.front()->get_shape().lens();
Khalique's avatar
Khalique committed
760

Khalique's avatar
Khalique committed
761
762
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
763
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
764

765
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
766
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
767
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
768
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
769
    }
Khalique's avatar
Khalique committed
770

Khalique's avatar
Khalique committed
771
772
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
773
774
775
776
777
778
779
    {
        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
780
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
781
782
    }

Khalique's avatar
Khalique committed
783
784
785
786
787
788
789
790
791
792
    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());
        }
793
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
794
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
795
796
797
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
Khalique's avatar
Khalique committed
798
799
800
801
802
803
804
805
806
807
808
809
        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());
    }
810
811
812
    // 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
813
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
814
815
    {
        if(args.size() != 1)
816
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
        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
853
854
        if(contains(attributes, "extra_shape"))
        {
855
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
856
857
        }

858
859
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
860
            if(args.size() != 1)
861
            {
862
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
863
864
            }

Shucai Xiao's avatar
Shucai Xiao committed
865
866
            if(contains(attributes, "shape"))
            {
867
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
868
                               "at the same time");
869
870
            }

871
872
873
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
874
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
875
            }
876

877
878
879
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
880
881
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
882
883
884
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
885
886
            if(!contains(attributes, "shape"))
            {
887
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
888
889
890
            }

            literal ls = parse_value(attributes.at("shape"));
891
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
892
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
893
            migraphx::shape s{type, dims};
894
895
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
896
897
898
        }
        else
        {
899
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
900
901
902
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
903
904
905
    instruction_ref parse_constant_of_shape(const std::string&,
                                            attribute_map attributes,
                                            std::vector<instruction_ref> args)
906
907
    {
        literal l_val{};
Shucai Xiao's avatar
Shucai Xiao committed
908
        if(contains(attributes, "value"))
909
910
        {
            l_val = parse_value(attributes.at("value"));
Shucai Xiao's avatar
Shucai Xiao committed
911
            if(l_val.get_shape().elements() != 1)
912
913
914
915
916
917
918
919
920
921
922
            {
                MIGRAPHX_THROW("ConstantOfShape: attribute value can contain only 1 elements!");
            }
        }
        else
        {
            l_val = literal({shape::float_type, {1}, {0}}, {0.0f});
        }

        // input is empty, output is a scalar
        auto type = l_val.get_shape().type();
Shucai Xiao's avatar
Shucai Xiao committed
923
        if(args.size() == 0)
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
        {
            return prog.add_literal(literal({type, {1}, {0}}, l_val.data()));
        }
        else
        {
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
                MIGRAPHX_THROW("ConstantOfShape: cannot handle dynamic shape as input");
            }

            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);

            literal l_out;
            l_val.visit([&](auto val) {
                using type = std::remove_cv_t<typename decltype(val)::value_type>;
                // l_val contains only one element
                std::vector<type> out_vec(s.elements(), *val.begin());
                l_out = literal(s, out_vec);
            });

            return prog.add_literal(l_out);
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
951
952
    instruction_ref
    parse_expand(const std::string&, attribute_map, std::vector<instruction_ref> args)
953
    {
Shucai Xiao's avatar
Shucai Xiao committed
954
955
        auto in_lens  = args[0]->get_shape().lens();
        auto ex_lens  = args[1]->get_shape().lens();
956
957
958
959
960
        auto out_lens = compute_broadcasted_lens(in_lens, ex_lens);

        return prog.add_instruction(op::multibroadcast{out_lens}, std::move(args[0]));
    }

Shucai Xiao's avatar
Shucai Xiao committed
961
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
962
963
964
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
965
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
966
967
968

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
969
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
970
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
971
972
973
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
974
975
976
977
978
979
980
981
982
        }

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

983
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
984
985
        if(direction == "bidirectional")
        {
986
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
987
988
989
        }
        else if(direction == "reverse")
        {
990
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
991
992
        }

993
        std::vector<std::string> vec_names{"tanh"};
994
995
996
997
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
998
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
999
1000
1001
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1002
1003
        }

1004
1005
1006
        auto name_it = std::find_if(vec_names.begin(), vec_names.end(), [&](auto& name) {
            return (map_actv_funcs.count(name) == 0);
        });
Shucai Xiao's avatar
Shucai Xiao committed
1007
        if(name_it != vec_names.end())
1008
1009
1010
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
1011

Shucai Xiao's avatar
Shucai Xiao committed
1012
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
1013
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
1014
        // if only one actv function is provided, we use it in both
1015
        // forward and reverse direction
1016
        if(dirct == op::rnn_direction::bidirectional)
1017
        {
Shucai Xiao's avatar
Shucai Xiao committed
1018
            if(vec_names.size() == 1)
1019
1020
1021
1022
1023
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
1024
1025
1026
        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];
1027
        });
Shucai Xiao's avatar
Shucai Xiao committed
1028

Shucai Xiao's avatar
Shucai Xiao committed
1029
1030
1031
1032
1033
1034
1035
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

1036
1037
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1038
        if(args.size() < 6)
1039
1040
1041
1042
1043
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
1044
1045
        // 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
1046
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1047

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

Shucai Xiao's avatar
Shucai Xiao committed
1051
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
1052
1053
    }

1054
    std::vector<instruction_ref>
1055
1056
1057
1058
1059
1060
1061
    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
1062
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1063
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1064
1065
1066
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
1067
1068
1069
1070
1071
1072
1073
1074
1075
        }

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

1076
        op::rnn_direction dirct = op::rnn_direction::forward;
1077
1078
        if(direction == "bidirectional")
        {
1079
            dirct = op::rnn_direction::bidirectional;
1080
1081
1082
        }
        else if(direction == "reverse")
        {
1083
            dirct = op::rnn_direction::reverse;
1084
1085
        }

1086
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
1087
1088
        if(contains(attributes, "activations"))
        {
1089
            auto names = attributes.at("activations").strings();
1090
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
1091
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1092
1093
1094
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1095
1096
        }

1097
        // need 4 activation functions
1098
        if(dirct == op::rnn_direction::bidirectional)
1099
        {
Shucai Xiao's avatar
Shucai Xiao committed
1100
            // 4 activation functions are used in the bidirectional
1101
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
1102
1103
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
1104
1105
            // 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
1106
1107
1108
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
1109
            if(vec_names.size() == 1)
1110
            {
1111
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
1112
            }
1113
            else if(vec_names.size() == 2)
1114
            {
1115
1116
1117
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
1118
            }
1119
            else if(vec_names.size() == 3)
1120
            {
1121
                vec_names.push_back(vec_names.at(2));
1122
1123
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
1124
        else
1125
        {
1126
            if(vec_names.size() == 1)
1127
            {
1128
                vec_names.push_back(vec_names.at(0));
1129
1130
1131
            }
        }

1132
1133
1134
        auto name_it = std::find_if(vec_names.begin(), vec_names.end(), [&](auto& name) {
            return (map_actv_funcs.count(name) == 0);
        });
Shucai Xiao's avatar
Shucai Xiao committed
1135
        if(name_it != vec_names.end())
1136
1137
1138
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1139

Shucai Xiao's avatar
Shucai Xiao committed
1140
1141
1142
        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
1143
        });
1144
1145
1146
1147
1148
1149
1150
1151

        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
1152
        if(contains(attributes, "linear_before_reset"))
1153
1154
1155
1156
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1157
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1158
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1159
1160
1161
1162
1163
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1164
1165
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1166
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1167
            std::move(args));
1168
1169

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

Shucai Xiao's avatar
Shucai Xiao committed
1172
        return {hidden_states, last_output};
1173
1174
    }

Shucai Xiao's avatar
Shucai Xiao committed
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
    std::vector<instruction_ref>
    parse_lstm(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"))
        {
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
            if(hidden_size != hidden_size_att)
            {
                MIGRAPHX_THROW("LSTM: hidden size mismatch in input and attribute");
            }
        }

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

Shucai Xiao's avatar
Shucai Xiao committed
1197
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1198
1199
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1200
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1201
1202
1203
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1204
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1205
        }
Shucai Xiao's avatar
Shucai Xiao committed
1206
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1207
        {
Shucai Xiao's avatar
Shucai Xiao committed
1208
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1209
1210
1211
1212
1213
1214
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1215
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
Shucai Xiao's avatar
Shucai Xiao committed
1216
1217
1218
1219
1220
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1221
1222
1223
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1224
1225
1226
        }

        // need 6 activation functions for bidirectional directions
Shucai Xiao's avatar
Shucai Xiao committed
1227
        if(dirct == op::rnn_direction::bidirectional)
Shucai Xiao's avatar
Shucai Xiao committed
1228
1229
1230
1231
1232
1233
        {
            // 6 activation functions are used in the bidirectional
            // scenario. No spec is provided in onnx::operator. we
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1st six times. If 2 actv functins are provided,
            // repeat 2nd once, then repeat all three once
Shucai Xiao's avatar
Shucai Xiao committed
1234
            // if 3 actv funcs are provide, repeat all three once.
Shucai Xiao's avatar
Shucai Xiao committed
1235
1236
1237
1238
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(vec_names.size())
            {
1239
            case 1:
Shucai Xiao's avatar
Shucai Xiao committed
1240
1241
1242
1243
1244
1245
                vec_names = {vec_names.at(0),
                             vec_names.at(0),
                             vec_names.at(0),
                             vec_names.at(0),
                             vec_names.at(0),
                             vec_names.at(0)};
1246
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1247
1248
1249

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1250
1251
1252
1253
1254
1255
                vec_names = {vec_names.at(0),
                             vec_names.at(1),
                             vec_names.at(1),
                             vec_names.at(0),
                             vec_names.at(1),
                             vec_names.at(1)};
Shucai Xiao's avatar
Shucai Xiao committed
1256
1257
1258
1259
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1260
1261
1262
1263
1264
1265
                vec_names = {vec_names.at(0),
                             vec_names.at(1),
                             vec_names.at(2),
                             vec_names.at(0),
                             vec_names.at(1),
                             vec_names.at(2)};
Shucai Xiao's avatar
Shucai Xiao committed
1266
1267
                break;

Shucai Xiao's avatar
Shucai Xiao committed
1268
1269
1270
1271
1272
1273
1274
            case 4:
                vec_names = {vec_names.at(0),
                             vec_names.at(1),
                             vec_names.at(2),
                             vec_names.at(3),
                             vec_names.at(3),
                             vec_names.at(3)};
1275
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1276

Shucai Xiao's avatar
Shucai Xiao committed
1277
1278
1279
1280
1281
1282
1283
            case 5:
                vec_names = {vec_names.at(0),
                             vec_names.at(1),
                             vec_names.at(2),
                             vec_names.at(3),
                             vec_names.at(4),
                             vec_names.at(4)};
1284
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1285

Shucai Xiao's avatar
Shucai Xiao committed
1286
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1287
1288
1289
1290
1291
1292
            }
        }
        else
        {
            switch(vec_names.size())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1293
            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
Shucai Xiao's avatar
Shucai Xiao committed
1294
1295
1296

            case 2:
                // repeat the 2nd actv func once, so we have 3 actv funcs
Shucai Xiao's avatar
Shucai Xiao committed
1297
                vec_names = {vec_names.at(0), vec_names.at(1), vec_names.at(1)};
Shucai Xiao's avatar
Shucai Xiao committed
1298
1299
                break;

Shucai Xiao's avatar
Shucai Xiao committed
1300
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1301
1302
1303
            }
        }

1304
1305
1306
        auto name_it = std::find_if(vec_names.begin(), vec_names.end(), [&](auto& name) {
            return (map_actv_funcs.count(name) == 0);
        });
Shucai Xiao's avatar
Shucai Xiao committed
1307
        if(name_it != vec_names.end())
1308
1309
1310
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332

        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];
        });

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

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

        // append undefined opeator to make 6 arguments
        if(args.size() < 8)
        {
            auto ins = prog.add_instruction(op::undefined{});
Shucai Xiao's avatar
Shucai Xiao committed
1333
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1334
1335
1336
1337
        }

        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1338
            op::lstm{hidden_size, vec_actv_funcs, dirct, clip, input_forget}, std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1339
1340

        // second output for last lstm output
Shucai Xiao's avatar
Shucai Xiao committed
1341
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1342
1343
1344
1345
1346
1347

        // third output for last cell output
        auto last_cell_output = prog.add_instruction(op::lstm_last_cell_output{}, hidden_states);

        return {hidden_states, last_output, last_cell_output};
    }
1348

Shucai Xiao's avatar
Shucai Xiao committed
1349
1350
1351
    instruction_ref parse_reduce_sum(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
    {
        std::size_t n_dim = args.front()->get_shape().lens().size();

        // default to reduce over all dimensions
        std::vector<std::size_t> axes(n_dim);
        std::iota(axes.begin(), axes.end(), 0);
        if(contains(attributes, "axes"))
        {
            axes.clear();
            auto&& attr_axes = attributes["axes"].ints();
Shucai Xiao's avatar
Shucai Xiao committed
1362
            axes             = std::vector<std::size_t>(attr_axes.begin(), attr_axes.end());
1363
1364
1365
        }

        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
1366
        if(contains(attributes, "keepdims"))
1367
1368
1369
1370
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1371
        if(keep_dims == 1)
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
        {
            return prog.add_instruction(op::reduce_sum{axes}, std::move(args));
        }
        else
        {
            auto ins = prog.add_instruction(op::reduce_sum{axes}, std::move(args));
            std::vector<int64_t> squeeze_axes{axes.begin(), axes.end()};
            return prog.add_instruction(op::squeeze{squeeze_axes}, ins);
        }
    }
1382

Shucai Xiao's avatar
Shucai Xiao committed
1383
1384
    instruction_ref
    parse_cast(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
1385
    {
Shucai Xiao's avatar
Shucai Xiao committed
1386
        if(!contains(attributes, "to"))
1387
1388
1389
1390
        {
            MIGRAPHX_THROW("PARSE_CAST: missing to type attribute!");
        }

Shucai Xiao's avatar
Shucai Xiao committed
1391
        int to_type        = parse_value(attributes.at("to")).at<int>();
1392
1393
1394
        shape::type_t type = get_type(to_type);
        return prog.add_instruction(op::convert{type}, std::move(args));
    }
Shucai Xiao's avatar
Shucai Xiao committed
1395

Paul's avatar
Paul committed
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
    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
1408
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
1409
1410
1411
1412
1413
1414
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
1415
1416
1417
1418
1419
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
1420
1421
1422
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
            // 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
1435
        }
Paul's avatar
Paul committed
1436
        for(auto&& output : graph.output())
Paul's avatar
Paul committed
1437
        {
Paul's avatar
Paul committed
1438
            this->parse_node(output.name());
Paul's avatar
Paul committed
1439
1440
1441
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1442
    void parse_undefined(const std::string& name)
1443
    {
Shucai Xiao's avatar
Shucai Xiao committed
1444
        auto ins           = prog.add_instruction(op::undefined{});
1445
1446
1447
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
1448
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
1449
    {
Paul's avatar
Paul committed
1450
        if(name.empty())
Paul's avatar
Paul committed
1451
            MIGRAPHX_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
1452
1453
1454
1455
1456
1457
1458
1459
        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
1460
1461
                    assert(name != input);
                    this->parse_node(input);
Paul's avatar
Paul committed
1462
                }
Shucai Xiao's avatar
Shucai Xiao committed
1463
                else if(input.empty())
Paul's avatar
Paul committed
1464
                {
1465
                    this->parse_undefined(input);
Paul's avatar
Paul committed
1466
                }
1467
                args.push_back(instructions.at(input));
Paul's avatar
Paul committed
1468
            }
Paul's avatar
Paul committed
1469
            std::vector<instruction_ref> result;
Paul's avatar
Paul committed
1470
1471
            if(ops.count(node.op_type()) == 0)
            {
1472
                result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
Paul's avatar
Paul committed
1473
1474
1475
            }
            else
            {
Paul's avatar
Paul committed
1476
                result = ops[node.op_type()](get_attributes(node), args);
Paul's avatar
Paul committed
1477
            }
Paul's avatar
Paul committed
1478
            // Even no output nodes produce output in migraphx
Paul's avatar
Paul committed
1479
            if(node.output().empty() and result.size() == 1)
Paul's avatar
Paul committed
1480
1481
            {
                instructions[name] = result.front();
Paul's avatar
Paul committed
1482
1483
1484
            }
            else
            {
Paul's avatar
Paul committed
1485
1486
1487
1488
1489
1490
                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
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
            }
        }
    }

    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
1508
        std::size_t n = 0;
Paul's avatar
Paul committed
1509
1510
        for(auto&& node : graph.node())
        {
Paul's avatar
Paul committed
1511
            if(node.output().empty())
Paul's avatar
Paul committed
1512
            {
Paul's avatar
Paul committed
1513
                if(node.name().empty())
Paul's avatar
Paul committed
1514
1515
1516
1517
1518
1519
1520
1521
1522
                {
                    result["migraphx_unamed_node_" + std::to_string(n)] = node;
                    n++;
                }
                else
                {
                    result[node.name()] = node;
                }
            }
Paul's avatar
Paul committed
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
            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
1548
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
1549
1550
1551
1552
1553
        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
1554
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
1555
1556
1557
1558
1559
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
1560
1561
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
1562
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
1563
1564
1565
            switch(t.data_type())
            {
            case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
1566
            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
Scott Thornton's avatar
Scott Thornton committed
1567
            case onnx::TensorProto::UINT8: throw std::runtime_error("");
1568
            case onnx::TensorProto::INT8: return create_literal(shape::int32_type, dims, s.data());
Khalique's avatar
Khalique committed
1569
1570
            case onnx::TensorProto::UINT16:
                return create_literal(shape::int32_type, dims, s.data());
1571
1572
1573
            case onnx::TensorProto::INT16: return create_literal(shape::int32_type, dims, s.data());
            case onnx::TensorProto::INT32: return create_literal(shape::int32_type, dims, s.data());
            case onnx::TensorProto::INT64: return create_literal(shape::int64_type, dims, s.data());
Scott Thornton's avatar
Scott Thornton committed
1574
            case onnx::TensorProto::STRING: throw std::runtime_error("");
1575
            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
Khalique's avatar
Khalique committed
1576
1577
1578
1579
            case onnx::TensorProto::FLOAT16:
                return create_literal(shape::half_type, dims, s.data());
            case onnx::TensorProto::DOUBLE:
                return create_literal(shape::double_type, dims, s.data());
Scott Thornton's avatar
Scott Thornton committed
1580
1581
1582
1583
1584
            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
1585
            MIGRAPHX_THROW("Invalid tensor type");
1586
        }
Paul's avatar
Paul committed
1587
1588
1589
1590
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Khalique's avatar
Khalique committed
1591
            return create_literal(shape::float_type, dims, t.float_data());
Paul's avatar
Paul committed
1592
1593
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Khalique's avatar
Khalique committed
1594
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1595
        case onnx::TensorProto::UINT16:
Khalique's avatar
Khalique committed
1596
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1597
        case onnx::TensorProto::INT16:
Khalique's avatar
Khalique committed
1598
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1599
        case onnx::TensorProto::INT32:
Khalique's avatar
Khalique committed
1600
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1601
        case onnx::TensorProto::INT64:
Khalique's avatar
Khalique committed
1602
            return create_literal(shape::int64_type, dims, t.int64_data());
Paul's avatar
Paul committed
1603
1604
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Khalique's avatar
Khalique committed
1605
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1606
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
1607
        {
Khalique's avatar
Khalique committed
1608
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
1609
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
1610
1611
1612
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
1613
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
1614
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
1615
        }
Paul's avatar
Paul committed
1616
        case onnx::TensorProto::DOUBLE:
Khalique's avatar
Khalique committed
1617
            return create_literal(shape::double_type, dims, t.double_data());
Paul's avatar
Paul committed
1618
1619
1620
1621
1622
        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
1623
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
1624
1625
    }

Khalique's avatar
Khalique committed
1626
    static literal
1627
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
1628
    {
Khalique's avatar
Khalique committed
1629
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1630
        if(dims.empty())
1631
            return literal{{shape_type}, data};
1632
1633
1634
        return literal{{shape_type, dims}, data};
    }

1635
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
1636
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
1637
1638
    {
        if(dims.empty())
1639
            return literal{{shape_type}, data.begin(), data.end()};
1640
        return literal{{shape_type, dims}, data.begin(), data.end()};
1641
1642
    }

Paul's avatar
Paul committed
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
    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
1662
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1663
1664
1665
1666
1667
1668
1669
1670
1671
        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
1672
        auto&& tensor_dims = t.tensor_type().shape().dim();
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
        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
1684
1685
        return {shape_type, dims};
    }
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707

    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
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
};

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
1731
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
1732
} // namespace migraphx