onnx.cpp 23.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
#include <migraph/fallthrough.hpp>
#include <migraph/program.hpp>
#include <migraph/operators.hpp>
#include <migraph/ranges.hpp>
16
#include <migraph/instruction.hpp>
Paul's avatar
Paul committed
17

Paul's avatar
Paul committed
18
namespace migraph {
Paul's avatar
Paul committed
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

struct unknown
{
    std::string op;
    std::string name() const { return "unknown:" + op; }
    shape compute_shape(std::vector<shape> input) const
    {
        if(input.empty())
            return {};
        else
            return input.front();
    }
    friend std::ostream& operator<<(std::ostream& os, const unknown& x)
    {
        os << x.name();
        return os;
    }
};

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
42
    using op_func = std::function<instruction_ref(attribute_map, std::vector<instruction_ref>)>;
Paul's avatar
Paul committed
43
44
45
46
47
48
49
50
    node_map nodes;
    std::unordered_map<std::string, instruction_ref> instructions;
    program prog = program();

    std::unordered_map<std::string, op_func> ops;

    onnx_parser()
    {
51
52
        add_generic_op("Add", op::add{});
        add_generic_op("Div", op::div{});
Shucai Xiao's avatar
Shucai Xiao committed
53
        add_generic_op("MatMul", op::dot{});
54
55
56
57
        add_generic_op("Mul", op::mul{});
        add_generic_op("Relu", op::activation{"relu"});
        add_generic_op("Sub", op::sub{});
        add_generic_op("Sum", op::add{});
Paul's avatar
Paul committed
58

Khalique's avatar
Khalique committed
59
        add_mem_op("ImageScaler", &onnx_parser::parse_imagescaler);
60
        add_mem_op("LeakyRelu", &onnx_parser::parse_leaky_relu);
Paul's avatar
Paul committed
61
62
        add_mem_op("Constant", &onnx_parser::parse_constant);
        add_mem_op("Conv", &onnx_parser::parse_conv);
Paul's avatar
Paul committed
63
64
        add_mem_op("MaxPool", &onnx_parser::parse_pooling);
        add_mem_op("AveragePool", &onnx_parser::parse_pooling);
65
66
        add_mem_op("GlobalMaxPool", &onnx_parser::parse_pooling);
        add_mem_op("GlobalAveragePool", &onnx_parser::parse_pooling);
Paul's avatar
Paul committed
67
        add_mem_op("Reshape", &onnx_parser::parse_reshape);
Paul's avatar
Paul committed
68
69
        add_mem_op("Flatten", &onnx_parser::parse_flatten);
        add_mem_op("Gemm", &onnx_parser::parse_gemm);
70
        add_mem_op("BatchNormalization", &onnx_parser::parse_batchnorm);
Paul's avatar
Paul committed
71
        add_mem_op("Softmax", &onnx_parser::parse_softmax);
72
73
74
        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
75
        add_mem_op("Concat", &onnx_parser::parse_concat);
Paul's avatar
Paul committed
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
    }

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

Paul's avatar
Paul committed
92
    template <class T>
Paul's avatar
Paul committed
93
94
95
96
97
98
99
100
101
102
103
    void add_generic_op(std::string name, T x)
    {
        ops.emplace(name, [this, x](attribute_map attributes, std::vector<instruction_ref> args) {
            if(args.size() == 2 and 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;
Scott Thornton's avatar
Scott Thornton committed
104
105
                    auto l =
                        prog.add_instruction(op::broadcast{axis, args[0]->get_shape()}, args[1]);
Paul's avatar
Paul committed
106
107
108
109
110
111
112
                    return prog.add_instruction(x, args[0], l);
                }
            }
            return prog.add_instruction(x, args);
        });
    }

Paul's avatar
Paul committed
113
    instruction_ref
Paul's avatar
Paul committed
114
    parse_softmax(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
115
116
    {
        auto dims = args.front()->get_shape().lens();
Scott Thornton's avatar
Scott Thornton committed
117
118
        auto r =
            prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1]), 1, 1}}, args.front());
119
120
        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
121
122
    }

Paul's avatar
Paul committed
123
    instruction_ref
Paul's avatar
Paul committed
124
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
125
    {
126
        op::convolution op;
Paul's avatar
Paul committed
127
128
129
        if(contains(attributes, "pads"))
        {
            copy(attributes["pads"].ints(), op.padding.begin());
Paul's avatar
Paul committed
130
        }
Paul's avatar
Paul committed
131
132
133
134
135
136
137
138
139
140
141
142
        if(contains(attributes, "strides"))
        {
            copy(attributes["strides"].ints(), op.stride.begin());
        }
        if(contains(attributes, "dilations"))
        {
            copy(attributes["dilations"].ints(), op.dilation.begin());
        }
        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
143
            auto l2       = prog.add_instruction(op::broadcast{axis, l1->get_shape()}, args[2]);
144
            return prog.add_instruction(op::add{}, l1, l2);
Paul's avatar
Paul committed
145
        }
Paul's avatar
Paul committed
146
147
        return prog.add_instruction(op, args);
    }
Paul's avatar
Paul committed
148

Paul's avatar
Paul committed
149
150
151
    instruction_ref parse_pooling(const std::string& name,
                                  attribute_map attributes,
                                  std::vector<instruction_ref> args)
Paul's avatar
Paul committed
152
    {
153
154
155
156
157
158
159
160
161
        op::pooling op{name == "MaxPool" or name == "GlobalMaxPool" ? "max" : "average"};
        if(name == "GlobalMaxPool" or name == "GlobalAveragePool")
        {
            auto lens        = args.front()->get_shape().lens();
            auto num_lengths = lens.size() - 2; // ignore N and C values in lens
            assert(num_lengths > 0);
            op.lengths = std::vector<std::size_t>(num_lengths);
            std::copy_n(lens.begin() + 2, num_lengths, op.lengths.begin());
        }
Paul's avatar
Paul committed
162
163
164
165
166
167
168
169
170
171
172
173
        if(contains(attributes, "pads"))
        {
            copy(attributes["pads"].ints(), op.padding.begin());
        }
        if(contains(attributes, "strides"))
        {
            copy(attributes["strides"].ints(), op.stride.begin());
        }
        if(contains(attributes, "kernel_shape"))
        {
            copy(attributes["kernel_shape"].ints(), op.lengths.begin());
        }
Paul's avatar
Paul committed
174
        return prog.add_instruction(op, std::move(args));
Paul's avatar
Paul committed
175
176
    }

Paul's avatar
Paul committed
177
    instruction_ref
Paul's avatar
Paul committed
178
    parse_reshape(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
179
    {
180
        op::reshape op;
Paul's avatar
Paul committed
181
182
183
184
185
186
187
        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
188
            literal s = args[1]->get_literal();
Paul's avatar
Paul committed
189
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
190
        }
Paul's avatar
Paul committed
191
192
193
        return prog.add_instruction(op, args[0]);
    }

Paul's avatar
Paul committed
194
    instruction_ref
Paul's avatar
Paul committed
195
    parse_flatten(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
196
197
    {
        uint64_t axis = 0;
Paul's avatar
Paul committed
198
199
200
201
        if(contains(attributes, "axis"))
        {
            axis = parse_value(attributes.at("axis")).at<int>();
        }
202
        return prog.add_instruction(op::flatten{axis}, args[0]);
Paul's avatar
Paul committed
203
204
    }

205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
    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
223
224
225
226
227
228
229
    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));
    }
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250

    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
251
252
253
    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
Paul's avatar
Paul committed
254
255
256
257
    {
        literal v = parse_value(attributes.at("value"));
        return prog.add_literal(v);
    }
Paul's avatar
Paul committed
258

Paul's avatar
Paul committed
259
    instruction_ref
Paul's avatar
Paul committed
260
    parse_gemm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
    {
        float alpha = 1.0f;
        float beta  = 0.0f;
        bool transa = false;
        bool transb = false;
        if(contains(attributes, "alpha"))
        {
            alpha = parse_value(attributes.at("alpha")).at<float>();
        }
        if(contains(attributes, "beta"))
        {
            alpha = parse_value(attributes.at("beta")).at<float>();
        }
        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};
283
284
        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
285
286
287
        if(args.size() == 3)
        {
            uint64_t axis = 1;
Shucai Xiao's avatar
Shucai Xiao committed
288
            auto l3       = prog.add_instruction(op::dot{alpha, beta}, l1, l2);
Scott Thornton's avatar
Scott Thornton committed
289
            auto l4       = prog.add_instruction(op::broadcast{axis, l3->get_shape()}, args[2]);
290
            return prog.add_instruction(op::add{}, l3, l4);
Paul's avatar
Paul committed
291
        }
Shucai Xiao's avatar
Shucai Xiao committed
292
        return prog.add_instruction(op::dot{alpha, beta}, l1, l2);
Paul's avatar
Paul committed
293
294
    }

295
    instruction_ref
Paul's avatar
Paul committed
296
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
297
    {
Scott Thornton's avatar
Scott Thornton committed
298
299
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
300
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Scott Thornton's avatar
Scott Thornton committed
301
        bool is_test                                      = false;
302
303
304
305
306
307
        if(contains(attributes, "epsilon"))
        {
            epsilon = parse_value(attributes.at("epsilon")).at<float>();
        }
        if(contains(attributes, "momentum"))
        {
308
            momentum = parse_value(attributes.at("momentum")).at<float>();
309
310
311
        }
        if(contains(attributes, "is_test"))
        {
wsttiger's avatar
wsttiger committed
312
            is_test = parse_value(attributes.at("is_test")).at<uint64_t>() > 0;
313
314
315
        }
        if(contains(attributes, "spatial"))
        {
316
            bn_mode = (parse_value(attributes.at("spatial")).at<uint64_t>() > 0)
317
318
                          ? op::batch_norm_inference::spatial
                          : op::batch_norm_inference::per_activation;
319
        }
Paul's avatar
Paul committed
320
        (void)is_test;
Paul's avatar
Paul committed
321
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
Paul's avatar
Paul committed
322
        return prog.add_instruction(op, std::move(args));
323
324
    }

325
326
327
328
    instruction_ref parse_leaky_relu(const std::string&,
                                     attribute_map attributes,
                                     std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
329
        float alpha = 0.01; // default alpha val for leaky relu
330
331
332
333
334
335
336
337
        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
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
    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
355

Khalique's avatar
Khalique committed
356
357
358
359
360
361
362
363
364
        auto scale_val = prog.add_literal(scale);
        auto bias_vals = prog.add_literal(
            migraph::literal{migraph::shape{migraph::shape::float_type, {bias.size()}}, bias});

        auto scale_tensor = prog.add_instruction(migraph::op::scalar{input_shape}, scale_val);
        auto img_scaled   = prog.add_instruction(migraph::op::mul{}, args.front(), scale_tensor);
        auto bias_bcast   = prog.add_instruction(migraph::op::broadcast{1, input_shape}, bias_vals);
        return prog.add_instruction(migraph::op::add{}, img_scaled, bias_bcast);
    }
Khalique's avatar
Khalique committed
365

Paul's avatar
Paul committed
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
    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);
385
386
387
388
389
        std::unordered_map<std::string, onnx::TensorProto> initializer_data;
        for(auto&& f : graph.initializer())
        {
            initializer_data[f.name()] = f;
        }
Paul's avatar
Paul committed
390
391
392
        for(auto&& input : graph.input())
        {
            const std::string& name = input.name();
393
394
395
396
397
398
399
400
401
402
403
404
            // 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
405
406
407
        }
        for(auto&& p : nodes)
        {
408
            this->parse_node(get_name(p.second));
Paul's avatar
Paul committed
409
410
411
        }
    }

Paul's avatar
Paul committed
412
    void parse_node(const std::string& name)
Paul's avatar
Paul committed
413
    {
Paul's avatar
Paul committed
414
        if(name.empty())
Paul's avatar
Paul committed
415
            MIGRAPH_THROW("Onnx node must have a name");
Paul's avatar
Paul committed
416
417
418
419
420
421
422
423
        if(instructions.count(name) == 0)
        {
            auto&& node = nodes.at(name);
            std::vector<instruction_ref> args;
            for(auto&& input : node.input())
            {
                if(nodes.count(input) > 0)
                {
424
                    auto&& iname = get_name(nodes.at(input));
Paul's avatar
Paul committed
425
                    assert(name != iname);
Paul's avatar
Paul committed
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
                    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;
    }

455
456
457
458
459
    static std::string get_name(const onnx::NodeProto& node)
    {
        if(node.name().empty())
        {
            std::string generated = "migraph_unnamed_node";
Paul's avatar
Paul committed
460
461
462
463
            return std::accumulate(node.output().begin(),
                                   node.output().end(),
                                   generated,
                                   [](auto x, auto y) { return x + "_" + y; });
464
465
466
467
        }
        return node.name();
    }

Paul's avatar
Paul committed
468
469
470
471
472
    static node_map get_nodes(const onnx::GraphProto& graph)
    {
        std::unordered_map<std::string, onnx::NodeProto> result;
        for(auto&& node : graph.node())
        {
473
            result[get_name(node)] = node;
Paul's avatar
Paul committed
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
            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
499
        case onnx::AttributeProto::FLOATS: return from_repeated(shape::float_type, attr.floats());
Paul's avatar
Paul committed
500
501
502
503
504
        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
505
        MIGRAPH_THROW("Invalid attribute type");
Paul's avatar
Paul committed
506
507
508
509
510
    }

    static literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
511
512
        if(t.has_raw_data())
        {
wsttiger's avatar
wsttiger committed
513
            const std::string& s = t.raw_data();
Scott Thornton's avatar
Scott Thornton committed
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
            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()};
            case onnx::TensorProto::FLOAT16: throw std::runtime_error("");
            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("");
            }
            MIGRAPH_THROW("Invalid tensor type");
534
        }
Paul's avatar
Paul committed
535
536
537
538
        switch(t.data_type())
        {
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
Paul's avatar
Paul committed
539
            return literal{{shape::float_type, dims}, t.float_data().begin(), t.float_data().end()};
Paul's avatar
Paul committed
540
541
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
Paul's avatar
Paul committed
542
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
543
        case onnx::TensorProto::UINT16:
Paul's avatar
Paul committed
544
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
545
        case onnx::TensorProto::INT16:
Paul's avatar
Paul committed
546
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
547
        case onnx::TensorProto::INT32:
Paul's avatar
Paul committed
548
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
549
        case onnx::TensorProto::INT64:
Paul's avatar
Paul committed
550
            return literal{{shape::int64_type, dims}, t.int64_data().begin(), t.int64_data().end()};
Paul's avatar
Paul committed
551
552
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
Paul's avatar
Paul committed
553
            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
Paul's avatar
Paul committed
554
555
556
557
558
559
560
561
562
        case onnx::TensorProto::FLOAT16: throw std::runtime_error("");
        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
563
        MIGRAPH_THROW("Invalid tensor type");
Paul's avatar
Paul committed
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
593
594
595
    }

    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");
        case onnx::TensorProto::FLOAT16:
            break; // throw std::runtime_error("Unsupported type FLOAT16");
        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
596
        auto&& tensor_dims = t.tensor_type().shape().dim();
597
598
599
600
601
602
603
604
605
        std::transform(
            tensor_dims.begin(), tensor_dims.end(), std::back_inserter(dims), [](auto&& d) {
                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
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
        return {shape_type, dims};
    }
};

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
631
} // namespace migraph