onnx.cpp 61 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

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

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

    void init_actv_func()
    {
106
107
108
109
110
111
        // 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
112
113
114
115
    }

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

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

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

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

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

        return out_lens;
    }

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
267
268
269
    instruction_ref parse_logsoftmax(const std::string&,
                                     const attribute_map& attributes,
                                     std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
270
271
272
273
274
275
276
277
278
279
    {
        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));
    }

280
    instruction_ref parse_argmax(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
281
282
                                 const attribute_map& attributes,
                                 std::vector<instruction_ref> args)
283
284
285
286
287
288
289
    {
        int axis = 0;
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
290
        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
291
        if(contains(attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
292
293
294
295
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
296
        if(keep_dims == 0)
297
298
299
300
301
302
303
304
        {
            auto ins = prog.add_instruction(op::argmax{axis}, std::move(args));
            return prog.add_instruction(op::squeeze{{static_cast<int64_t>(axis)}}, ins);
        }
        else
        {
            return prog.add_instruction(op::argmax{axis}, std::move(args));
        }
305
306
307
    }

    instruction_ref parse_argmin(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
308
309
                                 const attribute_map& attributes,
                                 std::vector<instruction_ref> args)
310
311
312
313
314
315
316
    {
        int axis = 0;
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
317
        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
318
        if(contains(attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
319
320
321
322
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
323
        if(keep_dims == 0)
324
325
326
327
328
329
330
331
        {
            auto ins = prog.add_instruction(op::argmin{axis}, std::move(args));
            return prog.add_instruction(op::squeeze{{static_cast<int64_t>(axis)}}, ins);
        }
        else
        {
            return prog.add_instruction(op::argmin{axis}, std::move(args));
        }
332
333
    }

Paul's avatar
Paul committed
334
    instruction_ref
Paul's avatar
Paul committed
335
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
336
    {
337
        op::convolution op;
338
        auto l0 = args[0];
Paul's avatar
Paul committed
339
340
        if(contains(attributes, "pads"))
        {
Scott Thornton's avatar
Scott Thornton committed
341
            if(contains(attributes, "auto_pad"))
342
            {
Paul's avatar
Paul committed
343
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
344
            }
345
346
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
347
            if(padding.size() != 4)
348
            {
Paul's avatar
Paul committed
349
                MIGRAPHX_THROW("padding should have 4 values");
350
            }
Scott Thornton's avatar
Scott Thornton committed
351
            if(padding[0] != padding[2] || padding[1] != padding[3])
352
            {
353
354
                // insert zeros for pad op (args[0] has 4 dims)
                padding = {0, 0, padding[0], padding[1], 0, 0, padding[2], padding[3]};
Khalique's avatar
Khalique committed
355
                l0      = prog.add_instruction(op::pad{padding}, l0);
356
            }
357
358
359
360
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
361
            }
Paul's avatar
Paul committed
362
        }
Paul's avatar
Paul committed
363
364
365
366
367
368
369
370
        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
371
        if(contains(attributes, "auto_pad"))
372
373
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
374
            if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
375
            {
Paul's avatar
Paul committed
376
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
377
378
            }

wsttiger's avatar
fixes  
wsttiger committed
379
            if(s.find("SAME") != std::string::npos)
380
            {
381
                op.padding_mode = op::padding_mode_t::same;
382
383
            }
        }
Khalique's avatar
Khalique committed
384
385
386
387
        if(contains(attributes, "group"))
        {
            op.group = parse_value(attributes.at("group")).at<int>();
        }
Paul's avatar
Paul committed
388
389
390
391
        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
392
            auto l2 = prog.add_instruction(op::broadcast{axis, l1->get_shape().lens()}, args[2]);
393
            return prog.add_instruction(op::add{}, l1, l2);
Paul's avatar
Paul committed
394
        }
395
        return prog.add_instruction(op, l0, args[1]);
Paul's avatar
Paul committed
396
    }
Paul's avatar
Paul committed
397

Paul's avatar
Paul committed
398
399
400
    instruction_ref parse_pooling(const std::string& name,
                                  attribute_map attributes,
                                  std::vector<instruction_ref> args)
Paul's avatar
Paul committed
401
    {
Khalique's avatar
Khalique committed
402
        op::pooling op{ends_with(name, "MaxPool") ? "max" : "average"};
403
        auto l0 = args[0];
Khalique's avatar
Khalique committed
404
        if(starts_with(name, "Global"))
405
        {
Khalique's avatar
Khalique committed
406
407
            auto lens  = args.front()->get_shape().lens();
            op.lengths = {lens[2], lens[3]};
408
        }
Paul's avatar
Paul committed
409
410
        if(contains(attributes, "pads"))
        {
411
412
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
413
            if(padding.size() != 4)
414
            {
Paul's avatar
Paul committed
415
                MIGRAPHX_THROW("padding should have 4 values");
416
            }
Scott Thornton's avatar
Scott Thornton committed
417
            if(padding[0] != padding[2] || padding[1] != padding[3])
418
            {
419
420
                // 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
421
422
                l0 = prog.add_instruction(op::pad{padding, std::numeric_limits<float>::lowest()},
                                          l0);
423
424
425
426
427
            }
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
428
            }
Paul's avatar
Paul committed
429
430
431
432
433
434
435
436
437
        }
        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
438
        if(contains(attributes, "auto_pad"))
439
440
        {
            auto s = attributes["auto_pad"].s();
441
            if(s.find("SAME_UPPER") == std::string::npos)
442
            {
443
                MIGRAPHX_THROW("auto_pad only supports SAME_UPPER for pooling");
444
            }
445
            op.padding_mode = op::padding_mode_t::same;
446
447
        }

448
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
449
450
    }

Paul's avatar
Paul committed
451
    instruction_ref
Paul's avatar
Paul committed
452
    parse_reshape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
453
    {
454
        op::reshape op;
Paul's avatar
Paul committed
455
456
457
458
459
460
461
        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
462
            auto s = args[1]->eval();
Paul's avatar
Paul committed
463
            if(s.empty())
Paul's avatar
Paul committed
464
                MIGRAPHX_THROW("Dynamic shape is not supported.");
Paul's avatar
Paul committed
465
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
466
        }
Paul's avatar
Paul committed
467
468
469
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
470
    instruction_ref
Paul's avatar
Paul committed
471
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
472
    {
473
        uint64_t axis = 1;
Paul's avatar
Paul committed
474
475
476
477
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
478
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
479
480
    }

481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
    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
499
500
501
502
503
504
505
    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));
    }
506

507
508
509
    instruction_ref
    parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
510
        int axis = 0;
511
512
513
514
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
515
        op::gather op{axis};
516
517
518
        return prog.add_instruction(op, std::move(args));
    }

519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
    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
539
540
541
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
542
    {
Shucai Xiao's avatar
Shucai Xiao committed
543
        literal v     = parse_value(attributes.at("value"));
544
545
546
        auto dim_size = attributes.at("value").t().dims_size();
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
547
        {
548
            migraphx::shape scalar_shape{v.get_shape().type()};
549
550
551
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

Paul's avatar
Paul committed
552
553
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
554

Paul's avatar
Paul committed
555
    instruction_ref
Paul's avatar
Paul committed
556
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
557
558
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
559
        float beta  = 1.0f;
Paul's avatar
Paul committed
560
561
562
563
564
565
566
567
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
568
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
569
570
571
572
573
574
575
576
577
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
578
579
580
581
582
583

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

584
585
        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
586
587
        if(args.size() == 3)
        {
588
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
589
            {
Shucai Xiao's avatar
Shucai Xiao committed
590
                auto out_lens   = l1->get_shape().lens();
591
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
592
                auto l3         = args[2];
Shucai Xiao's avatar
Shucai Xiao committed
593
594
                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
595
                {
596
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
597
                }
598
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
599
            }
Paul's avatar
Paul committed
600
        }
601
602

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

605
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
606
    parse_matmul(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
607
    {
Shucai Xiao's avatar
Shucai Xiao committed
608
609
        auto l0      = args[0];
        auto l1      = args[1];
610
611
612
613
614
        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
615
        if(l0_lens.size() == 1)
616
617
618
619
620
621
622
        {
            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
623
        if(l1_lens.size() == 1)
624
625
626
627
628
629
630
631
        {
            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
632
        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
633
634
635
636
637
638
        {
            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);
639
            l0_broadcasted_lens = output_lens;
640
            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
641
            l1_broadcasted_lens = output_lens;
642
            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
Shucai Xiao's avatar
Shucai Xiao committed
643
            if(l0_lens != l0_broadcasted_lens)
644
645
646
            {
                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
            }
Shucai Xiao's avatar
Shucai Xiao committed
647
            if(l1_lens != l1_broadcasted_lens)
648
649
650
651
652
            {
                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
653
        auto dot_res     = prog.add_instruction(op::dot{1.0f, 0.0f}, bl0, bl1);
654
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
655
        if(is_a_prepended)
656
657
658
659
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
660
        if(is_b_appended)
661
662
663
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
664

665
666
667
        return dot_res;
    }

668
    instruction_ref
Paul's avatar
Paul committed
669
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
670
    {
Scott Thornton's avatar
Scott Thornton committed
671
672
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
673
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Scott Thornton's avatar
Scott Thornton committed
674
        bool is_test                                      = false;
675
676
677
678
679
680
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
681
            momentum = parse_value(attributes.at("momentum")).at<float>();
682
683
684
        }
        if(contains(attributes, "is_test"))
        {
wsttiger's avatar
wsttiger committed
685
            is_test = parse_value(attributes.at("is_test")).at<uint64_t>() > 0;
686
687
688
        }
        if(contains(attributes, "spatial"))
        {
689
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
690
691
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
692
        }
Paul's avatar
Paul committed
693
        (void)is_test;
Paul's avatar
Paul committed
694
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
695
        return prog.add_instruction(op, std::move(args));
696
697
    }

698
699
700
701
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
702
        float alpha = 0.01; // default alpha val for leaky relu
703
704
705
706
707
708
709
710
        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
711
712
    instruction_ref
    parse_elu(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
713
714
715
716
717
718
719
720
721
722
    {
        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
723
724
    instruction_ref
    parse_lrn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
725
726
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
727
728
729
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
Khalique's avatar
Khalique committed
730
731
732
733
734
735
736
737
738
739
740
741
        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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
    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());
        }
758
        auto input_lens = args.front()->get_shape().lens();
Khalique's avatar
Khalique committed
759

Khalique's avatar
Khalique committed
760
761
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
762
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
763

764
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
765
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
766
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
767
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
768
    }
Khalique's avatar
Khalique committed
769

Khalique's avatar
Khalique committed
770
771
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
772
773
774
775
776
777
778
    {
        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
779
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
780
781
    }

Khalique's avatar
Khalique committed
782
783
784
785
786
787
788
789
790
791
    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());
        }
792
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
793
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
794
795
796
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
Khalique's avatar
Khalique committed
797
798
799
800
801
802
803
804
805
806
807
808
        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());
    }
809
810
811
    // 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
812
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
813
814
    {
        if(args.size() != 1)
815
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
        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
852
853
        if(contains(attributes, "extra_shape"))
        {
854
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
855
856
        }

857
858
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
859
            if(args.size() != 1)
860
            {
861
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
862
863
            }

Shucai Xiao's avatar
Shucai Xiao committed
864
865
            if(contains(attributes, "shape"))
            {
866
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
867
                               "at the same time");
868
869
            }

870
871
872
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
873
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
874
            }
875

876
877
878
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
879
880
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
881
882
883
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
884
885
            if(!contains(attributes, "shape"))
            {
886
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
887
888
889
            }

            literal ls = parse_value(attributes.at("shape"));
890
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
891
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
892
            migraphx::shape s{type, dims};
893
894
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
895
896
897
        }
        else
        {
898
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
899
900
901
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
902
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
903
904
905
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
906
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
907
908
909

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
910
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
911
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
912
913
914
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
915
916
917
918
919
920
921
922
923
        }

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

924
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
925
926
        if(direction == "bidirectional")
        {
927
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
928
929
930
        }
        else if(direction == "reverse")
        {
931
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
932
933
        }

934
        std::vector<std::string> vec_names{"tanh"};
935
936
937
938
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
939
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
940
941
942
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
943
944
        }

945
946
947
        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
948
        if(name_it != vec_names.end())
949
950
951
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
952

Shucai Xiao's avatar
Shucai Xiao committed
953
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
954
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
955
        // if only one actv function is provided, we use it in both
956
        // forward and reverse direction
957
        if(dirct == op::rnn_direction::bidirectional)
958
        {
Shucai Xiao's avatar
Shucai Xiao committed
959
            if(vec_names.size() == 1)
960
961
962
963
964
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
965
966
967
        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];
968
        });
Shucai Xiao's avatar
Shucai Xiao committed
969

Shucai Xiao's avatar
Shucai Xiao committed
970
971
972
973
974
975
976
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

977
978
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
979
        if(args.size() < 6)
980
981
982
983
984
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
985
986
        // 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
987
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
988

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

Shucai Xiao's avatar
Shucai Xiao committed
992
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
993
994
    }

995
    std::vector<instruction_ref>
996
997
998
999
1000
1001
1002
    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
1003
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1004
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1005
1006
1007
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
1008
1009
1010
1011
1012
1013
1014
1015
1016
        }

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

1017
        op::rnn_direction dirct = op::rnn_direction::forward;
1018
1019
        if(direction == "bidirectional")
        {
1020
            dirct = op::rnn_direction::bidirectional;
1021
1022
1023
        }
        else if(direction == "reverse")
        {
1024
            dirct = op::rnn_direction::reverse;
1025
1026
        }

1027
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
1028
1029
        if(contains(attributes, "activations"))
        {
1030
            auto names = attributes.at("activations").strings();
1031
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
1032
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1033
1034
1035
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1036
1037
        }

1038
        // need 4 activation functions
1039
        if(dirct == op::rnn_direction::bidirectional)
1040
        {
Shucai Xiao's avatar
Shucai Xiao committed
1041
            // 4 activation functions are used in the bidirectional
1042
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
1043
1044
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
1045
1046
            // 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
1047
1048
1049
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
1050
            if(vec_names.size() == 1)
1051
            {
1052
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
1053
            }
1054
            else if(vec_names.size() == 2)
1055
            {
1056
1057
1058
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
1059
            }
1060
            else if(vec_names.size() == 3)
1061
            {
1062
                vec_names.push_back(vec_names.at(2));
1063
1064
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
1065
        else
1066
        {
1067
            if(vec_names.size() == 1)
1068
            {
1069
                vec_names.push_back(vec_names.at(0));
1070
1071
1072
            }
        }

1073
1074
1075
        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
1076
        if(name_it != vec_names.end())
1077
1078
1079
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1080

Shucai Xiao's avatar
Shucai Xiao committed
1081
1082
1083
        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
1084
        });
1085
1086
1087
1088
1089
1090
1091
1092

        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
1093
        if(contains(attributes, "linear_before_reset"))
1094
1095
1096
1097
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1098
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1099
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1100
1101
1102
1103
1104
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1105
1106
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1107
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1108
            std::move(args));
1109
1110

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

Shucai Xiao's avatar
Shucai Xiao committed
1113
        return {hidden_states, last_output};
1114
1115
    }

Shucai Xiao's avatar
Shucai Xiao committed
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
    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
1138
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1139
1140
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1141
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1142
1143
1144
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1145
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1146
        }
Shucai Xiao's avatar
Shucai Xiao committed
1147
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1148
        {
Shucai Xiao's avatar
Shucai Xiao committed
1149
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1150
1151
1152
1153
1154
1155
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1156
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
Shucai Xiao's avatar
Shucai Xiao committed
1157
1158
1159
1160
1161
        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
1162
1163
1164
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1165
1166
1167
        }

        // need 6 activation functions for bidirectional directions
Shucai Xiao's avatar
Shucai Xiao committed
1168
        if(dirct == op::rnn_direction::bidirectional)
Shucai Xiao's avatar
Shucai Xiao committed
1169
1170
1171
1172
1173
1174
        {
            // 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
1175
            // if 3 actv funcs are provide, repeat all three once.
Shucai Xiao's avatar
Shucai Xiao committed
1176
1177
1178
1179
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(vec_names.size())
            {
1180
            case 1:
Shucai Xiao's avatar
Shucai Xiao committed
1181
1182
1183
1184
1185
1186
                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)};
1187
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1188
1189
1190

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1191
1192
1193
1194
1195
1196
                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
1197
1198
1199
1200
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1201
1202
1203
1204
1205
1206
                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
1207
1208
                break;

Shucai Xiao's avatar
Shucai Xiao committed
1209
1210
1211
1212
1213
1214
1215
            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)};
1216
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1217

Shucai Xiao's avatar
Shucai Xiao committed
1218
1219
1220
1221
1222
1223
1224
            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)};
1225
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1226

Shucai Xiao's avatar
Shucai Xiao committed
1227
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1228
1229
1230
1231
1232
1233
            }
        }
        else
        {
            switch(vec_names.size())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1234
            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
Shucai Xiao's avatar
Shucai Xiao committed
1235
1236
1237

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

Shucai Xiao's avatar
Shucai Xiao committed
1241
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1242
1243
1244
            }
        }

1245
1246
1247
        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
1248
        if(name_it != vec_names.end())
1249
1250
1251
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273

        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
1274
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1275
1276
1277
1278
        }

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

        // second output for last lstm output
Shucai Xiao's avatar
Shucai Xiao committed
1282
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1283
1284
1285
1286
1287
1288

        // 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};
    }
1289

Shucai Xiao's avatar
Shucai Xiao committed
1290
1291
1292
    instruction_ref parse_reduce_sum(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
    {
        std::size_t n_dim = args.front()->get_shape().lens().size();

        // default to reduce over all dimensions
        std::vector<std::size_t> axes(n_dim);
        std::iota(axes.begin(), axes.end(), 0);
        if(contains(attributes, "axes"))
        {
            axes.clear();
            auto&& attr_axes = attributes["axes"].ints();
Shucai Xiao's avatar
Shucai Xiao committed
1303
            axes             = std::vector<std::size_t>(attr_axes.begin(), attr_axes.end());
1304
1305
1306
        }

        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
1307
        if(contains(attributes, "keepdims"))
1308
1309
1310
1311
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1312
        if(keep_dims == 1)
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
        {
            return prog.add_instruction(op::reduce_sum{axes}, std::move(args));
        }
        else
        {
            auto ins = prog.add_instruction(op::reduce_sum{axes}, std::move(args));
            std::vector<int64_t> squeeze_axes{axes.begin(), axes.end()};
            return prog.add_instruction(op::squeeze{squeeze_axes}, ins);
        }
    }
Shucai Xiao's avatar
Shucai Xiao committed
1323

Paul's avatar
Paul committed
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
    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
1336
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
1337
1338
1339
1340
1341
1342
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
1343
1344
1345
1346
1347
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
1348
1349
1350
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
            // 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
1363
        }
Paul's avatar
Paul committed
1364
        for(auto&& output : graph.output())
Paul's avatar
Paul committed
1365
        {
Paul's avatar
Paul committed
1366
            this->parse_node(output.name());
Paul's avatar
Paul committed
1367
1368
1369
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1370
    void parse_undefined(const std::string& name)
1371
    {
Shucai Xiao's avatar
Shucai Xiao committed
1372
        auto ins           = prog.add_instruction(op::undefined{});
1373
1374
1375
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
1376
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
1377
    {
Paul's avatar
Paul committed
1378
        if(name.empty())
Paul's avatar
Paul committed
1379
            MIGRAPHX_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
1380
1381
1382
1383
1384
1385
1386
1387
        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
1388
1389
                    assert(name != input);
                    this->parse_node(input);
Paul's avatar
Paul committed
1390
                }
Shucai Xiao's avatar
Shucai Xiao committed
1391
                else if(input.empty())
Paul's avatar
Paul committed
1392
                {
1393
                    this->parse_undefined(input);
Paul's avatar
Paul committed
1394
                }
1395
                args.push_back(instructions.at(input));
Paul's avatar
Paul committed
1396
            }
Paul's avatar
Paul committed
1397
            std::vector<instruction_ref> result;
Paul's avatar
Paul committed
1398
1399
            if(ops.count(node.op_type()) == 0)
            {
1400
                result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
Paul's avatar
Paul committed
1401
1402
1403
            }
            else
            {
Paul's avatar
Paul committed
1404
                result = ops[node.op_type()](get_attributes(node), args);
Paul's avatar
Paul committed
1405
            }
Paul's avatar
Paul committed
1406
            // Even no output nodes produce output in migraphx
Paul's avatar
Paul committed
1407
            if(node.output().empty() and result.size() == 1)
Paul's avatar
Paul committed
1408
1409
            {
                instructions[name] = result.front();
Paul's avatar
Paul committed
1410
1411
1412
            }
            else
            {
Paul's avatar
Paul committed
1413
1414
1415
1416
1417
1418
                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
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
            }
        }
    }

    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
1436
        std::size_t n = 0;
Paul's avatar
Paul committed
1437
1438
        for(auto&& node : graph.node())
        {
Paul's avatar
Paul committed
1439
            if(node.output().empty())
Paul's avatar
Paul committed
1440
            {
Paul's avatar
Paul committed
1441
                if(node.name().empty())
Paul's avatar
Paul committed
1442
1443
1444
1445
1446
1447
1448
1449
1450
                {
                    result["migraphx_unamed_node_" + std::to_string(n)] = node;
                    n++;
                }
                else
                {
                    result[node.name()] = node;
                }
            }
Paul's avatar
Paul committed
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
            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
1476
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
1477
1478
1479
1480
1481
        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
1482
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
1483
1484
1485
1486
1487
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
1488
1489
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
1490
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
1491
1492
1493
            switch(t.data_type())
            {
            case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
1494
            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
Scott Thornton's avatar
Scott Thornton committed
1495
            case onnx::TensorProto::UINT8: throw std::runtime_error("");
1496
            case onnx::TensorProto::INT8: return create_literal(shape::int32_type, dims, s.data());
Khalique's avatar
Khalique committed
1497
1498
            case onnx::TensorProto::UINT16:
                return create_literal(shape::int32_type, dims, s.data());
1499
1500
1501
            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
1502
            case onnx::TensorProto::STRING: throw std::runtime_error("");
1503
            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
Khalique's avatar
Khalique committed
1504
1505
1506
1507
            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
1508
1509
1510
1511
1512
            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
1513
            MIGRAPHX_THROW("Invalid tensor type");
1514
        }
Paul's avatar
Paul committed
1515
1516
1517
1518
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Khalique's avatar
Khalique committed
1519
            return create_literal(shape::float_type, dims, t.float_data());
Paul's avatar
Paul committed
1520
1521
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Khalique's avatar
Khalique committed
1522
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1523
        case onnx::TensorProto::UINT16:
Khalique's avatar
Khalique committed
1524
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1525
        case onnx::TensorProto::INT16:
Khalique's avatar
Khalique committed
1526
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1527
        case onnx::TensorProto::INT32:
Khalique's avatar
Khalique committed
1528
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1529
        case onnx::TensorProto::INT64:
Khalique's avatar
Khalique committed
1530
            return create_literal(shape::int64_type, dims, t.int64_data());
Paul's avatar
Paul committed
1531
1532
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Khalique's avatar
Khalique committed
1533
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1534
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
1535
        {
Khalique's avatar
Khalique committed
1536
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
1537
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
1538
1539
1540
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
1541
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
1542
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
1543
        }
Paul's avatar
Paul committed
1544
        case onnx::TensorProto::DOUBLE:
Khalique's avatar
Khalique committed
1545
            return create_literal(shape::double_type, dims, t.double_data());
Paul's avatar
Paul committed
1546
1547
1548
1549
1550
        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
1551
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
1552
1553
    }

Khalique's avatar
Khalique committed
1554
    static literal
1555
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
1556
    {
Khalique's avatar
Khalique committed
1557
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1558
        if(dims.empty())
1559
            return literal{{shape_type}, data};
1560
1561
1562
        return literal{{shape_type, dims}, data};
    }

1563
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
1564
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
1565
1566
    {
        if(dims.empty())
1567
            return literal{{shape_type}, data.begin(), data.end()};
1568
        return literal{{shape_type, dims}, data.begin(), data.end()};
1569
1570
    }

Paul's avatar
Paul committed
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
    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
1590
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1591
1592
1593
1594
1595
1596
1597
1598
1599
        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
1600
        auto&& tensor_dims = t.tensor_type().shape().dim();
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
        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
1612
1613
        return {shape_type, dims};
    }
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635

    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
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
};

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