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

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

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

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

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

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

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

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

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

382
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
383
384
    }

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

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

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

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

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

Paul's avatar
Paul committed
486
487
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
488

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

513
514
515
516
517
        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));

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

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

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

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

599
600
601
        return dot_res;
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
897
898
899
        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];
900
        });
Shucai Xiao's avatar
Shucai Xiao committed
901

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

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

Shucai Xiao's avatar
Shucai Xiao committed
917
918
        // 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
919
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
920

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

Shucai Xiao's avatar
Shucai Xiao committed
924
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
925
926
    }

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

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

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

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

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

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

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

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
1043
        return {hidden_states, last_output};
1044
1045
    }

Shucai Xiao's avatar
Shucai Xiao committed
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
    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
1068
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1069
1070
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1071
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1072
1073
1074
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1075
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1076
        }
Shucai Xiao's avatar
Shucai Xiao committed
1077
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1078
        {
Shucai Xiao's avatar
Shucai Xiao committed
1079
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
        }
        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());
1092
            std::copy(names.begin(), names.end(), vec_names.begin());
Shucai Xiao's avatar
Shucai Xiao committed
1093
1094
1095
        }

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

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

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

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
1169
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1170
1171
1172
            }
        }

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

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

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

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

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

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

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

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

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

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

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

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

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

    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
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
};

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