onnx.cpp 91.3 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>
19
#include <migraphx/pad_calc.hpp>
Paul's avatar
Paul committed
20
21

namespace migraphx {
Paul's avatar
Paul committed
22
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
23

24
25
namespace onnx = onnx_for_migraphx;

Paul's avatar
Paul committed
26
27
28
struct onnx_parser
{
    using attribute_map = std::unordered_map<std::string, onnx::AttributeProto>;
29
30
31
32
33
34
    struct node_info
    {
        attribute_map attributes{};
        std::size_t num_outputs = 1;
    };
    using node_map = std::unordered_map<std::string, onnx::NodeProto>;
Paul's avatar
Paul committed
35
    using op_func =
36
        std::function<std::vector<instruction_ref>(node_info, std::vector<instruction_ref>)>;
Paul's avatar
Paul committed
37
38
    node_map nodes;
    std::unordered_map<std::string, instruction_ref> instructions;
39
40
41
42
    program prog                  = program();
    bool is_pytorch               = false;
    std::size_t default_dim_value = 1;
    std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims;
43
    bool skip_unknown_operators = false;
Paul's avatar
Paul committed
44
45

    std::unordered_map<std::string, op_func> ops;
46
    std::unordered_map<std::string, operation> map_actv_funcs;
Paul's avatar
Paul committed
47
48
49

    onnx_parser()
    {
50
        // sort onnx operator alphabetically through name
Khalique's avatar
Khalique committed
51
        add_generic_op("Abs", op::abs{});
52
53
54
55
56
57
58
59
60
        add_generic_op("Acos", op::acos{});
        add_generic_op("Acosh", op::acosh{});
        add_generic_op("Asin", op::asin{});
        add_generic_op("Asinh", op::asinh{});
        add_generic_op("Atan", op::atan{});
        add_generic_op("Atanh", op::atanh{});
        add_generic_op("Ceil", op::ceil{});
        add_generic_op("Cos", op::cos{});
        add_generic_op("Cosh", op::cosh{});
Shucai Xiao's avatar
Shucai Xiao committed
61
        add_generic_op("Erf", op::erf{});
62
        add_generic_op("Exp", op::exp{});
Khalique's avatar
Khalique committed
63
        add_generic_op("Dropout", op::identity{});
64
65
        add_generic_op("Log", op::log{});
        add_generic_op("Floor", op::floor{});
Khalique's avatar
Khalique committed
66
        add_generic_op("Identity", op::identity{});
kahmed10's avatar
kahmed10 committed
67
        add_generic_op("Reciprocal", op::recip{});
68
69
70
71
        add_generic_op("Relu", op::relu{});
        add_generic_op("Round", op::round{});
        add_generic_op("Sigmoid", op::sigmoid{});
        add_generic_op("Sign", op::sign{});
Shucai Xiao's avatar
Shucai Xiao committed
72
        add_generic_op("Sin", op::sin{});
73
        add_generic_op("Sinh", op::sinh{});
74
        add_generic_op("Sqrt", op::sqrt{});
75
76
        add_generic_op("Tan", op::tan{});
        add_generic_op("Tanh", op::tanh{});
Paul's avatar
Paul committed
77

Khalique's avatar
Khalique committed
78
79
80
        add_binary_op("Add", op::add{});
        add_binary_op("Div", op::div{});
        add_binary_op("Mul", op::mul{});
Shucai Xiao's avatar
Shucai Xiao committed
81
        add_binary_op("Pow", op::pow{});
Shucai Xiao's avatar
Shucai Xiao committed
82
        add_binary_op("PRelu", op::prelu{});
83
        add_binary_op("Sub", op::sub{});
Khalique's avatar
Khalique committed
84

Khalique's avatar
Khalique committed
85
86
87
        add_variadic_op("Sum", op::add{});
        add_variadic_op("Max", op::max{});
        add_variadic_op("Min", op::min{});
Paul's avatar
Paul committed
88

89
        add_mem_op("AveragePool", &onnx_parser::parse_pooling);
90
91
        add_mem_op("ArgMax", &onnx_parser::parse_arg_op<op::argmax>);
        add_mem_op("ArgMin", &onnx_parser::parse_arg_op<op::argmin>);
92
        add_mem_op("BatchNormalization", &onnx_parser::parse_batchnorm);
93
        add_mem_op("Cast", &onnx_parser::parse_cast);
Khalique's avatar
Khalique committed
94
        add_mem_op("Clip", &onnx_parser::parse_clip);
95
        add_mem_op("Concat", &onnx_parser::parse_concat);
Paul's avatar
Paul committed
96
        add_mem_op("Constant", &onnx_parser::parse_constant);
97
98
99
100
        add_mem_op("ConstantFill", &onnx_parser::parse_constant_fill);
        add_mem_op("ConstantOfShape", &onnx_parser::parse_constant_of_shape);
        add_mem_op("Conv", &onnx_parser::parse_conv<op::convolution>);
        add_mem_op("ConvInteger", &onnx_parser::parse_conv<op::quant_convolution>);
kahmed10's avatar
kahmed10 committed
101
        add_mem_op("ConvTranspose", &onnx_parser::parse_conv_transpose);
102
103
        add_mem_op("Elu", &onnx_parser::parse_elu);
        add_mem_op("Expand", &onnx_parser::parse_expand);
Paul's avatar
Paul committed
104
        add_mem_op("Flatten", &onnx_parser::parse_flatten);
105
        add_mem_op("Gather", &onnx_parser::parse_gather);
Paul's avatar
Paul committed
106
        add_mem_op("Gemm", &onnx_parser::parse_gemm);
107
108
109
110
        add_mem_op("GlobalAveragePool", &onnx_parser::parse_pooling);
        add_mem_op("GlobalMaxPool", &onnx_parser::parse_pooling);
        add_mem_op("GRU", &onnx_parser::parse_gru);
        add_mem_op("ImageScaler", &onnx_parser::parse_imagescaler);
kahmed10's avatar
kahmed10 committed
111
        add_mem_op("InstanceNormalization", &onnx_parser::parse_instancenorm);
112
        add_mem_op("LeakyRelu", &onnx_parser::parse_leaky_relu);
113
        add_mem_op("LogSoftmax", &onnx_parser::parse_softmax<op::logsoftmax>);
114
115
116
117
        add_mem_op("LRN", &onnx_parser::parse_lrn);
        add_mem_op("MatMul", &onnx_parser::parse_matmul<op::dot>);
        add_mem_op("MatMulInteger", &onnx_parser::parse_matmul<op::quant_dot>);
        add_mem_op("MaxPool", &onnx_parser::parse_pooling);
kahmed10's avatar
kahmed10 committed
118
        add_mem_op("OneHot", &onnx_parser::parse_onehot);
kahmed10's avatar
kahmed10 committed
119
        add_mem_op("Range", &onnx_parser::parse_range);
Shucai Xiao's avatar
Shucai Xiao committed
120
121
122
123
124
        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
125
        add_mem_op("ReduceMean", &onnx_parser::parse_reduce_oper<op::reduce_mean>);
Shucai Xiao's avatar
Shucai Xiao committed
126
        add_mem_op("ReduceMin", &onnx_parser::parse_reduce_oper<op::reduce_min>);
Shucai Xiao's avatar
Shucai Xiao committed
127
128
129
        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);
130
131
132
133
134
135
        add_mem_op("Reshape", &onnx_parser::parse_reshape);
        add_mem_op("RNN", &onnx_parser::parse_rnn);
        add_mem_op("Pad", &onnx_parser::parse_pad);
        add_mem_op("Shape", &onnx_parser::parse_shape);
        add_mem_op("Slice", &onnx_parser::parse_slice);
        add_mem_op("Softmax", &onnx_parser::parse_softmax<op::softmax>);
136
        add_mem_op("Split", &onnx_parser::parse_split);
137
        add_mem_op("Squeeze", &onnx_parser::parse_squeeze);
kahmed10's avatar
kahmed10 committed
138
        add_mem_op("Tile", &onnx_parser::parse_tile);
139
140
141
        add_mem_op("Transpose", &onnx_parser::parse_transpose);
        add_mem_op("Unsqueeze", &onnx_parser::parse_unsqueeze);
        add_mem_op("LSTM", &onnx_parser::parse_lstm);
142
143
144
145
146
147
148

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

    void init_actv_func()
    {
149
150
151
152
153
154
        // 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
155
156
157
158
    }

    template <class F>
    void add_op(std::string name, F f)
Paul's avatar
Paul committed
159
160
161
162
163
164
165
166
167
    {
        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
168
169
170
171
172
173
174
    {
        ops.emplace(name, f);
    }

    template <class F>
    void add_mem_op(std::string name, F f)
    {
Paul's avatar
Paul committed
175
        add_op(name, [=](auto&&... xs) {
Paul's avatar
Paul committed
176
177
178
            return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
        });
    }
Khalique's avatar
Khalique committed
179

180
    template <class T>
Khalique's avatar
Khalique committed
181
    void add_binary_op(std::string name, T x)
182
    {
183
        add_op(name, [this, x](node_info info, std::vector<instruction_ref> args) {
Scott Thornton's avatar
Scott Thornton committed
184
            if(args.size() != 2)
Paul's avatar
Paul committed
185
                MIGRAPHX_THROW("binary operators should have 2 operands");
186
            if(contains(info.attributes, "broadcast") and contains(info.attributes, "axis"))
187
            {
188
                uint64_t broadcasted = parse_value(info.attributes.at("broadcast")).at<uint64_t>();
189
190
                if(broadcasted != 0)
                {
191
                    uint64_t axis = parse_value(info.attributes.at("axis")).at<uint64_t>();
Shucai Xiao's avatar
Shucai Xiao committed
192
193
                    auto l = prog.add_instruction(op::broadcast{axis, args[0]->get_shape().lens()},
                                                  args[1]);
194
195
                    return prog.add_instruction(x, args[0], l);
                }
196
                return prog.add_instruction(x, args);
197
            }
Paul's avatar
Paul committed
198
            else
199
            {
Khalique's avatar
Khalique committed
200
                return add_broadcastable_binary_op(args[0], args[1], x);
201
202
203
204
            }
        });
    }

Shucai Xiao's avatar
Shucai Xiao committed
205
206
    std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
                                                      std::vector<std::size_t> s1)
207
208
209
210
211
212
213
214
215
216
217
218
219
    {
        // 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
220
        if(s0.size() > s1.size())
221
222
223
224
225
226
        {
            s0.swap(s1);
        }

        std::vector<std::size_t> out_lens(s1);
        auto offset = s1.size() - s0.size();
Shucai Xiao's avatar
Shucai Xiao committed
227
228
229
230
        std::transform(s0.begin(),
                       s0.end(),
                       s1.begin() + offset,
                       out_lens.begin() + offset,
231
                       [&](auto a, auto b) {
Shucai Xiao's avatar
Shucai Xiao committed
232
                           if(a != b and a != 1 and b != 1)
233
                           {
Shucai Xiao's avatar
Shucai Xiao committed
234
235
236
237
238
239
                               MIGRAPHX_THROW("COMPUTE_BROADCASTLEN: shape {" +
                                              to_string_range(s0) + "} and {" +
                                              to_string_range(s1) + "} mismatch!");
                           }
                           return std::max(a, b);
                       });
240
241
242
243

        return out_lens;
    }

Shucai Xiao's avatar
Shucai Xiao committed
244
245
    instruction_ref make_contiguous(instruction_ref ins)
    {
Shucai Xiao's avatar
Shucai Xiao committed
246
        if(ins->get_shape().standard())
Shucai Xiao's avatar
Shucai Xiao committed
247
248
249
250
251
252
253
        {
            return ins;
        }

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

Khalique's avatar
Khalique committed
254
255
256
    template <class T>
    instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
    {
Khalique's avatar
Khalique committed
257
        if(arg0->get_shape().lens() != arg1->get_shape().lens())
Khalique's avatar
Khalique committed
258
259
        {
            // Get lengths for both arguments
Shucai Xiao's avatar
Shucai Xiao committed
260
261
            auto s0       = arg0->get_shape().lens();
            auto s1       = arg1->get_shape().lens();
262
            auto out_lens = compute_broadcasted_lens(s0, s1);
263
264
265
266
267
268
269
270
271

            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
272
273
274
275
276
277
            return prog.add_instruction(x, l0, l1);
        }
        else
        {
            return prog.add_instruction(x, {arg0, arg1});
        }
278
279
    }

Paul's avatar
Paul committed
280
    template <class T>
Paul's avatar
Paul committed
281
282
    void add_generic_op(std::string name, T x)
    {
283
        add_op(name, [this, x](const node_info&, std::vector<instruction_ref> args) {
Paul's avatar
Paul committed
284
285
286
287
            return prog.add_instruction(x, args);
        });
    }

Khalique's avatar
Khalique committed
288
    template <class T>
Khalique's avatar
Khalique committed
289
    void add_variadic_op(std::string name, T x)
Khalique's avatar
Khalique committed
290
    {
291
        add_op(name, [this, x](const node_info&, std::vector<instruction_ref> args) {
Khalique's avatar
Khalique committed
292
            return std::accumulate(std::next(args.begin()),
Khalique's avatar
Khalique committed
293
294
295
296
297
                                   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
298
        });
Khalique's avatar
Khalique committed
299
300
    }

kahmed10's avatar
kahmed10 committed
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
    template <class T>
    std::vector<int64_t> to_int64_vector(const std::vector<T>& input_vector)
    {
        std::vector<int64_t> output_vector(input_vector.begin(), input_vector.end());
        return output_vector;
    }

    instruction_ref
    add_bias(const std::vector<instruction_ref>& args, instruction_ref curr_ins, uint64_t axis)
    {
        if(args.size() == 3)
        {
            auto bias_bcast =
                prog.add_instruction(op::broadcast{axis, curr_ins->get_shape().lens()}, args[2]);
            return prog.add_instruction(op::add{}, curr_ins, bias_bcast);
        }
        return curr_ins;
    }

320
321
    template <class Op>
    void check_asym_padding(instruction_ref& ins,
322
                            const std::vector<int64_t>& padding,
323
324
325
                            Op& op,
                            float pad_val = 0)
    {
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
        bool asym_padding = false;
        assert(padding.size() % 2 == 0);
        size_t pad_ndims = padding.size() / 2;

        auto left_pad_it  = padding.begin();
        auto right_pad_it = left_pad_it + pad_ndims;

        for(size_t i = 0; i < pad_ndims; i++)
        {
            if(padding[i] != padding[i + pad_ndims])
            {
                asym_padding = true;
                break;
            }
        }

        if(asym_padding)
343
        {
344
345
346
347
348
349
            std::vector<int64_t> asym_pads{0, 0, 0, 0}; // don't pad N and C
            // add left pads
            asym_pads.insert(asym_pads.begin() + 2, left_pad_it, right_pad_it);
            // add right pads
            asym_pads.insert(asym_pads.begin() + pad_ndims + 4, right_pad_it, padding.end());
            ins = prog.add_instruction(op::pad{asym_pads, pad_val}, ins);
350
351
352
        }
        else
        {
353
            op.padding = std::vector<size_t>(left_pad_it, right_pad_it);
354
355
356
        }
    }

357
358
    instruction_ref
    parse_clip(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
359
    {
kahmed10's avatar
kahmed10 committed
360
361
362
363
364
365
366
        auto input_lens = args[0]->get_shape().lens();
        instruction_ref min_arg;
        instruction_ref max_arg;
        bool min_used = false;
        bool max_used = false;

        if(args.size() == 3)
Khalique's avatar
Khalique committed
367
        {
kahmed10's avatar
kahmed10 committed
368
369
370
371
            min_arg  = args[1];
            max_arg  = args[2];
            min_used = true;
            max_used = true;
Khalique's avatar
Khalique committed
372
        }
kahmed10's avatar
kahmed10 committed
373
        else if(args.size() == 2)
Khalique's avatar
Khalique committed
374
        {
kahmed10's avatar
kahmed10 committed
375
376
377
378
379
380
381
382
383
384
385
386
387
            min_arg  = args[1];
            min_used = true;
        }
        // if using previous opset for attributes
        else if(contains(info.attributes, "min") and contains(info.attributes, "max"))
        {

            float min_val = parse_value(info.attributes.at("min")).at<float>();
            float max_val = parse_value(info.attributes.at("max")).at<float>();
            min_arg       = prog.add_literal(min_val);
            max_arg       = prog.add_literal(max_val);
            min_used      = true;
            max_used      = true;
Khalique's avatar
Khalique committed
388
        }
kahmed10's avatar
kahmed10 committed
389
390
391
392
393
394
395
396
397
398
399
400
401

        if(min_used)
            min_arg = prog.add_instruction(op::multibroadcast{input_lens}, min_arg);

        if(max_used)
            max_arg = prog.add_instruction(op::multibroadcast{input_lens}, max_arg);

        if(min_used and max_used)
            return prog.add_instruction(op::clip{}, args[0], min_arg, max_arg);
        if(min_used)
            return prog.add_instruction(op::max{}, args[0], min_arg);

        return prog.add_instruction(op::identity{}, args[0]);
Khalique's avatar
Khalique committed
402
403
    }

Shucai Xiao's avatar
Shucai Xiao committed
404
    template <class Op>
405
406
    instruction_ref
    parse_softmax(const std::string&, node_info info, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
407
    {
408
        int64_t axis = 1;
409
        if(contains(info.attributes, "axis"))
410
        {
411
            axis = parse_value(info.attributes.at("axis")).at<int>();
412
413
        }

414
        return prog.add_instruction(Op{axis}, std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
415
416
    }

Shucai Xiao's avatar
Shucai Xiao committed
417
    template <class Op>
418
419
    instruction_ref
    parse_arg_op(const std::string&, node_info info, std::vector<instruction_ref> args)
420
    {
421
        int64_t axis = 0;
422
        if(contains(info.attributes, "axis"))
423
        {
424
            axis = static_cast<int64_t>(parse_value(info.attributes.at("axis")).at<int>());
425
426
        }

Shucai Xiao's avatar
Shucai Xiao committed
427
        int keep_dims = 1;
428
        if(contains(info.attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
429
        {
430
            keep_dims = parse_value(info.attributes.at("keepdims")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
431
432
        }

Shucai Xiao's avatar
Shucai Xiao committed
433
        if(keep_dims == 0)
434
        {
435
            auto ins = prog.add_instruction(Op{axis}, std::move(args));
436
            return prog.add_instruction(op::squeeze{{axis}}, ins);
437
438
439
        }
        else
        {
440
            return prog.add_instruction(Op{axis}, std::move(args));
441
        }
442
443
    }

444
445
    template <class Op>
    instruction_ref process_auto_pad_attribute(instruction_ref ins,
446
                                               node_info info,
447
                                               Op& op,
448
449
                                               std::vector<std::size_t> k_lens,
                                               std::vector<std::size_t> dilation,
450
451
                                               const std::vector<std::size_t>& in_lens,
                                               float value = 0.0f)
452
    {
453
454
455
        size_t kdims = in_lens.size() - 2;
        assert(k_lens.size() == kdims and dilation.size() == kdims);

456
        if(!contains(info.attributes, "auto_pad"))
457
458
459
460
        {
            return ins;
        }

461
        auto auto_pad = info.attributes["auto_pad"].s();
462
463
        if(auto_pad.find("SAME") != std::string::npos)
        {
464
            op.padding_mode    = op::padding_mode_t::same;
465
            bool is_same_upper = (auto_pad.find("SAME_UPPER") != std::string::npos);
466
467
468
469
470
471
472
473
474
475
476
477
            std::vector<int64_t> padding(2 * kdims);

            for(size_t i = 0; i < padding.size() / 2; i++)
            {
                calculate_padding(i,
                                  padding,
                                  in_lens[i + 2],
                                  op.stride[i],
                                  dilation[i],
                                  k_lens[i],
                                  is_same_upper);
            }
478

479
            check_asym_padding(ins, padding, op, value);
480
481
482
483
484
        }

        return ins;
    }

kahmed10's avatar
kahmed10 committed
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
    void calc_reflect_indices(std::vector<int>& indices, const int64_t num_dims)
    {
        int k         = 0;
        bool reversed = false;
        // in reflect padding, if the num_pads > num_dims,
        // compute the extra pad indices periodically, ex. ( 1, 2, 3, 2, 1, 0)
        for(int& idx : indices)
        {
            if(k == num_dims - 1)
                reversed = true;
            if(k == 0)
                reversed = false;
            if(reversed)
                k--;
            else
                k++;
            idx = k;
        }
    }

    instruction_ref reflect_pad(const std::vector<int64_t>& pads, instruction_ref input)
    {
        size_t num_dims = pads.size() / 2;
        std::vector<int> ldims(pads.begin(), pads.begin() + num_dims);
        std::vector<int> rdims(pads.begin() + num_dims, pads.end());
        assert(ldims.size() == rdims.size());

        std::vector<int64_t> axes(num_dims);
        std::iota(axes.begin(), axes.end(), int64_t{0});

        // iterate over dimensions, starting from lowest dimension
        for(int64_t i = num_dims - 1; i >= 0; i--)
        {
            auto axis   = i;
            auto lcount = ldims.at(i);
            auto rcount = rdims.at(i);
            if(lcount == 0 and rcount == 0) // no padding for current dim
                continue;

            // calculate starts and ends for each iteration since shape may change
            std::vector<size_t> dims = input->get_shape().lens();
            std::vector<int64_t> starts(axes.size(), 0);
            std::vector<int64_t> ends(dims.begin(), dims.end());
            std::vector<instruction_ref> slices;

            auto starts_it = starts.begin() + i;
            auto ends_it   = ends.begin() + i;
            auto dims_it   = dims.begin() + i;

            std::vector<int> l_indices(lcount);
            std::vector<int> r_indices(rcount);

            // compute slice indices in a periodic fashion
            calc_reflect_indices(l_indices, *dims_it);
            calc_reflect_indices(r_indices, *dims_it);

            for(int idx : l_indices)
            {
                *starts_it = idx;
                *ends_it   = *starts_it + 1;
                slices.push_back(prog.add_instruction(op::slice{axes, starts, ends}, input));
            }
            // when padding on the left side, the outermost pad should be at the beginning
            std::reverse(slices.begin(), slices.end());
            slices.push_back(input);
            for(int idx : r_indices)
            {
                *starts_it = *dims_it - idx - 1;
                *ends_it   = *starts_it + 1;
                slices.push_back(prog.add_instruction(op::slice{axes, starts, ends}, input));
            }
            input = prog.add_instruction(op::concat{axis}, slices);
        }
        return input;
    }

561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
    void check_attr_sizes(size_t kdims, size_t attr_size, const std::string& error_msg)
    {
        if(kdims != attr_size)
        {
            MIGRAPHX_THROW(error_msg + " k-dims: " + to_string(kdims) +
                           " attribute size: " + to_string(attr_size));
        }
    }

    template <class Op>
    void recalc_conv_attributes(Op& op, size_t kdims)
    {
        if(op.padding.size() != kdims)
        {
            op.padding.resize(kdims);
            std::fill_n(op.padding.begin(), kdims, 0);
        }
        if(op.stride.size() != kdims)
        {
            op.stride.resize(kdims);
            std::fill_n(op.stride.begin(), kdims, 1);
        }
        if(op.dilation.size() != kdims)
        {
            op.dilation.resize(kdims);
            std::fill_n(op.dilation.begin(), kdims, 1);
        }
    }

590
    template <class Op>
Paul's avatar
Paul committed
591
    instruction_ref
592
    parse_conv(const std::string&, node_info info, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
593
    {
594
        Op op;
595
596
        auto l0      = args[0];
        auto weights = args[1];
597
598
599
600
        auto in_lens = l0->get_shape().lens();
        assert(in_lens.size() > 2);
        auto kdims = in_lens.size() - 2;

601
        std::vector<int64_t> padding;
602
        if(contains(info.attributes, "pads"))
Paul's avatar
Paul committed
603
        {
604
            if(contains(info.attributes, "auto_pad"))
605
            {
606
607
                auto s = info.attributes["auto_pad"].s();
                if(contains(info.attributes, "pads") and to_upper(s) != "NOTSET")
608
                {
609
610
                    MIGRAPHX_THROW(
                        "PARSE_CONV: auto_pad and padding cannot be specified simultaneously");
611
                }
612
            }
613
            op.padding.clear();
614
            copy(info.attributes["pads"].ints(), std::back_inserter(padding));
615
            check_attr_sizes(kdims, padding.size() / 2, "PARSE_CONV: inconsistent paddings");
616
            check_asym_padding(l0, padding, op);
Paul's avatar
Paul committed
617
        }
618
        if(contains(info.attributes, "strides"))
Paul's avatar
Paul committed
619
        {
620
621
622
            op.stride.clear();
            copy(info.attributes["strides"].ints(), std::back_inserter(op.stride));
            check_attr_sizes(kdims, op.stride.size(), "PARSE_CONV: inconsistent strides");
Paul's avatar
Paul committed
623
        }
624
        if(contains(info.attributes, "dilations"))
Paul's avatar
Paul committed
625
        {
626
627
628
            op.dilation.clear();
            copy(info.attributes["dilations"].ints(), std::back_inserter(op.dilation));
            check_attr_sizes(kdims, op.dilation.size(), "PARSE_CONV: inconsistent dilations");
Paul's avatar
Paul committed
629
        }
630
        if(contains(info.attributes, "auto_pad"))
631
        {
632
            auto weight_lens = weights->get_shape().lens();
633

634
            std::vector<std::size_t> k_lens(weight_lens.begin() + 2, weight_lens.end());
635
            l0 = process_auto_pad_attribute(l0, info, op, k_lens, op.dilation, in_lens);
636
        }
637
        if(contains(info.attributes, "group"))
Khalique's avatar
Khalique committed
638
        {
639
            op.group = parse_value(info.attributes.at("group")).at<int>();
Khalique's avatar
Khalique committed
640
        }
kahmed10's avatar
kahmed10 committed
641

642
643
        recalc_conv_attributes(op, kdims);

kahmed10's avatar
kahmed10 committed
644
645
646
647
        auto l1 = prog.add_instruction(op, l0, args[1]);
        return add_bias(args, l1, 1);
    }

648
649
    instruction_ref
    parse_conv_transpose(const std::string&, node_info info, std::vector<instruction_ref> args)
kahmed10's avatar
kahmed10 committed
650
651
652
653
654
    {
        op::deconvolution op;
        auto l0 = args[0];
        std::vector<std::int64_t> padding;
        bool asymm_padding = false;
655
        if(contains(info.attributes, "pads"))
kahmed10's avatar
kahmed10 committed
656
        {
657
            if(contains(info.attributes, "auto_pad"))
kahmed10's avatar
kahmed10 committed
658
            {
659
660
                auto s = info.attributes["auto_pad"].s();
                if(contains(info.attributes, "pads") and to_upper(s) != "NOTSET")
kahmed10's avatar
kahmed10 committed
661
662
663
664
                {
                    MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
                }
            }
665
            copy(info.attributes["pads"].ints(), std::back_inserter(padding));
kahmed10's avatar
kahmed10 committed
666
667
668
669
670
671
672
673
674
675
676
677
678
679
            if(padding.size() != 4)
            {
                MIGRAPHX_THROW("padding should have 4 values");
            }
            if(padding[0] != padding[2] || padding[1] != padding[3])
            {
                asymm_padding = true;
            }
            else
            {
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
            }
        }
680
        if(contains(info.attributes, "strides"))
kahmed10's avatar
kahmed10 committed
681
        {
682
            copy(info.attributes["strides"].ints(), op.stride.begin());
kahmed10's avatar
kahmed10 committed
683
        }
684
        if(contains(info.attributes, "dilations"))
Paul's avatar
Paul committed
685
        {
686
            copy(info.attributes["dilations"].ints(), op.dilation.begin());
Paul's avatar
Paul committed
687
        }
688
        if(contains(info.attributes, "auto_pad"))
kahmed10's avatar
kahmed10 committed
689
        {
690
691
            auto s = info.attributes["auto_pad"].s();
            if(contains(info.attributes, "pads") and to_upper(s) != "NOTSET")
kahmed10's avatar
kahmed10 committed
692
693
694
695
696
697
698
699
700
701
            {
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
            }

            if(s.find("SAME") != std::string::npos)
            {
                op.padding_mode = op::padding_mode_t::same;
            }
        }

702
        if(contains(info.attributes, "group"))
kahmed10's avatar
kahmed10 committed
703
        {
704
            op.group = parse_value(info.attributes.at("group")).at<int>();
kahmed10's avatar
kahmed10 committed
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
        }

        auto l1                   = prog.add_instruction(op, l0, args[1]);
        std::vector<int64_t> dims = to_int64_vector(l1->get_shape().lens());
        std::vector<int64_t> curr_shape{dims[2], dims[3]};
        if(asymm_padding)
        {
            op::slice slice_op;
            slice_op.axes   = {0, 1, 2, 3};
            slice_op.starts = {0, 0, 0 + padding[0], 0 + padding[1]};
            slice_op.ends   = {
                dims[0], dims[1], curr_shape[0] - padding[2], curr_shape[1] - padding[3]};

            l1 = prog.add_instruction(slice_op, l1);
        }

721
        if(contains(info.attributes, "output_padding"))
kahmed10's avatar
kahmed10 committed
722
723
        {
            std::vector<int64_t> output_padding;
724
            copy(info.attributes["output_padding"].ints(), std::back_inserter(output_padding));
kahmed10's avatar
kahmed10 committed
725
726
727
728
            output_padding = {0, 0, 0, 0, 0, 0, output_padding[0], output_padding[1]};
            l1             = prog.add_instruction(op::pad{output_padding}, l1);
        }

729
        if(contains(info.attributes, "output_shape"))
kahmed10's avatar
kahmed10 committed
730
731
        {
            std::vector<int64_t> output_shape;
732
            copy(info.attributes["output_shape"].ints(), std::back_inserter(output_shape));
kahmed10's avatar
kahmed10 committed
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
            dims       = to_int64_vector(l1->get_shape().lens());
            curr_shape = {dims[2], dims[3]};
            if(curr_shape != output_shape)
            {
                std::vector<int64_t> target_padding = {0,
                                                       0,
                                                       0,
                                                       0,
                                                       0,
                                                       0,
                                                       output_shape[0] - curr_shape[0],
                                                       output_shape[1] - curr_shape[1]};
                l1 = prog.add_instruction(op::pad{target_padding}, l1);
            }
        }

        return add_bias(args, l1, 1);
Paul's avatar
Paul committed
750
    }
Paul's avatar
Paul committed
751

752
753
    instruction_ref
    parse_pooling(const std::string& name, node_info info, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
754
    {
Khalique's avatar
Khalique committed
755
        op::pooling op{ends_with(name, "MaxPool") ? "max" : "average"};
756
757
758
759
760
        auto l0      = args[0];
        auto in_lens = l0->get_shape().lens();
        assert(in_lens.size() > 2);
        auto kdims = in_lens.size() - 2;

Khalique's avatar
Khalique committed
761
        if(starts_with(name, "Global"))
762
        {
763
            op.lengths = std::vector<size_t>(in_lens.begin() + 2, in_lens.end());
764
        }
765

766
        if(contains(info.attributes, "pads"))
Paul's avatar
Paul committed
767
        {
768
            if(contains(info.attributes, "auto_pad"))
769
            {
770
                auto s = info.attributes["auto_pad"].s();
771
772
773
774
775
776
                if(to_upper(s) != "NOTSET")
                {
                    MIGRAPHX_THROW(
                        "PARSE_POOLING: auto_pad and padding cannot be specified simultaneously");
                }
            }
777
778
            op.padding.clear();
            std::vector<int64_t> padding;
779
            copy(info.attributes["pads"].ints(), std::back_inserter(padding));
780
781
            check_attr_sizes(kdims, padding.size() / 2, "PARSE_POOLING: inconsistent paddings");

782
783
784
785
            float pad_val = 0;
            if(op.mode == "max")
                pad_val = std::numeric_limits<float>::lowest();
            check_asym_padding(l0, padding, op, pad_val);
786
            in_lens = l0->get_shape().lens();
Paul's avatar
Paul committed
787
        }
788

789
        if(contains(info.attributes, "strides"))
Paul's avatar
Paul committed
790
        {
791
792
793
            op.stride.clear();
            copy(info.attributes["strides"].ints(), std::back_inserter(op.stride));
            check_attr_sizes(kdims, op.stride.size(), "PARSE_POOLING: inconsistent strides");
Paul's avatar
Paul committed
794
        }
795
        if(contains(info.attributes, "kernel_shape"))
Paul's avatar
Paul committed
796
        {
797
798
799
            op.lengths.clear();
            copy(info.attributes["kernel_shape"].ints(), std::back_inserter(op.lengths));
            check_attr_sizes(kdims, op.lengths.size(), "PARSE_POOLING: inconsistent lengths");
Paul's avatar
Paul committed
800
        }
801

802
        if(contains(info.attributes, "auto_pad"))
803
        {
804
805
            op.padding.clear();
            float val = 0.0f;
806
807
808
809
810
811
            // MaxPool
            if(op.mode == "max")
            {
                val = std::numeric_limits<float>::lowest();
            }

812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
            l0      = process_auto_pad_attribute(l0, info, op, op.lengths, {1, 1}, in_lens, val);
            in_lens = l0->get_shape().lens();
        }

        if(op.padding.size() != kdims)
        {
            op.padding.resize(kdims);
            std::fill_n(op.padding.begin(), kdims, 0);
        }
        if(op.stride.size() != kdims)
        {
            op.stride.resize(kdims);
            std::fill_n(op.stride.begin(), kdims, 1);
        }

        for(size_t i = 0; i < kdims; i++)
        {
            if(op.lengths[i] > in_lens[i + 2] + 2 * op.padding[i])
                MIGRAPHX_THROW("PARSE_POOLING: kernel shape is too large");
831
832
        }

833
        return prog.add_instruction(op, l0);
Paul's avatar
Paul committed
834
835
    }

Paul's avatar
Paul committed
836
    instruction_ref
837
    parse_reshape(const std::string&, node_info info, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
838
    {
839
        op::reshape op;
Paul's avatar
Paul committed
840
841
        if(args.size() == 1)
        {
842
            literal s = parse_value(info.attributes.at("shape"));
843
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
844
845
846
        }
        if(args.size() == 2)
        {
Paul's avatar
Paul committed
847
            auto s = args[1]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
848
            check_arg_empty(s, "Reshape: dynamic shape is not supported");
Paul's avatar
Paul committed
849
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
850
        }
851

Shucai Xiao's avatar
Shucai Xiao committed
852
        return prog.add_instruction(op, make_contiguous(args[0]));
Paul's avatar
Paul committed
853
854
    }

Paul's avatar
Paul committed
855
    instruction_ref
856
    parse_flatten(const std::string&, node_info info, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
857
    {
858
        int64_t axis = 1;
859
        if(contains(info.attributes, "axis"))
Paul's avatar
Paul committed
860
        {
861
            axis = parse_value(info.attributes.at("axis")).at<int>();
Paul's avatar
Paul committed
862
        }
863
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
864
865
    }

866
    instruction_ref
867
    parse_squeeze(const std::string&, node_info info, std::vector<instruction_ref> args)
868
869
    {
        op::squeeze op;
870
        literal s = parse_value(info.attributes.at("axes"));
871
        s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
872
        return prog.add_instruction(op, make_contiguous(args[0]));
873
874
875
    }

    instruction_ref
876
    parse_unsqueeze(const std::string&, node_info info, std::vector<instruction_ref> args)
877
878
    {
        op::unsqueeze op;
879
        literal s = parse_value(info.attributes.at("axes"));
880
        s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
881
        return prog.add_instruction(op, make_contiguous(args[0]));
882
883
    }

Scott Thornton's avatar
Scott Thornton committed
884
    instruction_ref
885
    parse_concat(const std::string&, node_info info, std::vector<instruction_ref> args)
Scott Thornton's avatar
Scott Thornton committed
886
    {
Shucai Xiao's avatar
Shucai Xiao committed
887
        // change to hande axis to be negative values
888
        if(!contains(info.attributes, "axis"))
Shucai Xiao's avatar
Shucai Xiao committed
889
890
891
892
        {
            MIGRAPHX_THROW("PARSE_CONCAT: attribute axis is required!");
        }

893
        int axis = parse_value(info.attributes.at("axis")).at<int>();
Scott Thornton's avatar
Scott Thornton committed
894
895
896
        op::concat op{axis};
        return prog.add_instruction(op, std::move(args));
    }
897

898
    instruction_ref
899
    parse_gather(const std::string&, node_info info, std::vector<instruction_ref> args)
900
    {
901
        int axis = 0;
902
        if(contains(info.attributes, "axis"))
903
        {
904
            axis = parse_value(info.attributes.at("axis")).at<int>();
905
        }
906

907
        op::gather op{axis};
Shucai Xiao's avatar
Shucai Xiao committed
908
        return prog.add_instruction(op, make_contiguous(args[0]), make_contiguous(args[1]));
909
910
    }

911
    instruction_ref
912
    parse_slice(const std::string&, node_info info, std::vector<instruction_ref> args)
913
914
    {
        op::slice op;
Shucai Xiao's avatar
Shucai Xiao committed
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936

        // slice can have up to 5 inputs, we first check the 5th one
        // to decide whether MIGRAPHX can handle this slice
        if(args.size() == 5)
        {
            migraphx::argument step_arg = args.back()->eval();
            check_arg_empty(step_arg, "PARSE_SLICE: cannot handle variable steps for slice");
            std::vector<int> steps;
            step_arg.visit([&](auto s) { steps.assign(s.begin(), s.end()); });
            if(!std::all_of(steps.begin(), steps.end(), [](auto s) { return s == 1; }))
            {
                MIGRAPHX_THROW("PARSE_SLICE: cannot handle step other than 1");
            }
        }

        if(args.size() >= 4)
        {
            migraphx::argument axes_arg = args.at(3)->eval();
            check_arg_empty(axes_arg, "PARSE_SLICE: cannot handle variable axes for slice");
            axes_arg.visit([&](auto s) { op.axes.assign(s.begin(), s.end()); });
        }
        else if(contains(info.attributes, "axes"))
937
        {
938
            literal s = parse_value(info.attributes.at("axes"));
939
940
            s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        }
Shucai Xiao's avatar
Shucai Xiao committed
941
942

        if(args.size() >= 3)
Khalique's avatar
Khalique committed
943
        {
Shucai Xiao's avatar
Shucai Xiao committed
944
945
946
            migraphx::argument end_arg = args.at(2)->eval();
            check_arg_empty(end_arg, "PARSE_SLICE: cannot handle variable ends for slice");
            end_arg.visit([&](auto s) { op.ends.assign(s.begin(), s.end()); });
Khalique's avatar
Khalique committed
947
        }
Shucai Xiao's avatar
Shucai Xiao committed
948
        else if(contains(info.attributes, "ends"))
949
        {
950
951
            literal s = parse_value(info.attributes.at("ends"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.ends)); });
952
        }
Shucai Xiao's avatar
Shucai Xiao committed
953
954
955
956
957
958
959
960

        if(args.size() >= 2)
        {
            migraphx::argument start_arg = args.at(1)->eval();
            check_arg_empty(start_arg, "PARSE_SLICE: cannot handle variable starts for slice");
            start_arg.visit([&](auto s) { op.starts.assign(s.begin(), s.end()); });
        }
        else if(contains(info.attributes, "starts"))
961
        {
962
            literal s = parse_value(info.attributes.at("starts"));
963
964
            s.visit([&](auto v) { copy(v, std::back_inserter(op.starts)); });
        }
Shucai Xiao's avatar
Shucai Xiao committed
965

kahmed10's avatar
kahmed10 committed
966
967
968
969
970
971
972
        if(op.axes.empty())
        {
            std::vector<int64_t> axes(args[0]->get_shape().lens().size());
            std::iota(axes.begin(), axes.end(), int64_t{0});
            op.axes = axes;
        }

973
974
975
        return prog.add_instruction(op, args[0]);
    }

976
977
    instruction_ref
    parse_constant(const std::string&, node_info info, const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
978
    {
979
        literal v = parse_value(info.attributes.at("value"));
980
        // return empty literal
Shucai Xiao's avatar
Shucai Xiao committed
981
        if(v.get_shape().elements() == 0)
982
983
984
985
        {
            return prog.add_literal(literal{});
        }

986
        auto dim_size = info.attributes.at("value").t().dims_size();
987
988
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
989
        {
990
            migraphx::shape scalar_shape{v.get_shape().type()};
991
992
993
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

Paul's avatar
Paul committed
994
995
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
996

Paul's avatar
Paul committed
997
    instruction_ref
998
    parse_gemm(const std::string&, node_info info, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
999
1000
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
1001
        float beta  = 1.0f;
Paul's avatar
Paul committed
1002
1003
        bool transa = false;
        bool transb = false;
1004
        if(contains(info.attributes, "alpha"))
Paul's avatar
Paul committed
1005
        {
1006
            alpha = parse_value(info.attributes.at("alpha")).at<float>();
Paul's avatar
Paul committed
1007
        }
1008
        if(contains(info.attributes, "beta"))
Paul's avatar
Paul committed
1009
        {
1010
            beta = parse_value(info.attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
1011
        }
1012
        if(contains(info.attributes, "transA"))
Paul's avatar
Paul committed
1013
        {
1014
            transa = parse_value(info.attributes.at("transA")).at<bool>();
Paul's avatar
Paul committed
1015
        }
1016
        if(contains(info.attributes, "transB"))
Paul's avatar
Paul committed
1017
        {
1018
            transb = parse_value(info.attributes.at("transB")).at<bool>();
Paul's avatar
Paul committed
1019
        }
1020
1021
1022
1023
1024
1025

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

1026
1027
        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
1028
1029
        if(args.size() == 3)
        {
1030
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
1031
            {
Shucai Xiao's avatar
Shucai Xiao committed
1032
                auto out_lens   = l1->get_shape().lens();
1033
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
1034
                auto l3         = args[2];
Shucai Xiao's avatar
Shucai Xiao committed
1035
1036
                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
1037
                {
1038
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
1039
                }
1040
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
1041
            }
Paul's avatar
Paul committed
1042
        }
1043
1044

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

1047
    template <class Op>
1048
    instruction_ref
1049
    parse_matmul(const std::string&, const node_info&, std::vector<instruction_ref> args)
1050
    {
Shucai Xiao's avatar
Shucai Xiao committed
1051
1052
        auto l0      = args[0];
        auto l1      = args[1];
1053
1054
1055
1056
1057
        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
1058
        if(l0_lens.size() == 1)
1059
1060
1061
1062
1063
1064
1065
        {
            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
1066
        if(l1_lens.size() == 1)
1067
1068
1069
1070
1071
1072
1073
1074
        {
            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
1075
        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
1076
1077
1078
1079
1080
1081
        {
            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);
1082
            l0_broadcasted_lens = output_lens;
1083
            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
1084
            l1_broadcasted_lens = output_lens;
1085
            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
Shucai Xiao's avatar
Shucai Xiao committed
1086
            if(l0_lens != l0_broadcasted_lens)
1087
1088
1089
            {
                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
            }
Shucai Xiao's avatar
Shucai Xiao committed
1090
            if(l1_lens != l1_broadcasted_lens)
1091
1092
1093
1094
1095
            {
                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
            }
        }

1096
        auto dot_res     = prog.add_instruction(Op{1, 0}, bl0, bl1);
1097
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
1098
        if(is_a_prepended)
1099
1100
1101
1102
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
1103
        if(is_b_appended)
1104
1105
1106
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
1107

1108
1109
1110
        return dot_res;
    }

1111
    instruction_ref
1112
    parse_batchnorm(const std::string&, node_info info, std::vector<instruction_ref> args)
1113
    {
Scott Thornton's avatar
Scott Thornton committed
1114
1115
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
1116
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
1117
        if(contains(info.attributes, "epsilon"))
1118
        {
1119
            epsilon = parse_value(info.attributes.at("epsilon")).at<float>();
1120
        }
1121
        if(contains(info.attributes, "momentum"))
1122
        {
1123
            momentum = parse_value(info.attributes.at("momentum")).at<float>();
1124
        }
1125
        if(contains(info.attributes, "spatial"))
1126
        {
1127
            bn_mode = (parse_value(info.attributes.at("spatial")).at<uint64_t>() > 0)
1128
1129
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
1130
        }
Paul's avatar
Paul committed
1131
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
1132
        return prog.add_instruction(op, std::move(args));
1133
1134
    }

1135
1136
    instruction_ref
    parse_instancenorm(const std::string&, node_info info, std::vector<instruction_ref> args)
kahmed10's avatar
kahmed10 committed
1137
1138
1139
1140
1141
1142
    {
        // y = scale * ( x - mean ) / sqrt ( variance + epsilon ) + bias
        // mean = reduce_mean({H, W}, x)
        // variance = reduce_mean({H, W}, (x - mean)^2)

        float epsilon = 1e-5f;
1143
        if(contains(info.attributes, "epsilon"))
kahmed10's avatar
kahmed10 committed
1144
        {
1145
            epsilon = parse_value(info.attributes.at("epsilon")).at<float>();
kahmed10's avatar
kahmed10 committed
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
        }
        auto x     = args[0];
        auto scale = args[1];
        auto bias  = args[2];
        auto dims  = x->get_shape().lens();

        auto mean            = prog.add_instruction(op::reduce_mean{{2, 3}}, x);
        auto mean_bcast      = prog.add_instruction(op::multibroadcast{dims}, mean);
        auto l0              = prog.add_instruction(op::sqdiff{}, x, mean_bcast);
        auto variance        = prog.add_instruction(op::reduce_mean{{2, 3}}, l0);
        auto l1              = prog.add_instruction(op::sub{}, x, mean_bcast);
        auto epsilon_literal = prog.add_literal(epsilon);
        auto epsilon_bcast   = prog.add_instruction(op::multibroadcast{dims}, epsilon_literal);
        auto variance_bcast  = prog.add_instruction(op::multibroadcast{dims}, variance);
        auto l2              = prog.add_instruction(op::add{}, variance_bcast, epsilon_bcast);
        auto l3              = prog.add_instruction(op::rsqrt{}, l2);
        auto l4              = prog.add_instruction(op::mul{}, l1, l3);
        auto scale_bcast     = prog.add_instruction(op::broadcast{1, dims}, scale);
        ;
        auto bias_bcast = prog.add_instruction(op::broadcast{1, dims}, bias);
        auto l5         = prog.add_instruction(op::mul{}, l4, scale_bcast);
        return prog.add_instruction(op::add{}, l5, bias_bcast);
    }

1170
1171
    instruction_ref
    parse_leaky_relu(const std::string&, node_info info, std::vector<instruction_ref> args)
1172
    {
Khalique's avatar
Khalique committed
1173
        float alpha = 0.01; // default alpha val for leaky relu
1174
        if(contains(info.attributes, "alpha"))
1175
        {
1176
            alpha = parse_value(info.attributes.at("alpha")).at<float>();
1177
1178
1179
1180
1181
        }
        op::leaky_relu op{alpha};
        return prog.add_instruction(op, args.front());
    }

1182
    instruction_ref parse_elu(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1183
1184
    {
        float alpha = 1.0; // default alpha val for elu
1185
        if(contains(info.attributes, "alpha"))
Khalique's avatar
Khalique committed
1186
        {
1187
            alpha = parse_value(info.attributes.at("alpha")).at<float>();
Khalique's avatar
Khalique committed
1188
1189
1190
1191
1192
        }
        op::elu op{alpha};
        return prog.add_instruction(op, args.front());
    }

1193
    instruction_ref parse_lrn(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1194
1195
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
1196
1197
1198
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
1199
1200
1201
1202
1203
1204
1205
1206
        if(contains(info.attributes, "alpha"))
            alpha = parse_value(info.attributes.at("alpha")).at<float>();
        if(contains(info.attributes, "beta"))
            beta = parse_value(info.attributes.at("beta")).at<float>();
        if(contains(info.attributes, "bias"))
            bias = parse_value(info.attributes.at("bias")).at<float>();
        if(contains(info.attributes, "size"))
            size = parse_value(info.attributes.at("size")).at<int>();
Khalique's avatar
Khalique committed
1207
1208
1209
1210
        op::lrn op{alpha, beta, bias, size};
        return prog.add_instruction(op, args.front());
    }

1211
1212
    instruction_ref
    parse_imagescaler(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1213
1214
1215
    {
        float scale = 1.0;
        std::vector<float> bias{};
1216
        if(contains(info.attributes, "scale"))
Khalique's avatar
Khalique committed
1217
        {
1218
            scale = parse_value(info.attributes.at("scale")).at<float>();
Khalique's avatar
Khalique committed
1219
1220
        }

1221
        if(contains(info.attributes, "bias"))
Khalique's avatar
Khalique committed
1222
        {
1223
            auto&& bias_floats = info.attributes["bias"].floats();
Khalique's avatar
Khalique committed
1224
1225
            bias               = std::vector<float>(bias_floats.begin(), bias_floats.end());
        }
Shucai Xiao's avatar
Shucai Xiao committed
1226
1227
1228
        auto input_shape       = args.front()->get_shape();
        auto const& input_lens = input_shape.lens();
        auto input_type        = input_shape.type();
Khalique's avatar
Khalique committed
1229

Shucai Xiao's avatar
Shucai Xiao committed
1230
1231
        auto scale_val = prog.add_literal(literal{shape{input_type}, {scale}});
        auto bias_vals = prog.add_literal(literal{shape{input_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
1232

1233
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
1234
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
1235
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
1236
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
1237
    }
Khalique's avatar
Khalique committed
1238

Khalique's avatar
Khalique committed
1239
    instruction_ref
1240
    parse_transpose(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1241
1242
    {
        std::vector<int64_t> perm{};
1243
        if(contains(info.attributes, "perm"))
Khalique's avatar
Khalique committed
1244
        {
1245
            auto&& perm_vals = info.attributes["perm"].ints();
Khalique's avatar
Khalique committed
1246
1247
            perm             = std::vector<int64_t>(perm_vals.begin(), perm_vals.end());
        }
Paul's avatar
Paul committed
1248
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
1249
1250
    }

1251
    instruction_ref parse_pad(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1252
1253
    {
        std::vector<int64_t> pads{};
1254
1255
1256
1257
1258
1259
1260
        if(args.size() >= 2)
        {
            auto pad_arg = args.at(1)->eval();
            check_arg_empty(pad_arg, "PARSE_PAD: pad input must be constant");
            pad_arg.visit([&](auto v) { pads.assign(v.begin(), v.end()); });
        }
        else if(contains(info.attributes, "pads"))
Khalique's avatar
Khalique committed
1261
        {
1262
            auto&& pad_vals = info.attributes["pads"].ints();
Khalique's avatar
Khalique committed
1263
1264
            pads            = std::vector<int64_t>(pad_vals.begin(), pad_vals.end());
        }
1265
1266
1267
1268
1269
        else
        {
            MIGRAPHX_THROW("PARSE_PAD: pad must be available");
        }

1270
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
1271
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
1272
1273
1274
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
1275

kahmed10's avatar
kahmed10 committed
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
        if(contains(info.attributes, "mode"))
        {
            auto mode = info.attributes.at("mode").s();
            if(mode == "reflect")
                return reflect_pad(pads, args.front());
            if(mode != "constant")
            {
                MIGRAPHX_THROW(
                    "PARSE_PAD: migraphx currently only supports constant and reflect padding");
            }
        }

1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
        float value = 0.0f;
        // third input is the value
        if(args.size() == 3)
        {
            auto val_ins = args.at(2);
            if(!val_ins->can_eval())
            {
                MIGRAPHX_THROW("PARSE_PAD: input value must be constant");
            }
            auto val_arg = val_ins->eval();
            if(val_arg.get_shape().elements() != 1)
            {
                MIGRAPHX_THROW("PARSE_PAD: value should contain only one element");
            }
            value = val_arg.at<float>();
        }
        else if(contains(info.attributes, "value"))
Khalique's avatar
Khalique committed
1305
        {
1306
            value = parse_value(info.attributes.at("value")).at<float>();
Khalique's avatar
Khalique committed
1307
        }
1308

Khalique's avatar
Khalique committed
1309
1310
        return prog.add_instruction(migraphx::op::pad{pads, value}, args.front());
    }
1311
1312
1313
    // Use a literal instruction to replace the shape since, output of
    // shape operator are literals in migraphx
    instruction_ref
1314
    parse_shape(const std::string&, const node_info&, std::vector<instruction_ref> args)
1315
1316
    {
        if(args.size() != 1)
1317
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
        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
1330
1331
    instruction_ref
    parse_constant_fill(const std::string&, node_info info, std::vector<instruction_ref> args)
1332
1333
1334
1335
1336
    {
        int input_as_shape = 0;
        int dtype          = 1;
        float value        = 0.0f;

1337
        if(contains(info.attributes, "dtype"))
1338
        {
1339
            dtype = parse_value(info.attributes.at("dtype")).at<int>();
1340
        }
Shucai Xiao's avatar
Shucai Xiao committed
1341
        shape::type_t type = get_type(dtype);
1342

1343
        if(contains(info.attributes, "input_as_shape"))
1344
        {
1345
            input_as_shape = parse_value(info.attributes.at("input_as_shape")).at<int>();
1346
1347
        }

1348
        if(contains(info.attributes, "value"))
1349
        {
1350
            value = parse_value(info.attributes.at("value")).at<float>();
1351
1352
        }

1353
        if(contains(info.attributes, "extra_shape"))
Shucai Xiao's avatar
Shucai Xiao committed
1354
        {
1355
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
1356
1357
        }

1358
1359
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
1360
            if(args.size() != 1)
1361
            {
1362
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
1363
1364
            }

1365
            if(contains(info.attributes, "shape"))
Shucai Xiao's avatar
Shucai Xiao committed
1366
            {
1367
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
1368
                               "at the same time");
1369
1370
            }

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

1374
1375
1376
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
1377
1378
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
1379
1380
1381
        }
        else if(input_as_shape == 0)
        {
1382
            if(!contains(info.attributes, "shape"))
Shucai Xiao's avatar
Shucai Xiao committed
1383
            {
1384
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
1385
1386
            }

1387
            literal ls = parse_value(info.attributes.at("shape"));
1388
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
1389
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
1390
            migraphx::shape s{type, dims};
1391
1392
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
1393
1394
1395
        }
        else
        {
1396
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
1397
1398
1399
        }
    }

1400
1401
    instruction_ref
    parse_constant_of_shape(const std::string&, node_info info, std::vector<instruction_ref> args)
1402
1403
    {
        literal l_val{};
1404
        if(contains(info.attributes, "value"))
1405
        {
1406
            l_val = parse_value(info.attributes.at("value"));
Shucai Xiao's avatar
Shucai Xiao committed
1407
            if(l_val.get_shape().elements() != 1)
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
            {
                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();
1419

Shucai Xiao's avatar
Shucai Xiao committed
1420
        if(args.empty())
1421
        {
Shucai Xiao's avatar
Shucai Xiao committed
1422
            MIGRAPHX_THROW("ConstantOfShape : must have 1 input!");
1423
1424
1425
        }
        else
        {
1426
1427
            migraphx::shape s;
            // empty input tensor, output is a scalar
Shucai Xiao's avatar
Shucai Xiao committed
1428
            if(args[0]->get_shape().elements() == 0)
1429
            {
1430
                s = migraphx::shape{type, {1}, {0}};
1431
            }
1432
1433
1434
            else
            {
                migraphx::argument in = args[0]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
1435
                check_arg_empty(in, "ConstantOfShape: dynamic shape is not supported");
1436

1437
1438
1439
1440
                std::vector<std::size_t> dims;
                in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
                s = migraphx::shape{type, dims};
            }
1441

Shucai Xiao's avatar
Shucai Xiao committed
1442
            literal l_out{};
1443
            l_val.visit([&](auto val) {
Shucai Xiao's avatar
Shucai Xiao committed
1444
                using val_type = std::remove_cv_t<typename decltype(val)::value_type>;
1445
                // l_val contains only one element
1446
                std::vector<val_type> out_vec(s.elements(), val.front());
1447
1448
1449
1450
1451
1452
1453
                l_out = literal(s, out_vec);
            });

            return prog.add_literal(l_out);
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1454
    instruction_ref
1455
    parse_expand(const std::string&, const node_info&, std::vector<instruction_ref> args)
1456
    {
Shucai Xiao's avatar
Shucai Xiao committed
1457
        auto in_lens             = args[0]->get_shape().lens();
1458
        migraphx::argument arg_s = args[1]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
1459
        check_arg_empty(arg_s, "Expand: dynamic shape is not supported");
1460
1461
1462
        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
1463
        return prog.add_instruction(op::multibroadcast{out_lens}, args[0]);
1464
1465
    }

Shucai Xiao's avatar
Shucai Xiao committed
1466
    std::vector<instruction_ref>
1467
    parse_rnn(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1468
1469
    {
        migraphx::shape input_shape = args[0]->get_shape();
1470
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
1471

1472
        if(contains(info.attributes, "hidden_size"))
Shucai Xiao's avatar
Shucai Xiao committed
1473
        {
1474
            std::size_t hidden_size_att = parse_value(info.attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1475
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1476
1477
1478
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
1479
1480
1481
1482
        }

        // Handling of direction to be added later
        std::string direction{"forward"};
1483
        if(contains(info.attributes, "direction"))
Shucai Xiao's avatar
Shucai Xiao committed
1484
        {
1485
            direction = info.attributes.at("direction").s();
Shucai Xiao's avatar
Shucai Xiao committed
1486
1487
        }

1488
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1489
1490
        if(direction == "bidirectional")
        {
1491
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1492
1493
1494
        }
        else if(direction == "reverse")
        {
1495
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1496
1497
        }

1498
        std::vector<std::string> vec_names{"tanh"};
1499
        if(contains(info.attributes, "activations"))
1500
        {
1501
            auto names = info.attributes.at("activations").strings();
1502
            vec_names.clear();
1503
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1504
1505
1506
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1507
1508
        }

1509
1510
1511
        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
1512
        if(name_it != vec_names.end())
1513
1514
1515
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
1516

Shucai Xiao's avatar
Shucai Xiao committed
1517
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
1518
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
1519
        // if only one actv function is provided, we use it in both
1520
        // forward and reverse direction
1521
        if(dirct == op::rnn_direction::bidirectional)
1522
        {
Shucai Xiao's avatar
Shucai Xiao committed
1523
            if(vec_names.size() == 1)
1524
1525
1526
1527
1528
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
1529
        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1530
1531
1532
1533
        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
1534

Shucai Xiao's avatar
Shucai Xiao committed
1535
1536
        // To be added later
        float clip = 0.0;
1537
        if(contains(info.attributes, "clip"))
Shucai Xiao's avatar
Shucai Xiao committed
1538
        {
1539
            clip = parse_value(info.attributes.at("clip")).at<float>();
Shucai Xiao's avatar
Shucai Xiao committed
1540
1541
        }

1542
1543
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1544
        if(args.size() < 6)
1545
1546
1547
1548
1549
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
1550
1551
        // 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
1552
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1553

1554
        // second output for the last hidden state
Shucai Xiao's avatar
Shucai Xiao committed
1555
        auto last_output = prog.add_instruction(op::rnn_last_hs_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1556

Shucai Xiao's avatar
Shucai Xiao committed
1557
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
1558
1559
    }

1560
    std::vector<instruction_ref>
1561
    parse_gru(const std::string&, node_info info, std::vector<instruction_ref> args)
1562
1563
1564
1565
    {
        migraphx::shape input_shape = args[0]->get_shape();
        std::size_t hidden_size     = args[2]->get_shape().lens()[2];

1566
        if(contains(info.attributes, "hidden_size"))
1567
        {
1568
            std::size_t hidden_size_att = parse_value(info.attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1569
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1570
1571
1572
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
1573
1574
1575
1576
        }

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

1582
        op::rnn_direction dirct = op::rnn_direction::forward;
1583
1584
        if(direction == "bidirectional")
        {
1585
            dirct = op::rnn_direction::bidirectional;
1586
1587
1588
        }
        else if(direction == "reverse")
        {
1589
            dirct = op::rnn_direction::reverse;
1590
1591
        }

1592
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
1593
        if(contains(info.attributes, "activations"))
1594
        {
1595
            auto names = info.attributes.at("activations").strings();
1596
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
1597
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1598
1599
1600
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1601
1602
        }

1603
        // need 4 activation functions
1604
        if(dirct == op::rnn_direction::bidirectional)
1605
        {
Shucai Xiao's avatar
Shucai Xiao committed
1606
            // 4 activation functions are used in the bidirectional
1607
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
1608
1609
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
1610
1611
            // 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
1612
1613
1614
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
1615
            if(vec_names.size() == 1)
1616
            {
1617
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
1618
            }
1619
            else if(vec_names.size() == 2)
1620
            {
1621
1622
1623
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
1624
            }
1625
            else if(vec_names.size() == 3)
1626
            {
1627
                vec_names.push_back(vec_names.at(2));
1628
1629
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
1630
        else
1631
        {
1632
            if(vec_names.size() == 1)
1633
            {
1634
                vec_names.push_back(vec_names.at(0));
1635
1636
1637
            }
        }

1638
1639
1640
        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
1641
        if(name_it != vec_names.end())
1642
1643
1644
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1645

Shucai Xiao's avatar
Shucai Xiao committed
1646
        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1647
1648
1649
1650
        std::transform(vec_names.begin(),
                       vec_names.end(),
                       vec_actv_funcs.begin(),
                       [&](const auto& name) { return map_actv_funcs[name]; });
1651
1652

        float clip = 0.0;
1653
        if(contains(info.attributes, "clip"))
1654
        {
1655
            clip = parse_value(info.attributes.at("clip")).at<float>();
1656
1657
1658
        }

        int linear_before_reset = 0;
1659
        if(contains(info.attributes, "linear_before_reset"))
1660
        {
1661
            linear_before_reset = parse_value(info.attributes.at("linear_before_reset")).at<int>();
1662
1663
        }

Shucai Xiao's avatar
Shucai Xiao committed
1664
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1665
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1666
1667
1668
1669
1670
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1671
1672
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1673
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1674
            std::move(args));
1675
1676

        // second output for last gru output
Shucai Xiao's avatar
Shucai Xiao committed
1677
        auto last_output = prog.add_instruction(op::rnn_last_hs_output{}, hidden_states);
1678

Shucai Xiao's avatar
Shucai Xiao committed
1679
        return {hidden_states, last_output};
1680
1681
    }

Shucai Xiao's avatar
Shucai Xiao committed
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
    void lstm_actv_functions(op::rnn_direction dirct, std::vector<std::string>& actv_func_names)
    {
        // need 6 activation functions for bidirectional directions
        if(dirct == op::rnn_direction::bidirectional)
        {
            // 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
            // if 3 actv funcs are provide, repeat all three once.
            // the same algorithm is used for 4, 5, and 6 actv funcions
            // provided. This may need change later
            switch(actv_func_names.size())
            {
            case 1:
                actv_func_names = {actv_func_names.at(0),
                                   actv_func_names.at(0),
                                   actv_func_names.at(0),
                                   actv_func_names.at(0),
                                   actv_func_names.at(0),
                                   actv_func_names.at(0)};
                break;

            case 2:
                // repeat the 2nd actv func once, then repeat all three another time
                actv_func_names = {actv_func_names.at(0),
                                   actv_func_names.at(1),
                                   actv_func_names.at(1),
                                   actv_func_names.at(0),
                                   actv_func_names.at(1),
                                   actv_func_names.at(1)};
                break;

            case 3:
                // repeat all three actv funcs once
                actv_func_names = {actv_func_names.at(0),
                                   actv_func_names.at(1),
                                   actv_func_names.at(2),
                                   actv_func_names.at(0),
                                   actv_func_names.at(1),
                                   actv_func_names.at(2)};
                break;

            case 4:
                actv_func_names = {actv_func_names.at(0),
                                   actv_func_names.at(1),
                                   actv_func_names.at(2),
                                   actv_func_names.at(3),
                                   actv_func_names.at(3),
                                   actv_func_names.at(3)};
                break;

            case 5:
                actv_func_names = {actv_func_names.at(0),
                                   actv_func_names.at(1),
                                   actv_func_names.at(2),
                                   actv_func_names.at(3),
                                   actv_func_names.at(4),
                                   actv_func_names.at(4)};
                break;

            default: break;
            }
        }
        else
        {
            switch(actv_func_names.size())
            {
            case 1:
                actv_func_names = {
                    actv_func_names.at(0), actv_func_names.at(0), actv_func_names.at(0)};
                break;

            case 2:
                // repeat the 2nd actv func once, so we have 3 actv funcs
                actv_func_names = {
                    actv_func_names.at(0), actv_func_names.at(1), actv_func_names.at(1)};
                break;

            default: break;
            }
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1767
    std::vector<instruction_ref>
1768
    parse_lstm(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1769
1770
1771
1772
    {
        migraphx::shape input_shape = args[0]->get_shape();
        std::size_t hidden_size     = args[2]->get_shape().lens()[2];

1773
        if(contains(info.attributes, "hidden_size"))
Shucai Xiao's avatar
Shucai Xiao committed
1774
        {
1775
            std::size_t hidden_size_att = parse_value(info.attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1776
1777
1778
1779
1780
1781
1782
1783
            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"};
1784
        if(contains(info.attributes, "direction"))
Shucai Xiao's avatar
Shucai Xiao committed
1785
        {
1786
            direction = info.attributes.at("direction").s();
Shucai Xiao's avatar
Shucai Xiao committed
1787
1788
        }

Shucai Xiao's avatar
Shucai Xiao committed
1789
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1790
1791
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1792
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1793
1794
1795
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1796
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1797
        }
Shucai Xiao's avatar
Shucai Xiao committed
1798
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1799
        {
Shucai Xiao's avatar
Shucai Xiao committed
1800
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1801
1802
1803
1804
1805
1806
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1807
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
1808
        if(contains(info.attributes, "activations"))
Shucai Xiao's avatar
Shucai Xiao committed
1809
        {
1810
            auto names = info.attributes.at("activations").strings();
Shucai Xiao's avatar
Shucai Xiao committed
1811
1812
            vec_names.clear();
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1813
1814
1815
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1816
1817
        }

Shucai Xiao's avatar
Shucai Xiao committed
1818
        lstm_actv_functions(dirct, vec_names);
Shucai Xiao's avatar
Shucai Xiao committed
1819

1820
1821
1822
        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
1823
        if(name_it != vec_names.end())
1824
1825
1826
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1827
1828

        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1829
1830
1831
1832
        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
1833
1834

        float clip = 0.0;
1835
        if(contains(info.attributes, "clip"))
Shucai Xiao's avatar
Shucai Xiao committed
1836
        {
1837
            clip = parse_value(info.attributes.at("clip")).at<float>();
Shucai Xiao's avatar
Shucai Xiao committed
1838
1839
1840
        }

        int input_forget = 0;
1841
        if(contains(info.attributes, "input_forget"))
Shucai Xiao's avatar
Shucai Xiao committed
1842
        {
1843
            input_forget = parse_value(info.attributes.at("input_forget")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1844
1845
1846
1847
1848
1849
        }

        // 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
1850
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1851
1852
1853
1854
        }

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

Shucai Xiao's avatar
Shucai Xiao committed
1857
        auto last_output = prog.add_instruction(op::rnn_last_hs_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1858
1859

        // third output for last cell output
Shucai Xiao's avatar
Shucai Xiao committed
1860
        auto last_cell_output = prog.add_instruction(op::rnn_last_cell_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1861
1862
1863

        return {hidden_states, last_output, last_cell_output};
    }
1864

Shucai Xiao's avatar
Shucai Xiao committed
1865
    template <class T>
1866
1867
    instruction_ref
    parse_reduce_oper(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1868
1869
1870
1871
    {
        std::size_t n_dim = args.front()->get_shape().lens().size();

        // default to reduce over all dimensions
1872
        std::vector<int64_t> axes(n_dim);
Shucai Xiao's avatar
Shucai Xiao committed
1873
        std::iota(axes.begin(), axes.end(), 0);
1874
        if(contains(info.attributes, "axes"))
Shucai Xiao's avatar
Shucai Xiao committed
1875
1876
        {
            axes.clear();
1877
            auto&& attr_axes = info.attributes["axes"].ints();
1878
            axes             = std::vector<int64_t>(attr_axes.begin(), attr_axes.end());
Shucai Xiao's avatar
Shucai Xiao committed
1879
1880
1881
        }

        int keep_dims = 1;
1882
        if(contains(info.attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
1883
        {
1884
            keep_dims = parse_value(info.attributes.at("keepdims")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1885
1886
1887
1888
        }

        if(keep_dims == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
1889
            return prog.add_instruction(T{axes}, std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1890
1891
1892
        }
        else
        {
Shucai Xiao's avatar
Shucai Xiao committed
1893
            auto ins = prog.add_instruction(T{axes}, std::move(args));
1894
            return prog.add_instruction(op::squeeze{axes}, ins);
1895
1896
        }
    }
1897

Shucai Xiao's avatar
Shucai Xiao committed
1898
    instruction_ref
1899
    parse_reduce_l1(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1900
1901
    {
        auto abs_ins = prog.add_instruction(op::abs{}, args[0]);
1902
        return parse_reduce_oper<op::reduce_sum>({}, std::move(info), {abs_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1903
1904
1905
    }

    instruction_ref
1906
    parse_reduce_l2(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1907
1908
    {
        auto square_ins = prog.add_instruction(op::mul{}, args[0], args[0]);
1909
        auto sum_ins    = parse_reduce_oper<op::reduce_sum>({}, std::move(info), {square_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1910
1911
1912
        return prog.add_instruction(op::sqrt{}, sum_ins);
    }

1913
1914
    instruction_ref
    parse_reduce_log_sum(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1915
    {
1916
        auto sum_ins = parse_reduce_oper<op::reduce_sum>({}, std::move(info), std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1917
1918
1919
        return prog.add_instruction(op::log{}, sum_ins);
    }

1920
1921
    instruction_ref
    parse_reduce_log_sum_exp(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1922
1923
    {
        auto exp_ins = prog.add_instruction(op::exp{}, args[0]);
1924
        auto sum_ins = parse_reduce_oper<op::reduce_sum>({}, std::move(info), {exp_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1925
1926
1927
        return prog.add_instruction(op::log{}, sum_ins);
    }

1928
1929
    instruction_ref
    parse_reduce_sum_square(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1930
1931
    {
        auto square_ins = prog.add_instruction(op::mul{}, args[0], args[0]);
1932
        return parse_reduce_oper<op::reduce_sum>({}, std::move(info), {square_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1933
1934
    }

Shucai Xiao's avatar
Shucai Xiao committed
1935
    instruction_ref
1936
    parse_cast(const std::string&, node_info info, std::vector<instruction_ref> args)
1937
    {
1938
        if(!contains(info.attributes, "to"))
1939
1940
1941
1942
        {
            MIGRAPHX_THROW("PARSE_CAST: missing to type attribute!");
        }

1943
        int to_type        = parse_value(info.attributes.at("to")).at<int>();
1944
1945
1946
        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
1947

1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
    std::vector<instruction_ref>
    parse_split(const std::string&, node_info info, std::vector<instruction_ref> args)
    {
        int64_t axis = 0;
        if(contains(info.attributes, "axis"))
        {
            axis = parse_value(info.attributes.at("axis")).at<int>();
        }

        auto lens      = args[0]->get_shape().lens();
        int64_t n_rank = static_cast<int64_t>(lens.size());
        if((axis < -n_rank) || (axis >= n_rank))
        {
            MIGRAPHX_THROW("PARSE_SPLIT: axis attribute out of rank!");
        }
        int64_t tuned_axis = (axis < 0) ? axis + n_rank : axis;

        std::vector<int64_t> vec_splits;
        if(contains(info.attributes, "split"))
        {
            literal s = parse_value(info.attributes.at("split"));
            s.visit([&](auto v) { vec_splits.assign(v.begin(), v.end()); });

            if(std::accumulate(vec_splits.begin(), vec_splits.end(), int64_t(0)) !=
               static_cast<int64_t>(lens[tuned_axis]))
            {
                MIGRAPHX_THROW("PARSE_SPLIT: sum of split attribute unequal to dim size of axis!");
            }
        }
        // no split attribute, input is equally divided
        else
        {
            if((lens[tuned_axis] % info.num_outputs) != 0)
            {
                MIGRAPHX_THROW("PARSE_SPLIT: input cannot be equally divided into " +
                               to_string(info.num_outputs) + " splits!");
            }
            auto dl = lens[tuned_axis] / info.num_outputs;
            vec_splits.resize(info.num_outputs, dl);
        }

        std::vector<instruction_ref> ret_ins;
        int64_t start = 0;
        for(auto sl : vec_splits)
        {
            ret_ins.push_back(
                prog.add_instruction(op::slice{{axis}, {start}, {start + sl}}, args[0]));
            start += sl;
        }

        return ret_ins;
    }

kahmed10's avatar
kahmed10 committed
2001
2002
2003
2004
    instruction_ref
    parse_onehot(const std::string&, node_info info, std::vector<instruction_ref> args)
    {
        migraphx::argument depth_arg = args[1]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
2005
        check_arg_empty(depth_arg, "PARSE_ONEHOT: depth - dynamic shape not supported");
kahmed10's avatar
kahmed10 committed
2006
2007
2008
        size_t depth = depth_arg.at<size_t>();

        int64_t axis = -1;
Shucai Xiao's avatar
Shucai Xiao committed
2009
2010
2011
2012
        if(contains(info.attributes, "axis"))
        {
            axis = info.attributes.at("axis").i();
        }
kahmed10's avatar
kahmed10 committed
2013

Shucai Xiao's avatar
Shucai Xiao committed
2014
        std::vector<float> depth_input(depth * depth, 0.0f);
kahmed10's avatar
kahmed10 committed
2015
2016
        for(int i = 0; i < depth; i++)
        {
Shucai Xiao's avatar
Shucai Xiao committed
2017
            depth_input[depth * i + i] = 1.0f;
kahmed10's avatar
kahmed10 committed
2018
2019
        }

Shucai Xiao's avatar
Shucai Xiao committed
2020
2021
2022
2023
2024
2025
2026
2027
        auto type = args[2]->get_shape().type();
        shape s{type, {depth, depth}};
        auto l_val      = prog.add_literal({s, depth_input});
        auto gather_out = prog.add_instruction(op::gather{0}, {l_val, args[0]});

        // Finally, we need a transpose to move the inner most dim to the axis dim
        int n_rank = gather_out->get_shape().lens().size();
        if(axis < -n_rank or axis >= n_rank)
kahmed10's avatar
kahmed10 committed
2028
        {
Shucai Xiao's avatar
Shucai Xiao committed
2029
            MIGRAPHX_THROW("PARSE_ONEHOT: axis out of range");
kahmed10's avatar
kahmed10 committed
2030
        }
Shucai Xiao's avatar
Shucai Xiao committed
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
        int64_t tuned_axis = (axis < 0) ? axis + n_rank : axis;
        std::vector<int64_t> perm(n_rank - 1);
        std::iota(perm.begin(), perm.end(), 0);
        perm.insert(perm.begin() + tuned_axis, n_rank - 1);
        auto tr_out = prog.add_instruction(op::transpose{perm}, gather_out);
        auto lens   = tr_out->get_shape().lens();

        auto off_val       = prog.add_instruction(op::slice{{0}, {0}, {1}}, args[2]);
        auto on_val        = prog.add_instruction(op::slice{{0}, {1}, {2}}, args[2]);
        auto diff          = prog.add_instruction(op::sub{}, on_val, off_val);
        auto unsq_off_val  = prog.add_instruction(op::multibroadcast{lens}, off_val);
        auto unsq_diff_val = prog.add_instruction(op::multibroadcast{lens}, diff);
        auto l_mul         = prog.add_instruction(op::mul{}, tr_out, unsq_diff_val);
        return prog.add_instruction(op::add{}, l_mul, unsq_off_val);
kahmed10's avatar
kahmed10 committed
2045
2046
    }

kahmed10's avatar
kahmed10 committed
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
    instruction_ref
    parse_tile(const std::string&, const node_info&, std::vector<instruction_ref> args)
    {
        migraphx::argument arg_s = args[1]->eval();
        check_arg_empty(arg_s, "PARSE_TILE: dynamic shape is not supported");
        std::vector<std::int64_t> repeats;
        arg_s.visit([&](auto input) { repeats.assign(input.begin(), input.end()); });

        auto l0 = args[0];
        for(int i = 0; i < repeats.size(); i++)
        {
            auto l1 = l0;
            for(int j = 1; j < repeats[i]; j++)
            {
                l0 = prog.add_instruction(op::concat{i}, l0, l1);
            }
        }
        return l0;
    }

kahmed10's avatar
kahmed10 committed
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
    instruction_ref
    parse_range(const std::string&, const node_info&, std::vector<instruction_ref> args)
    {

        auto start_arg = args[0]->eval();
        check_arg_empty(start_arg, "PARSE_RANGE: start arg dynamic shape is not supported");
        auto limit_arg = args[1]->eval();
        check_arg_empty(limit_arg, "PARSE_RANGE: limit arg dynamic shape is not supported");
        auto delta_arg = args[2]->eval();
        check_arg_empty(delta_arg, "PARSE_RANGE: delta arg dynamic shape is not supported");

        assert(args[0]->get_shape().elements() == 1 and args[1]->get_shape().elements() == 1 and
               args[2]->get_shape().elements() == 1);

        instruction_ref l0;

        visit_all(start_arg, limit_arg, delta_arg)([&](auto start, auto limit, auto delta) {
            auto start_val = start.front();
            auto limit_val = limit.front();
            auto delta_val = delta.front();

            size_t num_elements = static_cast<size_t>(
                ceil(static_cast<double>(limit_val - start_val) / static_cast<double>(delta_val)));

            assert(num_elements > 0);

            using type = decltype(start_val);

            std::vector<type> range_vals(num_elements);

            std::generate(range_vals.begin(), range_vals.end(), [&]() {
                auto result = start_val;
                start_val += delta_val;
                return result;
            });

            l0 = prog.add_literal({shape{args[0]->get_shape().type(), {num_elements}}, range_vals});
        });
        return l0;
    }

Paul's avatar
Paul committed
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
    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
2120
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
2121
2122
2123
        }
    }

Paul Fultz II's avatar
Paul Fultz II committed
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
    void parse_from(const void* data, std::size_t size)
    {
        onnx::ModelProto model;
        if(model.ParseFromArray(data, size))
        {
            if(model.has_graph())
            {
                this->parse_graph(model.graph());
            }
        }
        else
        {
            MIGRAPHX_THROW("Failed reading onnx file.");
        }
    }

Paul's avatar
Paul committed
2140
2141
    void parse_graph(const onnx::GraphProto& graph)
    {
2142
        for(auto&& f : graph.initializer())
2143
2144
            instructions[f.name()] = prog.add_literal(parse_tensor(f));

Paul's avatar
Paul committed
2145
2146
2147
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
2148
2149
            // input not in initializer_data, so it is a real input
            if(!contains(instructions, name))
2150
            {
2151
2152
2153
2154
2155
2156
2157
                std::vector<std::size_t> dims;
                if(map_input_dims.count(name) > 0)
                {
                    dims = map_input_dims.at(name);
                }

                shape s            = parse_type(input.type(), dims);
2158
2159
                instructions[name] = prog.add_parameter(name, s);
            }
Paul's avatar
Paul committed
2160
        }
2161
2162

        for(auto&& node : graph.node())
Paul's avatar
Paul committed
2163
        {
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
            std::vector<instruction_ref> args;
            for(auto&& input : node.input())
            {
                if(input.empty())
                {
                    this->parse_undefined(input);
                }
                if(instructions.count(input) == 0)
                {
                    MIGRAPHX_THROW("PARSE_GRAPH: invalid onnx file. Input \"" + input +
                                   "\" is unavailable due to unordered nodes!");
                }
                args.push_back(instructions.at(input));
            }

            std::vector<instruction_ref> result;
            std::size_t output_num = static_cast<std::size_t>(node.output().size());
            if(ops.count(node.op_type()) == 0)
            {
2183
2184
2185
2186
                if(skip_unknown_operators)
                    result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
                else
                    MIGRAPHX_THROW("Unknown operator: " + node.op_type());
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
            }
            else
            {
                result = ops[node.op_type()]({get_attributes(node), output_num}, args);
            }

            output_num = std::min<std::size_t>(output_num, result.size());
            std::transform(node.output().begin(),
                           node.output().begin() + output_num,
                           result.begin(),
                           std::inserter(instructions, instructions.end()),
                           [](auto&& x, auto&& y) { return std::make_pair(x, y); });
Paul's avatar
Paul committed
2199
        }
Shucai Xiao's avatar
Shucai Xiao committed
2200

2201
        // Find instructions corresponding to the output
Shucai Xiao's avatar
Shucai Xiao committed
2202
        auto prog_output = graph.output();
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
        std::vector<std::string> all_output_names;
        std::vector<std::string> prog_output_names;
        std::transform(prog_output.begin(),
                       prog_output.end(),
                       std::back_inserter(all_output_names),
                       [](auto& node) { return node.name(); });
        std::copy_if(
            all_output_names.begin(),
            all_output_names.end(),
            std::back_inserter(prog_output_names),
            [&](const auto& name) { return !(name.empty() or instructions.count(name) == 0); });

        std::vector<instruction_ref> output_ins;
        std::transform(prog_output_names.begin(),
                       prog_output_names.end(),
                       std::back_inserter(output_ins),
                       [&](const auto& name) { return instructions[name]; });

        // add the return instuction
        prog.add_return(output_ins);
Paul's avatar
Paul committed
2223
2224
    }

Shucai Xiao's avatar
Shucai Xiao committed
2225
    void parse_undefined(const std::string& name)
2226
    {
Shucai Xiao's avatar
Shucai Xiao committed
2227
        auto ins           = prog.add_instruction(op::undefined{});
2228
2229
2230
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
    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;
    }

    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
2255
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
2256
        case onnx::AttributeProto::INTS: return from_repeated(shape::int64_type, attr.ints());
Paul's avatar
Paul committed
2257
2258
2259
2260
2261
        case onnx::AttributeProto::UNDEFINED:
        case onnx::AttributeProto::GRAPH:
        case onnx::AttributeProto::STRING:
        case onnx::AttributeProto::STRINGS:
        case onnx::AttributeProto::TENSORS:
2262
2263
        case onnx::AttributeProto::SPARSE_TENSOR:
        case onnx::AttributeProto::SPARSE_TENSORS:
Paul's avatar
Paul committed
2264
2265
        case onnx::AttributeProto::GRAPHS: return {};
        }
Paul's avatar
Paul committed
2266
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
2267
2268
2269
2270
2271
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
2272
2273
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
2274
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
2275
2276
            switch(t.data_type())
            {
2277
            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
Khalique's avatar
Khalique committed
2278
2279
2280
2281
            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());
2282
            case onnx::TensorProto::INT64: return create_literal(shape::int64_type, dims, s.data());
Paul's avatar
Paul committed
2283
2284
2285
2286
            case onnx::TensorProto::INT8:
            case onnx::TensorProto::UINT16:
            case onnx::TensorProto::INT16:
            case onnx::TensorProto::INT32:
2287
            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
Paul's avatar
Paul committed
2288
2289
2290
2291
2292
2293
            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
2294
2295
            case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
            }
Paul's avatar
Paul committed
2296
            MIGRAPHX_THROW("Invalid tensor type");
2297
        }
Paul's avatar
Paul committed
2298
2299
2300
2301
2302
2303
        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
2304
        case onnx::TensorProto::BOOL:
Khalique's avatar
Khalique committed
2305
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
2306
        case onnx::TensorProto::INT64:
Khalique's avatar
Khalique committed
2307
            return create_literal(shape::int64_type, dims, t.int64_data());
Paul's avatar
Paul committed
2308
2309
2310
2311
        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
2312
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
2313
        {
Khalique's avatar
Khalique committed
2314
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
2315
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
2316
2317
2318
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
2319
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
2320
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
2321
        }
Paul's avatar
Paul committed
2322
2323
2324
2325
2326
2327
        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
2328
2329
        case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
        }
Paul's avatar
Paul committed
2330
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
2331
2332
    }

Khalique's avatar
Khalique committed
2333
    static literal
2334
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
2335
    {
Khalique's avatar
Khalique committed
2336
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
2337
        if(dims.empty())
2338
            return literal{{shape_type}, data};
2339
2340
2341
        return literal{{shape_type, dims}, data};
    }

2342
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
2343
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
2344
2345
    {
        if(dims.empty())
2346
            return literal{{shape_type}, data.begin(), data.end()};
2347
        return literal{{shape_type, dims}, data.begin(), data.end()};
2348
2349
    }

2350
    shape parse_type(const onnx::TypeProto& t, const std::vector<std::size_t>& input_dims)
Paul's avatar
Paul committed
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
    {
        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
2361
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
2362
2363
2364
        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;
2365
        case onnx::TensorProto::UINT8: shape_type = shape::uint8_type; break;
Paul's avatar
Paul committed
2366
2367
2368
        case onnx::TensorProto::STRING:
        case onnx::TensorProto::BOOL:
        case onnx::TensorProto::UNDEFINED:
Paul's avatar
Paul committed
2369
2370
        case onnx::TensorProto::COMPLEX64:
        case onnx::TensorProto::COMPLEX128:
Paul's avatar
Paul committed
2371
            break; // throw std::runtime_error("Unsupported type");
Paul's avatar
Paul committed
2372
        }
2373
2374
2375
2376
2377
2378

        if(!input_dims.empty())
        {
            return {shape_type, input_dims};
        }

Paul's avatar
Paul committed
2379
        std::vector<std::size_t> dims;
Paul's avatar
Paul committed
2380
        auto&& tensor_dims = t.tensor_type().shape().dim();
2381
2382
2383
        std::transform(tensor_dims.begin(),
                       tensor_dims.end(),
                       std::back_inserter(dims),
2384
2385
                       [&](auto&& d) -> std::size_t {
                           if(d.has_dim_value())
2386
                           {
2387
                               if(static_cast<int>(d.dim_value()) <= 0)
2388
2389
2390
                               {
                                   return default_dim_value;
                               }
2391
                               return d.dim_value();
2392
                           }
2393
2394
2395
2396
                           else
                           {
                               return default_dim_value;
                           }
2397
                       });
2398

2399
2400
2401
        if(dims.empty())
            return {shape_type};

Paul's avatar
Paul committed
2402
2403
        return {shape_type, dims};
    }
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425

    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
2426
2427
2428

    void check_arg_empty(const argument& arg, const std::string& msg)
    {
Shucai Xiao's avatar
Shucai Xiao committed
2429
        if(arg.empty())
Shucai Xiao's avatar
Shucai Xiao committed
2430
2431
2432
2433
        {
            MIGRAPHX_THROW(msg);
        }
    }
Paul's avatar
Paul committed
2434
2435
};

Paul Fultz II's avatar
Paul Fultz II committed
2436
template <class... Ts>
2437
program parse_onnx_from(const onnx_options& options, Ts&&... xs)
Paul's avatar
Paul committed
2438
2439
{
    onnx_parser parser;
2440
2441
2442
    parser.map_input_dims         = options.map_input_dims;
    parser.default_dim_value      = options.default_dim_value;
    parser.skip_unknown_operators = options.skip_unknown_operators;
2443

2444
    if(options.print_program_on_error)
Paul's avatar
Paul committed
2445
    {
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
        // Log the program when it can't be parsed
        try
        {
            parser.parse_from(std::forward<Ts>(xs)...);
        }
        catch(...)
        {
            std::cerr << parser.prog << std::endl;
            throw;
        }
Paul's avatar
Paul committed
2456
    }
2457
    else
Paul's avatar
Paul committed
2458
    {
2459
        parser.parse_from(std::forward<Ts>(xs)...);
Paul's avatar
Paul committed
2460
2461
2462
2463
    }
    return std::move(parser.prog);
}

2464
program parse_onnx(const std::string& name, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2465
2466
2467
2468
2469
{
    std::fstream input(name.c_str(), std::ios::in | std::ios::binary);
    return parse_onnx_from(options, input);
}

2470
program parse_onnx_buffer(const std::string& buffer, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2471
2472
2473
2474
{
    return parse_onnx_from(options, buffer.data(), buffer.size());
}

2475
program parse_onnx_buffer(const void* data, std::size_t size, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2476
2477
2478
2479
{
    return parse_onnx_from(options, data, size);
}

Paul's avatar
Paul committed
2480
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
2481
} // namespace migraphx