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

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

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

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

    std::unordered_map<std::string, op_func> ops;
34
    std::unordered_map<std::string, operation> actv_funcs;
Paul's avatar
Paul committed
35
36
37

    onnx_parser()
    {
Shucai Xiao's avatar
Shucai Xiao committed
38
        add_generic_op("MatMul", op::dot{});
Khalique's avatar
Khalique committed
39
        add_generic_op("Relu", op::relu{});
Khalique's avatar
Khalique committed
40
41
        add_generic_op("Sigmoid", op::sigmoid{});
        add_generic_op("Abs", op::abs{});
Shucai Xiao's avatar
Shucai Xiao committed
42
43
        add_generic_op("Exp", op::exp{});
        add_generic_op("Log", op::log{});
Khalique's avatar
Khalique committed
44
45
        // disable dropout for inference
        add_generic_op("Dropout", op::identity{});
Khalique's avatar
Khalique committed
46
        add_generic_op("Identity", op::identity{});
Shucai Xiao's avatar
Shucai Xiao committed
47
48
49
        add_generic_op("Sin", op::sin{});
        add_generic_op("Cos", op::cos{});
        add_generic_op("Tan", op::tan{});
50
51
        add_generic_op("Sinh", op::sinh{});
        add_generic_op("Cosh", op::cosh{});
52
        add_generic_op("Tanh", op::tanh{});
53
54
55
        add_generic_op("Asin", op::asin{});
        add_generic_op("Acos", op::acos{});
        add_generic_op("Atan", op::atan{});
Paul's avatar
Paul committed
56

Khalique's avatar
Khalique committed
57
58
59
60
61
        add_binary_op("Add", op::add{});
        add_binary_op("Div", op::div{});
        add_binary_op("Mul", op::mul{});
        add_binary_op("Sub", op::sub{});

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

Khalique's avatar
Khalique committed
66
        add_mem_op("ImageScaler", &onnx_parser::parse_imagescaler);
67
        add_mem_op("LeakyRelu", &onnx_parser::parse_leaky_relu);
Khalique's avatar
Khalique committed
68
        add_mem_op("Elu", &onnx_parser::parse_elu);
Paul's avatar
Paul committed
69
70
        add_mem_op("Constant", &onnx_parser::parse_constant);
        add_mem_op("Conv", &onnx_parser::parse_conv);
Paul's avatar
Paul committed
71
72
        add_mem_op("MaxPool", &onnx_parser::parse_pooling);
        add_mem_op("AveragePool", &onnx_parser::parse_pooling);
73
74
        add_mem_op("GlobalMaxPool", &onnx_parser::parse_pooling);
        add_mem_op("GlobalAveragePool", &onnx_parser::parse_pooling);
Paul's avatar
Paul committed
75
        add_mem_op("Reshape", &onnx_parser::parse_reshape);
Paul's avatar
Paul committed
76
77
        add_mem_op("Flatten", &onnx_parser::parse_flatten);
        add_mem_op("Gemm", &onnx_parser::parse_gemm);
78
        add_mem_op("BatchNormalization", &onnx_parser::parse_batchnorm);
Paul's avatar
Paul committed
79
        add_mem_op("Softmax", &onnx_parser::parse_softmax);
80
81
82
        add_mem_op("Squeeze", &onnx_parser::parse_squeeze);
        add_mem_op("Unsqueeze", &onnx_parser::parse_unsqueeze);
        add_mem_op("Slice", &onnx_parser::parse_slice);
Scott Thornton's avatar
Scott Thornton committed
83
        add_mem_op("Concat", &onnx_parser::parse_concat);
84
85
86
        add_mem_op("Gather", &onnx_parser::parse_gather);
        add_mem_op("Shape", &onnx_parser::parse_shape);
        add_mem_op("ConstantFill", &onnx_parser::parse_constant_fill);
Khalique's avatar
Khalique committed
87
        add_mem_op("Transpose", &onnx_parser::parse_transpose);
Shucai Xiao's avatar
Shucai Xiao committed
88
        add_mem_op("RNN", &onnx_parser::parse_rnn);
89
        add_mem_op("GRU", &onnx_parser::parse_gru);
90
91
92
93
94
95
96
97
98
99

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

    void init_actv_func()
    {
        actv_funcs.insert(std::make_pair("tanh", op::tanh{}));
        actv_funcs.insert(std::make_pair("relu", op::relu{}));
        actv_funcs.insert(std::make_pair("sigmoid", op::sigmoid{}));
100
101
        actv_funcs.insert(std::make_pair("leakyrelu", op::leaky_relu{}));
        actv_funcs.insert(std::make_pair("elu", op::elu{}));
Paul's avatar
Paul committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
    }

    template <class F>
    void add_op(std::string name, F f)
    {
        ops.emplace(name, f);
    }

    template <class F>
    void add_mem_op(std::string name, F f)
    {
        ops.emplace(name, [=](auto&&... xs) {
            return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
        });
    }
Khalique's avatar
Khalique committed
117

118
    template <class T>
Khalique's avatar
Khalique committed
119
    void add_binary_op(std::string name, T x)
120
121
    {
        ops.emplace(name, [this, x](attribute_map attributes, std::vector<instruction_ref> args) {
Scott Thornton's avatar
Scott Thornton committed
122
            if(args.size() != 2)
Paul's avatar
Paul committed
123
                MIGRAPHX_THROW("binary operators should have 2 operands");
124
125
126
127
128
129
130
131
132
133
134
135
            if(contains(attributes, "broadcast"))
            {
                uint64_t broadcasted = parse_value(attributes.at("broadcast")).at<uint64_t>();
                if(broadcasted != 0)
                {
                    uint64_t axis = (contains(attributes, "axis"))
                                        ? parse_value(attributes.at("axis")).at<uint64_t>()
                                        : 0;
                    auto l =
                        prog.add_instruction(op::broadcast{axis, args[0]->get_shape()}, args[1]);
                    return prog.add_instruction(x, args[0], l);
                }
136
                return prog.add_instruction(x, args);
137
            }
Paul's avatar
Paul committed
138
            else
139
            {
Khalique's avatar
Khalique committed
140
                return add_broadcastable_binary_op(args[0], args[1], x);
141
142
143
144
            }
        });
    }

Khalique's avatar
Khalique committed
145
146
147
148
149
    template <class T>
    instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
    {
        if(arg0->get_shape() != arg1->get_shape())
        {
Khalique's avatar
Khalique committed
150
151
152
153
154
155
156
157
158
159
160
161
162
            // 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)
            //
Khalique's avatar
Khalique committed
163
164
165
166
167
168
169
170
            // Get lengths for both arguments
            const std::vector<std::size_t>* s0 = &arg0->get_shape().lens();
            const std::vector<std::size_t>* s1 = &arg1->get_shape().lens();

            // Make sure s0 is the smaller size
            if(s0->size() > s1->size())
                std::swap(s0, s1);

Khalique's avatar
Khalique committed
171
            std::vector<std::size_t> output_lens(*s1);
Khalique's avatar
Khalique committed
172
173
            auto offset = s1->size() - s0->size();
            std::transform(s0->begin(),
Khalique's avatar
Khalique committed
174
175
176
177
                           s0->end(),
                           s1->begin() + offset,
                           output_lens.begin() + offset,
                           [](auto a, auto b) { return std::max(a, b); });
Khalique's avatar
Khalique committed
178
179
180
181
182
183
184
185
186

            auto l0 = prog.add_instruction(op::multibroadcast{output_lens}, arg0);
            auto l1 = prog.add_instruction(op::multibroadcast{output_lens}, arg1);
            return prog.add_instruction(x, l0, l1);
        }
        else
        {
            return prog.add_instruction(x, {arg0, arg1});
        }
187
188
    }

Paul's avatar
Paul committed
189
    template <class T>
Paul's avatar
Paul committed
190
191
    void add_generic_op(std::string name, T x)
    {
Paul's avatar
Paul committed
192
        ops.emplace(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
Paul's avatar
Paul committed
193
194
195
196
            return prog.add_instruction(x, args);
        });
    }

Khalique's avatar
Khalique committed
197
    template <class T>
Khalique's avatar
Khalique committed
198
    void add_variadic_op(std::string name, T x)
Khalique's avatar
Khalique committed
199
    {
Khalique's avatar
Khalique committed
200
201
        ops.emplace(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
            return std::accumulate(std::next(args.begin()),
Khalique's avatar
Khalique committed
202
203
204
205
206
                                   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
207
        });
Khalique's avatar
Khalique committed
208
209
    }

Paul's avatar
Paul committed
210
    instruction_ref
Paul's avatar
Paul committed
211
    parse_softmax(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
212
213
    {
        auto dims = args.front()->get_shape().lens();
Scott Thornton's avatar
Scott Thornton committed
214
215
        auto r =
            prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1]), 1, 1}}, args.front());
216
217
        auto s = prog.add_instruction(op::softmax{}, r);
        return prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1])}}, s);
Paul's avatar
Paul committed
218
219
    }

Paul's avatar
Paul committed
220
    instruction_ref
Paul's avatar
Paul committed
221
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
222
    {
223
        op::convolution op;
Paul's avatar
Paul committed
224
225
        if(contains(attributes, "pads"))
        {
Scott Thornton's avatar
Scott Thornton committed
226
            if(contains(attributes, "auto_pad"))
227
            {
Paul's avatar
Paul committed
228
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
229
230
231
            }
            std::vector<std::size_t> padding(4);
            copy(attributes["pads"].ints(), padding.begin());
Scott Thornton's avatar
Scott Thornton committed
232
            if(padding.size() != 4)
233
            {
Paul's avatar
Paul committed
234
                MIGRAPHX_THROW("padding should have 4 values");
235
            }
Scott Thornton's avatar
Scott Thornton committed
236
            if(padding[0] != padding[2] || padding[1] != padding[3])
237
            {
Paul's avatar
Paul committed
238
                MIGRAPHX_THROW("migraphx does not support asymetric padding");
239
240
241
            }
            op.padding[0] = padding[0];
            op.padding[1] = padding[1];
Paul's avatar
Paul committed
242
        }
Paul's avatar
Paul committed
243
244
245
246
247
248
249
250
        if(contains(attributes, "strides"))
        {
            copy(attributes["strides"].ints(), op.stride.begin());
        }
        if(contains(attributes, "dilations"))
        {
            copy(attributes["dilations"].ints(), op.dilation.begin());
        }
Scott Thornton's avatar
Scott Thornton committed
251
        if(contains(attributes, "auto_pad"))
252
253
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
254
            if(contains(attributes, "pads") and to_upper(s) != "NOTSET")
255
            {
Paul's avatar
Paul committed
256
                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
257
258
            }

wsttiger's avatar
fixes  
wsttiger committed
259
            if(s.find("SAME") != std::string::npos)
260
261
262
263
            {
                op.padding_mode = op::convolution::same;
            }
        }
Khalique's avatar
Khalique committed
264
265
266
267
        if(contains(attributes, "group"))
        {
            op.group = parse_value(attributes.at("group")).at<int>();
        }
Paul's avatar
Paul committed
268
269
270
271
        if(args.size() == 3)
        {
            uint64_t axis = 1;
            auto l1       = prog.add_instruction(op, args[0], args[1]);
Scott Thornton's avatar
Scott Thornton committed
272
            auto l2       = prog.add_instruction(op::broadcast{axis, l1->get_shape()}, args[2]);
273
            return prog.add_instruction(op::add{}, l1, l2);
Paul's avatar
Paul committed
274
        }
Paul's avatar
Paul committed
275
276
        return prog.add_instruction(op, args);
    }
Paul's avatar
Paul committed
277

Paul's avatar
Paul committed
278
279
280
    instruction_ref parse_pooling(const std::string& name,
                                  attribute_map attributes,
                                  std::vector<instruction_ref> args)
Paul's avatar
Paul committed
281
    {
Khalique's avatar
Khalique committed
282
283
        op::pooling op{ends_with(name, "MaxPool") ? "max" : "average"};
        if(starts_with(name, "Global"))
284
        {
Khalique's avatar
Khalique committed
285
286
            auto lens  = args.front()->get_shape().lens();
            op.lengths = {lens[2], lens[3]};
287
        }
Paul's avatar
Paul committed
288
289
        if(contains(attributes, "pads"))
        {
290
291
            std::vector<std::size_t> padding(4);
            copy(attributes["pads"].ints(), padding.begin());
Scott Thornton's avatar
Scott Thornton committed
292
            if(padding.size() != 4)
293
            {
Paul's avatar
Paul committed
294
                MIGRAPHX_THROW("padding should have 4 values");
295
            }
Scott Thornton's avatar
Scott Thornton committed
296
            if(padding[0] != padding[2] || padding[1] != padding[3])
297
            {
Paul's avatar
Paul committed
298
                MIGRAPHX_THROW("migraphx does not support asymetric padding");
299
300
301
            }
            op.padding[0] = padding[0];
            op.padding[1] = padding[1];
Paul's avatar
Paul committed
302
303
304
305
306
307
308
309
310
        }
        if(contains(attributes, "strides"))
        {
            copy(attributes["strides"].ints(), op.stride.begin());
        }
        if(contains(attributes, "kernel_shape"))
        {
            copy(attributes["kernel_shape"].ints(), op.lengths.begin());
        }
Scott Thornton's avatar
Scott Thornton committed
311
        if(contains(attributes, "auto_pad"))
312
313
        {
            auto s = attributes["auto_pad"].s();
Scott Thornton's avatar
Scott Thornton committed
314
            if(to_upper(s) != "NOTSET")
315
            {
Paul's avatar
Paul committed
316
                MIGRAPHX_THROW("auto_pad is not supported for pooling");
317
318
319
            }
        }

Paul's avatar
Paul committed
320
        return prog.add_instruction(op, std::move(args));
Paul's avatar
Paul committed
321
322
    }

Paul's avatar
Paul committed
323
    instruction_ref
Paul's avatar
Paul committed
324
    parse_reshape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
325
    {
326
        op::reshape op;
Paul's avatar
Paul committed
327
328
329
330
331
332
333
        if(args.size() == 1)
        {
            literal s = parse_value(attributes.at("shape"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
        }
        if(args.size() == 2)
        {
Paul's avatar
Paul committed
334
            literal s = args[1]->get_literal();
Paul's avatar
Paul committed
335
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
336
        }
Paul's avatar
Paul committed
337
338
339
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
340
    instruction_ref
Paul's avatar
Paul committed
341
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
342
    {
343
        uint64_t axis = 1;
Paul's avatar
Paul committed
344
345
346
347
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
348
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
349
350
    }

351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
    instruction_ref
    parse_squeeze(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::squeeze op;
        literal s = parse_value(attributes.at("axes"));
        s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        return prog.add_instruction(op, args[0]);
    }

    instruction_ref
    parse_unsqueeze(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::unsqueeze op;
        literal s = parse_value(attributes.at("axes"));
        s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        return prog.add_instruction(op, args[0]);
    }

Scott Thornton's avatar
Scott Thornton committed
369
370
371
372
373
374
375
    instruction_ref
    parse_concat(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        std::size_t axis = parse_value(attributes.at("axis")).at<int>();
        op::concat op{axis};
        return prog.add_instruction(op, std::move(args));
    }
376

377
378
379
380
381
382
383
384
    instruction_ref
    parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        std::size_t axis = 0;
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
385
        op::gather op{axis};
386
387
388
        return prog.add_instruction(op, std::move(args));
    }

389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
    instruction_ref
    parse_slice(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::slice op;
        if(contains(attributes, "axes"))
        {
            literal s = parse_value(attributes.at("axes"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
        }
        {
            literal s = parse_value(attributes.at("ends"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.ends)); });
        }
        {
            literal s = parse_value(attributes.at("starts"));
            s.visit([&](auto v) { copy(v, std::back_inserter(op.starts)); });
        }
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
409
410
411
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
412
413
414
415
    {
        literal v = parse_value(attributes.at("value"));
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
416

Paul's avatar
Paul committed
417
    instruction_ref
Paul's avatar
Paul committed
418
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
419
420
    {
        float alpha = 1.0f;
Khalique's avatar
Khalique committed
421
        float beta  = 1.0f;
Paul's avatar
Paul committed
422
423
424
425
426
427
428
429
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
430
            beta = parse_value(attributes.at("beta")).at<float>();
Paul's avatar
Paul committed
431
432
433
434
435
436
437
438
439
440
        }
        if(contains(attributes, "transA"))
        {
            transa = parse_value(attributes.at("transA")).at<bool>();
        }
        if(contains(attributes, "transB"))
        {
            transb = parse_value(attributes.at("transB")).at<bool>();
        }
        std::vector<int64_t> perm = {1, 0};
441
442
        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
443
444
        if(args.size() == 3)
        {
Khalique's avatar
Khalique committed
445
            if(beta != 0.f)
446
            {
Khalique's avatar
Khalique committed
447
                auto l3 = prog.add_instruction(op::dot{alpha}, l1, l2);
Khalique's avatar
Khalique committed
448
                auto l4 = args[2];
Khalique's avatar
Khalique committed
449
                if(l4->get_shape().scalar()) // ignore args[2] (no C value added to alpha*A*B)
Khalique's avatar
Khalique committed
450
                    return l3;
Khalique's avatar
Khalique committed
451
                if(beta != 1.f)
Khalique's avatar
Khalique committed
452
453
                {
                    auto beta_val = prog.add_literal(beta);
Khalique's avatar
Khalique committed
454
455
                    auto l5 = prog.add_instruction(op::scalar{args[2]->get_shape()}, beta_val);
                    l4      = prog.add_instruction(op::mul{}, args[2], l5);
Khalique's avatar
Khalique committed
456
457
                }
                return add_broadcastable_binary_op(l3, l4, op::add{});
458
            }
Paul's avatar
Paul committed
459
        }
Shucai Xiao's avatar
Shucai Xiao committed
460
        return prog.add_instruction(op::dot{alpha, beta}, l1, l2);
Paul's avatar
Paul committed
461
462
    }

463
    instruction_ref
Paul's avatar
Paul committed
464
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
465
    {
Scott Thornton's avatar
Scott Thornton committed
466
467
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
468
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Scott Thornton's avatar
Scott Thornton committed
469
        bool is_test                                      = false;
470
471
472
473
474
475
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
476
            momentum = parse_value(attributes.at("momentum")).at<float>();
477
478
479
        }
        if(contains(attributes, "is_test"))
        {
wsttiger's avatar
wsttiger committed
480
            is_test = parse_value(attributes.at("is_test")).at<uint64_t>() > 0;
481
482
483
        }
        if(contains(attributes, "spatial"))
        {
484
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
485
486
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
487
        }
Paul's avatar
Paul committed
488
        (void)is_test;
Paul's avatar
Paul committed
489
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
490
        return prog.add_instruction(op, std::move(args));
491
492
    }

493
494
495
496
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
497
        float alpha = 0.01; // default alpha val for leaky relu
498
499
500
501
502
503
504
505
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        op::leaky_relu op{alpha};
        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
506
507
    instruction_ref
    parse_elu(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
508
509
510
511
512
513
514
515
516
517
    {
        float alpha = 1.0; // default alpha val for elu
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        op::elu op{alpha};
        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
    instruction_ref parse_imagescaler(const std::string&,
                                      attribute_map attributes,
                                      std::vector<instruction_ref> args)
    {
        float scale = 1.0;
        std::vector<float> bias{};
        if(contains(attributes, "scale"))
        {
            scale = parse_value(attributes.at("scale")).at<float>();
        }

        if(contains(attributes, "bias"))
        {
            auto&& bias_floats = attributes["bias"].floats();
            bias               = std::vector<float>(bias_floats.begin(), bias_floats.end());
        }
        auto input_shape = args.front()->get_shape();
Khalique's avatar
Khalique committed
535

Khalique's avatar
Khalique committed
536
537
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
Paul's avatar
Paul committed
538
            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
Khalique's avatar
Khalique committed
539

Paul's avatar
Paul committed
540
541
        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_shape}, scale_val);
        auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
Paul's avatar
Paul committed
542
        auto bias_bcast = prog.add_instruction(migraphx::op::broadcast{1, input_shape}, bias_vals);
Paul's avatar
Paul committed
543
        return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
Khalique's avatar
Khalique committed
544
    }
Khalique's avatar
Khalique committed
545

Khalique's avatar
Khalique committed
546
547
    instruction_ref
    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
548
549
550
551
552
553
554
    {
        std::vector<int64_t> perm{};
        if(contains(attributes, "perm"))
        {
            auto&& perm_vals = attributes["perm"].ints();
            perm             = std::vector<int64_t>(perm_vals.begin(), perm_vals.end());
        }
Paul's avatar
Paul committed
555
        return prog.add_instruction(migraphx::op::transpose{perm}, args.front());
Khalique's avatar
Khalique committed
556
557
    }

558
559
560
    // Use a literal instruction to replace the shape since, output of
    // shape operator are literals in migraphx
    instruction_ref
Shucai Xiao's avatar
Shucai Xiao committed
561
    parse_shape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
562
563
    {
        if(args.size() != 1)
564
            MIGRAPHX_THROW("Shape: operator should have 1 operand");
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
593
594
595
596
597
598
599
600
        std::vector<std::size_t> arg_shape = args[0]->get_shape().lens();
        std::vector<int64_t> vec_shape(arg_shape.size());
        migraphx::shape s(migraphx::shape::int64_type, {arg_shape.size()});
        std::transform(arg_shape.begin(), arg_shape.end(), vec_shape.begin(), [](auto i) {
            return int64_t(i);
        });
        return prog.add_literal(migraphx::literal{s, vec_shape});
    }

    // Use a literal instruction to replace the constantFill operator. In RNN, input shape
    // and value are fixed, so no need to do the actual computation for the constantFill
    // operator
    instruction_ref parse_constant_fill(const std::string&,
                                        attribute_map attributes,
                                        std::vector<instruction_ref> args)
    {
        int input_as_shape = 0;
        int dtype          = 1;
        float value        = 0.0f;

        if(contains(attributes, "dtype"))
        {
            dtype = parse_value(attributes.at("dtype")).at<int>();
        }
        migraphx::shape::type_t type = get_type(dtype);

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

        if(contains(attributes, "value"))
        {
            value = parse_value(attributes.at("value")).at<float>();
        }

Shucai Xiao's avatar
Shucai Xiao committed
601
602
        if(contains(attributes, "extra_shape"))
        {
603
            MIGRAPHX_THROW("ConstantFill: cannot handle extra shape attribute");
604
605
        }

606
607
        if(input_as_shape == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
608
            if(args.size() != 1)
609
            {
610
                MIGRAPHX_THROW("ConstantFill: need an input argument as output shape");
611
612
            }

Shucai Xiao's avatar
Shucai Xiao committed
613
614
            if(contains(attributes, "shape"))
            {
615
                MIGRAPHX_THROW("ConstantFill: cannot set the shape argument and pass in an input "
Shucai Xiao's avatar
Shucai Xiao committed
616
                               "at the same time");
617
618
            }

619
620
621
            migraphx::argument in = args[0]->eval();
            if(in.empty())
            {
622
                MIGRAPHX_THROW("ConstantFill: cannot handle dynamic shape as input");
623
            }
624

625
626
627
            std::vector<std::size_t> dims;
            in.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
            migraphx::shape s(type, dims);
628
629
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
630
631
632
        }
        else if(input_as_shape == 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
633
634
            if(!contains(attributes, "shape"))
            {
635
                MIGRAPHX_THROW("ConstantFill: attribute output shape is needed");
636
637
638
            }

            literal ls = parse_value(attributes.at("shape"));
639
            std::vector<std::size_t> dims;
Shucai Xiao's avatar
Shucai Xiao committed
640
            ls.visit([&](auto s) { dims.assign(s.begin(), s.end()); });
641
            migraphx::shape s{type, dims};
642
643
            std::vector<float> values(s.elements(), value);
            return prog.add_literal(migraphx::literal(s, values));
644
645
646
        }
        else
        {
647
            MIGRAPHX_THROW("ConstantFill: wrong value of attribute input_as_shape");
648
649
650
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
651
652
653
654
    instruction_ref
    parse_rnn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
655
        std::size_t hidden_size     = args[1]->get_shape().lens()[1];
Shucai Xiao's avatar
Shucai Xiao committed
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671

        if(contains(attributes, "hidden_size"))
        {
            hidden_size = parse_value(attributes.at("hidden_size")).at<int>();
        }
        else
        {
            MIGRAPHX_THROW("RNN: hidden size attribute missing");
        }

        std::string activation_func = {"tanh"};
        if(contains(attributes, "activations"))
        {
            activation_func = attributes.at("activations").strings(0);
        }

672
        if(actv_funcs.count(activation_func) == 0)
Shucai Xiao's avatar
Shucai Xiao committed
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
        {
            MIGRAPHX_THROW("RNN: activation function " + activation_func + " not supported");
        }

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

        op::rnn::rnn_direction_t dirct = op::rnn::forward;
        if(direction == "bidirectional")
        {
            dirct = op::rnn::bidirectional;
        }
        else if(direction == "reverse")
        {
            dirct = op::rnn::reverse;
        }

        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

701
702
        return prog.add_instruction(op::rnn{hidden_size, actv_funcs[activation_func], dirct, clip},
                                    std::move(args));
Shucai Xiao's avatar
Shucai Xiao committed
703
704
    }

705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
    instruction_ref
    parse_gru(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        migraphx::shape input_shape = args[0]->get_shape();
        std::size_t hidden_size     = args[2]->get_shape().lens()[2];

        if(contains(attributes, "hidden_size"))
        {
            hidden_size = parse_value(attributes.at("hidden_size")).at<int>();
        }
        else
        {
            MIGRAPHX_THROW("GRU: hidden size attribute missing");
        }

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

        op::gru::gru_direction_t dirct = op::gru::forward;
        if(direction == "bidirectional")
        {
            dirct = op::gru::bidirectional;
        }
        else if(direction == "reverse")
        {
            dirct = op::gru::reverse;
        }

        std::vector<std::string> act_funcs = {"sigmoid", "tanh"};
        if(contains(attributes, "activations"))
        {
            act_funcs[0] = attributes.at("activations").strings(0);
            act_funcs[1] = attributes.at("activations").strings(1);
        }

        if (act_funcs.size() != 2)
        {
            MIGRAPHX_THROW("GRU: wrong activation function attribute");
        }

        for (std::size_t i = 0; i < act_funcs.size(); ++i)
        {
            if(actv_funcs.count(act_funcs.at(i)) == 0)
            {
                MIGRAPHX_THROW("GRU: activation function " + act_funcs.at(i) + " not supported");
            }
        }

        // To be added later
        float clip = 0.0;
        if(contains(attributes, "clip"))
        {
            clip = parse_value(attributes.at("clip")).at<float>();
        }

        int linear_before_reset = 0;
        if (contains(attributes, "linear_before_reset"))
        {
            linear_before_reset = parse_value(attributes.at("linear_before_reset")).at<int>();
        }

        return prog.add_instruction(op::gru{hidden_size, 
                                    {actv_funcs[act_funcs.at(0)], actv_funcs[act_funcs.at(1)]}, 
                                    dirct, clip, linear_before_reset},
                                    std::move(args));
    }

Paul's avatar
Paul committed
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
    void parse_from(std::istream& is)
    {
        onnx::ModelProto model;
        if(model.ParseFromIstream(&is))
        {
            if(model.has_graph())
            {
                this->parse_graph(model.graph());
            }
        }
        else
        {
            throw std::runtime_error("Failed reading");
        }
    }

    void parse_graph(const onnx::GraphProto& graph)
    {
        nodes = get_nodes(graph);
795
796
797
798
799
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
800
801
802
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
803
804
805
806
807
808
809
810
811
812
813
814
            // Does the input have an initializer?
            if(contains(initializer_data, name))
            {
                auto t             = initializer_data[name];
                instructions[name] = prog.add_literal(parse_tensor(t));
            }
            else
            {
                // TODO: Get shape of input parameter
                shape s            = parse_type(input.type());
                instructions[name] = prog.add_parameter(name, s);
            }
Paul's avatar
Paul committed
815
816
817
        }
        for(auto&& p : nodes)
        {
818
            this->parse_node(get_name(p.second));
Paul's avatar
Paul committed
819
820
821
        }
    }

Paul's avatar
Paul committed
822
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
823
    {
Paul's avatar
Paul committed
824
        if(name.empty())
Paul's avatar
Paul committed
825
            MIGRAPHX_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
826
827
828
829
830
831
        if(instructions.count(name) == 0)
        {
            auto&& node = nodes.at(name);
            std::vector<instruction_ref> args;
            for(auto&& input : node.input())
            {
832
833
834
835
836
837
838
839
840
841
                // For RNN, LSTM, and GRU operators, one of the input arguments
                // is prim::Undefined, and it is ignored by protobuf. We use a
                // hack to ignore this argument for these three operators
                std::string op_type = node.op_type();
                if((op_type == "RNN" || op_type == "LSTM" || op_type == "GRU") &&
                   input.empty() == true)
                {
                    continue;
                }

Paul's avatar
Paul committed
842
843
                if(nodes.count(input) > 0)
                {
844
                    auto&& iname = get_name(nodes.at(input));
Paul's avatar
Paul committed
845
                    assert(name != iname);
Paul's avatar
Paul committed
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
                    this->parse_node(iname);
                    args.push_back(instructions.at(iname));
                }
                else
                {
                    args.push_back(instructions.at(input));
                }
            }
            if(ops.count(node.op_type()) == 0)
            {
                instructions[name] = prog.add_instruction(unknown{node.op_type()}, args);
            }
            else
            {
                instructions[name] = ops[node.op_type()](get_attributes(node), args);
            }
        }
    }

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

875
876
877
878
    static std::string get_name(const onnx::NodeProto& node)
    {
        if(node.name().empty())
        {
Paul's avatar
Paul committed
879
            std::string generated = "migraphx_unnamed_node";
Paul's avatar
Paul committed
880
881
882
883
            return std::accumulate(node.output().begin(),
                                   node.output().end(),
                                   generated,
                                   [](auto x, auto y) { return x + "_" + y; });
884
885
886
887
        }
        return node.name();
    }

Paul's avatar
Paul committed
888
889
890
891
892
    static node_map get_nodes(const onnx::GraphProto& graph)
    {
        std::unordered_map<std::string, onnx::NodeProto> result;
        for(auto&& node : graph.node())
        {
893
            result[get_name(node)] = node;
Paul's avatar
Paul committed
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
            for(auto&& output : node.output())
            {
                result[output] = node;
            }
        }
        return result;
    }

    template <class T>
    static literal from_repeated(shape::type_t t, const T& r)
    {
        std::size_t size = r.size();
        return literal{{t, {size}}, r.begin(), r.end()};
    }

    static literal parse_value(const onnx::AttributeProto& attr)
    {
        switch(attr.type())
        {
        case onnx::AttributeProto::UNDEFINED: return {};
        case onnx::AttributeProto::FLOAT: return literal{attr.f()};
        case onnx::AttributeProto::INT: return literal{attr.i()};
        case onnx::AttributeProto::STRING: return {};
        case onnx::AttributeProto::TENSOR: return parse_tensor(attr.t());
        case onnx::AttributeProto::GRAPH: return {};
Paul's avatar
Paul committed
919
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
920
921
922
923
924
        case onnx::AttributeProto::INTS: return from_repeated(shape::int64_type, attr.ints());
        case onnx::AttributeProto::STRINGS: return {};
        case onnx::AttributeProto::TENSORS: return {};
        case onnx::AttributeProto::GRAPHS: return {};
        }
Paul's avatar
Paul committed
925
        MIGRAPHX_THROW("Invalid attribute type");
Paul's avatar
Paul committed
926
927
928
929
930
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
Khalique's avatar
Khalique committed
931
        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
932
        if(dims.empty())
Khalique's avatar
Khalique committed
933
934
935
        {
            dims = {1};
        }
936
937
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
938
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
939
940
941
942
943
944
945
946
947
948
949
950
            switch(t.data_type())
            {
            case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
            case onnx::TensorProto::FLOAT: return literal{{shape::float_type, dims}, s.data()};
            case onnx::TensorProto::UINT8: throw std::runtime_error("");
            case onnx::TensorProto::INT8: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::UINT16: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT16: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT32: return literal{{shape::int32_type, dims}, s.data()};
            case onnx::TensorProto::INT64: return literal{{shape::int64_type, dims}, s.data()};
            case onnx::TensorProto::STRING: throw std::runtime_error("");
            case onnx::TensorProto::BOOL: return literal{{shape::int32_type, dims}, s.data()};
Paul's avatar
Paul committed
951
            case onnx::TensorProto::FLOAT16: return literal{{shape::half_type, dims}, s.data()};
Scott Thornton's avatar
Scott Thornton committed
952
953
954
955
956
957
            case onnx::TensorProto::DOUBLE: return literal{{shape::double_type, dims}, s.data()};
            case onnx::TensorProto::UINT32: throw std::runtime_error("");
            case onnx::TensorProto::UINT64: throw std::runtime_error("");
            case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
            case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
            }
Paul's avatar
Paul committed
958
            MIGRAPHX_THROW("Invalid tensor type");
959
        }
Paul's avatar
Paul committed
960
961
962
963
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Paul's avatar
Paul committed
964
            return literal{{shape::float_type, dims}, t.float_data().begin(), t.float_data().end()};
Paul's avatar
Paul committed
965
966
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Paul's avatar
Paul committed
967
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
968
        case onnx::TensorProto::UINT16:
Paul's avatar
Paul committed
969
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
970
        case onnx::TensorProto::INT16:
Paul's avatar
Paul committed
971
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
972
        case onnx::TensorProto::INT32:
Paul's avatar
Paul committed
973
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
974
        case onnx::TensorProto::INT64:
Paul's avatar
Paul committed
975
            return literal{{shape::int64_type, dims}, t.int64_data().begin(), t.int64_data().end()};
Paul's avatar
Paul committed
976
977
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Paul's avatar
Paul committed
978
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
979
980
        case onnx::TensorProto::FLOAT16:
            return literal{{shape::half_type, dims}, t.float_data().begin(), t.float_data().end()};
Paul's avatar
Paul committed
981
982
983
984
985
986
987
988
        case onnx::TensorProto::DOUBLE:
            return literal{
                {shape::double_type, dims}, t.double_data().begin(), t.double_data().end()};
        case onnx::TensorProto::UINT32: throw std::runtime_error("");
        case onnx::TensorProto::UINT64: throw std::runtime_error("");
        case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
        case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
        }
Paul's avatar
Paul committed
989
        MIGRAPHX_THROW("Invalid tensor type");
Paul's avatar
Paul committed
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
    }

    static shape parse_type(const onnx::TypeProto& t)
    {
        shape::type_t shape_type{};
        switch(t.tensor_type().elem_type())
        {
        case onnx::TensorProto::UNDEFINED:
            break; // throw std::runtime_error("Unsupported type UNDEFINED");
        case onnx::TensorProto::FLOAT: shape_type = shape::float_type; break;
        case onnx::TensorProto::UINT8:
            break; // throw std::runtime_error("Unsupported type UINT8");
        case onnx::TensorProto::INT8: shape_type = shape::int8_type; break;
        case onnx::TensorProto::UINT16: shape_type = shape::uint16_type; break;
        case onnx::TensorProto::INT16: shape_type = shape::int16_type; break;
        case onnx::TensorProto::INT32: shape_type = shape::int32_type; break;
        case onnx::TensorProto::INT64: shape_type = shape::int64_type; break;
        case onnx::TensorProto::STRING:
            break; // throw std::runtime_error("Unsupported type STRING");
        case onnx::TensorProto::BOOL:
            break; // throw std::runtime_error("Unsupported type BOOL");
Paul's avatar
Paul committed
1011
        case onnx::TensorProto::FLOAT16: shape_type = shape::half_type; break;
Paul's avatar
Paul committed
1012
1013
1014
1015
1016
1017
1018
1019
1020
        case onnx::TensorProto::DOUBLE: shape_type = shape::double_type; break;
        case onnx::TensorProto::UINT32: shape_type = shape::uint32_type; break;
        case onnx::TensorProto::UINT64: shape_type = shape::uint64_type; break;
        case onnx::TensorProto::COMPLEX64:
            break; // throw std::runtime_error("Unsupported type COMPLEX64");
        case onnx::TensorProto::COMPLEX128:
            break; // throw std::runtime_error("Unsupported type COMPLEX128");
        }
        std::vector<std::size_t> dims;
Paul's avatar
Paul committed
1021
        auto&& tensor_dims = t.tensor_type().shape().dim();
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
        std::transform(tensor_dims.begin(),
                       tensor_dims.end(),
                       std::back_inserter(dims),
                       [](auto&& d) -> std::size_t {
                           if(not d.has_dim_value())
                           {
                               long default_batch_size = 1; // FIXME
                               return default_batch_size;
                           }
                           return d.dim_value();
                       });
Paul's avatar
Paul committed
1033
1034
        return {shape_type, dims};
    }
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056

    shape::type_t get_type(int dtype)
    {
        switch(dtype)
        {
        case 1: return shape::float_type;
        case 2: return shape::uint8_type;
        case 3: return shape::int8_type;
        case 4: return shape::uint16_type;
        case 5: return shape::int16_type;
        case 6: return shape::int32_type;
        case 7: return shape::int64_type;
        case 10: return shape::half_type;
        case 11: return shape::double_type;
        case 12: return shape::uint32_type;
        case 13: return shape::uint64_type;
        default:
        {
            MIGRAPHX_THROW("Prototensor data type " + std::to_string(dtype) + " not supported");
        }
        }
    }
Paul's avatar
Paul committed
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
};

program parse_onnx(const std::string& name)
{
    std::fstream input(name.c_str(), std::ios::in | std::ios::binary);
    onnx_parser parser;
#ifndef NDEBUG
    // Log the program when it can't be parsed
    try
    {
        parser.parse_from(input);
    }
    catch(...)
    {
        std::cerr << parser.prog << std::endl;
        throw;
    }
#else
    parser.parse_from(input);
#endif
    return std::move(parser.prog);
}

Paul's avatar
Paul committed
1080
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
1081
} // namespace migraphx