onnx.cpp 56.8 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
105
106
107
        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
108
109
110
111
    }

    template <class F>
    void add_op(std::string name, F f)
Paul's avatar
Paul committed
112
113
114
115
116
117
118
119
120
    {
        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
121
122
123
124
125
126
127
    {
        ops.emplace(name, f);
    }

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

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

Shucai Xiao's avatar
Shucai Xiao committed
158
159
    std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
                                                      std::vector<std::size_t> s1)
160
161
162
163
164
165
166
167
168
169
170
171
172
    {
        // Example:
        // s0 = (3,2,4,5) and s1 = (2,1,1)
        //
        // In this case we need to broadcast (:,1,1) portion of
        // s1 plus broadcast the 1st dimension of s1
        // giving output_lens = (3,2,4,5)
        //
        // Another example:
        // s0 = (3,2,1,5) and s1 = (2,7,5)
        // In this case we need to broadcast the (:,:,1:,:) axis
        // of s0 plus the 1st dimension of s1 giving
        // output_lens = (3,2,7,5)
Shucai Xiao's avatar
Shucai Xiao committed
173
        if(s0.size() > s1.size())
174
175
176
177
178
179
        {
            s0.swap(s1);
        }

        std::vector<std::size_t> out_lens(s1);
        auto offset = s1.size() - s0.size();
Shucai Xiao's avatar
Shucai Xiao committed
180
181
182
183
184
        std::transform(s0.begin(),
                       s0.end(),
                       s1.begin() + offset,
                       out_lens.begin() + offset,
                       [](auto a, auto b) { return std::max(a, b); });
185
186
187
188

        return out_lens;
    }

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

Paul's avatar
Paul committed
208
    template <class T>
Paul's avatar
Paul committed
209
210
    void add_generic_op(std::string name, T x)
    {
Paul's avatar
Paul committed
211
        add_op(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
Paul's avatar
Paul committed
212
213
214
215
            return prog.add_instruction(x, args);
        });
    }

Khalique's avatar
Khalique committed
216
    template <class T>
Khalique's avatar
Khalique committed
217
    void add_variadic_op(std::string name, T x)
Khalique's avatar
Khalique committed
218
    {
Paul's avatar
Paul committed
219
        add_op(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
Khalique's avatar
Khalique committed
220
            return std::accumulate(std::next(args.begin()),
Khalique's avatar
Khalique committed
221
222
223
224
225
                                   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
226
        });
Khalique's avatar
Khalique committed
227
228
    }

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

Shucai Xiao's avatar
Shucai Xiao committed
253
254
255
    instruction_ref parse_logsoftmax(const std::string&,
                                     const attribute_map& attributes,
                                     std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
256
257
258
259
260
261
262
263
264
265
    {
        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
266
    instruction_ref
Paul's avatar
Paul committed
267
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
268
    {
269
        op::convolution op;
270
        auto l0 = args[0];
Paul's avatar
Paul committed
271
272
        if(contains(attributes, "pads"))
        {
Scott Thornton's avatar
Scott Thornton committed
273
            if(contains(attributes, "auto_pad"))
274
            {
Paul's avatar
Paul committed
275
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
276
            }
277
278
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
279
            if(padding.size() != 4)
280
            {
Paul's avatar
Paul committed
281
                MIGRAPHX_THROW("padding should have 4 values");
282
            }
Scott Thornton's avatar
Scott Thornton committed
283
            if(padding[0] != padding[2] || padding[1] != padding[3])
284
            {
285
286
                // 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
287
                l0      = prog.add_instruction(op::pad{padding}, l0);
288
            }
289
290
291
292
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
293
            }
Paul's avatar
Paul committed
294
        }
Paul's avatar
Paul committed
295
296
297
298
299
300
301
302
        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
303
        if(contains(attributes, "auto_pad"))
304
305
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
306
            if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
307
            {
Paul's avatar
Paul committed
308
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
309
310
            }

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

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

379
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
380
381
    }

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

Paul's avatar
Paul committed
401
    instruction_ref
Paul's avatar
Paul committed
402
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
403
    {
404
        uint64_t axis = 1;
Paul's avatar
Paul committed
405
406
407
408
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
409
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
410
411
    }

412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
    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
430
431
432
433
434
435
436
    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));
    }
437

438
439
440
    instruction_ref
    parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
441
        int axis = 0;
442
443
444
445
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
446
        op::gather op{axis};
447
448
449
        return prog.add_instruction(op, std::move(args));
    }

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

Paul's avatar
Paul committed
483
484
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
485

Paul's avatar
Paul committed
486
    instruction_ref
Paul's avatar
Paul committed
487
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
488
489
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
490
        float beta  = 1.0f;
Paul's avatar
Paul committed
491
492
493
494
495
496
497
498
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
499
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
500
501
502
503
504
505
506
507
508
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
509

510
511
512
513
514
        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));

515
516
        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
517
518
        if(args.size() == 3)
        {
519
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
520
            {
Shucai Xiao's avatar
Shucai Xiao committed
521
                auto out_lens   = l1->get_shape().lens();
522
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
523
524
525
                auto l3         = args[2];
                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
526
                {
527
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
528
                }
529
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
530
            }
Paul's avatar
Paul committed
531
        }
532

Shucai Xiao's avatar
Shucai Xiao committed
533
        return prog.add_instruction(op::dot{alpha, beta}, l1, l2);
Paul's avatar
Paul committed
534
535
    }

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

Shucai Xiao's avatar
Shucai Xiao committed
584
        auto dot_res     = prog.add_instruction(op::dot{1.0f, 0.0f}, bl0, bl1);
585
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
586
        if(is_a_prepended)
587
588
589
590
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
591
        if(is_b_appended)
592
593
594
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
595

596
597
598
        return dot_res;
    }

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

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

Khalique's avatar
Khalique committed
691
692
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
693
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
694

695
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
696
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
697
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
698
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
699
    }
Khalique's avatar
Khalique committed
700

Khalique's avatar
Khalique committed
701
702
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
703
704
705
706
707
708
709
    {
        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
710
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
711
712
    }

Khalique's avatar
Khalique committed
713
714
715
716
717
718
719
720
721
722
    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());
        }
723
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
724
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
725
726
727
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
Khalique's avatar
Khalique committed
728
729
730
731
732
733
734
735
736
737
738
739
        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());
    }
740
741
742
    // 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
743
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
744
745
    {
        if(args.size() != 1)
746
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
747
748
749
750
751
752
753
754
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
        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
783
784
        if(contains(attributes, "extra_shape"))
        {
785
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
786
787
        }

788
789
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
790
            if(args.size() != 1)
791
            {
792
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
793
794
            }

Shucai Xiao's avatar
Shucai Xiao committed
795
796
            if(contains(attributes, "shape"))
            {
797
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
798
                               "at the same time");
799
800
            }

801
802
803
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
804
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
805
            }
806

807
808
809
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
810
811
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
812
813
814
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
815
816
            if(!contains(attributes, "shape"))
            {
817
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
818
819
820
            }

            literal ls = parse_value(attributes.at("shape"));
821
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
822
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
823
            migraphx::shape s{type, dims};
824
825
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
826
827
828
        }
        else
        {
829
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
830
831
832
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
833
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
834
835
836
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
837
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
838
839
840

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
841
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
842
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
843
844
845
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
846
847
848
849
850
851
852
853
854
        }

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

855
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
856
857
        if(direction == "bidirectional")
        {
858
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
859
860
861
        }
        else if(direction == "reverse")
        {
862
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
863
864
        }

865
866
867
868
869
        std::vector<std::string> vec_names{"tanh"};
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
870
            vec_names.resize(names.size());
871
            std::copy(names.begin(), names.end(), vec_names.begin());
872
873
        }

874
875
        auto name_it = std::find_if(vec_names.begin(), vec_names.end(), [&](auto& name) {
            return (map_actv_funcs.count(name) == 0);
876
        });
Shucai Xiao's avatar
Shucai Xiao committed
877
        if(name_it != vec_names.end())
878
879
880
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
881

Shucai Xiao's avatar
Shucai Xiao committed
882
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
883
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
884
        // if only one actv function is provided, we use it in both
885
        // forward and reverse direction
886
        if(dirct == op::rnn_direction::bidirectional)
887
        {
Shucai Xiao's avatar
Shucai Xiao committed
888
            if(vec_names.size() == 1)
889
890
891
892
893
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
894
895
896
        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];
897
        });
Shucai Xiao's avatar
Shucai Xiao committed
898

Shucai Xiao's avatar
Shucai Xiao committed
899
900
901
902
903
904
905
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

906
907
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
908
        if(args.size() < 6)
909
910
911
912
913
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
914
915
        // 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
916
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
917

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

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

924
    std::vector<instruction_ref>
925
926
927
928
929
930
931
    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
932
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
933
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
934
935
936
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
937
938
939
940
941
942
943
944
945
        }

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

946
        op::rnn_direction dirct = op::rnn_direction::forward;
947
948
        if(direction == "bidirectional")
        {
949
            dirct = op::rnn_direction::bidirectional;
950
951
952
        }
        else if(direction == "reverse")
        {
953
            dirct = op::rnn_direction::reverse;
954
955
        }

956
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
957
958
        if(contains(attributes, "activations"))
        {
959
            auto names = attributes.at("activations").strings();
960
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
961
            vec_names.resize(names.size());
962
            std::copy(names.begin(), names.end(), vec_names.begin());
963
964
        }

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

1000
1001
        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
1002
        });
Shucai Xiao's avatar
Shucai Xiao committed
1003
1004
        if(name_it != vec_names.end())
        {
1005
1006
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1007

Shucai Xiao's avatar
Shucai Xiao committed
1008
1009
1010
        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
1011
        });
1012
1013
1014
1015
1016
1017
1018
1019

        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
1020
        if(contains(attributes, "linear_before_reset"))
1021
1022
1023
1024
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1025
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1026
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1027
1028
1029
1030
1031
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1032
1033
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1034
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1035
            std::move(args));
1036
1037

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

Shucai Xiao's avatar
Shucai Xiao committed
1040
        return {hidden_states, last_output};
1041
1042
    }

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

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

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

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1116
1117
1118
1119
1120
1121
                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
1122
1123
1124
1125
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1126
1127
1128
1129
1130
1131
                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
1132
1133
                break;

Shucai Xiao's avatar
Shucai Xiao committed
1134
1135
1136
1137
1138
1139
1140
            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)};
1141
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1142

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

Shucai Xiao's avatar
Shucai Xiao committed
1152
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1153
1154
1155
1156
1157
1158
            }
        }
        else
        {
            switch(vec_names.size())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1159
            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
Shucai Xiao's avatar
Shucai Xiao committed
1160
1161
1162

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

Shucai Xiao's avatar
Shucai Xiao committed
1166
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1167
1168
1169
            }
        }

1170
1171
1172
        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
1173
        if(name_it != vec_names.end())
1174
1175
1176
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198

        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
1199
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1200
1201
1202
1203
        }

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

        // second output for last lstm output
Shucai Xiao's avatar
Shucai Xiao committed
1207
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1208
1209
1210
1211
1212
1213
1214

        // 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
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
    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
1227
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
1228
1229
1230
1231
1232
1233
        }
    }

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

Shucai Xiao's avatar
Shucai Xiao committed
1261
    void parse_undefined(const std::string& name)
1262
    {
Shucai Xiao's avatar
Shucai Xiao committed
1263
        auto ins           = prog.add_instruction(op::undefined{});
1264
1265
1266
        instructions[name] = ins;
    }

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

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

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
Khalique's avatar
Khalique committed
1379
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1380
        if(dims.empty())
Khalique's avatar
Khalique committed
1381
1382
1383
        {
            dims = {1};
        }
1384
1385
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
1386
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
            switch(t.data_type())
            {
            case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
            case onnx::TensorProto::FLOAT: return literal{{shape::float_type, dims}, s.data()};
            case onnx::TensorProto::UINT8: throw std::runtime_error("");
            case onnx::TensorProto::INT8: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::UINT16: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT16: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT32: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT64: return literal{{shape::int64_type, dims}, s.data()};
            case onnx::TensorProto::STRING: throw std::runtime_error("");
            case onnx::TensorProto::BOOL: return literal{{shape::int32_type, dims}, s.data()};
Paul's avatar
Paul committed
1399
            case onnx::TensorProto::FLOAT16: return literal{{shape::half_type, dims}, s.data()};
Scott Thornton's avatar
Scott Thornton committed
1400
1401
1402
1403
1404
1405
            case onnx::TensorProto::DOUBLE: return literal{{shape::double_type, dims}, s.data()};
            case onnx::TensorProto::UINT32: throw std::runtime_error("");
            case onnx::TensorProto::UINT64: throw std::runtime_error("");
            case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
            case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
            }
Paul's avatar
Paul committed
1406
            MIGRAPHX_THROW("Invalid tensor type");
1407
        }
Paul's avatar
Paul committed
1408
1409
1410
1411
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Paul's avatar
Paul committed
1412
            return literal{{shape::float_type, dims}, t.float_data().begin(), t.float_data().end()};
Paul's avatar
Paul committed
1413
1414
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Paul's avatar
Paul committed
1415
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1416
        case onnx::TensorProto::UINT16:
Paul's avatar
Paul committed
1417
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1418
        case onnx::TensorProto::INT16:
Paul's avatar
Paul committed
1419
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1420
        case onnx::TensorProto::INT32:
Paul's avatar
Paul committed
1421
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1422
        case onnx::TensorProto::INT64:
Paul's avatar
Paul committed
1423
            return literal{{shape::int64_type, dims}, t.int64_data().begin(), t.int64_data().end()};
Paul's avatar
Paul committed
1424
1425
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Paul's avatar
Paul committed
1426
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1427
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
1428
        {
Khalique's avatar
Khalique committed
1429
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
1430
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
1431
1432
1433
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
1434
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
1435
            return literal{{shape::half_type, dims}, data_half.begin(), data_half.end()};
Khalique's avatar
Khalique committed
1436
        }
Paul's avatar
Paul committed
1437
1438
1439
1440
1441
1442
1443
1444
        case onnx::TensorProto::DOUBLE:
            return literal{
                {shape::double_type, dims}, t.double_data().begin(), t.double_data().end()};
        case onnx::TensorProto::UINT32: throw std::runtime_error("");
        case onnx::TensorProto::UINT64: throw std::runtime_error("");
        case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
        case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
        }
Paul's avatar
Paul committed
1445
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
    }

    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
1467
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1468
1469
1470
1471
1472
1473
1474
1475
1476
        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
1477
        auto&& tensor_dims = t.tensor_type().shape().dim();
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
        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
1489
1490
        return {shape_type, dims};
    }
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512

    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
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
};

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