onnx.cpp 58.1 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
        add_generic_op("Exp", op::exp{});
Shucai Xiao's avatar
Shucai Xiao committed
43
        add_generic_op("Erf", op::erf{});
Shucai Xiao's avatar
Shucai Xiao committed
44
        add_generic_op("Log", op::log{});
Khalique's avatar
Khalique committed
45
46
        // disable dropout for inference
        add_generic_op("Dropout", op::identity{});
Khalique's avatar
Khalique committed
47
        add_generic_op("Identity", op::identity{});
Shucai Xiao's avatar
Shucai Xiao committed
48
49
50
        add_generic_op("Sin", op::sin{});
        add_generic_op("Cos", op::cos{});
        add_generic_op("Tan", op::tan{});
51
52
        add_generic_op("Sinh", op::sinh{});
        add_generic_op("Cosh", op::cosh{});
53
        add_generic_op("Tanh", op::tanh{});
54
55
56
        add_generic_op("Asin", op::asin{});
        add_generic_op("Acos", op::acos{});
        add_generic_op("Atan", op::atan{});
Paul's avatar
Paul committed
57

Khalique's avatar
Khalique committed
58
59
60
61
62
        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
63
64
65
        add_variadic_op("Sum", op::add{});
        add_variadic_op("Max", op::max{});
        add_variadic_op("Min", op::min{});
Paul's avatar
Paul committed
66

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

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

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

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

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

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

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

        std::vector<std::size_t> out_lens(s1);
        auto offset = s1.size() - s0.size();
Shucai Xiao's avatar
Shucai Xiao committed
182
183
184
185
        std::transform(s0.begin(),
                       s0.end(),
                       s1.begin() + offset,
                       out_lens.begin() + offset,
186
                       [&](auto a, auto b) {
Shucai Xiao's avatar
Shucai Xiao committed
187
                           if(a != b and a != 1 and b != 1)
188
                           {
Shucai Xiao's avatar
Shucai Xiao committed
189
190
191
192
193
194
                               MIGRAPHX_THROW("COMPUTE_BROADCASTLEN: shape {" +
                                              to_string_range(s0) + "} and {" +
                                              to_string_range(s1) + "} mismatch!");
                           }
                           return std::max(a, b);
                       });
195
196
197
198

        return out_lens;
    }

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

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

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

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

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

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

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

392
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
393
394
    }

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

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

425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
    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
443
444
445
446
447
448
449
    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));
    }
450

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

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

Paul's avatar
Paul committed
496
497
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
498

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

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

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

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

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

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

609
610
611
        return dot_res;
    }

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

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

Khalique's avatar
Khalique committed
704
705
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
706
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
707

708
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
709
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
710
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
711
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
712
    }
Khalique's avatar
Khalique committed
713

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

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

801
802
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
803
            if(args.size() != 1)
804
            {
805
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
806
807
            }

Shucai Xiao's avatar
Shucai Xiao committed
808
809
            if(contains(attributes, "shape"))
            {
810
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
811
                               "at the same time");
812
813
            }

814
815
816
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
817
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
818
            }
819

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

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

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

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

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

868
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
869
870
        if(direction == "bidirectional")
        {
871
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
872
873
874
        }
        else if(direction == "reverse")
        {
875
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
876
877
        }

878
        std::vector<std::string> vec_names{"tanh"};
879
880
881
882
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
883
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
884
885
886
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
887
888
        }

889
890
891
        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
892
        if(name_it != vec_names.end())
893
894
895
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
896

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

Shucai Xiao's avatar
Shucai Xiao committed
909
910
911
        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];
912
        });
Shucai Xiao's avatar
Shucai Xiao committed
913

Shucai Xiao's avatar
Shucai Xiao committed
914
915
916
917
918
919
920
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

921
922
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
923
        if(args.size() < 6)
924
925
926
927
928
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
929
930
        // 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
931
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
932

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

Shucai Xiao's avatar
Shucai Xiao committed
936
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
937
938
    }

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

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

961
        op::rnn_direction dirct = op::rnn_direction::forward;
962
963
        if(direction == "bidirectional")
        {
964
            dirct = op::rnn_direction::bidirectional;
965
966
967
        }
        else if(direction == "reverse")
        {
968
            dirct = op::rnn_direction::reverse;
969
970
        }

971
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
972
973
        if(contains(attributes, "activations"))
        {
974
            auto names = attributes.at("activations").strings();
975
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
976
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
977
978
979
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
980
981
        }

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

1017
1018
1019
        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
1020
        if(name_it != vec_names.end())
1021
1022
1023
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1024

Shucai Xiao's avatar
Shucai Xiao committed
1025
1026
1027
        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
1028
        });
1029
1030
1031
1032
1033
1034
1035
1036

        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
1037
        if(contains(attributes, "linear_before_reset"))
1038
1039
1040
1041
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1042
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1043
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1044
1045
1046
1047
1048
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1049
1050
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1051
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1052
            std::move(args));
1053
1054

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

Shucai Xiao's avatar
Shucai Xiao committed
1057
        return {hidden_states, last_output};
1058
1059
    }

Shucai Xiao's avatar
Shucai Xiao committed
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
    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
1082
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1083
1084
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1085
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1086
1087
1088
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1089
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1090
        }
Shucai Xiao's avatar
Shucai Xiao committed
1091
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1092
        {
Shucai Xiao's avatar
Shucai Xiao committed
1093
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1094
1095
1096
1097
1098
1099
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1100
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
Shucai Xiao's avatar
Shucai Xiao committed
1101
1102
1103
1104
1105
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1106
1107
1108
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1109
1110
1111
        }

        // need 6 activation functions for bidirectional directions
Shucai Xiao's avatar
Shucai Xiao committed
1112
        if(dirct == op::rnn_direction::bidirectional)
Shucai Xiao's avatar
Shucai Xiao committed
1113
1114
1115
1116
1117
1118
        {
            // 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
1119
            // if 3 actv funcs are provide, repeat all three once.
Shucai Xiao's avatar
Shucai Xiao committed
1120
1121
1122
1123
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(vec_names.size())
            {
1124
            case 1:
Shucai Xiao's avatar
Shucai Xiao committed
1125
1126
1127
1128
1129
1130
                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)};
1131
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1132
1133
1134

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1135
1136
1137
1138
1139
1140
                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
1141
1142
1143
1144
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1145
1146
1147
1148
1149
1150
                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
1151
1152
                break;

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

Shucai Xiao's avatar
Shucai Xiao committed
1162
1163
1164
1165
1166
1167
1168
            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)};
1169
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1170

Shucai Xiao's avatar
Shucai Xiao committed
1171
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1172
1173
1174
1175
1176
1177
            }
        }
        else
        {
            switch(vec_names.size())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1178
            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
Shucai Xiao's avatar
Shucai Xiao committed
1179
1180
1181

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

Shucai Xiao's avatar
Shucai Xiao committed
1185
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1186
1187
1188
            }
        }

1189
1190
1191
        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
1192
        if(name_it != vec_names.end())
1193
1194
1195
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217

        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
1218
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1219
1220
1221
1222
        }

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

        // second output for last lstm output
Shucai Xiao's avatar
Shucai Xiao committed
1226
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1227
1228
1229
1230
1231
1232
1233

        // 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
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
    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
1246
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
1247
1248
1249
1250
1251
1252
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
1253
1254
1255
1256
1257
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
1258
1259
1260
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
            // 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
1273
        }
Paul's avatar
Paul committed
1274
        for(auto&& output : graph.output())
Paul's avatar
Paul committed
1275
        {
Paul's avatar
Paul committed
1276
            this->parse_node(output.name());
Paul's avatar
Paul committed
1277
1278
1279
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1280
    void parse_undefined(const std::string& name)
1281
    {
Shucai Xiao's avatar
Shucai Xiao committed
1282
        auto ins           = prog.add_instruction(op::undefined{});
1283
1284
1285
        instructions[name] = ins;
    }

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

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

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

Khalique's avatar
Khalique committed
1464
    static literal
1465
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
1466
    {
Khalique's avatar
Khalique committed
1467
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1468
        if(dims.empty())
1469
            return literal{{shape_type}, data};
1470
1471
1472
        return literal{{shape_type, dims}, data};
    }

1473
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
1474
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
1475
1476
    {
        if(dims.empty())
1477
            return literal{{shape_type}, data.begin(), data.end()};
1478
        return literal{{shape_type, dims}, data.begin(), data.end()};
1479
1480
    }

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

    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
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
};

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