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

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

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

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

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

    onnx_parser()
    {
Khalique's avatar
Khalique committed
39
        add_generic_op("Relu", op::relu{});
Khalique's avatar
Khalique committed
40
41
        add_generic_op("Sigmoid", op::sigmoid{});
        add_generic_op("Abs", op::abs{});
Shucai Xiao's avatar
Shucai Xiao committed
42
43
        add_generic_op("Exp", op::exp{});
        add_generic_op("Log", op::log{});
Khalique's avatar
Khalique committed
44
45
        // disable dropout for inference
        add_generic_op("Dropout", op::identity{});
Khalique's avatar
Khalique committed
46
        add_generic_op("Identity", op::identity{});
Shucai Xiao's avatar
Shucai Xiao committed
47
48
49
        add_generic_op("Sin", op::sin{});
        add_generic_op("Cos", op::cos{});
        add_generic_op("Tan", op::tan{});
50
51
        add_generic_op("Sinh", op::sinh{});
        add_generic_op("Cosh", op::cosh{});
52
        add_generic_op("Tanh", op::tanh{});
53
54
55
        add_generic_op("Asin", op::asin{});
        add_generic_op("Acos", op::acos{});
        add_generic_op("Atan", op::atan{});
Paul's avatar
Paul committed
56

Khalique's avatar
Khalique committed
57
58
59
60
61
        add_binary_op("Add", op::add{});
        add_binary_op("Div", op::div{});
        add_binary_op("Mul", op::mul{});
        add_binary_op("Sub", op::sub{});

Khalique's avatar
Khalique committed
62
63
64
        add_variadic_op("Sum", op::add{});
        add_variadic_op("Max", op::max{});
        add_variadic_op("Min", op::min{});
Paul's avatar
Paul committed
65

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

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

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

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

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

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

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

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

        return out_lens;
    }

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

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

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

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

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

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

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

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

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

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

415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
    instruction_ref
    parse_squeeze(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::squeeze op;
        literal s = parse_value(attributes.at("axes"));
        s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        return prog.add_instruction(op, args[0]);
    }

    instruction_ref
    parse_unsqueeze(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::unsqueeze op;
        literal s = parse_value(attributes.at("axes"));
        s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        return prog.add_instruction(op, args[0]);
    }

Scott Thornton's avatar
Scott Thornton committed
433
434
435
436
437
438
439
    instruction_ref
    parse_concat(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        std::size_t axis = parse_value(attributes.at("axis")).at<int>();
        op::concat op{axis};
        return prog.add_instruction(op, std::move(args));
    }
440

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

453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
    instruction_ref
    parse_slice(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::slice op;
        if(contains(attributes, "axes"))
        {
            literal s = parse_value(attributes.at("axes"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        }
        {
            literal s = parse_value(attributes.at("ends"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.ends)); });
        }
        {
            literal s = parse_value(attributes.at("starts"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.starts)); });
        }
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
473
474
475
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
476
    {
Shucai Xiao's avatar
Shucai Xiao committed
477
        literal v     = parse_value(attributes.at("value"));
478
479
480
        auto dim_size = attributes.at("value").t().dims_size();
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
481
        {
482
            migraphx::shape scalar_shape{v.get_shape().type()};
483
484
485
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

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

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

        std::vector<int64_t> perm(args[0]->get_shape().lens().size());
        std::iota(perm.begin(), perm.end(), int64_t{0});
        // swap the last two elements
        std::swap(*perm.rbegin(), *(perm.rbegin() + 1));

518
519
        auto l1 = (transa) ? prog.add_instruction(op::transpose{perm}, args[0]) : args[0];
        auto l2 = (transb) ? prog.add_instruction(op::transpose{perm}, args[1]) : args[1];
Paul's avatar
Paul committed
520
521
        if(args.size() == 3)
        {
522
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
523
            {
Shucai Xiao's avatar
Shucai Xiao committed
524
                auto out_lens   = l1->get_shape().lens();
525
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
526
                auto l3         = args[2];
Shucai Xiao's avatar
Shucai Xiao committed
527
528
                auto l3_lens    = l3->get_shape().lens();
                if(!std::equal(out_lens.begin(), out_lens.end(), l3_lens.begin(), l3_lens.end()))
Khalique's avatar
Khalique committed
529
                {
530
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
531
                }
532
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
533
            }
Paul's avatar
Paul committed
534
        }
535
536

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

539
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
540
    parse_matmul(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
541
    {
Shucai Xiao's avatar
Shucai Xiao committed
542
543
        auto l0      = args[0];
        auto l1      = args[1];
544
545
546
547
548
        auto l0_lens = l0->get_shape().lens();
        auto l1_lens = l1->get_shape().lens();

        // args[0] is a vector, prepend 1 to the shape
        bool is_a_prepended = false;
Shucai Xiao's avatar
Shucai Xiao committed
549
        if(l0_lens.size() == 1)
550
551
552
553
554
555
556
        {
            is_a_prepended = true;
            l0_lens.insert(l0_lens.begin(), 1);
            l0 = prog.add_instruction(op::unsqueeze{{0}}, args[0]);
        }

        bool is_b_appended = false;
Shucai Xiao's avatar
Shucai Xiao committed
557
        if(l1_lens.size() == 1)
558
559
560
561
562
563
564
565
        {
            is_b_appended = true;
            l1_lens.push_back(1);
            l1 = prog.add_instruction(op::unsqueeze{{1}}, args[1]);
        }

        instruction_ref bl0 = l0;
        instruction_ref bl1 = l1;
Shucai Xiao's avatar
Shucai Xiao committed
566
        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
567
568
569
570
571
572
        {
            auto l0_it = l0_lens.begin() + l0_lens.size() - 2;
            std::vector<std::size_t> l0_broadcasted_lens(l0_lens.begin(), l0_it);
            auto l1_it = l1_lens.begin() + l1_lens.size() - 2;
            std::vector<std::size_t> l1_broadcasted_lens(l1_lens.begin(), l1_it);
            auto output_lens = compute_broadcasted_lens(l0_broadcasted_lens, l1_broadcasted_lens);
573
            l0_broadcasted_lens = output_lens;
574
            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
575
            l1_broadcasted_lens = output_lens;
576
            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
Shucai Xiao's avatar
Shucai Xiao committed
577
            if(l0_lens != l0_broadcasted_lens)
578
579
580
            {
                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
            }
Shucai Xiao's avatar
Shucai Xiao committed
581
            if(l1_lens != l1_broadcasted_lens)
582
583
584
585
586
            {
                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
            }
        }

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

599
600
601
        return dot_res;
    }

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

632
633
634
635
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
636
        float alpha = 0.01; // default alpha val for leaky relu
637
638
639
640
641
642
643
644
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        op::leaky_relu op{alpha};
        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
645
646
    instruction_ref
    parse_elu(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
647
648
649
650
651
652
653
654
655
656
    {
        float alpha = 1.0; // default alpha val for elu
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        op::elu op{alpha};
        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
657
658
    instruction_ref
    parse_lrn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
659
660
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
661
662
663
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
Khalique's avatar
Khalique committed
664
665
666
667
668
669
670
671
672
673
674
675
        if(contains(attributes, "alpha"))
            alpha = parse_value(attributes.at("alpha")).at<float>();
        if(contains(attributes, "beta"))
            beta = parse_value(attributes.at("beta")).at<float>();
        if(contains(attributes, "bias"))
            bias = parse_value(attributes.at("bias")).at<float>();
        if(contains(attributes, "size"))
            size = parse_value(attributes.at("size")).at<int>();
        op::lrn op{alpha, beta, bias, size};
        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
    instruction_ref parse_imagescaler(const std::string&,
                                      attribute_map attributes,
                                      std::vector<instruction_ref> args)
    {
        float scale = 1.0;
        std::vector<float> bias{};
        if(contains(attributes, "scale"))
        {
            scale = parse_value(attributes.at("scale")).at<float>();
        }

        if(contains(attributes, "bias"))
        {
            auto&& bias_floats = attributes["bias"].floats();
            bias               = std::vector<float>(bias_floats.begin(), bias_floats.end());
        }
692
        auto input_lens = args.front()->get_shape().lens();
Khalique's avatar
Khalique committed
693

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

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

Khalique's avatar
Khalique committed
704
705
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
706
707
708
709
710
711
712
    {
        std::vector<int64_t> perm{};
        if(contains(attributes, "perm"))
        {
            auto&& perm_vals = attributes["perm"].ints();
            perm             = std::vector<int64_t>(perm_vals.begin(), perm_vals.end());
        }
Paul's avatar
Paul committed
713
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
714
715
    }

Khalique's avatar
Khalique committed
716
717
718
719
720
721
722
723
724
725
    instruction_ref
    parse_pad(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        std::vector<int64_t> pads{};
        float value = 0.0f;
        if(contains(attributes, "pads"))
        {
            auto&& pad_vals = attributes["pads"].ints();
            pads            = std::vector<int64_t>(pad_vals.begin(), pad_vals.end());
        }
726
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
727
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
728
729
730
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
Khalique's avatar
Khalique committed
731
732
733
734
735
736
737
738
739
740
741
742
        if(contains(attributes, "value"))
        {
            value = parse_value(attributes.at("value")).at<float>();
        }
        if(contains(attributes, "mode"))
        {
            auto mode = attributes.at("mode").s();
            if(mode != "constant")
                MIGRAPHX_THROW("migraphx currently only supports constant padding");
        }
        return prog.add_instruction(migraphx::op::pad{pads, value}, args.front());
    }
743
744
745
    // Use a literal instruction to replace the shape since, output of
    // shape operator are literals in migraphx
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
746
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
747
748
    {
        if(args.size() != 1)
749
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
        std::vector<std::size_t> arg_shape = args[0]->get_shape().lens();
        std::vector<int64_t> vec_shape(arg_shape.size());
        migraphx::shape s(migraphx::shape::int64_type, {arg_shape.size()});
        std::transform(arg_shape.begin(), arg_shape.end(), vec_shape.begin(), [](auto i) {
            return int64_t(i);
        });
        return prog.add_literal(migraphx::literal{s, vec_shape});
    }

    // Use a literal instruction to replace the constantFill operator. In RNN, input shape
    // and value are fixed, so no need to do the actual computation for the constantFill
    // operator
    instruction_ref parse_constant_fill(const std::string&,
                                        attribute_map attributes,
                                        std::vector<instruction_ref> args)
    {
        int input_as_shape = 0;
        int dtype          = 1;
        float value        = 0.0f;

        if(contains(attributes, "dtype"))
        {
            dtype = parse_value(attributes.at("dtype")).at<int>();
        }
        migraphx::shape::type_t type = get_type(dtype);

        if(contains(attributes, "input_as_shape"))
        {
            input_as_shape = parse_value(attributes.at("input_as_shape")).at<int>();
        }

        if(contains(attributes, "value"))
        {
            value = parse_value(attributes.at("value")).at<float>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
786
787
        if(contains(attributes, "extra_shape"))
        {
788
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
789
790
        }

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

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

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

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

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

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

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

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

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

868
        std::vector<std::string> vec_names{"tanh"};
869
870
871
872
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
873
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
874
875
876
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
877
878
        }

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

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

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

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
926
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
927
928
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

1090
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
Shucai Xiao's avatar
Shucai Xiao committed
1091
1092
1093
1094
1095
        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
1096
1097
1098
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1099
1100
1101
        }

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

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

            case 3:
                // repeat all three actv funcs once
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(2),
                             vec_names.at(0),
                             vec_names.at(1),
                             vec_names.at(2)};
Shucai Xiao's avatar
Shucai Xiao committed
1141
1142
                break;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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