onnx.cpp 63.4 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
192
193
194
195
196
197
                       [](auto a, auto b) {
                           if (a != b and a != 1 and b != 1)
                           {
                               MIGRAPHX_THROW("COMPUTE_BROADCASTED_LEN: input shapes mismatch!");
                           }
                           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
        }
    }

903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
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
951
952
953
954
955
956
957
    instruction_ref parse_constant_of_shape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        literal l_val{};
        if (contains(attributes, "value"))
        {
            l_val = parse_value(attributes.at("value"));
            if (l_val.get_shape().elements() != 1)
            {
                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();
        if (args.size() == 0)
        {
            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);
        }
    }

    instruction_ref parse_expand(const std::string&, attribute_map, std::vector<instruction_ref> args)
    {
        auto in_lens = args[0]->get_shape().lens();
        auto ex_lens = args[1]->get_shape().lens();
        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
958
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
959
960
961
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
962
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
963
964
965

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

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

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

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

1001
1002
1003
        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
1004
        if(name_it != vec_names.end())
1005
1006
1007
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
1008

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

Shucai Xiao's avatar
Shucai Xiao committed
1021
1022
1023
        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];
1024
        });
Shucai Xiao's avatar
Shucai Xiao committed
1025

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

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

Shucai Xiao's avatar
Shucai Xiao committed
1041
1042
        // 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
1043
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1044

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

Shucai Xiao's avatar
Shucai Xiao committed
1048
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
1049
1050
    }

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

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

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

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

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

1129
1130
1131
        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
1132
        if(name_it != vec_names.end())
1133
1134
1135
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1136

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

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
1169
        return {hidden_states, last_output};
1170
1171
    }

Shucai Xiao's avatar
Shucai Xiao committed
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
    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
1194
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1195
1196
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1197
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1198
1199
1200
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1201
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1202
        }
Shucai Xiao's avatar
Shucai Xiao committed
1203
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1204
        {
Shucai Xiao's avatar
Shucai Xiao committed
1205
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1206
1207
1208
1209
1210
1211
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1212
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
Shucai Xiao's avatar
Shucai Xiao committed
1213
1214
1215
1216
1217
        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
1218
1219
1220
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1221
1222
1223
        }

        // need 6 activation functions for bidirectional directions
Shucai Xiao's avatar
Shucai Xiao committed
1224
        if(dirct == op::rnn_direction::bidirectional)
Shucai Xiao's avatar
Shucai Xiao committed
1225
1226
1227
1228
1229
1230
        {
            // 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
1231
            // if 3 actv funcs are provide, repeat all three once.
Shucai Xiao's avatar
Shucai Xiao committed
1232
1233
1234
1235
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(vec_names.size())
            {
1236
            case 1:
Shucai Xiao's avatar
Shucai Xiao committed
1237
1238
1239
1240
1241
1242
                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)};
1243
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1244
1245
1246

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1247
1248
1249
1250
1251
1252
                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
1253
1254
1255
1256
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1257
1258
1259
1260
1261
1262
                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
1263
1264
                break;

Shucai Xiao's avatar
Shucai Xiao committed
1265
1266
1267
1268
1269
1270
1271
            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)};
1272
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1273

Shucai Xiao's avatar
Shucai Xiao committed
1274
1275
1276
1277
1278
1279
1280
            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)};
1281
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1282

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

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

Shucai Xiao's avatar
Shucai Xiao committed
1297
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1298
1299
1300
            }
        }

1301
1302
1303
        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
1304
        if(name_it != vec_names.end())
1305
1306
1307
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329

        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
1330
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1331
1332
1333
1334
        }

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

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

        // 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};
    }
1345

Shucai Xiao's avatar
Shucai Xiao committed
1346
1347
1348
    instruction_ref parse_reduce_sum(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
    {
        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
1359
            axes             = std::vector<std::size_t>(attr_axes.begin(), attr_axes.end());
1360
1361
1362
        }

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

Shucai Xiao's avatar
Shucai Xiao committed
1368
        if(keep_dims == 1)
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
        {
            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);
        }
    }
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390

    instruction_ref parse_cast(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        if (!contains(attributes, "to"))
        {
            MIGRAPHX_THROW("PARSE_CAST: missing to type attribute!");
        }

        int to_type = parse_value(attributes.at("to")).at<int>();
        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
1391

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

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

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

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

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

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

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

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

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

    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
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
};

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