onnx.cpp 65.5 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
        add_binary_op("Add", op::add{});
        add_binary_op("Div", op::div{});
        add_binary_op("Mul", op::mul{});
        add_binary_op("Sub", op::sub{});
Shucai Xiao's avatar
Shucai Xiao committed
61
        add_binary_op("Pow", op::pow{});
Khalique's avatar
Khalique committed
62

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

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

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

    void init_actv_func()
    {
111
112
113
114
115
116
        // 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
117
118
119
120
    }

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

    template <class F>
    void add_mem_op(std::string name, F f)
    {
Paul's avatar
Paul committed
137
        add_op(name, [=](auto&&... xs) {
Paul's avatar
Paul committed
138
139
140
            return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
        });
    }
Khalique's avatar
Khalique committed
141

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

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

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

        return out_lens;
    }

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
272
273
274
    instruction_ref parse_logsoftmax(const std::string&,
                                     const attribute_map& attributes,
                                     std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
275
276
277
278
279
280
281
282
283
284
    {
        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));
    }

285
    instruction_ref parse_argmax(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
286
287
                                 const attribute_map& attributes,
                                 std::vector<instruction_ref> args)
288
    {
289
        int64_t axis = 0;
290
291
        if(contains(attributes, "axis"))
        {
292
            axis = static_cast<int64_t>(parse_value(attributes.at("axis")).at<int>());
293
294
        }

Shucai Xiao's avatar
Shucai Xiao committed
295
        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
296
        if(contains(attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
297
298
299
300
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
301
        if(keep_dims == 0)
302
303
        {
            auto ins = prog.add_instruction(op::argmax{axis}, std::move(args));
304
            return prog.add_instruction(op::squeeze{{axis}}, ins);
305
306
307
308
309
        }
        else
        {
            return prog.add_instruction(op::argmax{axis}, std::move(args));
        }
310
311
312
    }

    instruction_ref parse_argmin(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
313
314
                                 const attribute_map& attributes,
                                 std::vector<instruction_ref> args)
315
    {
316
        int64_t axis = 0;
317
318
        if(contains(attributes, "axis"))
        {
319
            axis = static_cast<int64_t>(parse_value(attributes.at("axis")).at<int>());
320
321
        }

Shucai Xiao's avatar
Shucai Xiao committed
322
        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
323
        if(contains(attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
324
325
326
327
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
328
        if(keep_dims == 0)
329
330
        {
            auto ins = prog.add_instruction(op::argmin{axis}, std::move(args));
331
            return prog.add_instruction(op::squeeze{{axis}}, ins);
332
333
334
335
336
        }
        else
        {
            return prog.add_instruction(op::argmin{axis}, std::move(args));
        }
337
338
    }

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

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

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

453
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
454
455
    }

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

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

486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
    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
504
505
506
507
508
509
510
    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));
    }
511

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

524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
    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
544
545
546
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
547
    {
Shucai Xiao's avatar
Shucai Xiao committed
548
        literal v = parse_value(attributes.at("value"));
549
        // return empty literal
Shucai Xiao's avatar
Shucai Xiao committed
550
        if(v.get_shape().elements() == 0)
551
552
553
554
        {
            return prog.add_literal(literal{});
        }

555
556
557
        auto dim_size = attributes.at("value").t().dims_size();
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
558
        {
559
            migraphx::shape scalar_shape{v.get_shape().type()};
560
561
562
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

Paul's avatar
Paul committed
563
564
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
565

Paul's avatar
Paul committed
566
    instruction_ref
Paul's avatar
Paul committed
567
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
568
569
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
570
        float beta  = 1.0f;
Paul's avatar
Paul committed
571
572
573
574
575
576
577
578
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
579
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
580
581
582
583
584
585
586
587
588
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
589
590
591
592
593
594

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

595
596
        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
597
598
        if(args.size() == 3)
        {
599
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
600
            {
Shucai Xiao's avatar
Shucai Xiao committed
601
                auto out_lens   = l1->get_shape().lens();
602
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
603
                auto l3         = args[2];
Shucai Xiao's avatar
Shucai Xiao committed
604
605
                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
606
                {
607
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
608
                }
609
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
610
            }
Paul's avatar
Paul committed
611
        }
612
613

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

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

Shucai Xiao's avatar
Shucai Xiao committed
664
        auto dot_res     = prog.add_instruction(op::dot{1.0f, 0.0f}, bl0, bl1);
665
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
666
        if(is_a_prepended)
667
668
669
670
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
671
        if(is_b_appended)
672
673
674
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
675

676
677
678
        return dot_res;
    }

679
    instruction_ref
Paul's avatar
Paul committed
680
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
681
    {
Scott Thornton's avatar
Scott Thornton committed
682
683
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
684
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Scott Thornton's avatar
Scott Thornton committed
685
        bool is_test                                      = false;
686
687
688
689
690
691
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
692
            momentum = parse_value(attributes.at("momentum")).at<float>();
693
694
695
        }
        if(contains(attributes, "is_test"))
        {
wsttiger's avatar
wsttiger committed
696
            is_test = parse_value(attributes.at("is_test")).at<uint64_t>() > 0;
697
698
699
        }
        if(contains(attributes, "spatial"))
        {
700
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
701
702
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
703
        }
Paul's avatar
Paul committed
704
        (void)is_test;
Paul's avatar
Paul committed
705
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
706
        return prog.add_instruction(op, std::move(args));
707
708
    }

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

Khalique's avatar
Khalique committed
771
772
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
773
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
774

775
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
776
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
777
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
778
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
779
    }
Khalique's avatar
Khalique committed
780

Khalique's avatar
Khalique committed
781
782
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
783
784
785
786
787
788
789
    {
        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
790
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
791
792
    }

Khalique's avatar
Khalique committed
793
794
795
796
797
798
799
800
801
802
    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());
        }
803
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
804
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
805
806
807
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
Khalique's avatar
Khalique committed
808
809
810
811
812
813
814
815
816
817
818
819
        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());
    }
820
821
822
    // 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
823
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
824
825
    {
        if(args.size() != 1)
826
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
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
852
853
854
855
856
857
858
859
860
861
862
        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
863
864
        if(contains(attributes, "extra_shape"))
        {
865
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
866
867
        }

868
869
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
870
            if(args.size() != 1)
871
            {
872
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
873
874
            }

Shucai Xiao's avatar
Shucai Xiao committed
875
876
            if(contains(attributes, "shape"))
            {
877
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
878
                               "at the same time");
879
880
            }

881
882
883
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
884
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
885
            }
886

887
888
889
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
890
891
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
892
893
894
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
895
896
            if(!contains(attributes, "shape"))
            {
897
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
898
899
900
            }

            literal ls = parse_value(attributes.at("shape"));
901
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
902
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
903
            migraphx::shape s{type, dims};
904
905
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
906
907
908
        }
        else
        {
909
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
910
911
912
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
913
914
915
    instruction_ref parse_constant_of_shape(const std::string&,
                                            attribute_map attributes,
                                            std::vector<instruction_ref> args)
916
917
    {
        literal l_val{};
Shucai Xiao's avatar
Shucai Xiao committed
918
        if(contains(attributes, "value"))
919
920
        {
            l_val = parse_value(attributes.at("value"));
Shucai Xiao's avatar
Shucai Xiao committed
921
            if(l_val.get_shape().elements() != 1)
922
923
924
925
926
927
928
929
930
931
932
            {
                MIGRAPHX_THROW("ConstantOfShape: attribute value can contain only 1 elements!");
            }
        }
        else
        {
            l_val = literal({shape::float_type, {1}, {0}}, {0.0f});
        }

        // input is empty, output is a scalar
        auto type = l_val.get_shape().type();
933

Shucai Xiao's avatar
Shucai Xiao committed
934
        if(args.empty())
935
        {
936
            MIGRAPHX_THROW("Parse ConstantOfShape : must have 1 input!");
937
938
939
        }
        else
        {
940
941
            migraphx::shape s;
            // empty input tensor, output is a scalar
Shucai Xiao's avatar
Shucai Xiao committed
942
            if(args[0]->get_shape().elements() == 0)
943
            {
944
                s = migraphx::shape{type, {1}, {0}};
945
            }
946
947
948
949
950
951
952
            else
            {
                migraphx::argument in = args[0]->eval();
                if(in.empty())
                {
                    MIGRAPHX_THROW("Parse ConstantOfShape: cannot handle dynamic shape as input");
                }
953

954
955
956
957
                std::vector<std::size_t> dims;
                in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
                s = migraphx::shape{type, dims};
            }
958
959
960
961
962
963
964
965
966
967
968
969
970

            literal l_out;
            l_val.visit([&](auto val) {
                using type = std::remove_cv_t<typename decltype(val)::value_type>;
                // l_val contains only one element
                std::vector<type> out_vec(s.elements(), *val.begin());
                l_out = literal(s, out_vec);
            });

            return prog.add_literal(l_out);
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
971
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
972
    parse_expand(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
973
    {
Shucai Xiao's avatar
Shucai Xiao committed
974
        auto in_lens             = args[0]->get_shape().lens();
975
        migraphx::argument arg_s = args[1]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
976
        if(arg_s.empty())
977
978
979
980
981
982
        {
            MIGRAPHX_THROW("Parse Expand: cannot handle dynamic shape as input");
        }
        std::vector<std::size_t> dims;
        arg_s.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
        auto out_lens = compute_broadcasted_lens(in_lens, dims);
983

Shucai Xiao's avatar
Shucai Xiao committed
984
        return prog.add_instruction(op::multibroadcast{out_lens}, args[0]);
985
986
    }

Shucai Xiao's avatar
Shucai Xiao committed
987
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
988
989
990
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
991
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
992
993
994

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
995
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
996
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
997
998
999
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
1000
1001
1002
1003
1004
1005
1006
1007
1008
        }

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

1009
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1010
1011
        if(direction == "bidirectional")
        {
1012
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1013
1014
1015
        }
        else if(direction == "reverse")
        {
1016
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1017
1018
        }

1019
        std::vector<std::string> vec_names{"tanh"};
1020
1021
1022
1023
        if(contains(attributes, "activations"))
        {
            auto names = attributes.at("activations").strings();
            vec_names.clear();
1024
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1025
1026
1027
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1028
1029
        }

1030
1031
1032
        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
1033
        if(name_it != vec_names.end())
1034
1035
1036
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
1037

Shucai Xiao's avatar
Shucai Xiao committed
1038
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
1039
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
1040
        // if only one actv function is provided, we use it in both
1041
        // forward and reverse direction
1042
        if(dirct == op::rnn_direction::bidirectional)
1043
        {
Shucai Xiao's avatar
Shucai Xiao committed
1044
            if(vec_names.size() == 1)
1045
1046
1047
1048
1049
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
1050
1051
1052
        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];
1053
        });
Shucai Xiao's avatar
Shucai Xiao committed
1054

Shucai Xiao's avatar
Shucai Xiao committed
1055
1056
1057
1058
1059
1060
1061
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

1062
1063
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1064
        if(args.size() < 6)
1065
1066
1067
1068
1069
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
1070
1071
        // 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
1072
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1073

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

Shucai Xiao's avatar
Shucai Xiao committed
1077
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
1078
1079
    }

1080
    std::vector<instruction_ref>
1081
1082
1083
1084
1085
1086
1087
    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
1088
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1089
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1090
1091
1092
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
1093
1094
1095
1096
1097
1098
1099
1100
1101
        }

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

1102
        op::rnn_direction dirct = op::rnn_direction::forward;
1103
1104
        if(direction == "bidirectional")
        {
1105
            dirct = op::rnn_direction::bidirectional;
1106
1107
1108
        }
        else if(direction == "reverse")
        {
1109
            dirct = op::rnn_direction::reverse;
1110
1111
        }

1112
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
1113
1114
        if(contains(attributes, "activations"))
        {
1115
            auto names = attributes.at("activations").strings();
1116
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
1117
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1118
1119
1120
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1121
1122
        }

1123
        // need 4 activation functions
1124
        if(dirct == op::rnn_direction::bidirectional)
1125
        {
Shucai Xiao's avatar
Shucai Xiao committed
1126
            // 4 activation functions are used in the bidirectional
1127
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
1128
1129
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
1130
1131
            // 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
1132
1133
1134
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
1135
            if(vec_names.size() == 1)
1136
            {
1137
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
1138
            }
1139
            else if(vec_names.size() == 2)
1140
            {
1141
1142
1143
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
1144
            }
1145
            else if(vec_names.size() == 3)
1146
            {
1147
                vec_names.push_back(vec_names.at(2));
1148
1149
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
1150
        else
1151
        {
1152
            if(vec_names.size() == 1)
1153
            {
1154
                vec_names.push_back(vec_names.at(0));
1155
1156
1157
            }
        }

1158
1159
1160
        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
1161
        if(name_it != vec_names.end())
1162
1163
1164
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1165

Shucai Xiao's avatar
Shucai Xiao committed
1166
1167
1168
        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
1169
        });
1170
1171
1172
1173
1174
1175
1176
1177

        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
1178
        if(contains(attributes, "linear_before_reset"))
1179
1180
1181
1182
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1183
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1184
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1185
1186
1187
1188
1189
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1190
1191
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1192
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1193
            std::move(args));
1194
1195

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

Shucai Xiao's avatar
Shucai Xiao committed
1198
        return {hidden_states, last_output};
1199
1200
    }

Shucai Xiao's avatar
Shucai Xiao committed
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
    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
1223
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1224
1225
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1226
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1227
1228
1229
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1230
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1231
        }
Shucai Xiao's avatar
Shucai Xiao committed
1232
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1233
        {
Shucai Xiao's avatar
Shucai Xiao committed
1234
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1235
1236
1237
1238
1239
1240
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1241
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
Shucai Xiao's avatar
Shucai Xiao committed
1242
1243
1244
1245
1246
        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
1247
1248
1249
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1250
1251
1252
        }

        // need 6 activation functions for bidirectional directions
Shucai Xiao's avatar
Shucai Xiao committed
1253
        if(dirct == op::rnn_direction::bidirectional)
Shucai Xiao's avatar
Shucai Xiao committed
1254
1255
1256
1257
1258
1259
        {
            // 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
1260
            // if 3 actv funcs are provide, repeat all three once.
Shucai Xiao's avatar
Shucai Xiao committed
1261
1262
1263
1264
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(vec_names.size())
            {
1265
            case 1:
Shucai Xiao's avatar
Shucai Xiao committed
1266
1267
1268
1269
1270
1271
                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)};
1272
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1273
1274
1275

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1276
1277
1278
1279
1280
1281
                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
1282
1283
1284
1285
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1286
1287
1288
1289
1290
1291
                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
1292
1293
                break;

Shucai Xiao's avatar
Shucai Xiao committed
1294
1295
1296
1297
1298
1299
1300
            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)};
1301
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1302

Shucai Xiao's avatar
Shucai Xiao committed
1303
1304
1305
1306
1307
1308
1309
            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)};
1310
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1311

Shucai Xiao's avatar
Shucai Xiao committed
1312
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1313
1314
1315
1316
1317
1318
            }
        }
        else
        {
            switch(vec_names.size())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1319
            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
Shucai Xiao's avatar
Shucai Xiao committed
1320
1321
1322

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

Shucai Xiao's avatar
Shucai Xiao committed
1326
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1327
1328
1329
            }
        }

1330
1331
1332
        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
1333
        if(name_it != vec_names.end())
1334
1335
1336
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358

        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
1359
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1360
1361
1362
1363
        }

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

        // second output for last lstm output
Shucai Xiao's avatar
Shucai Xiao committed
1367
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1368
1369
1370
1371
1372
1373

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

Shucai Xiao's avatar
Shucai Xiao committed
1375
1376
1377
    instruction_ref parse_reduce_sum(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
1378
1379
1380
1381
    {
        std::size_t n_dim = args.front()->get_shape().lens().size();

        // default to reduce over all dimensions
Shucai Xiao's avatar
Shucai Xiao committed
1382
        std::vector<int64_t> axes(n_dim);
1383
1384
1385
1386
1387
        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
1388
            axes             = std::vector<int64_t>(attr_axes.begin(), attr_axes.end());
1389
1390
1391
        }

        int keep_dims = 1;
Shucai Xiao's avatar
Shucai Xiao committed
1392
        if(contains(attributes, "keepdims"))
1393
1394
1395
1396
        {
            keep_dims = parse_value(attributes.at("keepdims")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1397
        if(keep_dims == 1)
1398
1399
1400
1401
1402
1403
        {
            return prog.add_instruction(op::reduce_sum{axes}, std::move(args));
        }
        else
        {
            auto ins = prog.add_instruction(op::reduce_sum{axes}, std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1404
            return prog.add_instruction(op::squeeze{axes}, ins);
1405
1406
        }
    }
Shucai Xiao's avatar
Shucai Xiao committed
1407
1408

    instruction_ref parse_reduce_mean(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
1409
1410
                                      attribute_map attributes,
                                      std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
    {
        std::size_t n_dim = args.front()->get_shape().lens().size();

        // default to reduce over all dimensions
        std::vector<int64_t> axes(n_dim);
        std::iota(axes.begin(), axes.end(), 0);
        if(contains(attributes, "axes"))
        {
            axes.clear();
            auto&& attr_axes = attributes["axes"].ints();
            axes             = std::vector<int64_t>(attr_axes.begin(), attr_axes.end());
        }

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

        if(keep_dims == 1)
        {
            return prog.add_instruction(op::reduce_mean{axes}, std::move(args));
        }
        else
        {
            auto ins = prog.add_instruction(op::reduce_mean{axes}, std::move(args));
            return prog.add_instruction(op::squeeze{axes}, ins);
1438
1439
        }
    }
1440

Shucai Xiao's avatar
Shucai Xiao committed
1441
1442
    instruction_ref
    parse_cast(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
1443
    {
Shucai Xiao's avatar
Shucai Xiao committed
1444
        if(!contains(attributes, "to"))
1445
1446
1447
1448
        {
            MIGRAPHX_THROW("PARSE_CAST: missing to type attribute!");
        }

Shucai Xiao's avatar
Shucai Xiao committed
1449
        int to_type        = parse_value(attributes.at("to")).at<int>();
1450
1451
1452
        shape::type_t type = get_type(to_type);
        return prog.add_instruction(op::convert{type}, std::move(args));
    }
Shucai Xiao's avatar
Shucai Xiao committed
1453

Paul's avatar
Paul committed
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
    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
1466
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
1467
1468
1469
1470
1471
1472
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
1473
1474
1475
1476
1477
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
1478
1479
1480
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
            // 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
1493
        }
Paul's avatar
Paul committed
1494
        for(auto&& output : graph.output())
Paul's avatar
Paul committed
1495
        {
Paul's avatar
Paul committed
1496
            this->parse_node(output.name());
Paul's avatar
Paul committed
1497
1498
1499
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1500
    void parse_undefined(const std::string& name)
1501
    {
Shucai Xiao's avatar
Shucai Xiao committed
1502
        auto ins           = prog.add_instruction(op::undefined{});
1503
1504
1505
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
1506
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
1507
    {
Paul's avatar
Paul committed
1508
        if(name.empty())
Paul's avatar
Paul committed
1509
            MIGRAPHX_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
1510
1511
1512
1513
1514
1515
1516
1517
        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
1518
1519
                    assert(name != input);
                    this->parse_node(input);
Paul's avatar
Paul committed
1520
                }
Shucai Xiao's avatar
Shucai Xiao committed
1521
                else if(input.empty())
Paul's avatar
Paul committed
1522
                {
1523
                    this->parse_undefined(input);
Paul's avatar
Paul committed
1524
                }
1525
                args.push_back(instructions.at(input));
Paul's avatar
Paul committed
1526
            }
Paul's avatar
Paul committed
1527
            std::vector<instruction_ref> result;
Paul's avatar
Paul committed
1528
1529
            if(ops.count(node.op_type()) == 0)
            {
1530
                result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
Paul's avatar
Paul committed
1531
1532
1533
            }
            else
            {
Paul's avatar
Paul committed
1534
                result = ops[node.op_type()](get_attributes(node), args);
Paul's avatar
Paul committed
1535
            }
Paul's avatar
Paul committed
1536
            // Even no output nodes produce output in migraphx
Paul's avatar
Paul committed
1537
            if(node.output().empty() and result.size() == 1)
Paul's avatar
Paul committed
1538
1539
            {
                instructions[name] = result.front();
Paul's avatar
Paul committed
1540
1541
1542
            }
            else
            {
Paul's avatar
Paul committed
1543
1544
1545
1546
1547
1548
                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
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
            }
        }
    }

    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
1566
        std::size_t n = 0;
Paul's avatar
Paul committed
1567
1568
        for(auto&& node : graph.node())
        {
Paul's avatar
Paul committed
1569
            if(node.output().empty())
Paul's avatar
Paul committed
1570
            {
Paul's avatar
Paul committed
1571
                if(node.name().empty())
Paul's avatar
Paul committed
1572
1573
1574
1575
1576
1577
1578
1579
1580
                {
                    result["migraphx_unamed_node_" + std::to_string(n)] = node;
                    n++;
                }
                else
                {
                    result[node.name()] = node;
                }
            }
Paul's avatar
Paul committed
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
            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
1606
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
1607
1608
1609
1610
1611
        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
1612
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
1613
1614
1615
1616
1617
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
1618
1619
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
1620
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
1621
1622
1623
            switch(t.data_type())
            {
            case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
1624
            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
Scott Thornton's avatar
Scott Thornton committed
1625
            case onnx::TensorProto::UINT8: throw std::runtime_error("");
1626
            case onnx::TensorProto::INT8: return create_literal(shape::int32_type, dims, s.data());
Khalique's avatar
Khalique committed
1627
1628
            case onnx::TensorProto::UINT16:
                return create_literal(shape::int32_type, dims, s.data());
1629
1630
1631
            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
1632
            case onnx::TensorProto::STRING: throw std::runtime_error("");
1633
            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
Khalique's avatar
Khalique committed
1634
1635
1636
1637
            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
1638
1639
1640
1641
1642
            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
1643
            MIGRAPHX_THROW("Invalid tensor type");
1644
        }
Paul's avatar
Paul committed
1645
1646
1647
1648
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Khalique's avatar
Khalique committed
1649
            return create_literal(shape::float_type, dims, t.float_data());
Paul's avatar
Paul committed
1650
1651
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Khalique's avatar
Khalique committed
1652
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1653
        case onnx::TensorProto::UINT16:
Khalique's avatar
Khalique committed
1654
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1655
        case onnx::TensorProto::INT16:
Khalique's avatar
Khalique committed
1656
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1657
        case onnx::TensorProto::INT32:
Khalique's avatar
Khalique committed
1658
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1659
        case onnx::TensorProto::INT64:
Khalique's avatar
Khalique committed
1660
            return create_literal(shape::int64_type, dims, t.int64_data());
Paul's avatar
Paul committed
1661
1662
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Khalique's avatar
Khalique committed
1663
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1664
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
1665
        {
Khalique's avatar
Khalique committed
1666
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
1667
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
1668
1669
1670
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
1671
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
1672
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
1673
        }
Paul's avatar
Paul committed
1674
        case onnx::TensorProto::DOUBLE:
Khalique's avatar
Khalique committed
1675
            return create_literal(shape::double_type, dims, t.double_data());
Paul's avatar
Paul committed
1676
1677
1678
1679
1680
        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
1681
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
1682
1683
    }

Khalique's avatar
Khalique committed
1684
    static literal
1685
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
1686
    {
Khalique's avatar
Khalique committed
1687
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1688
        if(dims.empty())
1689
            return literal{{shape_type}, data};
1690
1691
1692
        return literal{{shape_type, dims}, data};
    }

1693
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
1694
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
1695
1696
    {
        if(dims.empty())
1697
            return literal{{shape_type}, data.begin(), data.end()};
1698
        return literal{{shape_type, dims}, data.begin(), data.end()};
1699
1700
    }

Paul's avatar
Paul committed
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
    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
1720
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1721
1722
1723
1724
1725
1726
1727
1728
1729
        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
1730
        auto&& tensor_dims = t.tensor_type().shape().dim();
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
        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
1742
1743
        return {shape_type, dims};
    }
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765

    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
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
};

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