tf.cpp 39 KB
Newer Older
Khalique's avatar
Khalique committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
#include <google/protobuf/text_format.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <graph.pb.h>
#include <iostream>
#include <fstream>
#include <unordered_map>
#include <unordered_set>
#include <functional>
#include <array>
#include <utility>
#include <vector>

#include <migraphx/fallthrough.hpp>
#include <migraphx/program.hpp>
#include <migraphx/operators.hpp>
#include <migraphx/ranges.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/config.hpp>
#include <migraphx/tf.hpp>

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {

struct tf_parser
{
    using attribute_map = std::unordered_map<std::string, tensorflow::AttrValue>;
Khalique's avatar
Khalique committed
27
    using node_map      = std::unordered_map<std::string, tensorflow::NodeDef>;
Khalique's avatar
Khalique committed
28
29
    // using input_node_map = std::unordered_map<std::string, std::unordered_set<std::string>>;
    using op_func = std::function<instruction_ref(attribute_map, std::vector<instruction_ref>)>;
Khalique's avatar
Khalique committed
30

Khalique's avatar
Khalique committed
31
32
33
34
35
36
37
38
    node_map nodes;
    std::vector<tensorflow::NodeDef> input_nodes;
    std::unordered_map<std::string, instruction_ref> instructions;
    program prog = program();
    bool is_nhwc = true;

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

Khalique's avatar
Khalique committed
39
    std::vector<size_t> parse_axes(const attribute_map& attributes, const std::string& s) const
40
    {
41
42
43
        auto attrs = attributes.at(s).list().i();
        std::vector<size_t> axes;
        copy(attrs.begin(), attrs.end(), std::back_inserter(axes));
Khalique's avatar
Khalique committed
44
        if(is_nhwc)
45
        {
Khalique's avatar
Khalique committed
46
47
48
            std::transform(axes.begin(), axes.end(), axes.begin(), [&](size_t axis) {
                return parse_axis(axis);
            });
49
50
51
52
        }
        return axes;
    }

Khalique's avatar
Khalique committed
53
54
55
56
57
58
    template <class T>
    std::vector<T> parse_axes(std::vector<T> axes) const
    {
        std::vector<T> new_axes;
        if(is_nhwc)
        {
Khalique's avatar
Khalique committed
59
60
61
62
            std::transform(axes.begin(),
                           axes.end(),
                           std::back_inserter(new_axes),
                           [&](size_t axis) { return parse_axis(axis); });
Khalique's avatar
Khalique committed
63
64
65
66
        }
        return new_axes;
    }

Khalique's avatar
Khalique committed
67
68
69
    // tf stores certain attributes such as strides, dilations, as a 4D input.
    // The first and last dims are equal to 1, and the relevant data is in dims 2 and 3.
    // This helper function reorders the data to store for the respective operator member variables.
70
    template <class T>
71
    void reorder_data(std::vector<T>& prev_data) const
72
73
    {
        std::vector<T> new_data(prev_data.size());
74
        for(size_t i = 0; i < new_data.size(); i++)
75
        {
Khalique's avatar
Khalique committed
76
            auto new_idx         = parse_axis(i);
77
            new_data.at(new_idx) = prev_data.at(i);
78
        }
79
80
81
82
        prev_data = new_data;
    }

    template <class T>
Khalique's avatar
Khalique committed
83
    T parse_axis(const T& dim) const
84
    {
Khalique's avatar
Khalique committed
85
        T new_dim = dim;
86
87
88
89
        if(is_nhwc)
        {
            switch(dim)
            {
Khalique's avatar
Khalique committed
90
91
92
93
94
            case 0: new_dim = 0; break;
            case 1: new_dim = 2; break;
            case 2: new_dim = 3; break;
            case 3: new_dim = 1; break;
            default: break;
95
96
            }
        }
Khalique's avatar
Khalique committed
97
        return new_dim;
98
99
    }

100
101
102
103
104
105
106
    std::vector<int64_t> get_axes(size_t num_axes) const
    {
        std::vector<int64_t> axes(num_axes);
        std::iota(axes.begin(), axes.end(), 0);
        return axes;
    }

Khalique's avatar
Khalique committed
107
108
109
110
111
    tf_parser()
    {
        add_generic_op("Identity", op::identity{});
        add_generic_op("Relu", op::relu{});

112
        add_binary_op("Add", op::add{});
Khalique's avatar
Khalique committed
113
        add_binary_op("Mul", op::mul{});
Khalique's avatar
Khalique committed
114

115
        add_mem_op("AvgPool", &tf_parser::parse_pooling);
116
117
        add_mem_op("BiasAdd", &tf_parser::parse_biasadd);
        add_mem_op("ConcatV2", &tf_parser::parse_concat);
Khalique's avatar
Khalique committed
118
119
120
        add_mem_op("Const", &tf_parser::parse_constant);
        add_mem_op("Conv2D", &tf_parser::parse_conv);
        add_mem_op("FusedBatchNorm", &tf_parser::parse_batchnorm);
121
        add_mem_op("MatMul", &tf_parser::parse_matmul);
122
        add_mem_op("MaxPool", &tf_parser::parse_pooling);
Khalique's avatar
Khalique committed
123
        add_mem_op("Mean", &tf_parser::parse_mean);
Khalique's avatar
Khalique committed
124
        add_mem_op("Pack", &tf_parser::parse_pack);
Khalique's avatar
Khalique committed
125
        add_mem_op("Pad", &tf_parser::parse_pad);
Khalique's avatar
Khalique committed
126
        add_mem_op("Relu6", &tf_parser::parse_relu6);
127
128
129
        add_mem_op("Reshape", &tf_parser::parse_reshape);
        add_mem_op("Softmax", &tf_parser::parse_softmax);
        add_mem_op("Squeeze", &tf_parser::parse_squeeze);
130
        add_mem_op("StridedSlice", &tf_parser::parse_stridedslice);
Khalique's avatar
Khalique committed
131
132
    }

133
134
135
136
137
138
139
140
141
142
143
144
145
    template <class F>
    void add_op(std::string name, F f)
    {
        ops.emplace(name, f);
    }

    // Multi output op
    template <class F>
    void add_multi_op(std::string name, F f)
    {
        ops.emplace(name, f);
    }

Khalique's avatar
Khalique committed
146
147
148
    template <class F>
    void add_mem_op(std::string name, F f)
    {
149
        add_op(name, [=](auto&&... xs) {
Khalique's avatar
Khalique committed
150
151
152
153
154
155
156
            return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
        });
    }

    template <class T>
    void add_binary_op(std::string name, T x)
    {
Paul's avatar
Paul committed
157
        add_op(name, [this, x](const attribute_map& attributes, std::vector<instruction_ref> args) {
Khalique's avatar
Khalique committed
158
159
            if(args.size() != 2)
                MIGRAPHX_THROW("binary operators should have 2 operands");
160
161
162
163
164
            auto l0 = args[1];
            if(contains(attributes, "data_format"))
            {
                if(is_nhwc)
                {
Khalique's avatar
Khalique committed
165
                    l0 = prog.add_instruction(op::transpose{{0, 3, 1, 2}}, args[1]);
166
167
168
                }
            }
            return add_broadcastable_binary_op(args[0], l0, x);
Khalique's avatar
Khalique committed
169
170
171
172
173
174
        });
    }

    template <class T>
    instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
    {
Khalique's avatar
Khalique committed
175
        if(arg0->get_shape().lens() != arg1->get_shape().lens())
Khalique's avatar
Khalique committed
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        {
            // 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)
            //
            // Get lengths for both arguments
191
192
            const std::vector<size_t>* s0 = &arg0->get_shape().lens();
            const std::vector<size_t>* s1 = &arg1->get_shape().lens();
Khalique's avatar
Khalique committed
193
194
195
196
197

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

198
            std::vector<size_t> output_lens(*s1);
Khalique's avatar
Khalique committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
            auto offset = s1->size() - s0->size();
            std::transform(s0->begin(),
                           s0->end(),
                           s1->begin() + offset,
                           output_lens.begin() + offset,
                           [](auto a, auto b) { return std::max(a, b); });

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

    template <class T>
    void add_generic_op(std::string name, T x)
    {
Paul's avatar
Paul committed
219
        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
Khalique's avatar
Khalique committed
220
221
222
223
224
225
226
            return prog.add_instruction(x, args);
        });
    }

    instruction_ref
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
227
228
229
        float epsilon                                     = 1e-5f;
        float momentum                                    = 0.9f;
        op::batch_norm_inference::bn_infer_mode_t bn_mode = op::batch_norm_inference::spatial;
Khalique's avatar
Khalique committed
230
231
232
233
234
235
236
237
        if(contains(attributes, "epsilon"))
        {
            epsilon = attributes.at("epsilon").f();
        }
        op::batch_norm_inference op{epsilon, momentum, bn_mode};
        return prog.add_instruction(op, std::move(args));
    }

238
    instruction_ref
Khalique's avatar
Khalique committed
239
    parse_biasadd(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
240
    {
241
        uint64_t axis = 1; // assume output of previous layer is in NCHW (broadcast on channel)
Shucai Xiao's avatar
Shucai Xiao committed
242
        auto l0 = prog.add_instruction(op::broadcast{axis, args[0]->get_shape().lens()}, args[1]);
243
        return prog.add_instruction(op::add{}, args[0], l0);
244
245
    }

Khalique's avatar
Khalique committed
246
247
248
249
    instruction_ref
    parse_concat(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        // get index for axis within args
250
        size_t axis_idx = attributes.at("N").i();
Khalique's avatar
Khalique committed
251
        size_t axis     = parse_axis(args[axis_idx]->eval().at<int64_t>());
Khalique's avatar
Khalique committed
252
        op::concat op{axis};
253
        // return only first N arguments (assuming last index is the axis value)
Khalique's avatar
Khalique committed
254
        return prog.add_instruction(
255
            op, std::vector<instruction_ref>(args.begin(), args.begin() + args.size() - 1));
Khalique's avatar
Khalique committed
256
257
258
259
260
261
    }

    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
    {
Khalique's avatar
Khalique committed
262
263
        literal v       = parse_tensor(attributes.at("value").tensor());
        auto l0         = prog.add_literal(v);
264
265
266
267
268
269
        size_t num_axes = l0->get_shape().lens().size();
        if(num_axes >= 4)
        {
            std::vector<int64_t> transpose_axes = get_axes(num_axes);
            reorder_data(transpose_axes);
            l0 = prog.add_instruction(op::transpose{transpose_axes}, l0);
Khalique's avatar
Khalique committed
270
        }
271
        return l0;
Khalique's avatar
Khalique committed
272
273
274
275
276
277
278
279
280
281
282
    }

    instruction_ref
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::convolution op;
        if(contains(attributes, "padding"))
        {
            const std::string& pad_mode = attributes.at("padding").s();
            if(pad_mode.find("SAME") != std::string::npos)
            {
283
                op.padding_mode = op::padding_mode_t::same;
Khalique's avatar
Khalique committed
284
            }
Khalique's avatar
Khalique committed
285
            else if(pad_mode.find("EXPLICIT") != std::string::npos)
Khalique's avatar
Khalique committed
286
            {
287
                std::vector<size_t> padding;
288
                copy(attributes.at("explicit_paddings").list().i(), std::back_inserter(padding));
Khalique's avatar
Khalique committed
289
290
291
292
293
294
295
296
297
298
299
300
301
302
                if(padding.size() != 4)
                {
                    MIGRAPHX_THROW("padding should have 4 values");
                }
                if(padding[0] != padding[2] || padding[1] != padding[3])
                {
                    MIGRAPHX_THROW("migraphx does not support asymetric padding");
                }
                op.padding[0] = padding[0];
                op.padding[1] = padding[1];
            }
        }
        if(contains(attributes, "strides"))
        {
303
            std::vector<size_t> stride;
304
            copy(attributes.at("strides").list().i(), std::back_inserter(stride));
305
            reorder_data(stride);
306
307
            if(stride.size() != 4)
            {
308
                MIGRAPHX_THROW("strides should have 4 values");
309
            }
310
311
            op.stride[0] = stride[2];
            op.stride[1] = stride[3];
Khalique's avatar
Khalique committed
312
313
314
        }
        if(contains(attributes, "dilations"))
        {
315
            std::vector<size_t> dilation;
316
            copy(attributes.at("dilations").list().i(), std::back_inserter(dilation));
317
            reorder_data(dilation);
318
319
320
321
            if(dilation.size() != 4)
            {
                MIGRAPHX_THROW("dilation should have 4 values");
            }
322
323
            op.dilation[0] = dilation[2];
            op.dilation[1] = dilation[3];
Khalique's avatar
Khalique committed
324
        }
Khalique's avatar
Khalique committed
325
        auto weights = args[1];
326
        // check if weights are from a constant
Khalique's avatar
Khalique committed
327
328

        if(weights->name() != "@param")
329
        {
Khalique's avatar
Khalique committed
330
331
332
333
334
335
336
337
            if(is_nhwc)
            {
                weights = prog.add_instruction(op::transpose{{1, 3, 0, 2}}, args[1]);
            }
            else
            {
                weights = prog.add_instruction(op::transpose{{3, 2, 0, 1}}, args[1]);
            }
338
        }
Khalique's avatar
Khalique committed
339

Khalique's avatar
Khalique committed
340
        return prog.add_instruction(op, {args[0], weights});
Khalique's avatar
Khalique committed
341
342
    }

Khalique's avatar
Khalique committed
343
344
    instruction_ref
    parse_matmul(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
345
346
347
    {
        bool transa = false;
        bool transb = false;
Khalique's avatar
Khalique committed
348

349
350
351
352
353
354
355
356
357
358
359
360
        if(contains(attributes, "transpose_a"))
        {
            transa = attributes.at("transpose_a").b();
        }
        if(contains(attributes, "transpose_b"))
        {
            transb = attributes.at("transpose_a").b();
        }

        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
Khalique's avatar
Khalique committed
361
        std::iter_swap(perm.end() - 1, perm.end() - 2);
362
363
364
365
366
367
368

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

        return prog.add_instruction(op::dot{}, l1, l2);
    }

Khalique's avatar
Khalique committed
369
370
    instruction_ref
    parse_mean(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
371
    {
Khalique's avatar
Khalique committed
372
373

        auto axes      = parse_axes(args[1]->eval().get<int32_t>().to_vector());
Khalique's avatar
Khalique committed
374
        bool keep_dims = attributes.at("keep_dims").b();
Khalique's avatar
Khalique committed
375
        std::vector<int32_t> hw_axes{2, 3};
Khalique's avatar
Khalique committed
376
377
378
379
380
381
382
383
384
385
386
        if(axes == hw_axes and keep_dims)
        {
            op::pooling op{"average"};
            std::vector<size_t> input_dims{args[0]->get_shape().lens()};
            op.lengths[0] = input_dims[2];
            op.lengths[1] = input_dims[3];
            return prog.add_instruction(op, args.front());
        }
        MIGRAPHX_THROW("MIGraphX does not support mean outside of GlobalAvgPool transformation");
    }

Khalique's avatar
Khalique committed
387
388
389
    instruction_ref parse_pack(const std::string&,
                               const attribute_map& attributes,
                               std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
390
391
392
393
394
395
    {
        // reinterpret as unsqueeze with concat
        std::vector<instruction_ref> unsqueezed_args;
        int64_t axis = 0;
        if(contains(attributes, "axis"))
            axis = attributes.at("axis").i();
396
397
398
        size_t input_size = args.front()->get_shape().lens().size();
        if(axis > input_size)
        {
Khalique's avatar
Khalique committed
399
400
            MIGRAPHX_THROW("TF_PARSER: axis value of " + to_string(axis) +
                           " must be smaller than input size " + to_string(input_size));
401
402
403
404
405
        }
        // check if input arg needs axis to be converted to NCHW
        if(input_size >= 4)
            axis = parse_axis(axis);

Khalique's avatar
Khalique committed
406
407
408
409
410
        std::transform(
            args.begin(),
            args.end(),
            std::back_inserter(unsqueezed_args),
            [&](instruction_ref arg) { return prog.add_instruction(op::unsqueeze{{axis}}, arg); });
Khalique's avatar
Khalique committed
411
412
413
        return prog.add_instruction(op::concat{static_cast<size_t>(axis)}, unsqueezed_args);
    }

Khalique's avatar
Khalique committed
414
415
416
417
418
    instruction_ref
    parse_pad(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
    {
        size_t ndims = args.front()->get_shape().lens().size();

Khalique's avatar
Khalique committed
419
420
        // in tf, the paddings are arranged as a 2d shape (ndims, 2),
        // the last dim contains the left padding and right padding respectively
Khalique's avatar
Khalique committed
421
422
        std::vector<std::pair<int32_t, int32_t>> pad_per_dim(ndims);
        auto tf_padding = args[1]->eval().get<int32_t>().to_vector();
Khalique's avatar
Khalique committed
423
        for(size_t i = 0; i < 2 * ndims; i += 2)
Khalique's avatar
Khalique committed
424
        {
Khalique's avatar
Khalique committed
425
426
            pad_per_dim[i / 2].first  = tf_padding[i];
            pad_per_dim[i / 2].second = tf_padding[i + 1];
Khalique's avatar
Khalique committed
427
428
429
430
        }
        reorder_data(pad_per_dim);

        op::pad op;
Khalique's avatar
Khalique committed
431
432
        std::vector<int64_t> pads(ndims * 2);
        for(size_t i = 0; i < ndims; i++)
Khalique's avatar
Khalique committed
433
        {
Khalique's avatar
Khalique committed
434
435
            pads[i]         = pad_per_dim[i].first;
            pads[i + ndims] = pad_per_dim[i].second;
Khalique's avatar
Khalique committed
436
437
438
439
440
        }
        op.pads = pads;
        return prog.add_instruction(op, args.front());
    }

441
442
443
444
445
    instruction_ref parse_pooling(const std::string& name,
                                  attribute_map attributes,
                                  std::vector<instruction_ref> args)
    {
        op::pooling op{starts_with(name, "Max") ? "max" : "average"};
Khalique's avatar
Khalique committed
446

447
448
449
450
451
452
453
454
455
456
457
458
459
460
        if(contains(attributes, "padding"))
        {
            const std::string& pad_mode = attributes.at("padding").s();
            if(pad_mode.find("SAME") != std::string::npos)
            {
                op.padding_mode = op::padding_mode_t::same;
            }
            else if(pad_mode.find("VALID") != std::string::npos)
            {
                op.padding_mode = op::padding_mode_t::valid;
            }
        }
        if(contains(attributes, "strides"))
        {
461
            std::vector<size_t> stride;
462
            copy(attributes.at("strides").list().i(), std::back_inserter(stride));
463
            reorder_data(stride);
464
465
466
467
            if(stride.size() != 4)
            {
                MIGRAPHX_THROW("strides should have 4 values");
            }
468
469
            op.stride[0] = stride[2];
            op.stride[1] = stride[3];
470
471
472
        }
        if(contains(attributes, "ksize"))
        {
473
            std::vector<size_t> ksize;
474
            copy(attributes.at("ksize").list().i(), std::back_inserter(ksize));
475
            reorder_data(ksize);
476
477
478
            if(ksize.size() != 4)
            {
                MIGRAPHX_THROW("ksize should have 4 values");
Khalique's avatar
Khalique committed
479
            }
480
481
            op.lengths[0] = ksize[2];
            op.lengths[1] = ksize[3];
482
        }
483
        return prog.add_instruction(op, args[0]);
484
    }
Khalique's avatar
Khalique committed
485

Khalique's avatar
Khalique committed
486
487
488
489
490
491
492
493
494
    instruction_ref
    parse_relu6(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
    {
        op::clip op;
        op.max_val = 6.0;
        op.min_val = 0.0;
        return prog.add_instruction(op, std::move(args));
    }

495
    instruction_ref
Khalique's avatar
Khalique committed
496
    parse_reshape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
497
498
499
500
    {
        op::reshape op;
        if(args.size() != 2)
            MIGRAPHX_THROW("reshape needs 2 arguments (input, new_shape)");
Khalique's avatar
Khalique committed
501
        auto s = args[1]->eval();
502
503
504
505
        s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
        return prog.add_instruction(op, args[0]);
    }

Khalique's avatar
Khalique committed
506
507
508
509
510
511
512
513
514
    void parse_from(std::istream& is)
    {
        tensorflow::GraphDef graph;
        if(graph.ParseFromIstream(&is))
        {
            this->parse_graph(graph);
        }
        else
        {
515
            throw std::runtime_error("Failed reading tf file");
Khalique's avatar
Khalique committed
516
517
518
        }
    }

519
520
521
522
523
524
525
526
527
528
    instruction_ref
    parse_softmax(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
    {
        auto dims = args.front()->get_shape().lens();
        auto r =
            prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1]), 1, 1}}, args.front());
        auto s = prog.add_instruction(op::softmax{}, r);
        return prog.add_instruction(op::reshape{{long(dims[0]), long(dims[1])}}, s);
    }

Khalique's avatar
Khalique committed
529
530
531
    instruction_ref parse_squeeze(const std::string&,
                                  const attribute_map& attributes,
                                  std::vector<instruction_ref> args)
532
533
    {
        op::squeeze op;
534
        auto axes = parse_axes(attributes, "squeeze_dims");
535
        copy(axes, std::back_inserter(op.axes));
536
        auto args0_dims = args[0]->get_shape().lens();
537
538
        if(op.axes.empty()) // no squeeze_dims provided, remove any dim that equals 1
        {
539
            for(size_t i = 0; i < args0_dims.size(); i++)
540
            {
541
                if(args0_dims.at(i) == 1)
542
543
544
545
                {
                    op.axes.push_back(i);
                }
            }
546
        }
547
        return prog.add_instruction(op, args[0]);
548
549
    }

Khalique's avatar
Khalique committed
550
551
552
    instruction_ref parse_stridedslice(const std::string&,
                                       const attribute_map& attributes,
                                       std::vector<instruction_ref> args)
553
554
    {
        op::slice op;
Khalique's avatar
Khalique committed
555
556
557
558
559
560
561
562
        auto starts     = args[1]->eval().get<int32_t>().to_vector();
        auto ends       = args[2]->eval().get<int32_t>().to_vector();
        size_t num_axes = args[0]->get_shape().lens().size();
        if(num_axes >= 4)
        {
            reorder_data(starts);
            reorder_data(ends);
        }
563

Khalique's avatar
Khalique committed
564
565
566
567
        op.starts = std::vector<int64_t>(starts.begin(), starts.end());
        op.ends   = std::vector<int64_t>(ends.begin(), ends.end());
        op.axes   = std::vector<int64_t>(num_axes);
        std::iota(op.axes.begin(), op.axes.end(), 0);
568
        uint32_t shrink_axis_mask = 0;
Khalique's avatar
Khalique committed
569
        uint32_t bitwise_compare  = 1;
570
571
572
        std::vector<int64_t> squeeze_axes;

        if(contains(attributes, "shrink_axis_mask"))
573
            shrink_axis_mask = static_cast<uint32_t>(attributes.at("shrink_axis_mask").i());
574

Khalique's avatar
Khalique committed
575
        for(size_t i = 0; i < num_axes; i++)
576
        {
577
            // the LSB corresponds to axis 0 when determining which axes to squeeze
Khalique's avatar
Khalique committed
578
            if(((shrink_axis_mask >> i) & bitwise_compare) == 1)
579
580
                squeeze_axes.push_back(i);
        }
Khalique's avatar
Khalique committed
581
582
583
584
585
        if(num_axes >= 4)
        {
            squeeze_axes = parse_axes(squeeze_axes);
        }

586
587
588
589
        auto l0 = prog.add_instruction(op, args[0]);
        return prog.add_instruction(op::squeeze{squeeze_axes}, l0);
    }

Khalique's avatar
Khalique committed
590
591
592
593
594
    void parse_graph(const tensorflow::GraphDef& graph)
    {
        nodes = get_nodes(graph, input_nodes);
        for(auto&& input : input_nodes)
        {
Khalique's avatar
Khalique committed
595
            const std::string& name   = input.name();
Khalique's avatar
Khalique committed
596
            attribute_map input_attrs = get_attributes(input);
Khalique's avatar
Khalique committed
597
598
            shape::type_t shape_type  = parse_type(input_attrs.at("dtype").type());
            std::vector<size_t> dims  = parse_dims(input_attrs.at("shape").shape());
599
            if(is_nhwc and dims.size() >= 4)
600
            {
601
                reorder_data(dims);
602
            }
Khalique's avatar
Khalique committed
603
604
            shape s            = shape{shape_type, dims};
            instructions[name] = prog.add_parameter(name, s);
Khalique's avatar
Khalique committed
605
606
607
        }
        for(auto&& p : nodes)
        {
608
            this->parse_node(p.first);
Khalique's avatar
Khalique committed
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
        }
    }

    void parse_node(const std::string& name)
    {
        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)
                {
                    auto&& iname = get_name(nodes.at(input));
                    assert(name != iname);
                    this->parse_node(iname);
                    args.push_back(instructions.at(iname));
                }
                else
                {
                    args.push_back(instructions.at(input));
                }
            }
            if(ops.count(node.op()) == 0)
            {
635
                instructions[name] = prog.add_instruction(op::unknown{node.op()}, args);
Khalique's avatar
Khalique committed
636
637
638
639
640
641
642
643
644
645
646
            }
            else
            {
                instructions[name] = ops[node.op()](get_attributes(node), args);
            }
        }
    }

    static attribute_map get_attributes(const tensorflow::NodeDef& node)
    {
        attribute_map result;
Khalique's avatar
Khalique committed
647
        for(auto&& attr : node.attr())
Khalique's avatar
Khalique committed
648
649
650
651
652
653
        {
            result[attr.first] = attr.second;
        }
        return result;
    }

Khalique's avatar
Khalique committed
654
    static std::string get_name(const tensorflow::NodeDef& node) { return node.name(); }
Khalique's avatar
Khalique committed
655

Khalique's avatar
Khalique committed
656
657
    static node_map get_nodes(const tensorflow::GraphDef& graph,
                              std::vector<tensorflow::NodeDef>& input_nodes)
Khalique's avatar
Khalique committed
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
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
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
    {
        node_map result;
        for(auto&& node : graph.node())
        {
            auto node_name = get_name(node);
            // assume each node in graph has an associated name
            if(node_name.empty())
                MIGRAPHX_THROW("tf node with no name found");
            result[node_name] = node;
            if(node.op() == "Placeholder")
            {
                input_nodes.push_back(node);
            }
        }
        return result;
    }

    static shape::type_t parse_type(const tensorflow::DataType t)
    {
        shape::type_t shape_type{};
        switch(t)
        {
        case tensorflow::DataType::DT_INVALID:
            break; // throw std::runtime_error("Unsupported type UNDEFINED");
        case tensorflow::DataType::DT_FLOAT: shape_type = shape::float_type; break;
        case tensorflow::DataType::DT_DOUBLE: shape_type = shape::double_type; break;
        case tensorflow::DataType::DT_INT32: shape_type = shape::int32_type; break;
        case tensorflow::DataType::DT_UINT8:
            break; // throw std::runtime_error("Unsupported type UINT8");
        case tensorflow::DataType::DT_INT16: shape_type = shape::int16_type; break;
        case tensorflow::DataType::DT_INT8: shape_type = shape::int8_type; break;
        case tensorflow::DataType::DT_STRING:
            break; // throw std::runtime_error("Unsupported type STRING");
        case tensorflow::DataType::DT_COMPLEX64:
            break; // throw std::runtime_error("Unsupported type COMPLEX64");
        case tensorflow::DataType::DT_INT64: shape_type = shape::int64_type; break;
        case tensorflow::DataType::DT_BOOL:
            break; // throw std::runtime_error("Unsupported type BOOL");
        case tensorflow::DataType::DT_QINT8:
            break; // throw std::runtime_error("Unsupported type QINT8");
        case tensorflow::DataType::DT_QUINT8:
            break; // throw std::runtime_error("Unsupported type QUINT8");
        case tensorflow::DataType::DT_QINT32:
            break; // throw std::runtime_error("Unsupported type QINT32");
        case tensorflow::DataType::DT_BFLOAT16:
            break; // throw std::runtime_error("Unsupported type BFLOAT16");
        case tensorflow::DataType::DT_QINT16:
            break; // throw std::runtime_error("Unsupported type QINT16");
        case tensorflow::DataType::DT_QUINT16:
            break; // throw std::runtime_error("Unsupported type QUINT16");
        case tensorflow::DataType::DT_UINT16: shape_type = shape::uint16_type; break;
        case tensorflow::DataType::DT_COMPLEX128:
            break; // throw std::runtime_error("Unsupported type COMPLEX128");
        case tensorflow::DataType::DT_HALF: shape_type = shape::half_type; break;
        case tensorflow::DataType::DT_RESOURCE:
            break; // throw std::runtime_error("Unsupported type RESOURCE");
        case tensorflow::DataType::DT_VARIANT:
            break; // throw std::runtime_error("Unsupported type VARIANT");
        case tensorflow::DataType::DT_UINT32: shape_type = shape::uint32_type; break;
Khalique's avatar
Khalique committed
717
718
719
        case tensorflow::DataType::DT_UINT64:
            shape_type = shape::uint64_type;
            break;
Khalique's avatar
Khalique committed
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

        // tf pb should not use these types
        case tensorflow::DataType::DT_FLOAT_REF: break;
        case tensorflow::DataType::DT_DOUBLE_REF: break;
        case tensorflow::DataType::DT_INT32_REF: break;
        case tensorflow::DataType::DT_UINT8_REF: break;
        case tensorflow::DataType::DT_INT16_REF: break;
        case tensorflow::DataType::DT_INT8_REF: break;
        case tensorflow::DataType::DT_STRING_REF: break;
        case tensorflow::DataType::DT_COMPLEX64_REF: break;
        case tensorflow::DataType::DT_INT64_REF: break;
        case tensorflow::DataType::DT_BOOL_REF: break;
        case tensorflow::DataType::DT_QINT8_REF: break;
        case tensorflow::DataType::DT_QUINT8_REF: break;
        case tensorflow::DataType::DT_QINT32_REF: break;
        case tensorflow::DataType::DT_BFLOAT16_REF: break;
        case tensorflow::DataType::DT_QINT16_REF: break;
        case tensorflow::DataType::DT_QUINT16_REF: break;
        case tensorflow::DataType::DT_UINT16_REF: break;
        case tensorflow::DataType::DT_COMPLEX128_REF: break;
        case tensorflow::DataType::DT_HALF_REF: break;
        case tensorflow::DataType::DT_RESOURCE_REF: break;
        case tensorflow::DataType::DT_VARIANT_REF: break;
        case tensorflow::DataType::DT_UINT32_REF: break;
        case tensorflow::DataType::DT_UINT64_REF: break;
        case tensorflow::DataType::DataType_INT_MAX_SENTINEL_DO_NOT_USE_: break;
        case tensorflow::DataType::DataType_INT_MIN_SENTINEL_DO_NOT_USE_: break;
Khalique's avatar
Khalique committed
747
748
749
750
        }
        return shape_type;
    }

Khalique's avatar
Khalique committed
751
    static literal parse_tensor(const tensorflow::TensorProto& t)
Khalique's avatar
Khalique committed
752
753
    {
        std::vector<size_t> dims = parse_dims(t.tensor_shape());
754
        size_t shape_size = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<size_t>());
Khalique's avatar
Khalique committed
755
756
        if(!t.tensor_content().empty()) // has raw data
        {
Khalique's avatar
Khalique committed
757
            const std::string& s = t.tensor_content();
Khalique's avatar
Khalique committed
758
759
760
            switch(t.dtype())
            {
            case tensorflow::DataType::DT_INVALID: throw std::runtime_error("");
Khalique's avatar
Khalique committed
761
762
            case tensorflow::DataType::DT_FLOAT:
                return literal{{shape::float_type, dims}, s.data()};
Khalique's avatar
Khalique committed
763
            case tensorflow::DataType::DT_UINT8: throw std::runtime_error("");
764
            case tensorflow::DataType::DT_INT8: return literal{{shape::int8_type, dims}, s.data()};
Khalique's avatar
Khalique committed
765
            case tensorflow::DataType::DT_UINT16:
766
                return literal{{shape::uint16_type, dims}, s.data()};
Khalique's avatar
Khalique committed
767
            case tensorflow::DataType::DT_INT16:
768
                return literal{{shape::int16_type, dims}, s.data()};
Khalique's avatar
Khalique committed
769
770
771
772
            case tensorflow::DataType::DT_INT32:
                return literal{{shape::int32_type, dims}, s.data()};
            case tensorflow::DataType::DT_INT64:
                return literal{{shape::int64_type, dims}, s.data()};
Khalique's avatar
Khalique committed
773
            case tensorflow::DataType::DT_STRING: throw std::runtime_error("");
774
            case tensorflow::DataType::DT_BOOL: return literal{{shape::int8_type, dims}, s.data()};
Khalique's avatar
Khalique committed
775
            case tensorflow::DataType::DT_HALF: return literal{{shape::half_type, dims}, s.data()};
Khalique's avatar
Khalique committed
776
777
            case tensorflow::DataType::DT_DOUBLE:
                return literal{{shape::double_type, dims}, s.data()};
Khalique's avatar
Khalique committed
778
779
780
781
            case tensorflow::DataType::DT_UINT32: throw std::runtime_error("");
            case tensorflow::DataType::DT_UINT64: throw std::runtime_error("");
            case tensorflow::DataType::DT_COMPLEX64: throw std::runtime_error("");
            case tensorflow::DataType::DT_COMPLEX128: throw std::runtime_error("");
Khalique's avatar
Khalique committed
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
            case tensorflow::DataType::DT_QINT8: throw std::runtime_error("");
            case tensorflow::DataType::DT_QUINT8: throw std::runtime_error("");
            case tensorflow::DataType::DT_QINT32: throw std::runtime_error("");
            case tensorflow::DataType::DT_BFLOAT16: throw std::runtime_error("");
            case tensorflow::DataType::DT_QINT16: throw std::runtime_error("");
            case tensorflow::DataType::DT_QUINT16: throw std::runtime_error("");
            case tensorflow::DataType::DT_RESOURCE: throw std::runtime_error("");
            case tensorflow::DataType::DT_VARIANT: throw std::runtime_error("");
            case tensorflow::DataType::DT_FLOAT_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_DOUBLE_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_INT32_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_UINT8_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_INT16_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_INT8_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_STRING_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_COMPLEX64_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_INT64_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_BOOL_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_QINT8_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_QUINT8_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_QINT32_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_BFLOAT16_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_QINT16_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_QUINT16_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_UINT16_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_COMPLEX128_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_HALF_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_RESOURCE_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_VARIANT_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_UINT32_REF: throw std::runtime_error("");
            case tensorflow::DataType::DT_UINT64_REF: throw std::runtime_error("");
Khalique's avatar
Khalique committed
813
814
815
816
            case tensorflow::DataType::DataType_INT_MAX_SENTINEL_DO_NOT_USE_:
                throw std::runtime_error("");
            case tensorflow::DataType::DataType_INT_MIN_SENTINEL_DO_NOT_USE_:
                throw std::runtime_error("");
Khalique's avatar
Khalique committed
817
818
819
820
821
822
823
            }
            MIGRAPHX_THROW("Invalid tensor type");
        }
        switch(t.dtype())
        {
        case tensorflow::DataType::DT_INVALID: throw std::runtime_error("");
        case tensorflow::DataType::DT_FLOAT:
Khalique's avatar
Khalique committed
824
825
            return create_literal(
                shape::float_type, dims, get_data_vals(t.float_val(), shape_size));
Khalique's avatar
Khalique committed
826
827
        case tensorflow::DataType::DT_UINT8: throw std::runtime_error("");
        case tensorflow::DataType::DT_INT8:
828
            return create_literal(shape::int8_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
829
        case tensorflow::DataType::DT_UINT16:
830
            return create_literal(shape::uint16_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
831
        case tensorflow::DataType::DT_INT16:
832
            return create_literal(shape::int16_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
833
        case tensorflow::DataType::DT_INT32:
834
            return create_literal(shape::int32_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
835
        case tensorflow::DataType::DT_INT64:
Khalique's avatar
Khalique committed
836
837
            return create_literal(
                shape::int64_type, dims, get_data_vals(t.int64_val(), shape_size));
Khalique's avatar
Khalique committed
838
839
        case tensorflow::DataType::DT_STRING: throw std::runtime_error("");
        case tensorflow::DataType::DT_BOOL:
840
            return create_literal(shape::int32_type, dims, get_data_vals(t.bool_val(), shape_size));
Khalique's avatar
Khalique committed
841
        case tensorflow::DataType::DT_HALF:
Khalique's avatar
Khalique committed
842
        {
843
844
            std::vector<int> data_int32 = get_data_vals(t.half_val(), shape_size);
            std::vector<uint16_t> data_uint16(data_int32.begin(), data_int32.end());
Khalique's avatar
Khalique committed
845
846
847
848
849
            std::vector<half> data_half;
            std::transform(data_uint16.begin(),
                           data_uint16.end(),
                           std::back_inserter(data_half),
                           [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
850
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
851
        }
Khalique's avatar
Khalique committed
852
        case tensorflow::DataType::DT_DOUBLE:
Khalique's avatar
Khalique committed
853
            return literal{{shape::double_type, dims}, get_data_vals(t.double_val(), shape_size)};
Khalique's avatar
Khalique committed
854
855
856
857
        case tensorflow::DataType::DT_UINT32: throw std::runtime_error("");
        case tensorflow::DataType::DT_UINT64: throw std::runtime_error("");
        case tensorflow::DataType::DT_COMPLEX64: throw std::runtime_error("");
        case tensorflow::DataType::DT_COMPLEX128: throw std::runtime_error("");
Khalique's avatar
Khalique committed
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
        case tensorflow::DataType::DT_QINT8: throw std::runtime_error("");
        case tensorflow::DataType::DT_QUINT8: throw std::runtime_error("");
        case tensorflow::DataType::DT_QINT32: throw std::runtime_error("");
        case tensorflow::DataType::DT_BFLOAT16: throw std::runtime_error("");
        case tensorflow::DataType::DT_QINT16: throw std::runtime_error("");
        case tensorflow::DataType::DT_QUINT16: throw std::runtime_error("");
        case tensorflow::DataType::DT_RESOURCE: throw std::runtime_error("");
        case tensorflow::DataType::DT_VARIANT: throw std::runtime_error("");
        case tensorflow::DataType::DT_FLOAT_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_DOUBLE_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_INT32_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_UINT8_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_INT16_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_INT8_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_STRING_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_COMPLEX64_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_INT64_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_BOOL_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_QINT8_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_QUINT8_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_QINT32_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_BFLOAT16_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_QINT16_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_QUINT16_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_UINT16_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_COMPLEX128_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_HALF_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_RESOURCE_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_VARIANT_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_UINT32_REF: throw std::runtime_error("");
        case tensorflow::DataType::DT_UINT64_REF: throw std::runtime_error("");
Khalique's avatar
Khalique committed
889
890
891
892
        case tensorflow::DataType::DataType_INT_MAX_SENTINEL_DO_NOT_USE_:
            throw std::runtime_error("");
        case tensorflow::DataType::DataType_INT_MIN_SENTINEL_DO_NOT_USE_:
            throw std::runtime_error("");
Khalique's avatar
Khalique committed
893
894
895
896
        }
        MIGRAPHX_THROW("Invalid tensor type");
    }

897
    template <class T>
Khalique's avatar
Khalique committed
898
    static std::vector<T> get_data_vals(const google::protobuf::RepeatedField<T>& data,
899
                                        const size_t& shape_size)
900
901
902
903
904
905
906
907
908
909
910
911
    {
        std::vector<T> data_vals(shape_size);
        // check if shape has enough data values given existing fields
        if(data.size() == 1)
        {
            std::fill(data_vals.begin(), data_vals.end(), data[0]);
        }
        else
            copy(data.begin(), data.end(), std::back_inserter(data_vals));
        return data_vals;
    }

Khalique's avatar
Khalique committed
912
913
914
915
    static std::vector<size_t> parse_dims(const tensorflow::TensorShapeProto& s)
    {
        std::vector<size_t> dims;
        auto input_dims = s.dim();
Khalique's avatar
Khalique committed
916
917
918
        std::transform(input_dims.begin(),
                       input_dims.end(),
                       std::back_inserter(dims),
Paul's avatar
Paul committed
919
                       [](const tensorflow::TensorShapeProto_Dim& dim) { return dim.size(); });
Khalique's avatar
Khalique committed
920
921
        return dims;
    }
922
923

    template <class T>
Khalique's avatar
Khalique committed
924
    static literal
925
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, std::vector<T> data)
926
    {
Khalique's avatar
Khalique committed
927
        // assume if explicit value is mentioned in protobuf and dim size <= 1, treat as scalar
928
        if(dims.empty() or (dims.size() == 1 and dims.front() == 1))
929
            return literal{{shape_type}, data};
930
931
        return literal{{shape_type, dims}, data};
    }
Khalique's avatar
Khalique committed
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
};

program parse_tf(const std::string& name, bool is_nhwc)
{
    std::fstream input(name.c_str(), std::ios::in | std::ios::binary);
    tf_parser parser;
    parser.is_nhwc = is_nhwc;

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

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx