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

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

    void init_actv_func()
    {
102
103
104
105
106
        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
107
108
109
110
    }

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
157
158
    std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
                                                      std::vector<std::size_t> s1)
159
160
161
162
163
164
165
166
167
168
169
170
171
    {
        // 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
172
        if(s0.size() > s1.size())
173
174
175
176
177
178
        {
            s0.swap(s1);
        }

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

        return out_lens;
    }

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

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

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

Paul's avatar
Paul committed
228
    instruction_ref
Paul's avatar
Paul committed
229
    parse_softmax(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
230
231
    {
        auto dims = args.front()->get_shape().lens();
Scott Thornton's avatar
Scott Thornton committed
232
233
        auto r =
            prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1]), 1, 1}}, args.front());
234
235
        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
236
237
    }

Shucai Xiao's avatar
Shucai Xiao committed
238
239
240
    instruction_ref parse_logsoftmax(const std::string&,
                                     const attribute_map& attributes,
                                     std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
241
242
243
244
245
246
247
248
249
250
    {
        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
251
    instruction_ref
Paul's avatar
Paul committed
252
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
253
    {
254
        op::convolution op;
255
        auto l0 = args[0];
Paul's avatar
Paul committed
256
257
        if(contains(attributes, "pads"))
        {
Scott Thornton's avatar
Scott Thornton committed
258
            if(contains(attributes, "auto_pad"))
259
            {
Paul's avatar
Paul committed
260
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
261
            }
262
263
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
264
            if(padding.size() != 4)
265
            {
Paul's avatar
Paul committed
266
                MIGRAPHX_THROW("padding should have 4 values");
267
            }
Scott Thornton's avatar
Scott Thornton committed
268
            if(padding[0] != padding[2] || padding[1] != padding[3])
269
            {
270
271
                // 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
272
                l0      = prog.add_instruction(op::pad{padding}, l0);
273
            }
274
275
276
277
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
278
            }
Paul's avatar
Paul committed
279
        }
Paul's avatar
Paul committed
280
281
282
283
284
285
286
287
        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
288
        if(contains(attributes, "auto_pad"))
289
290
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
291
            if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
292
            {
Paul's avatar
Paul committed
293
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
294
295
            }

wsttiger's avatar
fixes  
wsttiger committed
296
            if(s.find("SAME") != std::string::npos)
297
            {
298
                op.padding_mode = op::padding_mode_t::same;
299
300
            }
        }
Khalique's avatar
Khalique committed
301
302
303
304
        if(contains(attributes, "group"))
        {
            op.group = parse_value(attributes.at("group")).at<int>();
        }
Paul's avatar
Paul committed
305
306
307
308
        if(args.size() == 3)
        {
            uint64_t axis = 1;
            auto l1       = prog.add_instruction(op, args[0], args[1]);
Scott Thornton's avatar
Scott Thornton committed
309
            auto l2       = prog.add_instruction(op::broadcast{axis, l1->get_shape()}, args[2]);
310
            return prog.add_instruction(op::add{}, l1, l2);
Paul's avatar
Paul committed
311
        }
312
        return prog.add_instruction(op, l0, args[1]);
Paul's avatar
Paul committed
313
    }
Paul's avatar
Paul committed
314

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

364
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
365
366
    }

Paul's avatar
Paul committed
367
    instruction_ref
Paul's avatar
Paul committed
368
    parse_reshape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
369
    {
370
        op::reshape op;
Paul's avatar
Paul committed
371
372
373
374
375
376
377
        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
378
            auto s = args[1]->eval();
Paul's avatar
Paul committed
379
            if(s.empty())
Paul's avatar
Paul committed
380
                MIGRAPHX_THROW("Dynamic shape is not supported.");
Paul's avatar
Paul committed
381
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
382
        }
Paul's avatar
Paul committed
383
384
385
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
386
    instruction_ref
Paul's avatar
Paul committed
387
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
388
    {
389
        uint64_t axis = 1;
Paul's avatar
Paul committed
390
391
392
393
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
394
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
395
396
    }

397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
    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
415
416
417
418
419
420
421
    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));
    }
422

423
424
425
    instruction_ref
    parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
426
        int axis = 0;
427
428
429
430
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
431
        op::gather op{axis};
432
433
434
        return prog.add_instruction(op, std::move(args));
    }

435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
    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
455
456
457
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
458
    {
Shucai Xiao's avatar
Shucai Xiao committed
459
        literal v     = parse_value(attributes.at("value"));
460
461
462
        auto dim_size = attributes.at("value").t().dims_size();
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
463
        {
464
            migraphx::shape scalar_shape{v.get_shape().type()};
465
466
467
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

Paul's avatar
Paul committed
468
469
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
470

Paul's avatar
Paul committed
471
    instruction_ref
Paul's avatar
Paul committed
472
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
473
474
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
475
        float beta  = 1.0f;
Paul's avatar
Paul committed
476
477
478
479
480
481
482
483
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
484
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
485
486
487
488
489
490
491
492
493
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
494

495
496
497
498
499
        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));

500
501
        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
502
503
        if(args.size() == 3)
        {
504
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
505
            {
Shucai Xiao's avatar
Shucai Xiao committed
506
                auto out_lens   = l1->get_shape().lens();
507
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
508
509
510
                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
511
                {
512
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
513
                }
514
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
515
            }
Paul's avatar
Paul committed
516
        }
517

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

521
522
523
    instruction_ref
    parse_matmul(const std::string&, attribute_map, std::vector<instruction_ref> args)
    {
Shucai Xiao's avatar
Shucai Xiao committed
524
525
        auto l0      = args[0];
        auto l1      = args[1];
526
527
528
529
530
        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
531
        if(l0_lens.size() == 1)
532
533
534
535
536
537
538
        {
            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
539
        if(l1_lens.size() == 1)
540
541
542
543
544
545
546
547
        {
            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
548
        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
549
550
551
552
553
554
555
556
        {
            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);
            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
Shucai Xiao's avatar
Shucai Xiao committed
557
            if(l0_lens != l0_broadcasted_lens)
558
559
560
            {
                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
            }
Shucai Xiao's avatar
Shucai Xiao committed
561
            if(l1_lens != l1_broadcasted_lens)
562
563
564
565
566
            {
                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
567
        auto dot_res     = prog.add_instruction(op::dot{1.0f, 0.0f}, bl0, bl1);
568
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
569
        if(is_a_prepended)
570
571
572
573
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
574
        if(is_b_appended)
575
576
577
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
578

579
580
581
        return dot_res;
    }

582
    instruction_ref
Paul's avatar
Paul committed
583
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
584
    {
Scott Thornton's avatar
Scott Thornton committed
585
586
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
587
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Scott Thornton's avatar
Scott Thornton committed
588
        bool is_test                                      = false;
589
590
591
592
593
594
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
595
            momentum = parse_value(attributes.at("momentum")).at<float>();
596
597
598
        }
        if(contains(attributes, "is_test"))
        {
wsttiger's avatar
wsttiger committed
599
            is_test = parse_value(attributes.at("is_test")).at<uint64_t>() > 0;
600
601
602
        }
        if(contains(attributes, "spatial"))
        {
603
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
604
605
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
606
        }
Paul's avatar
Paul committed
607
        (void)is_test;
Paul's avatar
Paul committed
608
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
609
        return prog.add_instruction(op, std::move(args));
610
611
    }

612
613
614
615
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
616
        float alpha = 0.01; // default alpha val for leaky relu
617
618
619
620
621
622
623
624
        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
625
626
    instruction_ref
    parse_elu(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
627
628
629
630
631
632
633
634
635
636
    {
        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
637
638
    instruction_ref
    parse_lrn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
639
640
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
641
642
643
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
Khalique's avatar
Khalique committed
644
645
646
647
648
649
650
651
652
653
654
655
        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
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
    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());
        }
        auto input_shape = args.front()->get_shape();
Khalique's avatar
Khalique committed
673

Khalique's avatar
Khalique committed
674
675
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
676
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
677

Paul's avatar
Paul committed
678
679
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_shape}, scale_val);
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Paul's avatar
Paul committed
680
        auto bias_bcast = prog.add_instruction(migraphx::op::broadcast{1, input_shape}, bias_vals);
Paul's avatar
Paul committed
681
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
682
    }
Khalique's avatar
Khalique committed
683

Khalique's avatar
Khalique committed
684
685
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
686
687
688
689
690
691
692
    {
        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
693
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
694
695
    }

Khalique's avatar
Khalique committed
696
697
698
699
700
701
702
703
704
705
    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());
        }
706
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
707
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
708
709
710
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
Khalique's avatar
Khalique committed
711
712
713
714
715
716
717
718
719
720
721
722
        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());
    }
723
724
725
    // 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
726
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
727
728
    {
        if(args.size() != 1)
729
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
        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
766
767
        if(contains(attributes, "extra_shape"))
        {
768
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
769
770
        }

771
772
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
773
            if(args.size() != 1)
774
            {
775
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
776
777
            }

Shucai Xiao's avatar
Shucai Xiao committed
778
779
            if(contains(attributes, "shape"))
            {
780
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
781
                               "at the same time");
782
783
            }

784
785
786
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
787
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
788
            }
789

790
791
792
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
793
794
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
795
796
797
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
798
799
            if(!contains(attributes, "shape"))
            {
800
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
801
802
803
            }

            literal ls = parse_value(attributes.at("shape"));
804
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
805
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
806
            migraphx::shape s{type, dims};
807
808
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
809
810
811
        }
        else
        {
812
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
813
814
815
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
816
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
817
818
819
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
820
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
821
822
823

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
824
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
825
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
826
827
828
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
829
830
831
832
833
834
835
836
837
        }

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

838
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
839
840
        if(direction == "bidirectional")
        {
841
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
842
843
844
        }
        else if(direction == "reverse")
        {
845
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
846
847
        }

848
849
850
851
852
        std::vector<std::string> vec_names{"tanh"};
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
853
            vec_names.resize(names.size());
854
            std::copy(names.begin(), names.end(), vec_names.begin());
855
856
        }

857
858
        auto name_it = std::find_if(vec_names.begin(), vec_names.end(), [&](auto& name) {
            return (map_actv_funcs.count(name) == 0);
859
        });
Shucai Xiao's avatar
Shucai Xiao committed
860
        if(name_it != vec_names.end())
861
862
863
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
864

Shucai Xiao's avatar
Shucai Xiao committed
865
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
866
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
867
        // if only one actv function is provided, we use it in both
868
        // forward and reverse direction
869
        if(dirct == op::rnn_direction::bidirectional)
870
        {
Shucai Xiao's avatar
Shucai Xiao committed
871
            if(vec_names.size() == 1)
872
873
874
875
876
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
877
878
879
        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];
880
        });
Shucai Xiao's avatar
Shucai Xiao committed
881

Shucai Xiao's avatar
Shucai Xiao committed
882
883
884
885
886
887
888
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

889
890
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
891
        if(args.size() < 6)
892
893
894
895
896
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
897
898
        // 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
899
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
900

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

Shucai Xiao's avatar
Shucai Xiao committed
904
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
905
906
    }

907
    std::vector<instruction_ref>
908
909
910
911
912
913
914
    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
915
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
916
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
917
918
919
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
920
921
922
923
924
925
926
927
928
        }

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

929
        op::rnn_direction dirct = op::rnn_direction::forward;
930
931
        if(direction == "bidirectional")
        {
932
            dirct = op::rnn_direction::bidirectional;
933
934
935
        }
        else if(direction == "reverse")
        {
936
            dirct = op::rnn_direction::reverse;
937
938
        }

939
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
940
941
        if(contains(attributes, "activations"))
        {
942
            auto names = attributes.at("activations").strings();
943
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
944
            vec_names.resize(names.size());
945
            std::copy(names.begin(), names.end(), vec_names.begin());
946
947
        }

948
        // need 4 activation functions
949
        if(dirct == op::rnn_direction::bidirectional)
950
        {
Shucai Xiao's avatar
Shucai Xiao committed
951
            // 4 activation functions are used in the bidirectional
952
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
953
954
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
955
956
            // 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
957
958
959
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
960
            if(vec_names.size() == 1)
961
            {
962
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
963
            }
964
            else if(vec_names.size() == 2)
965
            {
966
967
968
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
969
            }
970
            else if(vec_names.size() == 3)
971
            {
972
                vec_names.push_back(vec_names.at(2));
973
974
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
975
        else
976
        {
977
            if(vec_names.size() == 1)
978
            {
979
                vec_names.push_back(vec_names.at(0));
980
981
982
            }
        }

983
984
        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
985
        });
Shucai Xiao's avatar
Shucai Xiao committed
986
987
        if(name_it != vec_names.end())
        {
988
989
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
990

Shucai Xiao's avatar
Shucai Xiao committed
991
992
993
        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
994
        });
995
996
997
998
999
1000
1001
1002

        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
1003
        if(contains(attributes, "linear_before_reset"))
1004
1005
1006
1007
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1008
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1009
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1010
1011
1012
1013
1014
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1015
1016
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1017
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1018
            std::move(args));
1019
1020

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

Shucai Xiao's avatar
Shucai Xiao committed
1023
        return {hidden_states, last_output};
1024
1025
    }

Shucai Xiao's avatar
Shucai Xiao committed
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
    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
1048
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1049
1050
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1051
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1052
1053
1054
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1055
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1056
        }
Shucai Xiao's avatar
Shucai Xiao committed
1057
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1058
        {
Shucai Xiao's avatar
Shucai Xiao committed
1059
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
        }
        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());
1072
            std::copy(names.begin(), names.end(), vec_names.begin());
Shucai Xiao's avatar
Shucai Xiao committed
1073
1074
1075
        }

        // need 6 activation functions for bidirectional directions
Shucai Xiao's avatar
Shucai Xiao committed
1076
        if(dirct == op::rnn_direction::bidirectional)
Shucai Xiao's avatar
Shucai Xiao committed
1077
1078
1079
1080
1081
1082
        {
            // 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
1083
            // if 3 actv funcs are provide, repeat all three once.
Shucai Xiao's avatar
Shucai Xiao committed
1084
1085
1086
1087
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(vec_names.size())
            {
1088
            case 1:
Shucai Xiao's avatar
Shucai Xiao committed
1089
1090
1091
1092
1093
1094
                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)};
1095
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1096
1097
1098

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1099
1100
1101
1102
1103
1104
                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
1105
1106
1107
1108
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1109
1110
1111
1112
1113
1114
                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
1115
1116
                break;

Shucai Xiao's avatar
Shucai Xiao committed
1117
1118
1119
1120
1121
1122
1123
            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)};
1124
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1125

Shucai Xiao's avatar
Shucai Xiao committed
1126
1127
1128
1129
1130
1131
1132
            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)};
1133
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1134

Shucai Xiao's avatar
Shucai Xiao committed
1135
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1136
1137
1138
1139
1140
1141
            }
        }
        else
        {
            switch(vec_names.size())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1142
            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
Shucai Xiao's avatar
Shucai Xiao committed
1143
1144
1145

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

Shucai Xiao's avatar
Shucai Xiao committed
1149
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1150
1151
1152
            }
        }

1153
1154
1155
        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
1156
        if(name_it != vec_names.end())
1157
1158
1159
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181

        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
1182
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1183
1184
1185
1186
        }

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

        // second output for last lstm output
Shucai Xiao's avatar
Shucai Xiao committed
1190
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1191
1192
1193
1194
1195
1196
1197

        // 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
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
    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
1210
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
1211
1212
1213
1214
1215
1216
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
1217
1218
1219
1220
1221
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
1222
1223
1224
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
            // 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
1237
        }
Paul's avatar
Paul committed
1238
        for(auto&& output : graph.output())
Paul's avatar
Paul committed
1239
        {
Paul's avatar
Paul committed
1240
            this->parse_node(output.name());
Paul's avatar
Paul committed
1241
1242
1243
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1244
    void parse_undefined(const std::string& name)
1245
    {
Shucai Xiao's avatar
Shucai Xiao committed
1246
        auto ins           = prog.add_instruction(op::undefined{});
1247
1248
1249
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
1250
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
1251
    {
Paul's avatar
Paul committed
1252
        if(name.empty())
Paul's avatar
Paul committed
1253
            MIGRAPHX_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
1254
1255
1256
1257
1258
1259
1260
1261
        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
1262
1263
                    assert(name != input);
                    this->parse_node(input);
Paul's avatar
Paul committed
1264
                }
Shucai Xiao's avatar
Shucai Xiao committed
1265
                else if(input.empty())
Paul's avatar
Paul committed
1266
                {
1267
                    this->parse_undefined(input);
Paul's avatar
Paul committed
1268
                }
1269
                args.push_back(instructions.at(input));
Paul's avatar
Paul committed
1270
            }
Paul's avatar
Paul committed
1271
            std::vector<instruction_ref> result;
Paul's avatar
Paul committed
1272
1273
            if(ops.count(node.op_type()) == 0)
            {
1274
                result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
Paul's avatar
Paul committed
1275
1276
1277
            }
            else
            {
Paul's avatar
Paul committed
1278
                result = ops[node.op_type()](get_attributes(node), args);
Paul's avatar
Paul committed
1279
            }
Paul's avatar
Paul committed
1280
            // Even no output nodes produce output in migraphx
Paul's avatar
Paul committed
1281
            if(node.output().empty() and result.size() == 1)
Paul's avatar
Paul committed
1282
1283
            {
                instructions[name] = result.front();
Paul's avatar
Paul committed
1284
1285
1286
            }
            else
            {
Paul's avatar
Paul committed
1287
1288
1289
1290
1291
1292
                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
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
            }
        }
    }

    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
1310
        std::size_t n = 0;
Paul's avatar
Paul committed
1311
1312
        for(auto&& node : graph.node())
        {
Paul's avatar
Paul committed
1313
            if(node.output().empty())
Paul's avatar
Paul committed
1314
            {
Paul's avatar
Paul committed
1315
                if(node.name().empty())
Paul's avatar
Paul committed
1316
1317
1318
1319
1320
1321
1322
1323
1324
                {
                    result["migraphx_unamed_node_" + std::to_string(n)] = node;
                    n++;
                }
                else
                {
                    result[node.name()] = node;
                }
            }
Paul's avatar
Paul committed
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
            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
1350
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
1351
1352
1353
1354
1355
        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
1356
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
1357
1358
1359
1360
1361
    }

    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
1362
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1363
        if(dims.empty())
Khalique's avatar
Khalique committed
1364
1365
1366
        {
            dims = {1};
        }
1367
1368
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
1369
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
            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
1382
            case onnx::TensorProto::FLOAT16: return literal{{shape::half_type, dims}, s.data()};
Scott Thornton's avatar
Scott Thornton committed
1383
1384
1385
1386
1387
1388
            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
1389
            MIGRAPHX_THROW("Invalid tensor type");
1390
        }
Paul's avatar
Paul committed
1391
1392
1393
1394
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Paul's avatar
Paul committed
1395
            return literal{{shape::float_type, dims}, t.float_data().begin(), t.float_data().end()};
Paul's avatar
Paul committed
1396
1397
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Paul's avatar
Paul committed
1398
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1399
        case onnx::TensorProto::UINT16:
Paul's avatar
Paul committed
1400
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1401
        case onnx::TensorProto::INT16:
Paul's avatar
Paul committed
1402
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1403
        case onnx::TensorProto::INT32:
Paul's avatar
Paul committed
1404
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1405
        case onnx::TensorProto::INT64:
Paul's avatar
Paul committed
1406
            return literal{{shape::int64_type, dims}, t.int64_data().begin(), t.int64_data().end()};
Paul's avatar
Paul committed
1407
1408
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Paul's avatar
Paul committed
1409
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
1410
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
1411
        {
Khalique's avatar
Khalique committed
1412
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
1413
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
1414
1415
1416
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
1417
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
1418
            return literal{{shape::half_type, dims}, data_half.begin(), data_half.end()};
Khalique's avatar
Khalique committed
1419
        }
Paul's avatar
Paul committed
1420
1421
1422
1423
1424
1425
1426
1427
        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
1428
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
    }

    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
1450
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1451
1452
1453
1454
1455
1456
1457
1458
1459
        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
1460
        auto&& tensor_dims = t.tensor_type().shape().dim();
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
        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
1472
1473
        return {shape_type, dims};
    }
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495

    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
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
};

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