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

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

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

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

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

    onnx_parser()
    {
Khalique's avatar
Khalique committed
40
        add_generic_op("Relu", op::relu{});
Khalique's avatar
Khalique committed
41
42
        add_generic_op("Sigmoid", op::sigmoid{});
        add_generic_op("Abs", op::abs{});
Shucai Xiao's avatar
Shucai Xiao committed
43
        add_generic_op("Exp", op::exp{});
Shucai Xiao's avatar
Shucai Xiao committed
44
        add_generic_op("Erf", op::erf{});
Shucai Xiao's avatar
Shucai Xiao committed
45
        add_generic_op("Log", op::log{});
Khalique's avatar
Khalique committed
46
47
        // disable dropout for inference
        add_generic_op("Dropout", op::identity{});
Khalique's avatar
Khalique committed
48
        add_generic_op("Identity", op::identity{});
Shucai Xiao's avatar
Shucai Xiao committed
49
50
51
        add_generic_op("Sin", op::sin{});
        add_generic_op("Cos", op::cos{});
        add_generic_op("Tan", op::tan{});
52
53
        add_generic_op("Sinh", op::sinh{});
        add_generic_op("Cosh", op::cosh{});
54
        add_generic_op("Tanh", op::tanh{});
55
56
57
        add_generic_op("Asin", op::asin{});
        add_generic_op("Acos", op::acos{});
        add_generic_op("Atan", op::atan{});
58
        add_generic_op("Sqrt", op::sqrt{});
59
        add_generic_op("Round", op::round{});
60
        add_generic_op("Sign", op::sign{});
Shucai Xiao's avatar
Shucai Xiao committed
61
62
        add_generic_op("Ceil", op::ceil{});
        add_generic_op("Floor", op::floor{});
Paul's avatar
Paul committed
63

Khalique's avatar
Khalique committed
64
65
66
67
        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
68
        add_binary_op("Pow", op::pow{});
Khalique's avatar
Khalique committed
69

Khalique's avatar
Khalique committed
70
71
72
        add_variadic_op("Sum", op::add{});
        add_variadic_op("Max", op::max{});
        add_variadic_op("Min", op::min{});
Paul's avatar
Paul committed
73

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

        add_mem_op("ReduceL1", &onnx_parser::parse_reduce_l1);
        add_mem_op("ReduceL2", &onnx_parser::parse_reduce_l2);
        add_mem_op("ReduceLogSum", &onnx_parser::parse_reduce_log_sum);
        add_mem_op("ReduceLogSumExp", &onnx_parser::parse_reduce_log_sum_exp);
        add_mem_op("ReduceMax", &onnx_parser::parse_reduce_oper<op::reduce_max>);
Shucai Xiao's avatar
Shucai Xiao committed
115
        add_mem_op("ReduceMean", &onnx_parser::parse_reduce_oper<op::reduce_mean>);
Shucai Xiao's avatar
Shucai Xiao committed
116
        add_mem_op("ReduceMin", &onnx_parser::parse_reduce_oper<op::reduce_min>);
Shucai Xiao's avatar
Shucai Xiao committed
117
118
119
        add_mem_op("ReduceProd", &onnx_parser::parse_reduce_oper<op::reduce_prod>);
        add_mem_op("ReduceSum", &onnx_parser::parse_reduce_oper<op::reduce_sum>);
        add_mem_op("ReduceSumSquare", &onnx_parser::parse_reduce_sum_square);
120
121
122
123
124
125
126

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

    void init_actv_func()
    {
127
128
129
130
131
132
        // 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
133
134
135
136
    }

    template <class F>
    void add_op(std::string name, F f)
Paul's avatar
Paul committed
137
138
139
140
141
142
143
144
145
    {
        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
146
147
148
149
150
151
152
    {
        ops.emplace(name, f);
    }

    template <class F>
    void add_mem_op(std::string name, F f)
    {
Paul's avatar
Paul committed
153
        add_op(name, [=](auto&&... xs) {
Paul's avatar
Paul committed
154
155
156
            return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
        });
    }
Khalique's avatar
Khalique committed
157

158
    template <class T>
Khalique's avatar
Khalique committed
159
    void add_binary_op(std::string name, T x)
160
    {
Paul's avatar
Paul committed
161
        add_op(name, [this, x](attribute_map attributes, std::vector<instruction_ref> args) {
Scott Thornton's avatar
Scott Thornton committed
162
            if(args.size() != 2)
Paul's avatar
Paul committed
163
                MIGRAPHX_THROW("binary operators should have 2 operands");
164
            if(contains(attributes, "broadcast") and contains(attributes, "axis"))
165
166
167
168
            {
                uint64_t broadcasted = parse_value(attributes.at("broadcast")).at<uint64_t>();
                if(broadcasted != 0)
                {
169
                    uint64_t axis = parse_value(attributes.at("axis")).at<uint64_t>();
Shucai Xiao's avatar
Shucai Xiao committed
170
171
                    auto l = prog.add_instruction(op::broadcast{axis, args[0]->get_shape().lens()},
                                                  args[1]);
172
173
                    return prog.add_instruction(x, args[0], l);
                }
174
                return prog.add_instruction(x, args);
175
            }
Paul's avatar
Paul committed
176
            else
177
            {
Khalique's avatar
Khalique committed
178
                return add_broadcastable_binary_op(args[0], args[1], x);
179
180
181
182
            }
        });
    }

Shucai Xiao's avatar
Shucai Xiao committed
183
184
    std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
                                                      std::vector<std::size_t> s1)
185
186
187
188
189
190
191
192
193
194
195
196
197
    {
        // 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
198
        if(s0.size() > s1.size())
199
200
201
202
203
204
        {
            s0.swap(s1);
        }

        std::vector<std::size_t> out_lens(s1);
        auto offset = s1.size() - s0.size();
Shucai Xiao's avatar
Shucai Xiao committed
205
206
207
208
        std::transform(s0.begin(),
                       s0.end(),
                       s1.begin() + offset,
                       out_lens.begin() + offset,
209
                       [&](auto a, auto b) {
Shucai Xiao's avatar
Shucai Xiao committed
210
                           if(a != b and a != 1 and b != 1)
211
                           {
Shucai Xiao's avatar
Shucai Xiao committed
212
213
214
215
216
217
                               MIGRAPHX_THROW("COMPUTE_BROADCASTLEN: shape {" +
                                              to_string_range(s0) + "} and {" +
                                              to_string_range(s1) + "} mismatch!");
                           }
                           return std::max(a, b);
                       });
218
219
220
221

        return out_lens;
    }

Shucai Xiao's avatar
Shucai Xiao committed
222
223
    instruction_ref make_contiguous(instruction_ref ins)
    {
Shucai Xiao's avatar
Shucai Xiao committed
224
        if(ins->get_shape().standard())
Shucai Xiao's avatar
Shucai Xiao committed
225
226
227
228
229
230
231
        {
            return ins;
        }

        return prog.add_instruction(op::contiguous{}, ins);
    }

Khalique's avatar
Khalique committed
232
233
234
    template <class T>
    instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
    {
Khalique's avatar
Khalique committed
235
        if(arg0->get_shape().lens() != arg1->get_shape().lens())
Khalique's avatar
Khalique committed
236
237
        {
            // Get lengths for both arguments
Shucai Xiao's avatar
Shucai Xiao committed
238
239
            auto s0       = arg0->get_shape().lens();
            auto s1       = arg1->get_shape().lens();
240
            auto out_lens = compute_broadcasted_lens(s0, s1);
241
242
243
244
245
246
247
248
249

            auto l0 = arg0;
            if(arg0->get_shape().lens() != out_lens)
                l0 = prog.add_instruction(op::multibroadcast{out_lens}, arg0);

            auto l1 = arg1;
            if(arg1->get_shape().lens() != out_lens)
                l1 = prog.add_instruction(op::multibroadcast{out_lens}, arg1);

Khalique's avatar
Khalique committed
250
251
252
253
254
255
            return prog.add_instruction(x, l0, l1);
        }
        else
        {
            return prog.add_instruction(x, {arg0, arg1});
        }
256
257
    }

Paul's avatar
Paul committed
258
    template <class T>
Paul's avatar
Paul committed
259
260
    void add_generic_op(std::string name, T x)
    {
Paul's avatar
Paul committed
261
        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
Paul's avatar
Paul committed
262
263
264
265
            return prog.add_instruction(x, args);
        });
    }

Khalique's avatar
Khalique committed
266
    template <class T>
Khalique's avatar
Khalique committed
267
    void add_variadic_op(std::string name, T x)
Khalique's avatar
Khalique committed
268
    {
Paul's avatar
Paul committed
269
        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
Khalique's avatar
Khalique committed
270
            return std::accumulate(std::next(args.begin()),
Khalique's avatar
Khalique committed
271
272
273
274
275
                                   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
276
        });
Khalique's avatar
Khalique committed
277
278
    }

Khalique's avatar
Khalique committed
279
280
281
    instruction_ref parse_clip(const std::string&,
                               const attribute_map& attributes,
                               std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
282
283
284
285
286
287
288
289
290
291
292
293
294
    {
        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));
    }

Shucai Xiao's avatar
Shucai Xiao committed
295
    template <class Op>
296
    instruction_ref parse_softmax(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
297
298
                                  const attribute_map& attributes,
                                  std::vector<instruction_ref> args)
Paul's avatar
Paul committed
299
    {
300
        int64_t axis = 1;
301
302
303
304
305
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }

306
        return prog.add_instruction(Op{axis}, std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
307
308
    }

Shucai Xiao's avatar
Shucai Xiao committed
309
    template <class Op>
310
    instruction_ref parse_arg_op(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
311
312
                                 const attribute_map& attributes,
                                 std::vector<instruction_ref> args)
313
    {
314
        int64_t axis = 0;
315
316
        if(contains(attributes, "axis"))
        {
317
            axis = static_cast<int64_t>(parse_value(attributes.at("axis")).at<int>());
318
319
        }

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

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

337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
    template <class Op>
    instruction_ref process_auto_pad_attribute(instruction_ref ins,
                                               attribute_map& attributes,
                                               Op& op,
                                               const std::vector<std::size_t>& in_lens)
    {
        if(!contains(attributes, "auto_pad"))
        {
            return ins;
        }

        auto auto_pad = attributes["auto_pad"].s();
        if(auto_pad.find("SAME") != std::string::npos)
        {
            // calculate the padding
            std::array<std::size_t, 2> out_lens;
            out_lens[0] = (in_lens[2] + op.stride[0] - 1) / op.stride[0];
            out_lens[1] = (in_lens[3] + op.stride[1] - 1) / op.stride[1];

            std::array<std::size_t, 2> explicit_pads;
            explicit_pads[0] = (out_lens[0] - 1) * op.stride[0] + op.lengths[0] - in_lens[2];
            explicit_pads[1] = (out_lens[1] - 1) * op.stride[1] + op.lengths[1] - in_lens[3];
            op.padding[0]    = explicit_pads[0] / 2;
            op.padding[1]    = explicit_pads[1] / 2;
            explicit_pads[0] -= 2 * op.padding[0];
            explicit_pads[1] -= 2 * op.padding[1];
            std::vector<std::int64_t> pads(8, 0);
            if(explicit_pads[0] != 0 or explicit_pads[1] != 0)
            {
                if(auto_pad == "SAME_UPPER")
                {
                    pads[6] = explicit_pads[0];
                    pads[7] = explicit_pads[1];
                }
                else if(auto_pad == "SAME_LOWER")
                {
                    pads[2] = explicit_pads[0];
                    pads[3] = explicit_pads[1];
                }

                // MaxPool
                if(op.mode == "max")
                {
                    ins = prog.add_instruction(op::pad{pads, std::numeric_limits<float>::lowest()},
                                               ins);
                }
                // AveragePool
                else
                {
                    ins = prog.add_instruction(op::pad{pads}, ins);
                }
            }

            op.padding_mode = op::padding_mode_t::same;
        }

        return ins;
    }

Paul's avatar
Paul committed
396
    instruction_ref
Paul's avatar
Paul committed
397
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
398
    {
399
        op::convolution op;
400
        auto l0 = args[0];
Paul's avatar
Paul committed
401
402
        if(contains(attributes, "pads"))
        {
Scott Thornton's avatar
Scott Thornton committed
403
            if(contains(attributes, "auto_pad"))
404
            {
405
406
407
408
409
                auto s = attributes["auto_pad"].s();
                if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
                {
                    MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
                }
410
            }
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
                l0      = prog.add_instruction(op::pad{padding}, l0);
422
            }
423
424
425
426
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
427
            }
Paul's avatar
Paul committed
428
        }
Paul's avatar
Paul committed
429
430
431
432
433
434
435
436
        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
437
        if(contains(attributes, "auto_pad"))
438
439
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
440
            if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
441
            {
Paul's avatar
Paul committed
442
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
443
444
            }

wsttiger's avatar
fixes  
wsttiger committed
445
            if(s.find("SAME") != std::string::npos)
446
            {
447
                op.padding_mode = op::padding_mode_t::same;
448
449
            }
        }
Khalique's avatar
Khalique committed
450
451
452
453
        if(contains(attributes, "group"))
        {
            op.group = parse_value(attributes.at("group")).at<int>();
        }
Paul's avatar
Paul committed
454
455
456
        if(args.size() == 3)
        {
            uint64_t axis = 1;
Khalique's avatar
Khalique committed
457
            auto l1       = prog.add_instruction(op, l0, args[1]);
Shucai Xiao's avatar
Shucai Xiao committed
458
            auto l2 = prog.add_instruction(op::broadcast{axis, l1->get_shape().lens()}, args[2]);
459
            return prog.add_instruction(op::add{}, l1, l2);
Paul's avatar
Paul committed
460
        }
461
        return prog.add_instruction(op, l0, args[1]);
Paul's avatar
Paul committed
462
    }
Paul's avatar
Paul committed
463

Paul's avatar
Paul committed
464
465
466
    instruction_ref parse_pooling(const std::string& name,
                                  attribute_map attributes,
                                  std::vector<instruction_ref> args)
Paul's avatar
Paul committed
467
    {
Khalique's avatar
Khalique committed
468
        op::pooling op{ends_with(name, "MaxPool") ? "max" : "average"};
469
        auto l0 = args[0];
Khalique's avatar
Khalique committed
470
        if(starts_with(name, "Global"))
471
        {
Khalique's avatar
Khalique committed
472
473
            auto lens  = args.front()->get_shape().lens();
            op.lengths = {lens[2], lens[3]};
474
        }
475

Paul's avatar
Paul committed
476
477
        if(contains(attributes, "pads"))
        {
478
479
480
481
482
483
484
485
486
487
            if(contains(attributes, "auto_pad"))
            {
                auto s = attributes["auto_pad"].s();
                if(to_upper(s) != "NOTSET")
                {
                    MIGRAPHX_THROW(
                        "PARSE_POOLING: auto_pad and padding cannot be specified simultaneously");
                }
            }

488
489
            std::vector<std::int64_t> padding;
            copy(attributes["pads"].ints(), std::back_inserter(padding));
Scott Thornton's avatar
Scott Thornton committed
490
            if(padding.size() != 4)
491
            {
492
                MIGRAPHX_THROW("PARSE_POOLING: padding should have 4 values");
493
            }
Scott Thornton's avatar
Scott Thornton committed
494
            if(padding[0] != padding[2] || padding[1] != padding[3])
495
            {
496
497
                // insert zeros for pad op (args[0] has 4 dims)
                padding = {0, 0, padding[0], padding[1], 0, 0, padding[2], padding[3]};
498
499
500
501
502
503
504
505
506
507
508
                // MaxPool
                if(op.mode == "max")
                {
                    l0 = prog.add_instruction(
                        op::pad{padding, std::numeric_limits<float>::lowest()}, l0);
                }
                // AveragePool
                else
                {
                    l0 = prog.add_instruction(op::pad{padding}, l0);
                }
509
510
511
512
513
            }
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
514
            }
Paul's avatar
Paul committed
515
        }
516

Paul's avatar
Paul committed
517
518
519
520
521
522
523
524
        if(contains(attributes, "strides"))
        {
            copy(attributes["strides"].ints(), op.stride.begin());
        }
        if(contains(attributes, "kernel_shape"))
        {
            copy(attributes["kernel_shape"].ints(), op.lengths.begin());
        }
525

Scott Thornton's avatar
Scott Thornton committed
526
        if(contains(attributes, "auto_pad"))
527
        {
528
529
            auto in_lens = args[0]->get_shape().lens();
            l0           = process_auto_pad_attribute(l0, attributes, op, in_lens);
530
531
        }

532
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
533
534
    }

Paul's avatar
Paul committed
535
    instruction_ref
Paul's avatar
Paul committed
536
    parse_reshape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
537
    {
538
        op::reshape op;
Paul's avatar
Paul committed
539
540
        if(args.size() == 1)
        {
541
542
            literal s = parse_value(attributes.at("shape"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
543
544
545
        }
        if(args.size() == 2)
        {
Paul's avatar
Paul committed
546
            auto s = args[1]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
547
            check_arg_empty(s, "Reshape: dynamic shape is not supported");
Paul's avatar
Paul committed
548
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
549
        }
550

Shucai Xiao's avatar
Shucai Xiao committed
551
        return prog.add_instruction(op, make_contiguous(args[0]));
Paul's avatar
Paul committed
552
553
    }

Paul's avatar
Paul committed
554
    instruction_ref
Paul's avatar
Paul committed
555
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
556
    {
557
        int64_t axis = 1;
Paul's avatar
Paul committed
558
559
560
561
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
562
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
563
564
    }

565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
    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
583
584
585
    instruction_ref
    parse_concat(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
Shucai Xiao's avatar
Shucai Xiao committed
586
587
588
589
590
591
592
        // change to hande axis to be negative values
        if(!contains(attributes, "axis"))
        {
            MIGRAPHX_THROW("PARSE_CONCAT: attribute axis is required!");
        }

        int axis = parse_value(attributes.at("axis")).at<int>();
Scott Thornton's avatar
Scott Thornton committed
593
594
595
        op::concat op{axis};
        return prog.add_instruction(op, std::move(args));
    }
596

597
598
599
    instruction_ref
    parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
600
        int axis = 0;
601
602
603
604
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
605

606
        op::gather op{axis};
Shucai Xiao's avatar
Shucai Xiao committed
607
        return prog.add_instruction(op, make_contiguous(args[0]), make_contiguous(args[1]));
608
609
    }

610
611
612
613
    instruction_ref
    parse_slice(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::slice op;
Khalique's avatar
Khalique committed
614
        std::vector<size_t> dims = args[0]->get_shape().lens();
Khalique's avatar
Khalique committed
615
        size_t num_dims          = dims.size();
616
617
618
619
620
        if(contains(attributes, "axes"))
        {
            literal s = parse_value(attributes.at("axes"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        }
Khalique's avatar
Khalique committed
621
622
623
624
625
        else
        {
            op.axes = std::vector<int64_t>(num_dims);
            std::iota(op.axes.begin(), op.axes.end(), 0);
        }
Khalique's avatar
Khalique committed
626

Khalique's avatar
Khalique committed
627
        if(contains(attributes, "ends"))
628
        {
Paul's avatar
Paul committed
629
            op.ends = get_indices(attributes.at("ends"));
630
        }
Khalique's avatar
Khalique committed
631
        if(contains(attributes, "starts"))
632
633
634
635
636
637
638
        {
            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
639
640
641
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
642
    {
Shucai Xiao's avatar
Shucai Xiao committed
643
        literal v = parse_value(attributes.at("value"));
644
        // return empty literal
Shucai Xiao's avatar
Shucai Xiao committed
645
        if(v.get_shape().elements() == 0)
646
647
648
649
        {
            return prog.add_literal(literal{});
        }

650
651
652
        auto dim_size = attributes.at("value").t().dims_size();
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
653
        {
654
            migraphx::shape scalar_shape{v.get_shape().type()};
655
656
657
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

Paul's avatar
Paul committed
658
659
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
660

Paul's avatar
Paul committed
661
    instruction_ref
Paul's avatar
Paul committed
662
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
663
664
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
665
        float beta  = 1.0f;
Paul's avatar
Paul committed
666
667
668
669
670
671
672
673
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
674
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
675
676
677
678
679
680
681
682
683
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
684
685
686
687
688
689

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

690
691
        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
692
693
        if(args.size() == 3)
        {
694
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
695
            {
Shucai Xiao's avatar
Shucai Xiao committed
696
                auto out_lens   = l1->get_shape().lens();
697
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
698
                auto l3         = args[2];
Shucai Xiao's avatar
Shucai Xiao committed
699
700
                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
701
                {
702
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
703
                }
704
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
705
            }
Paul's avatar
Paul committed
706
        }
707
708

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

711
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
712
    parse_matmul(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
713
    {
Shucai Xiao's avatar
Shucai Xiao committed
714
715
        auto l0      = args[0];
        auto l1      = args[1];
716
717
718
719
720
        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
721
        if(l0_lens.size() == 1)
722
723
724
725
726
727
728
        {
            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
729
        if(l1_lens.size() == 1)
730
731
732
733
734
735
736
737
        {
            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
738
        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
739
740
741
742
743
744
        {
            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);
745
            l0_broadcasted_lens = output_lens;
746
            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
747
            l1_broadcasted_lens = output_lens;
748
            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
Shucai Xiao's avatar
Shucai Xiao committed
749
            if(l0_lens != l0_broadcasted_lens)
750
751
752
            {
                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
            }
Shucai Xiao's avatar
Shucai Xiao committed
753
            if(l1_lens != l1_broadcasted_lens)
754
755
756
757
758
            {
                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
759
        auto dot_res     = prog.add_instruction(op::dot{1.0f, 0.0f}, bl0, bl1);
760
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
761
        if(is_a_prepended)
762
763
764
765
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
766
        if(is_b_appended)
767
768
769
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
770

771
772
773
        return dot_res;
    }

774
    instruction_ref
Paul's avatar
Paul committed
775
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
776
    {
Scott Thornton's avatar
Scott Thornton committed
777
778
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
779
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
780
781
782
783
784
785
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
786
            momentum = parse_value(attributes.at("momentum")).at<float>();
787
788
789
        }
        if(contains(attributes, "spatial"))
        {
790
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
791
792
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
793
        }
Paul's avatar
Paul committed
794
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
795
        return prog.add_instruction(op, std::move(args));
796
797
    }

798
799
800
801
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
802
        float alpha = 0.01; // default alpha val for leaky relu
803
804
805
806
807
808
809
810
        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
811
812
    instruction_ref
    parse_elu(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
813
814
815
816
817
818
819
820
821
822
    {
        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
823
824
    instruction_ref
    parse_lrn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
825
826
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
827
828
829
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
Khalique's avatar
Khalique committed
830
831
832
833
834
835
836
837
838
839
840
841
        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
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
    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());
        }
858
        auto input_lens = args.front()->get_shape().lens();
Khalique's avatar
Khalique committed
859

Khalique's avatar
Khalique committed
860
861
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
862
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
863

864
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
865
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
866
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
867
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
868
    }
Khalique's avatar
Khalique committed
869

Khalique's avatar
Khalique committed
870
871
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
872
873
874
875
876
877
878
    {
        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
879
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
880
881
    }

Khalique's avatar
Khalique committed
882
883
884
885
886
887
888
889
890
891
    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());
        }
892
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
893
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
894
895
896
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
Khalique's avatar
Khalique committed
897
898
899
900
901
902
903
904
905
906
907
908
        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());
    }
909
910
911
    // 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
912
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
913
914
    {
        if(args.size() != 1)
915
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
        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>();
        }
Shucai Xiao's avatar
Shucai Xiao committed
940
        shape::type_t type = get_type(dtype);
941
942
943
944
945
946
947
948
949
950
951

        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
952
953
        if(contains(attributes, "extra_shape"))
        {
954
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
955
956
        }

957
958
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
959
            if(args.size() != 1)
960
            {
961
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
962
963
            }

Shucai Xiao's avatar
Shucai Xiao committed
964
965
            if(contains(attributes, "shape"))
            {
966
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
967
                               "at the same time");
968
969
            }

970
            migraphx::argument in = args[0]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
971
            check_arg_empty(in, "ConstantFill: dynamic shape is not supported");
972

973
974
975
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
976
977
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
978
979
980
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
981
982
            if(!contains(attributes, "shape"))
            {
983
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
984
985
986
            }

            literal ls = parse_value(attributes.at("shape"));
987
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
988
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
989
            migraphx::shape s{type, dims};
990
991
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
992
993
994
        }
        else
        {
995
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
996
997
998
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
999
1000
1001
    instruction_ref parse_constant_of_shape(const std::string&,
                                            attribute_map attributes,
                                            std::vector<instruction_ref> args)
1002
1003
    {
        literal l_val{};
Shucai Xiao's avatar
Shucai Xiao committed
1004
        if(contains(attributes, "value"))
1005
1006
        {
            l_val = parse_value(attributes.at("value"));
Shucai Xiao's avatar
Shucai Xiao committed
1007
            if(l_val.get_shape().elements() != 1)
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
            {
                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();
1019

Shucai Xiao's avatar
Shucai Xiao committed
1020
        if(args.empty())
1021
        {
Shucai Xiao's avatar
Shucai Xiao committed
1022
            MIGRAPHX_THROW("ConstantOfShape : must have 1 input!");
1023
1024
1025
        }
        else
        {
1026
1027
            migraphx::shape s;
            // empty input tensor, output is a scalar
Shucai Xiao's avatar
Shucai Xiao committed
1028
            if(args[0]->get_shape().elements() == 0)
1029
            {
1030
                s = migraphx::shape{type, {1}, {0}};
1031
            }
1032
1033
1034
            else
            {
                migraphx::argument in = args[0]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
1035
                check_arg_empty(in, "ConstantOfShape: dynamic shape is not supported");
1036

1037
1038
1039
1040
                std::vector<std::size_t> dims;
                in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
                s = migraphx::shape{type, dims};
            }
1041

Shucai Xiao's avatar
Shucai Xiao committed
1042
            literal l_out{};
1043
            l_val.visit([&](auto val) {
Shucai Xiao's avatar
Shucai Xiao committed
1044
                using val_type = std::remove_cv_t<typename decltype(val)::value_type>;
1045
                // l_val contains only one element
1046
                std::vector<val_type> out_vec(s.elements(), val.front());
1047
1048
1049
1050
1051
1052
1053
                l_out = literal(s, out_vec);
            });

            return prog.add_literal(l_out);
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1054
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
1055
    parse_expand(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
1056
    {
Shucai Xiao's avatar
Shucai Xiao committed
1057
        auto in_lens             = args[0]->get_shape().lens();
1058
        migraphx::argument arg_s = args[1]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
1059
        check_arg_empty(arg_s, "Expand: dynamic shape is not supported");
1060
1061
1062
        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);
Shucai Xiao's avatar
Shucai Xiao committed
1063
        return prog.add_instruction(op::multibroadcast{out_lens}, args[0]);
1064
1065
    }

Shucai Xiao's avatar
Shucai Xiao committed
1066
    std::vector<instruction_ref>
Shucai Xiao's avatar
Shucai Xiao committed
1067
1068
1069
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
1070
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
1071
1072
1073

        if(contains(attributes, "hidden_size"))
        {
Shucai Xiao's avatar
Shucai Xiao committed
1074
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1075
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1076
1077
1078
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
1079
1080
1081
1082
1083
1084
1085
1086
1087
        }

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

1088
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1089
1090
        if(direction == "bidirectional")
        {
1091
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1092
1093
1094
        }
        else if(direction == "reverse")
        {
1095
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1096
1097
        }

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

1109
1110
1111
        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
1112
        if(name_it != vec_names.end())
1113
1114
1115
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
1116

Shucai Xiao's avatar
Shucai Xiao committed
1117
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
1118
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
1119
        // if only one actv function is provided, we use it in both
1120
        // forward and reverse direction
1121
        if(dirct == op::rnn_direction::bidirectional)
1122
        {
Shucai Xiao's avatar
Shucai Xiao committed
1123
            if(vec_names.size() == 1)
1124
1125
1126
1127
1128
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
1129
        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1130
1131
1132
1133
        std::transform(vec_names.begin(),
                       vec_names.end(),
                       vec_actv_funcs.begin(),
                       [&](const auto& fn) { return map_actv_funcs[fn]; });
Shucai Xiao's avatar
Shucai Xiao committed
1134

Shucai Xiao's avatar
Shucai Xiao committed
1135
1136
1137
1138
1139
1140
1141
        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

1142
1143
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1144
        if(args.size() < 6)
1145
1146
1147
1148
1149
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
1150
1151
        // 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
1152
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1153

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

Shucai Xiao's avatar
Shucai Xiao committed
1157
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
1158
1159
    }

1160
    std::vector<instruction_ref>
1161
1162
1163
1164
1165
1166
1167
    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
1168
            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1169
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1170
1171
1172
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
1173
1174
1175
1176
1177
1178
1179
1180
1181
        }

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

1182
        op::rnn_direction dirct = op::rnn_direction::forward;
1183
1184
        if(direction == "bidirectional")
        {
1185
            dirct = op::rnn_direction::bidirectional;
1186
1187
1188
        }
        else if(direction == "reverse")
        {
1189
            dirct = op::rnn_direction::reverse;
1190
1191
        }

1192
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
1193
1194
        if(contains(attributes, "activations"))
        {
1195
            auto names = attributes.at("activations").strings();
1196
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
1197
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1198
1199
1200
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1201
1202
        }

1203
        // need 4 activation functions
1204
        if(dirct == op::rnn_direction::bidirectional)
1205
        {
Shucai Xiao's avatar
Shucai Xiao committed
1206
            // 4 activation functions are used in the bidirectional
1207
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
1208
1209
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
1210
1211
            // 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
1212
1213
1214
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
1215
            if(vec_names.size() == 1)
1216
            {
1217
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
1218
            }
1219
            else if(vec_names.size() == 2)
1220
            {
1221
1222
1223
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
1224
            }
1225
            else if(vec_names.size() == 3)
1226
            {
1227
                vec_names.push_back(vec_names.at(2));
1228
1229
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
1230
        else
1231
        {
1232
            if(vec_names.size() == 1)
1233
            {
1234
                vec_names.push_back(vec_names.at(0));
1235
1236
1237
            }
        }

1238
1239
1240
        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
1241
        if(name_it != vec_names.end())
1242
1243
1244
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1245

Shucai Xiao's avatar
Shucai Xiao committed
1246
        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1247
1248
1249
1250
        std::transform(vec_names.begin(),
                       vec_names.end(),
                       vec_actv_funcs.begin(),
                       [&](const auto& name) { return map_actv_funcs[name]; });
1251
1252
1253
1254
1255
1256
1257
1258

        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
1259
        if(contains(attributes, "linear_before_reset"))
1260
1261
1262
1263
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
1264
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1265
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1266
1267
1268
1269
1270
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1271
1272
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1273
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1274
            std::move(args));
1275
1276

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

Shucai Xiao's avatar
Shucai Xiao committed
1279
        return {hidden_states, last_output};
1280
1281
    }

Shucai Xiao's avatar
Shucai Xiao committed
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
    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
1304
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1305
1306
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1307
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1308
1309
1310
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1311
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1312
        }
Shucai Xiao's avatar
Shucai Xiao committed
1313
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1314
        {
Shucai Xiao's avatar
Shucai Xiao committed
1315
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1316
1317
1318
1319
1320
1321
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1322
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
Shucai Xiao's avatar
Shucai Xiao committed
1323
1324
1325
1326
1327
        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
1328
1329
1330
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1331
1332
1333
        }

        // need 6 activation functions for bidirectional directions
Shucai Xiao's avatar
Shucai Xiao committed
1334
        if(dirct == op::rnn_direction::bidirectional)
Shucai Xiao's avatar
Shucai Xiao committed
1335
1336
1337
1338
1339
1340
        {
            // 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
1341
            // if 3 actv funcs are provide, repeat all three once.
Shucai Xiao's avatar
Shucai Xiao committed
1342
1343
1344
1345
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(vec_names.size())
            {
1346
            case 1:
Shucai Xiao's avatar
Shucai Xiao committed
1347
1348
1349
1350
1351
1352
                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)};
1353
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1354
1355
1356

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
Shucai Xiao's avatar
Shucai Xiao committed
1357
1358
1359
1360
1361
1362
                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
1363
1364
1365
1366
                break;

            case 3:
                // repeat all three actv funcs once
Shucai Xiao's avatar
Shucai Xiao committed
1367
1368
1369
1370
1371
1372
                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
1373
1374
                break;

Shucai Xiao's avatar
Shucai Xiao committed
1375
1376
1377
1378
1379
1380
1381
            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)};
1382
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1383

Shucai Xiao's avatar
Shucai Xiao committed
1384
1385
1386
1387
1388
1389
1390
            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)};
1391
                break;
Shucai Xiao's avatar
Shucai Xiao committed
1392

Shucai Xiao's avatar
Shucai Xiao committed
1393
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1394
1395
1396
1397
1398
1399
            }
        }
        else
        {
            switch(vec_names.size())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1400
            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
Shucai Xiao's avatar
Shucai Xiao committed
1401
1402
1403

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

Shucai Xiao's avatar
Shucai Xiao committed
1407
            default: break;
Shucai Xiao's avatar
Shucai Xiao committed
1408
1409
1410
            }
        }

1411
1412
1413
        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
1414
        if(name_it != vec_names.end())
1415
1416
1417
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1418
1419

        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1420
1421
1422
1423
        std::transform(vec_names.begin(),
                       vec_names.end(),
                       vec_actv_funcs.begin(),
                       [&](const auto& name) { return map_actv_funcs[name]; });
Shucai Xiao's avatar
Shucai Xiao committed
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440

        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
1441
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1442
1443
1444
1445
        }

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

        // second output for last lstm output
Shucai Xiao's avatar
Shucai Xiao committed
1449
        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1450
1451
1452
1453
1454
1455

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

Shucai Xiao's avatar
Shucai Xiao committed
1457
    template <class T>
Shucai Xiao's avatar
Shucai Xiao committed
1458
    instruction_ref parse_reduce_oper(const std::string&,
Shucai Xiao's avatar
Shucai Xiao committed
1459
1460
                                      attribute_map attributes,
                                      std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1461
1462
1463
1464
    {
        std::size_t n_dim = args.front()->get_shape().lens().size();

        // default to reduce over all dimensions
1465
        std::vector<int64_t> axes(n_dim);
Shucai Xiao's avatar
Shucai Xiao committed
1466
1467
1468
1469
1470
        std::iota(axes.begin(), axes.end(), 0);
        if(contains(attributes, "axes"))
        {
            axes.clear();
            auto&& attr_axes = attributes["axes"].ints();
1471
            axes             = std::vector<int64_t>(attr_axes.begin(), attr_axes.end());
Shucai Xiao's avatar
Shucai Xiao committed
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
        }

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

        if(keep_dims == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
1482
            return prog.add_instruction(T{axes}, std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1483
1484
1485
        }
        else
        {
Shucai Xiao's avatar
Shucai Xiao committed
1486
            auto ins = prog.add_instruction(T{axes}, std::move(args));
1487
            return prog.add_instruction(op::squeeze{axes}, ins);
1488
1489
        }
    }
1490

Shucai Xiao's avatar
Shucai Xiao committed
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
    instruction_ref
    parse_reduce_l1(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        auto abs_ins = prog.add_instruction(op::abs{}, args[0]);
        return parse_reduce_oper<op::reduce_sum>({}, std::move(attributes), {abs_ins});
    }

    instruction_ref
    parse_reduce_l2(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        auto square_ins = prog.add_instruction(op::mul{}, args[0], args[0]);
        auto sum_ins = parse_reduce_oper<op::reduce_sum>({}, std::move(attributes), {square_ins});
        return prog.add_instruction(op::sqrt{}, sum_ins);
    }

    instruction_ref parse_reduce_log_sum(const std::string&,
                                         attribute_map attributes,
                                         std::vector<instruction_ref> args)
    {
        auto sum_ins =
            parse_reduce_oper<op::reduce_sum>({}, std::move(attributes), std::move(args));
        return prog.add_instruction(op::log{}, sum_ins);
    }

    instruction_ref parse_reduce_log_sum_exp(const std::string&,
                                             attribute_map attributes,
                                             std::vector<instruction_ref> args)
    {
        auto exp_ins = prog.add_instruction(op::exp{}, args[0]);
        auto sum_ins = parse_reduce_oper<op::reduce_sum>({}, std::move(attributes), {exp_ins});
        return prog.add_instruction(op::log{}, sum_ins);
    }

    instruction_ref parse_reduce_sum_square(const std::string&,
                                            attribute_map attributes,
                                            std::vector<instruction_ref> args)
    {
        auto square_ins = prog.add_instruction(op::mul{}, args[0], args[0]);
        return parse_reduce_oper<op::reduce_sum>({}, std::move(attributes), {square_ins});
    }

Shucai Xiao's avatar
Shucai Xiao committed
1532
1533
    instruction_ref
    parse_cast(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
1534
    {
Shucai Xiao's avatar
Shucai Xiao committed
1535
        if(!contains(attributes, "to"))
1536
1537
1538
1539
        {
            MIGRAPHX_THROW("PARSE_CAST: missing to type attribute!");
        }

Shucai Xiao's avatar
Shucai Xiao committed
1540
        int to_type        = parse_value(attributes.at("to")).at<int>();
1541
1542
1543
        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
1544

Paul's avatar
Paul committed
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
    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
1557
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
1558
1559
1560
1561
1562
1563
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
1564
        for(auto&& f : graph.initializer())
1565
1566
            instructions[f.name()] = prog.add_literal(parse_tensor(f));

Paul's avatar
Paul committed
1567
1568
1569
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
1570
1571
            // input not in initializer_data, so it is a real input
            if(!contains(instructions, name))
1572
1573
            {
                // TODO: Get shape of input parameter
1574
                shape s            = parse_type(input.type(), batch_size);
1575
1576
                instructions[name] = prog.add_parameter(name, s);
            }
Paul's avatar
Paul committed
1577
        }
Paul's avatar
Paul committed
1578
        for(auto&& output : graph.output())
Paul's avatar
Paul committed
1579
        {
Paul's avatar
Paul committed
1580
            this->parse_node(output.name());
Paul's avatar
Paul committed
1581
        }
Shucai Xiao's avatar
Shucai Xiao committed
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594

        // For now, the last output with a valid name is considered
        // as the program output, and add an identity instruction at
        // the program end
        auto prog_output = graph.output();
        auto oit         = std::find_if(prog_output.rbegin(), prog_output.rend(), [](auto& node) {
            return !node.name().empty();
        });

        if(instructions.count(oit->name()) > 0)
        {
            prog.add_instruction(op::identity{}, instructions[oit->name()]);
        }
Paul's avatar
Paul committed
1595
1596
    }

Shucai Xiao's avatar
Shucai Xiao committed
1597
    void parse_undefined(const std::string& name)
1598
    {
Shucai Xiao's avatar
Shucai Xiao committed
1599
        auto ins           = prog.add_instruction(op::undefined{});
1600
1601
1602
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
1603
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
1604
    {
Paul's avatar
Paul committed
1605
        if(name.empty())
Paul's avatar
Paul committed
1606
            MIGRAPHX_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
1607
1608
1609
1610
1611
1612
        if(instructions.count(name) == 0)
        {
            auto&& node = nodes.at(name);
            std::vector<instruction_ref> args;
            for(auto&& input : node.input())
            {
Shucai Xiao's avatar
Shucai Xiao committed
1613
                if(input.empty())
Paul's avatar
Paul committed
1614
                {
Shucai Xiao's avatar
Shucai Xiao committed
1615
                    this->parse_undefined(input);
Paul's avatar
Paul committed
1616
                }
Shucai Xiao's avatar
Shucai Xiao committed
1617
                else if(nodes.count(input) > 0)
Paul's avatar
Paul committed
1618
                {
Shucai Xiao's avatar
Shucai Xiao committed
1619
1620
                    assert(name != input);
                    this->parse_node(input);
Paul's avatar
Paul committed
1621
                }
1622
                args.push_back(instructions.at(input));
Paul's avatar
Paul committed
1623
            }
Paul's avatar
Paul committed
1624
            std::vector<instruction_ref> result;
Paul's avatar
Paul committed
1625
1626
            if(ops.count(node.op_type()) == 0)
            {
1627
                result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
Paul's avatar
Paul committed
1628
1629
1630
            }
            else
            {
Paul's avatar
Paul committed
1631
                result = ops[node.op_type()](get_attributes(node), args);
Paul's avatar
Paul committed
1632
            }
Paul's avatar
Paul committed
1633
            // Even no output nodes produce output in migraphx
Paul's avatar
Paul committed
1634
            if(node.output().empty() and result.size() == 1)
Paul's avatar
Paul committed
1635
1636
            {
                instructions[name] = result.front();
Paul's avatar
Paul committed
1637
1638
1639
            }
            else
            {
Shucai Xiao's avatar
Shucai Xiao committed
1640
1641
1642
1643
                assert(node.output().size() <= result.size());
                std::transform(node.output().begin(),
                               node.output().end(),
                               result.begin(),
Paul's avatar
Paul committed
1644
                               std::inserter(instructions, instructions.end()),
Shucai Xiao's avatar
Shucai Xiao committed
1645
                               [](auto&& x, auto&& y) { return std::make_pair(x, y); });
Paul's avatar
Paul committed
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
            }
        }
    }

    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
1663
        std::size_t n = 0;
Paul's avatar
Paul committed
1664
1665
        for(auto&& node : graph.node())
        {
Paul's avatar
Paul committed
1666
            if(node.output().empty())
Paul's avatar
Paul committed
1667
            {
Paul's avatar
Paul committed
1668
                if(node.name().empty())
Paul's avatar
Paul committed
1669
1670
1671
1672
1673
1674
1675
1676
1677
                {
                    result["migraphx_unamed_node_" + std::to_string(n)] = node;
                    n++;
                }
                else
                {
                    result[node.name()] = node;
                }
            }
Paul's avatar
Paul committed
1678
1679
1680
1681
1682
1683
1684
1685
            for(auto&& output : node.output())
            {
                result[output] = node;
            }
        }
        return result;
    }

Paul's avatar
Paul committed
1686
1687
1688
1689
1690
1691
    static std::vector<int64_t> get_indices(const onnx::AttributeProto& attr)
    {
        std::vector<int64_t> result;
        literal s = parse_value(attr);
        s.visit([&](auto v) { copy(v, std::back_inserter(result)); });
        // Clamp large indices to -1
Paul's avatar
Paul committed
1692
1693
1694
1695
1696
        std::replace_if(
            result.begin(),
            result.end(),
            [](auto x) { return x > int64_t{std::numeric_limits<std::int32_t>::max()} / 2; },
            -1);
Paul's avatar
Paul committed
1697
1698
1699
        return result;
    }

Paul's avatar
Paul committed
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
    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::FLOAT: return literal{attr.f()};
        case onnx::AttributeProto::INT: return literal{attr.i()};
        case onnx::AttributeProto::TENSOR: return parse_tensor(attr.t());
Paul's avatar
Paul committed
1714
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
1715
        case onnx::AttributeProto::INTS: return from_repeated(shape::int64_type, attr.ints());
Paul's avatar
Paul committed
1716
1717
1718
1719
1720
        case onnx::AttributeProto::UNDEFINED:
        case onnx::AttributeProto::GRAPH:
        case onnx::AttributeProto::STRING:
        case onnx::AttributeProto::STRINGS:
        case onnx::AttributeProto::TENSORS:
Paul's avatar
Paul committed
1721
1722
        case onnx::AttributeProto::GRAPHS: return {};
        }
Paul's avatar
Paul committed
1723
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
1724
1725
1726
1727
1728
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
1729
1730
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
1731
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
1732
1733
            switch(t.data_type())
            {
1734
            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
Khalique's avatar
Khalique committed
1735
1736
1737
1738
            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());
1739
            case onnx::TensorProto::INT64: return create_literal(shape::int64_type, dims, s.data());
Paul's avatar
Paul committed
1740
1741
1742
1743
            case onnx::TensorProto::INT8:
            case onnx::TensorProto::UINT16:
            case onnx::TensorProto::INT16:
            case onnx::TensorProto::INT32:
1744
            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
Paul's avatar
Paul committed
1745
1746
1747
1748
1749
1750
            case onnx::TensorProto::UINT8:
            case onnx::TensorProto::STRING:
            case onnx::TensorProto::UNDEFINED:
            case onnx::TensorProto::UINT32:
            case onnx::TensorProto::UINT64:
            case onnx::TensorProto::COMPLEX64:
Scott Thornton's avatar
Scott Thornton committed
1751
1752
            case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
            }
Paul's avatar
Paul committed
1753
            MIGRAPHX_THROW("Invalid tensor type");
1754
        }
Paul's avatar
Paul committed
1755
1756
1757
1758
1759
1760
        switch(t.data_type())
        {
        case onnx::TensorProto::INT8:
        case onnx::TensorProto::UINT16:
        case onnx::TensorProto::INT16:
        case onnx::TensorProto::INT32:
Paul's avatar
Paul committed
1761
        case onnx::TensorProto::BOOL:
Khalique's avatar
Khalique committed
1762
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
1763
        case onnx::TensorProto::INT64:
Khalique's avatar
Khalique committed
1764
            return create_literal(shape::int64_type, dims, t.int64_data());
Paul's avatar
Paul committed
1765
1766
1767
1768
        case onnx::TensorProto::DOUBLE:
            return create_literal(shape::double_type, dims, t.double_data());
        case onnx::TensorProto::FLOAT:
            return create_literal(shape::float_type, dims, t.float_data());
Paul's avatar
Paul committed
1769
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
1770
        {
Khalique's avatar
Khalique committed
1771
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
1772
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
1773
1774
1775
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
1776
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
1777
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
1778
        }
Paul's avatar
Paul committed
1779
1780
1781
1782
1783
1784
        case onnx::TensorProto::UNDEFINED:
        case onnx::TensorProto::UINT8:
        case onnx::TensorProto::STRING:
        case onnx::TensorProto::UINT32:
        case onnx::TensorProto::UINT64:
        case onnx::TensorProto::COMPLEX64:
Paul's avatar
Paul committed
1785
1786
        case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
        }
Paul's avatar
Paul committed
1787
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
1788
1789
    }

Khalique's avatar
Khalique committed
1790
    static literal
1791
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
1792
    {
Khalique's avatar
Khalique committed
1793
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
1794
        if(dims.empty())
1795
            return literal{{shape_type}, data};
1796
1797
1798
        return literal{{shape_type, dims}, data};
    }

1799
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
1800
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
1801
1802
    {
        if(dims.empty())
1803
            return literal{{shape_type}, data.begin(), data.end()};
1804
        return literal{{shape_type, dims}, data.begin(), data.end()};
1805
1806
    }

1807
    static shape parse_type(const onnx::TypeProto& t, const unsigned int batch_size)
Paul's avatar
Paul committed
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
    {
        shape::type_t shape_type{};
        switch(t.tensor_type().elem_type())
        {
        case onnx::TensorProto::FLOAT: shape_type = shape::float_type; break;
        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;
Paul's avatar
Paul committed
1818
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1819
1820
1821
        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;
Paul's avatar
Paul committed
1822
1823
1824
1825
        case onnx::TensorProto::UINT8:
        case onnx::TensorProto::STRING:
        case onnx::TensorProto::BOOL:
        case onnx::TensorProto::UNDEFINED:
Paul's avatar
Paul committed
1826
1827
        case onnx::TensorProto::COMPLEX64:
        case onnx::TensorProto::COMPLEX128:
Paul's avatar
Paul committed
1828
            break; // throw std::runtime_error("Unsupported type");
Paul's avatar
Paul committed
1829
1830
        }
        std::vector<std::size_t> dims;
Paul's avatar
Paul committed
1831
        auto&& tensor_dims = t.tensor_type().shape().dim();
1832
1833
1834
        std::transform(tensor_dims.begin(),
                       tensor_dims.end(),
                       std::back_inserter(dims),
1835
1836
                       [&](auto&& d) -> std::size_t {
                           if(d.has_dim_value())
1837
                           {
1838
1839
1840
                               if(static_cast<int>(d.dim_value()) <= 0)
                                   return batch_size;
                               return d.dim_value();
1841
                           }
1842
                           return batch_size;
1843
                       });
1844
1845
1846
        if(dims.empty())
            return {shape_type};

Paul's avatar
Paul committed
1847
1848
        return {shape_type, dims};
    }
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870

    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");
        }
        }
    }
Shucai Xiao's avatar
Shucai Xiao committed
1871
1872
1873

    void check_arg_empty(const argument& arg, const std::string& msg)
    {
Shucai Xiao's avatar
Shucai Xiao committed
1874
        if(arg.empty())
Shucai Xiao's avatar
Shucai Xiao committed
1875
1876
1877
1878
        {
            MIGRAPHX_THROW(msg);
        }
    }
Paul's avatar
Paul committed
1879
1880
};

1881
program parse_onnx(const std::string& name, onnx_options options)
Paul's avatar
Paul committed
1882
1883
1884
{
    std::fstream input(name.c_str(), std::ios::in | std::ios::binary);
    onnx_parser parser;
1885
    parser.batch_size = options.batch_size;
Paul's avatar
Paul committed
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
#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
1903
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
1904
} // namespace migraphx