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

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

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

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

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

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

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

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

        return out_lens;
    }

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

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

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

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

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

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

kahmed10's avatar
kahmed10 committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
    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;
    }

321
322
    template <class Op>
    void check_asym_padding(instruction_ref& ins,
323
                            const std::vector<int64_t>& padding,
324
325
326
                            Op& op,
                            float pad_val = 0)
    {
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
        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)
344
        {
345
346
347
348
349
350
            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);
351
352
353
        }
        else
        {
354
            op.padding = std::vector<size_t>(left_pad_it, right_pad_it);
355
356
357
        }
    }

358
359
    instruction_ref
    parse_clip(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
360
    {
kahmed10's avatar
kahmed10 committed
361
362
363
364
365
366
367
        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
368
        {
kahmed10's avatar
kahmed10 committed
369
370
371
372
            min_arg  = args[1];
            max_arg  = args[2];
            min_used = true;
            max_used = true;
Khalique's avatar
Khalique committed
373
        }
kahmed10's avatar
kahmed10 committed
374
        else if(args.size() == 2)
Khalique's avatar
Khalique committed
375
        {
kahmed10's avatar
kahmed10 committed
376
377
378
379
380
381
382
383
384
385
386
387
388
            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
389
        }
kahmed10's avatar
kahmed10 committed
390
391
392
393
394
395
396
397
398
399
400
401
402

        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
403
404
    }

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

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

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

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

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

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

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

462
        auto auto_pad = info.attributes["auto_pad"].s();
463
464
        if(auto_pad.find("SAME") != std::string::npos)
        {
465
            op.padding_mode    = op::padding_mode_t::same;
466
            bool is_same_upper = (auto_pad.find("SAME_UPPER") != std::string::npos);
467
468
469
470
471
472
473
474
475
476
477
478
            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);
            }
479

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

        return ins;
    }

kahmed10's avatar
kahmed10 committed
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
561
    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;
    }

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

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

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

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

643
644
        recalc_conv_attributes(op, kdims);

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

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

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

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

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

730
        if(contains(info.attributes, "output_shape"))
kahmed10's avatar
kahmed10 committed
731
732
        {
            std::vector<int64_t> output_shape;
733
            copy(info.attributes["output_shape"].ints(), std::back_inserter(output_shape));
kahmed10's avatar
kahmed10 committed
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
            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
751
    }
Paul's avatar
Paul committed
752

753
754
    instruction_ref
    parse_pooling(const std::string& name, node_info info, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
755
    {
Khalique's avatar
Khalique committed
756
        op::pooling op{ends_with(name, "MaxPool") ? "max" : "average"};
757
758
759
760
761
        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
762
        if(starts_with(name, "Global"))
763
        {
764
            op.lengths = std::vector<size_t>(in_lens.begin() + 2, in_lens.end());
765
        }
766

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

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

790
        if(contains(info.attributes, "strides"))
Paul's avatar
Paul committed
791
        {
792
793
794
            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
795
        }
796
        if(contains(info.attributes, "kernel_shape"))
Paul's avatar
Paul committed
797
        {
798
799
800
            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
801
        }
802

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

813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
            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");
832
833
        }

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

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

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

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

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

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

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

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

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

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

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

        // 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"))
938
        {
939
            literal s = parse_value(info.attributes.at("axes"));
940
941
            s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        }
Shucai Xiao's avatar
Shucai Xiao committed
942
943

        if(args.size() >= 3)
Khalique's avatar
Khalique committed
944
        {
Shucai Xiao's avatar
Shucai Xiao committed
945
946
947
            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
948
        }
Shucai Xiao's avatar
Shucai Xiao committed
949
        else if(contains(info.attributes, "ends"))
950
        {
951
952
            literal s = parse_value(info.attributes.at("ends"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.ends)); });
953
        }
Shucai Xiao's avatar
Shucai Xiao committed
954
955
956
957
958
959
960
961

        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"))
962
        {
963
            literal s = parse_value(info.attributes.at("starts"));
964
965
            s.visit([&](auto v) { copy(v, std::back_inserter(op.starts)); });
        }
Shucai Xiao's avatar
Shucai Xiao committed
966

kahmed10's avatar
kahmed10 committed
967
968
969
970
971
972
973
        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;
        }

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

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

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

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

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

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

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

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

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

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

1109
1110
1111
        return dot_res;
    }

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

1136
1137
    instruction_ref
    parse_instancenorm(const std::string&, node_info info, std::vector<instruction_ref> args)
kahmed10's avatar
kahmed10 committed
1138
1139
1140
1141
1142
1143
    {
        // 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;
1144
        if(contains(info.attributes, "epsilon"))
kahmed10's avatar
kahmed10 committed
1145
        {
1146
            epsilon = parse_value(info.attributes.at("epsilon")).at<float>();
kahmed10's avatar
kahmed10 committed
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
        }
        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);
    }

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

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

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

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

1222
        if(contains(info.attributes, "bias"))
Khalique's avatar
Khalique committed
1223
        {
1224
            auto&& bias_floats = info.attributes["bias"].floats();
Khalique's avatar
Khalique committed
1225
1226
            bias               = std::vector<float>(bias_floats.begin(), bias_floats.end());
        }
Shucai Xiao's avatar
Shucai Xiao committed
1227
1228
1229
        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
1230

Shucai Xiao's avatar
Shucai Xiao committed
1231
1232
        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
1233

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

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

1252
    instruction_ref parse_pad(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1253
1254
    {
        std::vector<int64_t> pads{};
1255
1256
1257
1258
1259
1260
1261
        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
1262
        {
1263
            auto&& pad_vals = info.attributes["pads"].ints();
Khalique's avatar
Khalique committed
1264
1265
            pads            = std::vector<int64_t>(pad_vals.begin(), pad_vals.end());
        }
1266
1267
1268
1269
1270
        else
        {
            MIGRAPHX_THROW("PARSE_PAD: pad must be available");
        }

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

kahmed10's avatar
kahmed10 committed
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
        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");
            }
        }

1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
        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
1306
        {
1307
            value = parse_value(info.attributes.at("value")).at<float>();
Khalique's avatar
Khalique committed
1308
        }
1309

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            return prog.add_literal(l_out);
        }
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return {hidden_states, last_output, last_cell_output};
    }
1865

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

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

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

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

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

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

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

1921
1922
    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
1923
1924
    {
        auto exp_ins = prog.add_instruction(op::exp{}, args[0]);
1925
        auto sum_ins = parse_reduce_oper<op::reduce_sum>({}, std::move(info), {exp_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1926
1927
1928
        return prog.add_instruction(op::log{}, sum_ins);
    }

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

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

1944
        int to_type        = parse_value(info.attributes.at("to")).at<int>();
1945
1946
1947
        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
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
2001
    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
2002
2003
2004
2005
    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
2006
        check_arg_empty(depth_arg, "PARSE_ONEHOT: depth - dynamic shape not supported");
kahmed10's avatar
kahmed10 committed
2007
2008
2009
        size_t depth = depth_arg.at<size_t>();

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

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

Shucai Xiao's avatar
Shucai Xiao committed
2021
2022
2023
2024
2025
2026
2027
2028
        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
2029
        {
Shucai Xiao's avatar
Shucai Xiao committed
2030
            MIGRAPHX_THROW("PARSE_ONEHOT: axis out of range");
kahmed10's avatar
kahmed10 committed
2031
        }
Shucai Xiao's avatar
Shucai Xiao committed
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
        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
2046
2047
    }

kahmed10's avatar
kahmed10 committed
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
    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
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
2108
    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;
    }

2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
    enum class reduce_mode_t
    {
        sum  = 0,
        mean = 1,
        max  = 2
    };

    instruction_ref parse_embedding_bag(const node_info& info, std::vector<instruction_ref> args)
    {
        if(args[2]->get_shape().elements() != 1)
            MIGRAPHX_THROW("PARSE_EMBEDDING_BAG: MIGraphX only supports offsets of size 1");
        reduce_mode_t reduce_mode = reduce_mode_t::sum;
        if(contains(info.attributes, "mode"))
        {
            reduce_mode = static_cast<reduce_mode_t>(info.attributes.at("mode").i());
        }

        auto l0 = prog.add_instruction(op::gather{}, args[0], args[1]);
        switch(reduce_mode)
        {
        case reduce_mode_t::sum: l0 = prog.add_instruction(op::reduce_sum{{0}}, l0); break;
        case reduce_mode_t::mean: l0 = prog.add_instruction(op::reduce_mean{{0}}, l0); break;
        case reduce_mode_t::max: l0 = prog.add_instruction(op::reduce_max{{0}}, l0); break;
        }
        return l0;
    }

    instruction_ref
    parse_aten(const std::string&, const node_info& info, std::vector<instruction_ref> args)
    {
        if(contains(info.attributes, "operator"))
        {
            auto op_name = info.attributes.at("operator").s();
            if(op_name.find("embedding_bag") != std::string::npos)
            {
                return parse_embedding_bag(info, std::move(args));
            }
        }
        MIGRAPHX_THROW("PARSE_ATEN: unsupported custom operator");
    }

Paul's avatar
Paul committed
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
    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
2162
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
2163
2164
2165
        }
    }

Paul Fultz II's avatar
Paul Fultz II committed
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
    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
2182
2183
    void parse_graph(const onnx::GraphProto& graph)
    {
2184
        for(auto&& f : graph.initializer())
2185
2186
            instructions[f.name()] = prog.add_literal(parse_tensor(f));

Paul's avatar
Paul committed
2187
2188
2189
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
2190
2191
            // input not in initializer_data, so it is a real input
            if(!contains(instructions, name))
2192
            {
2193
2194
2195
2196
2197
2198
2199
                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);
2200
2201
                instructions[name] = prog.add_parameter(name, s);
            }
Paul's avatar
Paul committed
2202
        }
2203
2204

        for(auto&& node : graph.node())
Paul's avatar
Paul committed
2205
        {
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
            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)
            {
2225
2226
2227
2228
                if(skip_unknown_operators)
                    result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
                else
                    MIGRAPHX_THROW("Unknown operator: " + node.op_type());
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
            }
            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
2241
        }
Shucai Xiao's avatar
Shucai Xiao committed
2242

2243
        // Find instructions corresponding to the output
Shucai Xiao's avatar
Shucai Xiao committed
2244
        auto prog_output = graph.output();
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
        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
2265
2266
    }

Shucai Xiao's avatar
Shucai Xiao committed
2267
    void parse_undefined(const std::string& name)
2268
    {
Shucai Xiao's avatar
Shucai Xiao committed
2269
        auto ins           = prog.add_instruction(op::undefined{});
2270
2271
2272
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
    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
2297
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
2298
        case onnx::AttributeProto::INTS: return from_repeated(shape::int64_type, attr.ints());
Paul's avatar
Paul committed
2299
2300
2301
2302
2303
        case onnx::AttributeProto::UNDEFINED:
        case onnx::AttributeProto::GRAPH:
        case onnx::AttributeProto::STRING:
        case onnx::AttributeProto::STRINGS:
        case onnx::AttributeProto::TENSORS:
2304
2305
        case onnx::AttributeProto::SPARSE_TENSOR:
        case onnx::AttributeProto::SPARSE_TENSORS:
Paul's avatar
Paul committed
2306
2307
        case onnx::AttributeProto::GRAPHS: return {};
        }
Paul's avatar
Paul committed
2308
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
2309
2310
2311
2312
2313
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
2314
2315
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
2316
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
2317
2318
            switch(t.data_type())
            {
2319
            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
Khalique's avatar
Khalique committed
2320
2321
2322
2323
            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());
2324
            case onnx::TensorProto::INT64: return create_literal(shape::int64_type, dims, s.data());
Paul's avatar
Paul committed
2325
2326
2327
2328
            case onnx::TensorProto::INT8:
            case onnx::TensorProto::UINT16:
            case onnx::TensorProto::INT16:
            case onnx::TensorProto::INT32:
2329
            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
Paul's avatar
Paul committed
2330
2331
2332
2333
2334
2335
            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
2336
2337
            case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
            }
Paul's avatar
Paul committed
2338
            MIGRAPHX_THROW("Invalid tensor type");
2339
        }
Paul's avatar
Paul committed
2340
2341
2342
2343
2344
2345
        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
2346
        case onnx::TensorProto::BOOL:
Khalique's avatar
Khalique committed
2347
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
2348
        case onnx::TensorProto::INT64:
Khalique's avatar
Khalique committed
2349
            return create_literal(shape::int64_type, dims, t.int64_data());
Paul's avatar
Paul committed
2350
2351
2352
2353
        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
2354
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
2355
        {
Khalique's avatar
Khalique committed
2356
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
2357
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
2358
2359
2360
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
2361
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
2362
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
2363
        }
Paul's avatar
Paul committed
2364
2365
2366
2367
2368
2369
        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
2370
2371
        case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
        }
Paul's avatar
Paul committed
2372
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
2373
2374
    }

Khalique's avatar
Khalique committed
2375
    static literal
2376
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
2377
    {
Khalique's avatar
Khalique committed
2378
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
2379
        if(dims.empty())
2380
            return literal{{shape_type}, data};
2381
2382
2383
        return literal{{shape_type, dims}, data};
    }

2384
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
2385
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
2386
2387
    {
        if(dims.empty())
2388
            return literal{{shape_type}, data.begin(), data.end()};
2389
        return literal{{shape_type, dims}, data.begin(), data.end()};
2390
2391
    }

2392
    shape parse_type(const onnx::TypeProto& t, const std::vector<std::size_t>& input_dims)
Paul's avatar
Paul committed
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
    {
        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
2403
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
2404
2405
2406
        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;
2407
        case onnx::TensorProto::UINT8: shape_type = shape::uint8_type; break;
Paul's avatar
Paul committed
2408
2409
2410
        case onnx::TensorProto::STRING:
        case onnx::TensorProto::BOOL:
        case onnx::TensorProto::UNDEFINED:
Paul's avatar
Paul committed
2411
2412
        case onnx::TensorProto::COMPLEX64:
        case onnx::TensorProto::COMPLEX128:
Paul's avatar
Paul committed
2413
            break; // throw std::runtime_error("Unsupported type");
Paul's avatar
Paul committed
2414
        }
2415
2416
2417
2418
2419
2420

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

Paul's avatar
Paul committed
2421
        std::vector<std::size_t> dims;
Paul's avatar
Paul committed
2422
        auto&& tensor_dims = t.tensor_type().shape().dim();
2423
2424
2425
        std::transform(tensor_dims.begin(),
                       tensor_dims.end(),
                       std::back_inserter(dims),
2426
2427
                       [&](auto&& d) -> std::size_t {
                           if(d.has_dim_value())
2428
                           {
2429
                               if(static_cast<int>(d.dim_value()) <= 0)
2430
2431
2432
                               {
                                   return default_dim_value;
                               }
2433
                               return d.dim_value();
2434
                           }
2435
2436
2437
2438
                           else
                           {
                               return default_dim_value;
                           }
2439
                       });
2440

2441
2442
2443
        if(dims.empty())
            return {shape_type};

Paul's avatar
Paul committed
2444
2445
        return {shape_type, dims};
    }
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467

    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
2468
2469
2470

    void check_arg_empty(const argument& arg, const std::string& msg)
    {
Shucai Xiao's avatar
Shucai Xiao committed
2471
        if(arg.empty())
Shucai Xiao's avatar
Shucai Xiao committed
2472
2473
2474
2475
        {
            MIGRAPHX_THROW(msg);
        }
    }
Paul's avatar
Paul committed
2476
2477
};

Paul Fultz II's avatar
Paul Fultz II committed
2478
template <class... Ts>
2479
program parse_onnx_from(const onnx_options& options, Ts&&... xs)
Paul's avatar
Paul committed
2480
2481
{
    onnx_parser parser;
2482
2483
2484
    parser.map_input_dims         = options.map_input_dims;
    parser.default_dim_value      = options.default_dim_value;
    parser.skip_unknown_operators = options.skip_unknown_operators;
2485

2486
    if(options.print_program_on_error)
Paul's avatar
Paul committed
2487
    {
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
        // 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
2498
    }
2499
    else
Paul's avatar
Paul committed
2500
    {
2501
        parser.parse_from(std::forward<Ts>(xs)...);
Paul's avatar
Paul committed
2502
2503
2504
2505
    }
    return std::move(parser.prog);
}

2506
program parse_onnx(const std::string& name, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2507
2508
2509
2510
2511
{
    std::fstream input(name.c_str(), std::ios::in | std::ios::binary);
    return parse_onnx_from(options, input);
}

2512
program parse_onnx_buffer(const std::string& buffer, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2513
2514
2515
2516
{
    return parse_onnx_from(options, buffer.data(), buffer.size());
}

2517
program parse_onnx_buffer(const void* data, std::size_t size, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2518
2519
2520
2521
{
    return parse_onnx_from(options, data, size);
}

Paul's avatar
Paul committed
2522
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
2523
} // namespace migraphx