onnx.cpp 57.7 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

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

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

    void init_actv_func()
    {
103
104
        // Support name format of all lower case or the first letter capital
        map_actv_funcs.insert(std::make_pair("tanh", op::tanh{}));
105
        map_actv_funcs.insert(std::make_pair("Tanh", op::tanh{}));
106
        map_actv_funcs.insert(std::make_pair("relu", op::relu{}));
107
        map_actv_funcs.insert(std::make_pair("Relu", op::relu{}));
108
        map_actv_funcs.insert(std::make_pair("sigmoid", op::sigmoid{}));
109
        map_actv_funcs.insert(std::make_pair("Sigmoid", op::sigmoid{}));
110
        map_actv_funcs.insert(std::make_pair("leakyrelu", op::leaky_relu{}));
111
        map_actv_funcs.insert(std::make_pair("LeakyRelu", op::leaky_relu{}));
112
        map_actv_funcs.insert(std::make_pair("elu", op::elu{}));
113
        map_actv_funcs.insert(std::make_pair("Elu", op::elu{}));
Paul's avatar
Paul committed
114
115
116
117
    }

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

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

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

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

        std::vector<std::size_t> out_lens(s1);
        auto offset = s1.size() - s0.size();
Shucai Xiao's avatar
Shucai Xiao committed
186
187
188
189
190
        std::transform(s0.begin(),
                       s0.end(),
                       s1.begin() + offset,
                       out_lens.begin() + offset,
                       [](auto a, auto b) { return std::max(a, b); });
191
192
193
194

        return out_lens;
    }

Khalique's avatar
Khalique committed
195
196
197
    template <class T>
    instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
    {
Khalique's avatar
Khalique committed
198
        if(arg0->get_shape().lens() != arg1->get_shape().lens())
Khalique's avatar
Khalique committed
199
200
        {
            // Get lengths for both arguments
Shucai Xiao's avatar
Shucai Xiao committed
201
202
            auto s0       = arg0->get_shape().lens();
            auto s1       = arg1->get_shape().lens();
203
            auto out_lens = compute_broadcasted_lens(s0, s1);
Shucai Xiao's avatar
Shucai Xiao committed
204
205
            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
206
207
208
209
210
211
            return prog.add_instruction(x, l0, l1);
        }
        else
        {
            return prog.add_instruction(x, {arg0, arg1});
        }
212
213
    }

Paul's avatar
Paul committed
214
    template <class T>
Paul's avatar
Paul committed
215
216
    void add_generic_op(std::string name, T x)
    {
Paul's avatar
Paul committed
217
        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
Paul's avatar
Paul committed
218
219
220
221
            return prog.add_instruction(x, args);
        });
    }

Khalique's avatar
Khalique committed
222
    template <class T>
Khalique's avatar
Khalique committed
223
    void add_variadic_op(std::string name, T x)
Khalique's avatar
Khalique committed
224
    {
Paul's avatar
Paul committed
225
        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
Khalique's avatar
Khalique committed
226
            return std::accumulate(std::next(args.begin()),
Khalique's avatar
Khalique committed
227
228
229
230
231
                                   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
232
        });
Khalique's avatar
Khalique committed
233
234
    }

Khalique's avatar
Khalique committed
235
236
237
    instruction_ref parse_clip(const std::string&,
                               const attribute_map& attributes,
                               std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
238
239
240
241
242
243
244
245
246
247
248
249
250
    {
        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
251
    instruction_ref
Paul's avatar
Paul committed
252
    parse_softmax(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
253
254
    {
        auto dims = args.front()->get_shape().lens();
Scott Thornton's avatar
Scott Thornton committed
255
256
        auto r =
            prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1]), 1, 1}}, args.front());
257
258
        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
259
260
    }

Shucai Xiao's avatar
Shucai Xiao committed
261
262
263
    instruction_ref parse_logsoftmax(const std::string&,
                                     const attribute_map& attributes,
                                     std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
264
265
266
267
268
269
270
271
272
273
    {
        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));
    }

Paul's avatar
Paul committed
274
    instruction_ref
Paul's avatar
Paul committed
275
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
276
    {
277
        op::convolution op;
278
        auto l0 = args[0];
Paul's avatar
Paul committed
279
280
        if(contains(attributes, "pads"))
        {
Scott Thornton's avatar
Scott Thornton committed
281
            if(contains(attributes, "auto_pad"))
282
            {
Paul's avatar
Paul committed
283
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
284
            }
285
286
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
287
            if(padding.size() != 4)
288
            {
Paul's avatar
Paul committed
289
                MIGRAPHX_THROW("padding should have 4 values");
290
            }
Scott Thornton's avatar
Scott Thornton committed
291
            if(padding[0] != padding[2] || padding[1] != padding[3])
292
            {
293
294
                // 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
295
                l0      = prog.add_instruction(op::pad{padding}, l0);
296
            }
297
298
299
300
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
301
            }
Paul's avatar
Paul committed
302
        }
Paul's avatar
Paul committed
303
304
305
306
307
308
309
310
        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
311
        if(contains(attributes, "auto_pad"))
312
313
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
314
            if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
315
            {
Paul's avatar
Paul committed
316
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
317
318
            }

wsttiger's avatar
fixes  
wsttiger committed
319
            if(s.find("SAME") != std::string::npos)
320
            {
321
                op.padding_mode = op::padding_mode_t::same;
322
323
            }
        }
Khalique's avatar
Khalique committed
324
325
326
327
        if(contains(attributes, "group"))
        {
            op.group = parse_value(attributes.at("group")).at<int>();
        }
Paul's avatar
Paul committed
328
329
330
331
        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
332
            auto l2 = prog.add_instruction(op::broadcast{axis, l1->get_shape().lens()}, args[2]);
333
            return prog.add_instruction(op::add{}, l1, l2);
Paul's avatar
Paul committed
334
        }
335
        return prog.add_instruction(op, l0, args[1]);
Paul's avatar
Paul committed
336
    }
Paul's avatar
Paul committed
337

Paul's avatar
Paul committed
338
339
340
    instruction_ref parse_pooling(const std::string& name,
                                  attribute_map attributes,
                                  std::vector<instruction_ref> args)
Paul's avatar
Paul committed
341
    {
Khalique's avatar
Khalique committed
342
        op::pooling op{ends_with(name, "MaxPool") ? "max" : "average"};
343
        auto l0 = args[0];
Khalique's avatar
Khalique committed
344
        if(starts_with(name, "Global"))
345
        {
Khalique's avatar
Khalique committed
346
347
            auto lens  = args.front()->get_shape().lens();
            op.lengths = {lens[2], lens[3]};
348
        }
Paul's avatar
Paul committed
349
350
        if(contains(attributes, "pads"))
        {
351
352
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
353
            if(padding.size() != 4)
354
            {
Paul's avatar
Paul committed
355
                MIGRAPHX_THROW("padding should have 4 values");
356
            }
Scott Thornton's avatar
Scott Thornton committed
357
            if(padding[0] != padding[2] || padding[1] != padding[3])
358
            {
359
360
                // 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
361
                l0      = prog.add_instruction(op::pad{padding}, l0);
362
363
364
365
366
            }
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
367
            }
Paul's avatar
Paul committed
368
369
370
371
372
373
374
375
376
        }
        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
377
        if(contains(attributes, "auto_pad"))
378
379
        {
            auto s = attributes["auto_pad"].s();
380
            if(s.find("SAME_UPPER") == std::string::npos)
381
            {
382
                MIGRAPHX_THROW("auto_pad only supports SAME_UPPER for pooling");
383
            }
384
            op.padding_mode = op::padding_mode_t::same;
385
386
        }

387
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
388
389
    }

Paul's avatar
Paul committed
390
    instruction_ref
Paul's avatar
Paul committed
391
    parse_reshape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
392
    {
393
        op::reshape op;
Paul's avatar
Paul committed
394
395
396
397
398
399
400
        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
401
            auto s = args[1]->eval();
Paul's avatar
Paul committed
402
            if(s.empty())
Paul's avatar
Paul committed
403
                MIGRAPHX_THROW("Dynamic shape is not supported.");
Paul's avatar
Paul committed
404
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
405
        }
Paul's avatar
Paul committed
406
407
408
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
409
    instruction_ref
Paul's avatar
Paul committed
410
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
411
    {
412
        uint64_t axis = 1;
Paul's avatar
Paul committed
413
414
415
416
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
417
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
418
419
    }

420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
    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
438
439
440
441
442
443
444
    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));
    }
445

446
447
448
    instruction_ref
    parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
449
        int axis = 0;
450
451
452
453
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
454
        op::gather op{axis};
455
456
457
        return prog.add_instruction(op, std::move(args));
    }

458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
    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
478
479
480
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
481
    {
Shucai Xiao's avatar
Shucai Xiao committed
482
        literal v     = parse_value(attributes.at("value"));
483
484
485
        auto dim_size = attributes.at("value").t().dims_size();
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
486
        {
487
            migraphx::shape scalar_shape{v.get_shape().type()};
488
489
490
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

Paul's avatar
Paul committed
491
492
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
493

Paul's avatar
Paul committed
494
    instruction_ref
Paul's avatar
Paul committed
495
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
496
497
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
498
        float beta  = 1.0f;
Paul's avatar
Paul committed
499
500
501
502
503
504
505
506
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
507
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
508
509
510
511
512
513
514
515
516
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
517
518
519
520
521
522

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

523
524
        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
525
526
        if(args.size() == 3)
        {
527
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
528
            {
Shucai Xiao's avatar
Shucai Xiao committed
529
                auto out_lens   = l1->get_shape().lens();
530
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
531
                auto l3         = args[2];
Shucai Xiao's avatar
Shucai Xiao committed
532
533
                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
534
                {
535
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
536
                }
537
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
538
            }
Paul's avatar
Paul committed
539
        }
540
541

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

544
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
545
    parse_matmul(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
546
    {
Shucai Xiao's avatar
Shucai Xiao committed
547
548
        auto l0      = args[0];
        auto l1      = args[1];
549
550
551
552
553
        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
554
        if(l0_lens.size() == 1)
555
556
557
558
559
560
561
        {
            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
562
        if(l1_lens.size() == 1)
563
564
565
566
567
568
569
570
        {
            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
571
        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
572
573
574
575
576
577
        {
            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);
578
            l0_broadcasted_lens = output_lens;
579
            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
580
            l1_broadcasted_lens = output_lens;
581
            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
Shucai Xiao's avatar
Shucai Xiao committed
582
            if(l0_lens != l0_broadcasted_lens)
583
584
585
            {
                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
            }
Shucai Xiao's avatar
Shucai Xiao committed
586
            if(l1_lens != l1_broadcasted_lens)
587
588
589
590
591
            {
                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
592
        auto dot_res     = prog.add_instruction(op::dot{1.0f, 0.0f}, bl0, bl1);
593
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
594
        if(is_a_prepended)
595
596
597
598
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
599
        if(is_b_appended)
600
601
602
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
603

604
605
606
        return dot_res;
    }

607
    instruction_ref
Paul's avatar
Paul committed
608
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
609
    {
Scott Thornton's avatar
Scott Thornton committed
610
611
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
612
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Scott Thornton's avatar
Scott Thornton committed
613
        bool is_test                                      = false;
614
615
616
617
618
619
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
620
            momentum = parse_value(attributes.at("momentum")).at<float>();
621
622
623
        }
        if(contains(attributes, "is_test"))
        {
wsttiger's avatar
wsttiger committed
624
            is_test = parse_value(attributes.at("is_test")).at<uint64_t>() > 0;
625
626
627
        }
        if(contains(attributes, "spatial"))
        {
628
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
629
630
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
631
        }
Paul's avatar
Paul committed
632
        (void)is_test;
Paul's avatar
Paul committed
633
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
634
        return prog.add_instruction(op, std::move(args));
635
636
    }

637
638
639
640
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
641
        float alpha = 0.01; // default alpha val for leaky relu
642
643
644
645
646
647
648
649
        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
650
651
    instruction_ref
    parse_elu(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
652
653
654
655
656
657
658
659
660
661
    {
        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
662
663
    instruction_ref
    parse_lrn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
664
665
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
666
667
668
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
Khalique's avatar
Khalique committed
669
670
671
672
673
674
675
676
677
678
679
680
        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
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
    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());
        }
697
        auto input_lens = args.front()->get_shape().lens();
Khalique's avatar
Khalique committed
698

Khalique's avatar
Khalique committed
699
700
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
701
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
702

703
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
704
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
705
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
706
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
707
    }
Khalique's avatar
Khalique committed
708

Khalique's avatar
Khalique committed
709
710
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
711
712
713
714
715
716
717
    {
        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
718
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
719
720
    }

Khalique's avatar
Khalique committed
721
722
723
724
725
726
727
728
729
730
    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());
        }
731
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
732
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
733
734
735
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
Khalique's avatar
Khalique committed
736
737
738
739
740
741
742
743
744
745
746
747
        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());
    }
748
749
750
    // 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
751
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
752
753
    {
        if(args.size() != 1)
754
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
        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
791
792
        if(contains(attributes, "extra_shape"))
        {
793
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
794
795
        }

796
797
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
798
            if(args.size() != 1)
799
            {
800
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
801
802
            }

Shucai Xiao's avatar
Shucai Xiao committed
803
804
            if(contains(attributes, "shape"))
            {
805
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
806
                               "at the same time");
807
808
            }

809
810
811
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
812
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
813
            }
814

815
816
817
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
818
819
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
820
821
822
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
823
824
            if(!contains(attributes, "shape"))
            {
825
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
826
827
828
            }

            literal ls = parse_value(attributes.at("shape"));
829
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
830
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
831
            migraphx::shape s{type, dims};
832
833
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
834
835
836
        }
        else
        {
837
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
838
839
840
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
841
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
842
843
844
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
845
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
846
847
848

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
849
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
850
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
851
852
853
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
854
855
856
857
858
859
860
861
862
        }

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

863
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
864
865
        if(direction == "bidirectional")
        {
866
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
867
868
869
        }
        else if(direction == "reverse")
        {
870
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
871
872
        }

873
        std::vector<std::string> vec_names{"tanh"};
874
875
876
877
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
878
            vec_names.resize(names.size());
879
            std::copy(names.begin(), names.end(), vec_names.begin());
880
881
        }

882
883
884
        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
885
        if(name_it != vec_names.end())
886
887
888
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
889

Shucai Xiao's avatar
Shucai Xiao committed
890
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
891
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
892
        // if only one actv function is provided, we use it in both
893
        // forward and reverse direction
894
        if(dirct == op::rnn_direction::bidirectional)
895
        {
Shucai Xiao's avatar
Shucai Xiao committed
896
            if(vec_names.size() == 1)
897
898
899
900
901
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
902
903
904
        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];
905
        });
Shucai Xiao's avatar
Shucai Xiao committed
906

Shucai Xiao's avatar
Shucai Xiao committed
907
908
909
910
911
912
913
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

914
915
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
916
        if(args.size() < 6)
917
918
919
920
921
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
922
923
        // 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
924
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
925

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

Shucai Xiao's avatar
Shucai Xiao committed
929
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
930
931
    }

932
    std::vector<instruction_ref>
933
934
935
936
937
938
939
    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
940
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
941
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
942
943
944
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
945
946
947
948
949
950
951
952
953
        }

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

954
        op::rnn_direction dirct = op::rnn_direction::forward;
955
956
        if(direction == "bidirectional")
        {
957
            dirct = op::rnn_direction::bidirectional;
958
959
960
        }
        else if(direction == "reverse")
        {
961
            dirct = op::rnn_direction::reverse;
962
963
        }

964
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
965
966
        if(contains(attributes, "activations"))
        {
967
            auto names = attributes.at("activations").strings();
968
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
969
            vec_names.resize(names.size());
970
            std::copy(names.begin(), names.end(), vec_names.begin());
971
972
        }

973
        // need 4 activation functions
974
        if(dirct == op::rnn_direction::bidirectional)
975
        {
Shucai Xiao's avatar
Shucai Xiao committed
976
            // 4 activation functions are used in the bidirectional
977
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
978
979
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
980
981
            // 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
982
983
984
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
985
            if(vec_names.size() == 1)
986
            {
987
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
988
            }
989
            else if(vec_names.size() == 2)
990
            {
991
992
993
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
994
            }
995
            else if(vec_names.size() == 3)
996
            {
997
                vec_names.push_back(vec_names.at(2));
998
999
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
1000
        else
1001
        {
1002
            if(vec_names.size() == 1)
1003
            {
1004
                vec_names.push_back(vec_names.at(0));
1005
1006
1007
            }
        }

1008
1009
1010
        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
1011
        if(name_it != vec_names.end())
1012
1013
1014
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1015

Shucai Xiao's avatar
Shucai Xiao committed
1016
1017
1018
        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
1019
        });
1020
1021
1022
1023
1024
1025
1026
1027

        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
1028
        if(contains(attributes, "linear_before_reset"))
1029
1030
1031
1032
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1033
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1034
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1035
1036
1037
1038
1039
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1040
1041
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1042
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1043
            std::move(args));
1044
1045

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

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

Shucai Xiao's avatar
Shucai Xiao committed
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
    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
1073
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1074
1075
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1076
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1077
1078
1079
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1080
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1081
        }
Shucai Xiao's avatar
Shucai Xiao committed
1082
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1083
        {
Shucai Xiao's avatar
Shucai Xiao committed
1084
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1085
1086
1087
1088
1089
1090
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1091
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
Shucai Xiao's avatar
Shucai Xiao committed
1092
1093
1094
1095
1096
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
            vec_names.resize(names.size());
1097
            std::copy(names.begin(), names.end(), vec_names.begin());
Shucai Xiao's avatar
Shucai Xiao committed
1098
1099
1100
        }

        // need 6 activation functions for bidirectional directions
Shucai Xiao's avatar
Shucai Xiao committed
1101
        if(dirct == op::rnn_direction::bidirectional)
Shucai Xiao's avatar
Shucai Xiao committed
1102
1103
1104
1105
1106
1107
        {
            // 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
1108
            // if 3 actv funcs are provide, repeat all three once.
Shucai Xiao's avatar
Shucai Xiao committed
1109
1110
1111
1112
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(vec_names.size())
            {
1113
            case 1:
Shucai Xiao's avatar
Shucai Xiao committed
1114
1115
1116
1117
1118
1119
                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)};
1120
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1121
1122
1123

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1124
1125
1126
1127
1128
1129
                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
1130
1131
1132
1133
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1134
1135
1136
1137
1138
1139
                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
1140
1141
                break;

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

Shucai Xiao's avatar
Shucai Xiao committed
1151
1152
1153
1154
1155
1156
1157
            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)};
1158
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1159

Shucai Xiao's avatar
Shucai Xiao committed
1160
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1161
1162
1163
1164
1165
1166
            }
        }
        else
        {
            switch(vec_names.size())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1167
            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
Shucai Xiao's avatar
Shucai Xiao committed
1168
1169
1170

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

Shucai Xiao's avatar
Shucai Xiao committed
1174
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1175
1176
1177
            }
        }

1178
1179
1180
        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
1181
        if(name_it != vec_names.end())
1182
1183
1184
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206

        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
1207
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1208
1209
1210
1211
        }

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

        // second output for last lstm output
Shucai Xiao's avatar
Shucai Xiao committed
1215
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1216
1217
1218
1219
1220
1221
1222

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

Paul's avatar
Paul committed
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
    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
1235
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
1236
1237
1238
1239
1240
1241
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
1242
1243
1244
1245
1246
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
1247
1248
1249
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
            // 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
1262
        }
Paul's avatar
Paul committed
1263
        for(auto&& output : graph.output())
Paul's avatar
Paul committed
1264
        {
Paul's avatar
Paul committed
1265
            this->parse_node(output.name());
Paul's avatar
Paul committed
1266
1267
1268
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1269
    void parse_undefined(const std::string& name)
1270
    {
Shucai Xiao's avatar
Shucai Xiao committed
1271
        auto ins           = prog.add_instruction(op::undefined{});
1272
1273
1274
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
1275
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
1276
    {
Paul's avatar
Paul committed
1277
        if(name.empty())
Paul's avatar
Paul committed
1278
            MIGRAPHX_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
1279
1280
1281
1282
1283
1284
1285
1286
        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
1287
1288
                    assert(name != input);
                    this->parse_node(input);
Paul's avatar
Paul committed
1289
                }
Shucai Xiao's avatar
Shucai Xiao committed
1290
                else if(input.empty())
Paul's avatar
Paul committed
1291
                {
1292
                    this->parse_undefined(input);
Paul's avatar
Paul committed
1293
                }
1294
                args.push_back(instructions.at(input));
Paul's avatar
Paul committed
1295
            }
Paul's avatar
Paul committed
1296
            std::vector<instruction_ref> result;
Paul's avatar
Paul committed
1297
1298
            if(ops.count(node.op_type()) == 0)
            {
1299
                result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
Paul's avatar
Paul committed
1300
1301
1302
            }
            else
            {
Paul's avatar
Paul committed
1303
                result = ops[node.op_type()](get_attributes(node), args);
Paul's avatar
Paul committed
1304
            }
Paul's avatar
Paul committed
1305
            // Even no output nodes produce output in migraphx
Paul's avatar
Paul committed
1306
            if(node.output().empty() and result.size() == 1)
Paul's avatar
Paul committed
1307
1308
            {
                instructions[name] = result.front();
Paul's avatar
Paul committed
1309
1310
1311
            }
            else
            {
Paul's avatar
Paul committed
1312
1313
1314
1315
1316
1317
                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
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
            }
        }
    }

    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
1335
        std::size_t n = 0;
Paul's avatar
Paul committed
1336
1337
        for(auto&& node : graph.node())
        {
Paul's avatar
Paul committed
1338
            if(node.output().empty())
Paul's avatar
Paul committed
1339
            {
Paul's avatar
Paul committed
1340
                if(node.name().empty())
Paul's avatar
Paul committed
1341
1342
1343
1344
1345
1346
1347
1348
1349
                {
                    result["migraphx_unamed_node_" + std::to_string(n)] = node;
                    n++;
                }
                else
                {
                    result[node.name()] = node;
                }
            }
Paul's avatar
Paul committed
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
            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
1375
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
1376
1377
1378
1379
1380
        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
1381
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
1382
1383
1384
1385
1386
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
1387
1388
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
1389
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
1390
1391
1392
            switch(t.data_type())
            {
            case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
1393
            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
Scott Thornton's avatar
Scott Thornton committed
1394
            case onnx::TensorProto::UINT8: throw std::runtime_error("");
1395
            case onnx::TensorProto::INT8: return create_literal(shape::int32_type, dims, s.data());
Khalique's avatar
Khalique committed
1396
1397
            case onnx::TensorProto::UINT16:
                return create_literal(shape::int32_type, dims, s.data());
1398
1399
1400
            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
1401
            case onnx::TensorProto::STRING: throw std::runtime_error("");
1402
            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
Khalique's avatar
Khalique committed
1403
1404
1405
1406
            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
1407
1408
1409
1410
1411
            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
1412
            MIGRAPHX_THROW("Invalid tensor type");
1413
        }
Paul's avatar
Paul committed
1414
1415
1416
1417
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Khalique's avatar
Khalique committed
1418
            return create_literal(shape::float_type, dims, t.float_data());
Paul's avatar
Paul committed
1419
1420
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Khalique's avatar
Khalique committed
1421
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1422
        case onnx::TensorProto::UINT16:
Khalique's avatar
Khalique committed
1423
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1424
        case onnx::TensorProto::INT16:
Khalique's avatar
Khalique committed
1425
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1426
        case onnx::TensorProto::INT32:
Khalique's avatar
Khalique committed
1427
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1428
        case onnx::TensorProto::INT64:
Khalique's avatar
Khalique committed
1429
            return create_literal(shape::int64_type, dims, t.int64_data());
Paul's avatar
Paul committed
1430
1431
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Khalique's avatar
Khalique committed
1432
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1433
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
1434
        {
Khalique's avatar
Khalique committed
1435
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
1436
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
1437
1438
1439
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
1440
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
1441
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
1442
        }
Paul's avatar
Paul committed
1443
        case onnx::TensorProto::DOUBLE:
Khalique's avatar
Khalique committed
1444
            return create_literal(shape::double_type, dims, t.double_data());
Paul's avatar
Paul committed
1445
1446
1447
1448
1449
        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
1450
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
1451
1452
    }

Khalique's avatar
Khalique committed
1453
    static literal
1454
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
1455
    {
Khalique's avatar
Khalique committed
1456
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1457
        if(dims.empty())
1458
            return literal{{shape_type}, data};
1459
1460
1461
        return literal{{shape_type, dims}, data};
    }

1462
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
1463
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
1464
1465
    {
        if(dims.empty())
1466
            return literal{{shape_type}, data.begin(), data.end()};
1467
        return literal{{shape_type, dims}, data.begin(), data.end()};
1468
1469
    }

Paul's avatar
Paul committed
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
    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
1489
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1490
1491
1492
1493
1494
1495
1496
1497
1498
        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
1499
        auto&& tensor_dims = t.tensor_type().shape().dim();
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
        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
1511
1512
        return {shape_type, dims};
    }
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534

    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
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
};

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