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

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

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

24
25
namespace onnx = onnx_for_migraphx;

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

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

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

Khalique's avatar
Khalique committed
79
80
81
        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
82
        add_binary_op("Pow", op::pow{});
Shucai Xiao's avatar
Shucai Xiao committed
83
        add_binary_op("PRelu", op::prelu{});
84
        add_binary_op("Sub", op::sub{});
Khalique's avatar
Khalique committed
85

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

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

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

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

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

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

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

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

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

        return out_lens;
    }

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

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

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

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

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

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

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

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

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

        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
405
406
    }

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

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

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

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

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

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

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

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

482
            check_asym_padding(ins, padding, op, value);
483
484
485
486
487
        }

        return ins;
    }

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

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

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

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

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

645
646
        recalc_conv_attributes(op, kdims);

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Shucai Xiao's avatar
Shucai Xiao committed
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
    instruction_ref
    parse_gather_elements(const std::string&, node_info info, std::vector<instruction_ref> args)
    {
        int axis = 0;
        if(contains(info.attributes, "axis"))
        {
            axis = parse_value(info.attributes.at("axis")).at<int>();
        }

        // standardize input data and index
        auto arg_data = make_contiguous(args[0]);
        auto arg_ind  = make_contiguous(args[1]);

        auto data_s = arg_data->get_shape();
        auto ind_s  = arg_ind->get_shape();

        if(data_s.lens().size() != ind_s.lens().size())
        {
            MIGRAPHX_THROW("PARSE_GATHER_ELEMENTS: input data and index must have the same rank!");
        }

        int n_rank     = static_cast<int>(data_s.lens().size());
        int tuned_axis = (axis < 0) ? (axis + n_rank) : axis;

        auto axis_stride      = data_s.strides()[tuned_axis];
        int64_t data_elem_num = static_cast<int64_t>(data_s.elements());
        // reshape the input data as one dimension and used as input data
        // to the gather operator
        arg_data = prog.add_instruction(op::reshape{{data_elem_num}}, arg_data);

        std::size_t elem_num = ind_s.elements();
        std::vector<int> ind_index(elem_num);
        std::iota(ind_index.begin(), ind_index.end(), 0);

        // convert index in input indices to that in input data
        std::vector<int> data_indices(elem_num);
        std::transform(ind_index.begin(), ind_index.end(), data_indices.begin(), [&](auto i) {
            return data_s.index(ind_s.multi(i));
        });

        std::vector<int> vec_axis_ind(elem_num);
        std::transform(ind_index.begin(), ind_index.end(), vec_axis_ind.begin(), [&](auto i) {
            return ind_s.multi(i)[tuned_axis];
        });

        auto l_shape_idx =
            prog.add_literal(literal(ind_s, data_indices.begin(), data_indices.end()));
        auto l_dim_idx = prog.add_literal(literal(ind_s, vec_axis_ind.begin(), vec_axis_ind.end()));
        auto l_stride  = prog.add_literal(literal{{ind_s.type(), {1}}, {axis_stride}});
        l_stride       = prog.add_instruction(op::multibroadcast{ind_s.lens()}, l_stride);
        auto dim_diff  = prog.add_instruction(op::sub{}, arg_ind, l_dim_idx);
        auto delta     = prog.add_instruction(op::mul{}, dim_diff, l_stride);
        auto ind       = prog.add_instruction(op::add{}, l_shape_idx, delta);

        op::gather op{0};
        return prog.add_instruction(op, arg_data, ind);
    }

972
    instruction_ref
973
    parse_slice(const std::string&, node_info info, std::vector<instruction_ref> args)
974
975
    {
        op::slice op;
Shucai Xiao's avatar
Shucai Xiao committed
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997

        // 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"))
998
        {
999
            literal s = parse_value(info.attributes.at("axes"));
1000
1001
            s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        }
Shucai Xiao's avatar
Shucai Xiao committed
1002
1003

        if(args.size() >= 3)
Khalique's avatar
Khalique committed
1004
        {
Shucai Xiao's avatar
Shucai Xiao committed
1005
1006
1007
            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
1008
        }
Shucai Xiao's avatar
Shucai Xiao committed
1009
        else if(contains(info.attributes, "ends"))
1010
        {
1011
1012
            literal s = parse_value(info.attributes.at("ends"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.ends)); });
1013
        }
Shucai Xiao's avatar
Shucai Xiao committed
1014
1015
1016
1017
1018
1019
1020
1021

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

kahmed10's avatar
kahmed10 committed
1027
1028
1029
1030
1031
1032
1033
        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;
        }

1034
1035
1036
        return prog.add_instruction(op, args[0]);
    }

1037
1038
    instruction_ref
    parse_constant(const std::string&, node_info info, const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
1039
    {
1040
        literal v = parse_value(info.attributes.at("value"));
1041
        // return empty literal
Shucai Xiao's avatar
Shucai Xiao committed
1042
        if(v.get_shape().elements() == 0)
1043
1044
1045
1046
        {
            return prog.add_literal(literal{});
        }

1047
        auto dim_size = info.attributes.at("value").t().dims_size();
1048
1049
        // if dim_size is 0, it is a scalar
        if(dim_size == 0)
1050
        {
1051
            migraphx::shape scalar_shape{v.get_shape().type()};
1052
1053
1054
            return prog.add_literal(migraphx::literal{scalar_shape, v.data()});
        }

Paul's avatar
Paul committed
1055
1056
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
1057

Paul's avatar
Paul committed
1058
    instruction_ref
1059
    parse_gemm(const std::string&, node_info info, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
1060
1061
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
1062
        float beta  = 1.0f;
Paul's avatar
Paul committed
1063
1064
        bool transa = false;
        bool transb = false;
1065
        if(contains(info.attributes, "alpha"))
Paul's avatar
Paul committed
1066
        {
1067
            alpha = parse_value(info.attributes.at("alpha")).at<float>();
Paul's avatar
Paul committed
1068
        }
1069
        if(contains(info.attributes, "beta"))
Paul's avatar
Paul committed
1070
        {
1071
            beta = parse_value(info.attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
1072
        }
1073
        if(contains(info.attributes, "transA"))
Paul's avatar
Paul committed
1074
        {
1075
            transa = parse_value(info.attributes.at("transA")).at<bool>();
Paul's avatar
Paul committed
1076
        }
1077
        if(contains(info.attributes, "transB"))
Paul's avatar
Paul committed
1078
        {
1079
            transb = parse_value(info.attributes.at("transB")).at<bool>();
Paul's avatar
Paul committed
1080
        }
1081
1082
1083
1084
1085
1086

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

1087
1088
        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
1089
1090
        if(args.size() == 3)
        {
1091
            if(beta != 0.f && args[2]->get_shape().elements() > 0)
1092
            {
Shucai Xiao's avatar
Shucai Xiao committed
1093
                auto out_lens   = l1->get_shape().lens();
1094
                out_lens.back() = l2->get_shape().lens().back();
Shucai Xiao's avatar
Shucai Xiao committed
1095
                auto l3         = args[2];
Shucai Xiao's avatar
Shucai Xiao committed
1096
1097
                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
1098
                {
1099
                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
Khalique's avatar
Khalique committed
1100
                }
1101
                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
1102
            }
Paul's avatar
Paul committed
1103
        }
1104
1105

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

1108
    template <class Op>
1109
    instruction_ref
1110
    parse_matmul(const std::string&, const node_info&, std::vector<instruction_ref> args)
1111
    {
Shucai Xiao's avatar
Shucai Xiao committed
1112
1113
        auto l0      = args[0];
        auto l1      = args[1];
1114
1115
1116
1117
1118
        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
1119
        if(l0_lens.size() == 1)
1120
1121
1122
1123
1124
1125
1126
        {
            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
1127
        if(l1_lens.size() == 1)
1128
1129
1130
1131
1132
1133
1134
1135
        {
            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
1136
        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
1137
1138
1139
1140
1141
1142
        {
            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);
1143
            l0_broadcasted_lens = output_lens;
1144
            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
1145
            l1_broadcasted_lens = output_lens;
1146
            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
Shucai Xiao's avatar
Shucai Xiao committed
1147
            if(l0_lens != l0_broadcasted_lens)
1148
1149
1150
            {
                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
            }
Shucai Xiao's avatar
Shucai Xiao committed
1151
            if(l1_lens != l1_broadcasted_lens)
1152
1153
1154
1155
1156
            {
                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
            }
        }

1157
        auto dot_res     = prog.add_instruction(Op{1, 0}, bl0, bl1);
1158
        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
Shucai Xiao's avatar
Shucai Xiao committed
1159
        if(is_a_prepended)
1160
1161
1162
1163
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
            --num_axis;
        }
Shucai Xiao's avatar
Shucai Xiao committed
1164
        if(is_b_appended)
1165
1166
1167
        {
            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
        }
Shucai Xiao's avatar
Shucai Xiao committed
1168

1169
1170
1171
        return dot_res;
    }

1172
    instruction_ref
1173
    parse_batchnorm(const std::string&, node_info info, std::vector<instruction_ref> args)
1174
    {
Scott Thornton's avatar
Scott Thornton committed
1175
1176
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
1177
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
1178
        if(contains(info.attributes, "epsilon"))
1179
        {
1180
            epsilon = parse_value(info.attributes.at("epsilon")).at<float>();
1181
        }
1182
        if(contains(info.attributes, "momentum"))
1183
        {
1184
            momentum = parse_value(info.attributes.at("momentum")).at<float>();
1185
        }
1186
        if(contains(info.attributes, "spatial"))
1187
        {
1188
            bn_mode = (parse_value(info.attributes.at("spatial")).at<uint64_t>() > 0)
1189
1190
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
1191
        }
Paul's avatar
Paul committed
1192
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
1193
        return prog.add_instruction(op, std::move(args));
1194
1195
    }

1196
1197
    instruction_ref
    parse_instancenorm(const std::string&, node_info info, std::vector<instruction_ref> args)
kahmed10's avatar
kahmed10 committed
1198
1199
1200
1201
1202
1203
    {
        // 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;
1204
        if(contains(info.attributes, "epsilon"))
kahmed10's avatar
kahmed10 committed
1205
        {
1206
            epsilon = parse_value(info.attributes.at("epsilon")).at<float>();
kahmed10's avatar
kahmed10 committed
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
        }
        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);
    }

1231
1232
    instruction_ref
    parse_leaky_relu(const std::string&, node_info info, std::vector<instruction_ref> args)
1233
    {
Khalique's avatar
Khalique committed
1234
        float alpha = 0.01; // default alpha val for leaky relu
1235
        if(contains(info.attributes, "alpha"))
1236
        {
1237
            alpha = parse_value(info.attributes.at("alpha")).at<float>();
1238
1239
1240
1241
1242
        }
        op::leaky_relu op{alpha};
        return prog.add_instruction(op, args.front());
    }

1243
    instruction_ref parse_elu(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1244
1245
    {
        float alpha = 1.0; // default alpha val for elu
1246
        if(contains(info.attributes, "alpha"))
Khalique's avatar
Khalique committed
1247
        {
1248
            alpha = parse_value(info.attributes.at("alpha")).at<float>();
Khalique's avatar
Khalique committed
1249
1250
1251
1252
1253
        }
        op::elu op{alpha};
        return prog.add_instruction(op, args.front());
    }

1254
    instruction_ref parse_lrn(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1255
1256
    {
        float alpha = 0.0001;
Khalique's avatar
Khalique committed
1257
1258
1259
        float beta  = 0.75;
        float bias  = 1.0;
        int size    = 1;
1260
1261
1262
1263
1264
1265
1266
1267
        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
1268
1269
1270
1271
        op::lrn op{alpha, beta, bias, size};
        return prog.add_instruction(op, args.front());
    }

1272
1273
    instruction_ref
    parse_imagescaler(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1274
1275
1276
    {
        float scale = 1.0;
        std::vector<float> bias{};
1277
        if(contains(info.attributes, "scale"))
Khalique's avatar
Khalique committed
1278
        {
1279
            scale = parse_value(info.attributes.at("scale")).at<float>();
Khalique's avatar
Khalique committed
1280
1281
        }

1282
        if(contains(info.attributes, "bias"))
Khalique's avatar
Khalique committed
1283
        {
1284
            auto&& bias_floats = info.attributes["bias"].floats();
Khalique's avatar
Khalique committed
1285
1286
            bias               = std::vector<float>(bias_floats.begin(), bias_floats.end());
        }
Shucai Xiao's avatar
Shucai Xiao committed
1287
1288
1289
        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
1290

Shucai Xiao's avatar
Shucai Xiao committed
1291
1292
        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
1293

1294
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
Paul's avatar
Paul committed
1295
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Shucai Xiao's avatar
Shucai Xiao committed
1296
        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
Paul's avatar
Paul committed
1297
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
1298
    }
Khalique's avatar
Khalique committed
1299

Khalique's avatar
Khalique committed
1300
    instruction_ref
1301
    parse_transpose(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1302
1303
    {
        std::vector<int64_t> perm{};
1304
        if(contains(info.attributes, "perm"))
Khalique's avatar
Khalique committed
1305
        {
1306
            auto&& perm_vals = info.attributes["perm"].ints();
Khalique's avatar
Khalique committed
1307
1308
            perm             = std::vector<int64_t>(perm_vals.begin(), perm_vals.end());
        }
Paul's avatar
Paul committed
1309
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
1310
1311
    }

1312
    instruction_ref parse_pad(const std::string&, node_info info, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
1313
1314
    {
        std::vector<int64_t> pads{};
1315
1316
1317
1318
1319
1320
1321
        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
1322
        {
1323
            auto&& pad_vals = info.attributes["pads"].ints();
Khalique's avatar
Khalique committed
1324
1325
            pads            = std::vector<int64_t>(pad_vals.begin(), pad_vals.end());
        }
1326
1327
1328
1329
1330
        else
        {
            MIGRAPHX_THROW("PARSE_PAD: pad must be available");
        }

1331
        // check if padding is actually being done (at least one value is nonzero)
Khalique's avatar
Khalique committed
1332
        if(std::all_of(pads.begin(), pads.end(), [](const int& i) { return i == 0; }))
1333
1334
1335
        {
            return prog.add_instruction(migraphx::op::identity{}, args.front());
        }
1336

kahmed10's avatar
kahmed10 committed
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
        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");
            }
        }

1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
        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
1366
        {
1367
            value = parse_value(info.attributes.at("value")).at<float>();
Khalique's avatar
Khalique committed
1368
        }
1369

Khalique's avatar
Khalique committed
1370
1371
        return prog.add_instruction(migraphx::op::pad{pads, value}, args.front());
    }
1372
1373
1374
    // Use a literal instruction to replace the shape since, output of
    // shape operator are literals in migraphx
    instruction_ref
1375
    parse_shape(const std::string&, const node_info&, std::vector<instruction_ref> args)
1376
1377
    {
        if(args.size() != 1)
1378
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
        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
1391
1392
    instruction_ref
    parse_constant_fill(const std::string&, node_info info, std::vector<instruction_ref> args)
1393
1394
1395
1396
1397
    {
        int input_as_shape = 0;
        int dtype          = 1;
        float value        = 0.0f;

1398
        if(contains(info.attributes, "dtype"))
1399
        {
1400
            dtype = parse_value(info.attributes.at("dtype")).at<int>();
1401
        }
Shucai Xiao's avatar
Shucai Xiao committed
1402
        shape::type_t type = get_type(dtype);
1403

1404
        if(contains(info.attributes, "input_as_shape"))
1405
        {
1406
            input_as_shape = parse_value(info.attributes.at("input_as_shape")).at<int>();
1407
1408
        }

1409
        if(contains(info.attributes, "value"))
1410
        {
1411
            value = parse_value(info.attributes.at("value")).at<float>();
1412
1413
        }

1414
        if(contains(info.attributes, "extra_shape"))
Shucai Xiao's avatar
Shucai Xiao committed
1415
        {
1416
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
1417
1418
        }

1419
1420
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
1421
            if(args.size() != 1)
1422
            {
1423
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
1424
1425
            }

1426
            if(contains(info.attributes, "shape"))
Shucai Xiao's avatar
Shucai Xiao committed
1427
            {
1428
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
1429
                               "at the same time");
1430
1431
            }

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

1435
1436
1437
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
1438
1439
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
1440
1441
1442
        }
        else if(input_as_shape == 0)
        {
1443
            if(!contains(info.attributes, "shape"))
Shucai Xiao's avatar
Shucai Xiao committed
1444
            {
1445
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
1446
1447
            }

1448
            literal ls = parse_value(info.attributes.at("shape"));
1449
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
1450
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
1451
            migraphx::shape s{type, dims};
1452
1453
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
1454
1455
1456
        }
        else
        {
1457
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
1458
1459
1460
        }
    }

1461
1462
    instruction_ref
    parse_constant_of_shape(const std::string&, node_info info, std::vector<instruction_ref> args)
1463
1464
    {
        literal l_val{};
1465
        if(contains(info.attributes, "value"))
1466
        {
1467
            l_val = parse_value(info.attributes.at("value"));
Shucai Xiao's avatar
Shucai Xiao committed
1468
            if(l_val.get_shape().elements() != 1)
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
            {
                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();
1480

Shucai Xiao's avatar
Shucai Xiao committed
1481
        if(args.empty())
1482
        {
Shucai Xiao's avatar
Shucai Xiao committed
1483
            MIGRAPHX_THROW("ConstantOfShape : must have 1 input!");
1484
1485
1486
        }
        else
        {
1487
1488
            migraphx::shape s;
            // empty input tensor, output is a scalar
Shucai Xiao's avatar
Shucai Xiao committed
1489
            if(args[0]->get_shape().elements() == 0)
1490
            {
1491
                s = migraphx::shape{type, {1}, {0}};
1492
            }
1493
1494
1495
            else
            {
                migraphx::argument in = args[0]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
1496
                check_arg_empty(in, "ConstantOfShape: dynamic shape is not supported");
1497

1498
1499
1500
1501
                std::vector<std::size_t> dims;
                in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
                s = migraphx::shape{type, dims};
            }
1502

Shucai Xiao's avatar
Shucai Xiao committed
1503
            literal l_out{};
1504
            l_val.visit([&](auto val) {
Shucai Xiao's avatar
Shucai Xiao committed
1505
                using val_type = std::remove_cv_t<typename decltype(val)::value_type>;
1506
                // l_val contains only one element
1507
                std::vector<val_type> out_vec(s.elements(), val.front());
1508
1509
1510
1511
1512
1513
1514
                l_out = literal(s, out_vec);
            });

            return prog.add_literal(l_out);
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
1515
    instruction_ref
1516
    parse_expand(const std::string&, const node_info&, std::vector<instruction_ref> args)
1517
    {
Shucai Xiao's avatar
Shucai Xiao committed
1518
        auto in_lens             = args[0]->get_shape().lens();
1519
        migraphx::argument arg_s = args[1]->eval();
Shucai Xiao's avatar
Shucai Xiao committed
1520
        check_arg_empty(arg_s, "Expand: dynamic shape is not supported");
1521
1522
1523
        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
1524
        return prog.add_instruction(op::multibroadcast{out_lens}, args[0]);
1525
1526
    }

Shucai Xiao's avatar
Shucai Xiao committed
1527
    std::vector<instruction_ref>
1528
    parse_rnn(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1529
1530
    {
        migraphx::shape input_shape = args[0]->get_shape();
1531
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
1532

1533
        if(contains(info.attributes, "hidden_size"))
Shucai Xiao's avatar
Shucai Xiao committed
1534
        {
1535
            std::size_t hidden_size_att = parse_value(info.attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1536
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1537
1538
1539
            {
                MIGRAPHX_THROW("RNN: hidden size mismatch in input and attribute");
            }
Shucai Xiao's avatar
Shucai Xiao committed
1540
1541
1542
1543
        }

        // Handling of direction to be added later
        std::string direction{"forward"};
1544
        if(contains(info.attributes, "direction"))
Shucai Xiao's avatar
Shucai Xiao committed
1545
        {
1546
            direction = info.attributes.at("direction").s();
Shucai Xiao's avatar
Shucai Xiao committed
1547
1548
        }

1549
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1550
1551
        if(direction == "bidirectional")
        {
1552
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1553
1554
1555
        }
        else if(direction == "reverse")
        {
1556
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1557
1558
        }

1559
        std::vector<std::string> vec_names{"tanh"};
1560
        if(contains(info.attributes, "activations"))
1561
        {
1562
            auto names = info.attributes.at("activations").strings();
1563
            vec_names.clear();
1564
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1565
1566
1567
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1568
1569
        }

1570
1571
1572
        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
1573
        if(name_it != vec_names.end())
1574
1575
1576
        {
            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
        }
1577

Shucai Xiao's avatar
Shucai Xiao committed
1578
        // bidirectional case should have two activation functions.
Shucai Xiao's avatar
Shucai Xiao committed
1579
        // one is for forward, and the other is for reverse.
Shucai Xiao's avatar
Shucai Xiao committed
1580
        // if only one actv function is provided, we use it in both
1581
        // forward and reverse direction
1582
        if(dirct == op::rnn_direction::bidirectional)
1583
        {
Shucai Xiao's avatar
Shucai Xiao committed
1584
            if(vec_names.size() == 1)
1585
1586
1587
1588
1589
            {
                vec_names.push_back(vec_names.at(0));
            }
        }

Shucai Xiao's avatar
Shucai Xiao committed
1590
        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1591
1592
1593
1594
        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
1595

Shucai Xiao's avatar
Shucai Xiao committed
1596
1597
        // To be added later
        float clip = 0.0;
1598
        if(contains(info.attributes, "clip"))
Shucai Xiao's avatar
Shucai Xiao committed
1599
        {
1600
            clip = parse_value(info.attributes.at("clip")).at<float>();
Shucai Xiao's avatar
Shucai Xiao committed
1601
1602
        }

1603
1604
        // if the number of arguments is less than 6, append
        // undefined operator to have 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1605
        if(args.size() < 6)
1606
1607
1608
1609
1610
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), (6 - args.size()), ins);
        }

Shucai Xiao's avatar
Shucai Xiao committed
1611
1612
        // 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
1613
                                                  std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1614

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

Shucai Xiao's avatar
Shucai Xiao committed
1618
        return {hidden_states, last_output};
Shucai Xiao's avatar
Shucai Xiao committed
1619
1620
    }

1621
    std::vector<instruction_ref>
1622
    parse_gru(const std::string&, node_info info, std::vector<instruction_ref> args)
1623
1624
1625
1626
    {
        migraphx::shape input_shape = args[0]->get_shape();
        std::size_t hidden_size     = args[2]->get_shape().lens()[2];

1627
        if(contains(info.attributes, "hidden_size"))
1628
        {
1629
            std::size_t hidden_size_att = parse_value(info.attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1630
            if(hidden_size != hidden_size_att)
Shucai Xiao's avatar
Shucai Xiao committed
1631
1632
1633
            {
                MIGRAPHX_THROW("GRU: hidden size mismatch in input and attribute");
            }
1634
1635
1636
1637
        }

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

1643
        op::rnn_direction dirct = op::rnn_direction::forward;
1644
1645
        if(direction == "bidirectional")
        {
1646
            dirct = op::rnn_direction::bidirectional;
1647
1648
1649
        }
        else if(direction == "reverse")
        {
1650
            dirct = op::rnn_direction::reverse;
1651
1652
        }

1653
        std::vector<std::string> vec_names = {"sigmoid", "tanh"};
1654
        if(contains(info.attributes, "activations"))
1655
        {
1656
            auto names = info.attributes.at("activations").strings();
1657
            vec_names.clear();
Shucai Xiao's avatar
Shucai Xiao committed
1658
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1659
1660
1661
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
1662
1663
        }

1664
        // need 4 activation functions
1665
        if(dirct == op::rnn_direction::bidirectional)
1666
        {
Shucai Xiao's avatar
Shucai Xiao committed
1667
            // 4 activation functions are used in the bidirectional
1668
            // scenario. No spec is provided in onnx::operator. we
Shucai Xiao's avatar
Shucai Xiao committed
1669
1670
            // use the algorithm that: if 1 actv function is provided,
            // repeat 1 four times. If 2 actv functins are provided,
1671
1672
            // 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
1673
1674
1675
            // assume the 3rd one is repeated once and used by the
            // reverse direction.
            // This may need change later
1676
            if(vec_names.size() == 1)
1677
            {
1678
                vec_names.insert(vec_names.end(), 3, vec_names.at(0));
1679
            }
1680
            else if(vec_names.size() == 2)
1681
            {
1682
1683
1684
                // repeat the activation functions
                vec_names.push_back(vec_names.at(0));
                vec_names.push_back(vec_names.at(1));
1685
            }
1686
            else if(vec_names.size() == 3)
1687
            {
1688
                vec_names.push_back(vec_names.at(2));
1689
1690
            }
        }
Shucai Xiao's avatar
Shucai Xiao committed
1691
        else
1692
        {
1693
            if(vec_names.size() == 1)
1694
            {
1695
                vec_names.push_back(vec_names.at(0));
1696
1697
1698
            }
        }

1699
1700
1701
        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
1702
        if(name_it != vec_names.end())
1703
1704
1705
        {
            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
        }
1706

Shucai Xiao's avatar
Shucai Xiao committed
1707
        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1708
1709
1710
1711
        std::transform(vec_names.begin(),
                       vec_names.end(),
                       vec_actv_funcs.begin(),
                       [&](const auto& name) { return map_actv_funcs[name]; });
1712
1713

        float clip = 0.0;
1714
        if(contains(info.attributes, "clip"))
1715
        {
1716
            clip = parse_value(info.attributes.at("clip")).at<float>();
1717
1718
1719
        }

        int linear_before_reset = 0;
1720
        if(contains(info.attributes, "linear_before_reset"))
1721
        {
1722
            linear_before_reset = parse_value(info.attributes.at("linear_before_reset")).at<int>();
1723
1724
        }

Shucai Xiao's avatar
Shucai Xiao committed
1725
        // append undefined opeator to make 6 arguments
Shucai Xiao's avatar
Shucai Xiao committed
1726
        if(args.size() < 6)
Shucai Xiao's avatar
Shucai Xiao committed
1727
1728
1729
1730
1731
        {
            auto ins = prog.add_instruction(op::undefined{});
            args.insert(args.end(), 6 - args.size(), ins);
        }

1732
1733
        // first output for concatenation of hidden states
        auto hidden_states = prog.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1734
            op::gru{hidden_size, vec_actv_funcs, dirct, clip, linear_before_reset},
Shucai Xiao's avatar
Shucai Xiao committed
1735
            std::move(args));
1736
1737

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

Shucai Xiao's avatar
Shucai Xiao committed
1740
        return {hidden_states, last_output};
1741
1742
    }

Shucai Xiao's avatar
Shucai Xiao committed
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
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
    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
1828
    std::vector<instruction_ref>
1829
    parse_lstm(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1830
1831
1832
1833
    {
        migraphx::shape input_shape = args[0]->get_shape();
        std::size_t hidden_size     = args[2]->get_shape().lens()[2];

1834
        if(contains(info.attributes, "hidden_size"))
Shucai Xiao's avatar
Shucai Xiao committed
1835
        {
1836
            std::size_t hidden_size_att = parse_value(info.attributes.at("hidden_size")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1837
1838
1839
1840
1841
1842
1843
1844
            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"};
1845
        if(contains(info.attributes, "direction"))
Shucai Xiao's avatar
Shucai Xiao committed
1846
        {
1847
            direction = info.attributes.at("direction").s();
Shucai Xiao's avatar
Shucai Xiao committed
1848
1849
        }

Shucai Xiao's avatar
Shucai Xiao committed
1850
        op::rnn_direction dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1851
1852
        if(direction == "bidirectional")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1853
            dirct = op::rnn_direction::bidirectional;
Shucai Xiao's avatar
Shucai Xiao committed
1854
1855
1856
        }
        else if(direction == "reverse")
        {
Shucai Xiao's avatar
Shucai Xiao committed
1857
            dirct = op::rnn_direction::reverse;
Shucai Xiao's avatar
Shucai Xiao committed
1858
        }
Shucai Xiao's avatar
Shucai Xiao committed
1859
        else if(direction == "forward")
Shucai Xiao's avatar
Shucai Xiao committed
1860
        {
Shucai Xiao's avatar
Shucai Xiao committed
1861
            dirct = op::rnn_direction::forward;
Shucai Xiao's avatar
Shucai Xiao committed
1862
1863
1864
1865
1866
1867
        }
        else
        {
            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
        }

1868
        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
1869
        if(contains(info.attributes, "activations"))
Shucai Xiao's avatar
Shucai Xiao committed
1870
        {
1871
            auto names = info.attributes.at("activations").strings();
Shucai Xiao's avatar
Shucai Xiao committed
1872
1873
            vec_names.clear();
            vec_names.resize(names.size());
Shucai Xiao's avatar
Shucai Xiao committed
1874
1875
1876
            std::transform(names.begin(), names.end(), vec_names.begin(), [](auto name) {
                return to_lower(name);
            });
Shucai Xiao's avatar
Shucai Xiao committed
1877
1878
        }

Shucai Xiao's avatar
Shucai Xiao committed
1879
        lstm_actv_functions(dirct, vec_names);
Shucai Xiao's avatar
Shucai Xiao committed
1880

1881
1882
1883
        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
1884
        if(name_it != vec_names.end())
1885
1886
1887
        {
            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
        }
Shucai Xiao's avatar
Shucai Xiao committed
1888
1889

        std::vector<operation> vec_actv_funcs(vec_names.size());
Paul's avatar
Paul committed
1890
1891
1892
1893
        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
1894
1895

        float clip = 0.0;
1896
        if(contains(info.attributes, "clip"))
Shucai Xiao's avatar
Shucai Xiao committed
1897
        {
1898
            clip = parse_value(info.attributes.at("clip")).at<float>();
Shucai Xiao's avatar
Shucai Xiao committed
1899
1900
1901
        }

        int input_forget = 0;
1902
        if(contains(info.attributes, "input_forget"))
Shucai Xiao's avatar
Shucai Xiao committed
1903
        {
1904
            input_forget = parse_value(info.attributes.at("input_forget")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1905
1906
1907
1908
1909
1910
        }

        // 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
1911
            args.insert(args.end(), 8 - args.size(), ins);
Shucai Xiao's avatar
Shucai Xiao committed
1912
1913
1914
1915
        }

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

Shucai Xiao's avatar
Shucai Xiao committed
1918
        auto last_output = prog.add_instruction(op::rnn_last_hs_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1919
1920

        // third output for last cell output
Shucai Xiao's avatar
Shucai Xiao committed
1921
        auto last_cell_output = prog.add_instruction(op::rnn_last_cell_output{}, hidden_states);
Shucai Xiao's avatar
Shucai Xiao committed
1922
1923
1924

        return {hidden_states, last_output, last_cell_output};
    }
1925

Shucai Xiao's avatar
Shucai Xiao committed
1926
    template <class T>
1927
1928
    instruction_ref
    parse_reduce_oper(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1929
1930
1931
1932
    {
        std::size_t n_dim = args.front()->get_shape().lens().size();

        // default to reduce over all dimensions
1933
        std::vector<int64_t> axes(n_dim);
Shucai Xiao's avatar
Shucai Xiao committed
1934
        std::iota(axes.begin(), axes.end(), 0);
1935
        if(contains(info.attributes, "axes"))
Shucai Xiao's avatar
Shucai Xiao committed
1936
1937
        {
            axes.clear();
1938
            auto&& attr_axes = info.attributes["axes"].ints();
1939
            axes             = std::vector<int64_t>(attr_axes.begin(), attr_axes.end());
Shucai Xiao's avatar
Shucai Xiao committed
1940
1941
1942
        }

        int keep_dims = 1;
1943
        if(contains(info.attributes, "keepdims"))
Shucai Xiao's avatar
Shucai Xiao committed
1944
        {
1945
            keep_dims = parse_value(info.attributes.at("keepdims")).at<int>();
Shucai Xiao's avatar
Shucai Xiao committed
1946
1947
1948
1949
        }

        if(keep_dims == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
1950
            return prog.add_instruction(T{axes}, std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1951
1952
1953
        }
        else
        {
Shucai Xiao's avatar
Shucai Xiao committed
1954
            auto ins = prog.add_instruction(T{axes}, std::move(args));
1955
            return prog.add_instruction(op::squeeze{axes}, ins);
1956
1957
        }
    }
1958

Shucai Xiao's avatar
Shucai Xiao committed
1959
    instruction_ref
1960
    parse_reduce_l1(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1961
1962
    {
        auto abs_ins = prog.add_instruction(op::abs{}, args[0]);
1963
        return parse_reduce_oper<op::reduce_sum>({}, std::move(info), {abs_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1964
1965
1966
    }

    instruction_ref
1967
    parse_reduce_l2(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1968
1969
    {
        auto square_ins = prog.add_instruction(op::mul{}, args[0], args[0]);
1970
        auto sum_ins    = parse_reduce_oper<op::reduce_sum>({}, std::move(info), {square_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1971
1972
1973
        return prog.add_instruction(op::sqrt{}, sum_ins);
    }

1974
1975
    instruction_ref
    parse_reduce_log_sum(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1976
    {
1977
        auto sum_ins = parse_reduce_oper<op::reduce_sum>({}, std::move(info), std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
1978
1979
1980
        return prog.add_instruction(op::log{}, sum_ins);
    }

1981
1982
    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
1983
1984
    {
        auto exp_ins = prog.add_instruction(op::exp{}, args[0]);
1985
        auto sum_ins = parse_reduce_oper<op::reduce_sum>({}, std::move(info), {exp_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1986
1987
1988
        return prog.add_instruction(op::log{}, sum_ins);
    }

1989
1990
    instruction_ref
    parse_reduce_sum_square(const std::string&, node_info info, std::vector<instruction_ref> args)
Shucai Xiao's avatar
Shucai Xiao committed
1991
1992
    {
        auto square_ins = prog.add_instruction(op::mul{}, args[0], args[0]);
1993
        return parse_reduce_oper<op::reduce_sum>({}, std::move(info), {square_ins});
Shucai Xiao's avatar
Shucai Xiao committed
1994
1995
    }

Shucai Xiao's avatar
Shucai Xiao committed
1996
    instruction_ref
1997
    parse_cast(const std::string&, node_info info, std::vector<instruction_ref> args)
1998
    {
1999
        if(!contains(info.attributes, "to"))
2000
2001
2002
2003
        {
            MIGRAPHX_THROW("PARSE_CAST: missing to type attribute!");
        }

2004
        int to_type        = parse_value(info.attributes.at("to")).at<int>();
2005
2006
2007
        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
2008

2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
    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
2062
2063
2064
2065
    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
2066
        check_arg_empty(depth_arg, "PARSE_ONEHOT: depth - dynamic shape not supported");
kahmed10's avatar
kahmed10 committed
2067
2068
2069
        size_t depth = depth_arg.at<size_t>();

        int64_t axis = -1;
Shucai Xiao's avatar
Shucai Xiao committed
2070
2071
2072
2073
        if(contains(info.attributes, "axis"))
        {
            axis = info.attributes.at("axis").i();
        }
kahmed10's avatar
kahmed10 committed
2074

Shucai Xiao's avatar
Shucai Xiao committed
2075
        std::vector<float> depth_input(depth * depth, 0.0f);
kahmed10's avatar
kahmed10 committed
2076
2077
        for(int i = 0; i < depth; i++)
        {
Shucai Xiao's avatar
Shucai Xiao committed
2078
            depth_input[depth * i + i] = 1.0f;
kahmed10's avatar
kahmed10 committed
2079
2080
        }

Shucai Xiao's avatar
Shucai Xiao committed
2081
2082
2083
2084
2085
2086
2087
2088
        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
2089
        {
Shucai Xiao's avatar
Shucai Xiao committed
2090
            MIGRAPHX_THROW("PARSE_ONEHOT: axis out of range");
kahmed10's avatar
kahmed10 committed
2091
        }
Shucai Xiao's avatar
Shucai Xiao committed
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
        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
2106
2107
    }

kahmed10's avatar
kahmed10 committed
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
    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
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
    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;
    }

2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
    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
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
    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
2222
            MIGRAPHX_THROW("Failed reading onnx file.");
Paul's avatar
Paul committed
2223
2224
2225
        }
    }

Paul Fultz II's avatar
Paul Fultz II committed
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
    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
2242
2243
    void parse_graph(const onnx::GraphProto& graph)
    {
2244
        for(auto&& f : graph.initializer())
2245
2246
            instructions[f.name()] = prog.add_literal(parse_tensor(f));

Paul's avatar
Paul committed
2247
2248
2249
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
2250
2251
            // input not in initializer_data, so it is a real input
            if(!contains(instructions, name))
2252
            {
2253
2254
2255
2256
2257
2258
2259
                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);
2260
2261
                instructions[name] = prog.add_parameter(name, s);
            }
Paul's avatar
Paul committed
2262
        }
2263
2264

        for(auto&& node : graph.node())
Paul's avatar
Paul committed
2265
        {
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
            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)
            {
2285
2286
2287
2288
                if(skip_unknown_operators)
                    result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
                else
                    MIGRAPHX_THROW("Unknown operator: " + node.op_type());
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
            }
            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
2301
        }
Shucai Xiao's avatar
Shucai Xiao committed
2302

2303
        // Find instructions corresponding to the output
Shucai Xiao's avatar
Shucai Xiao committed
2304
        auto prog_output = graph.output();
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
        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
2325
2326
    }

Shucai Xiao's avatar
Shucai Xiao committed
2327
    void parse_undefined(const std::string& name)
2328
    {
Shucai Xiao's avatar
Shucai Xiao committed
2329
        auto ins           = prog.add_instruction(op::undefined{});
2330
2331
2332
        instructions[name] = ins;
    }

Paul's avatar
Paul committed
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
    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
2357
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
2358
        case onnx::AttributeProto::INTS: return from_repeated(shape::int64_type, attr.ints());
Paul's avatar
Paul committed
2359
2360
2361
2362
2363
        case onnx::AttributeProto::UNDEFINED:
        case onnx::AttributeProto::GRAPH:
        case onnx::AttributeProto::STRING:
        case onnx::AttributeProto::STRINGS:
        case onnx::AttributeProto::TENSORS:
2364
2365
        case onnx::AttributeProto::SPARSE_TENSOR:
        case onnx::AttributeProto::SPARSE_TENSORS:
Paul's avatar
Paul committed
2366
2367
        case onnx::AttributeProto::GRAPHS: return {};
        }
Paul's avatar
Paul committed
2368
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
2369
2370
2371
2372
2373
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
2374
2375
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
2376
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
2377
2378
            switch(t.data_type())
            {
2379
            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
Khalique's avatar
Khalique committed
2380
2381
2382
2383
            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());
2384
            case onnx::TensorProto::INT64: return create_literal(shape::int64_type, dims, s.data());
Paul's avatar
Paul committed
2385
2386
2387
2388
            case onnx::TensorProto::INT8:
            case onnx::TensorProto::UINT16:
            case onnx::TensorProto::INT16:
            case onnx::TensorProto::INT32:
2389
            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
Paul's avatar
Paul committed
2390
2391
2392
2393
2394
2395
            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
2396
2397
            case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
            }
Paul's avatar
Paul committed
2398
            MIGRAPHX_THROW("Invalid tensor type");
2399
        }
Paul's avatar
Paul committed
2400
2401
2402
2403
2404
2405
        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
2406
        case onnx::TensorProto::BOOL:
Khalique's avatar
Khalique committed
2407
            return create_literal(shape::int32_type, dims, t.int32_data());
Paul's avatar
Paul committed
2408
        case onnx::TensorProto::INT64:
Khalique's avatar
Khalique committed
2409
            return create_literal(shape::int64_type, dims, t.int64_data());
Paul's avatar
Paul committed
2410
2411
2412
2413
        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
2414
        case onnx::TensorProto::FLOAT16:
Khalique's avatar
Khalique committed
2415
        {
Khalique's avatar
Khalique committed
2416
            std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
2417
            std::vector<half> data_half;
Khalique's avatar
Khalique committed
2418
2419
2420
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
2421
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
2422
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
2423
        }
Paul's avatar
Paul committed
2424
2425
2426
2427
2428
2429
        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
2430
2431
        case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
        }
Paul's avatar
Paul committed
2432
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
2433
2434
    }

Khalique's avatar
Khalique committed
2435
    static literal
2436
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
2437
    {
Khalique's avatar
Khalique committed
2438
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
2439
        if(dims.empty())
2440
            return literal{{shape_type}, data};
2441
2442
2443
        return literal{{shape_type, dims}, data};
    }

2444
    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
Khalique's avatar
Khalique committed
2445
    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
2446
2447
    {
        if(dims.empty())
2448
            return literal{{shape_type}, data.begin(), data.end()};
2449
        return literal{{shape_type, dims}, data.begin(), data.end()};
2450
2451
    }

2452
    shape parse_type(const onnx::TypeProto& t, const std::vector<std::size_t>& input_dims)
Paul's avatar
Paul committed
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
    {
        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
2463
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
2464
2465
2466
        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;
2467
        case onnx::TensorProto::UINT8: shape_type = shape::uint8_type; break;
Paul's avatar
Paul committed
2468
2469
2470
        case onnx::TensorProto::STRING:
        case onnx::TensorProto::BOOL:
        case onnx::TensorProto::UNDEFINED:
Paul's avatar
Paul committed
2471
2472
        case onnx::TensorProto::COMPLEX64:
        case onnx::TensorProto::COMPLEX128:
Paul's avatar
Paul committed
2473
            break; // throw std::runtime_error("Unsupported type");
Paul's avatar
Paul committed
2474
        }
2475
2476
2477
2478
2479
2480

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

Paul's avatar
Paul committed
2481
        std::vector<std::size_t> dims;
Paul's avatar
Paul committed
2482
        auto&& tensor_dims = t.tensor_type().shape().dim();
2483
2484
2485
        std::transform(tensor_dims.begin(),
                       tensor_dims.end(),
                       std::back_inserter(dims),
2486
2487
                       [&](auto&& d) -> std::size_t {
                           if(d.has_dim_value())
2488
                           {
2489
                               if(static_cast<int>(d.dim_value()) <= 0)
2490
2491
2492
                               {
                                   return default_dim_value;
                               }
2493
                               return d.dim_value();
2494
                           }
2495
2496
2497
2498
                           else
                           {
                               return default_dim_value;
                           }
2499
                       });
2500

2501
2502
2503
        if(dims.empty())
            return {shape_type};

Paul's avatar
Paul committed
2504
2505
        return {shape_type, dims};
    }
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527

    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
2528
2529
2530

    void check_arg_empty(const argument& arg, const std::string& msg)
    {
Shucai Xiao's avatar
Shucai Xiao committed
2531
        if(arg.empty())
Shucai Xiao's avatar
Shucai Xiao committed
2532
2533
2534
2535
        {
            MIGRAPHX_THROW(msg);
        }
    }
Paul's avatar
Paul committed
2536
2537
};

Paul Fultz II's avatar
Paul Fultz II committed
2538
template <class... Ts>
2539
program parse_onnx_from(const onnx_options& options, Ts&&... xs)
Paul's avatar
Paul committed
2540
2541
{
    onnx_parser parser;
2542
2543
2544
    parser.map_input_dims         = options.map_input_dims;
    parser.default_dim_value      = options.default_dim_value;
    parser.skip_unknown_operators = options.skip_unknown_operators;
2545

2546
    if(options.print_program_on_error)
Paul's avatar
Paul committed
2547
    {
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
        // 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
2558
    }
2559
    else
Paul's avatar
Paul committed
2560
    {
2561
        parser.parse_from(std::forward<Ts>(xs)...);
Paul's avatar
Paul committed
2562
2563
2564
2565
    }
    return std::move(parser.prog);
}

2566
program parse_onnx(const std::string& name, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2567
2568
2569
2570
2571
{
    std::fstream input(name.c_str(), std::ios::in | std::ios::binary);
    return parse_onnx_from(options, input);
}

2572
program parse_onnx_buffer(const std::string& buffer, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2573
2574
2575
2576
{
    return parse_onnx_from(options, buffer.data(), buffer.size());
}

2577
program parse_onnx_buffer(const void* data, std::size_t size, const onnx_options& options)
Paul Fultz II's avatar
Paul Fultz II committed
2578
2579
2580
2581
{
    return parse_onnx_from(options, data, size);
}

Paul's avatar
Paul committed
2582
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
2583
} // namespace migraphx