tf.cpp 51.4 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
#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>
Khalique's avatar
Khalique committed
20
#include <migraphx/pad_calc.hpp>
Khalique's avatar
Khalique committed
21
22
23
24
25
26
27

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {

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

Khalique's avatar
Khalique committed
32
33
34
35
36
37
38
39
    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;

Paul's avatar
Paul committed
40
    bool should_transpose(instruction_ref ins) const
Paul's avatar
Paul committed
41
42
43
44
45
46
    {
        return is_nhwc and ins->get_shape().lens().size() == 4;
    }

    instruction_ref to_nhwc(instruction_ref ins)
    {
Paul's avatar
Paul committed
47
        if(should_transpose(ins))
Paul's avatar
Paul committed
48
49
50
51
52
53
            return prog.add_instruction(op::transpose{{0, 2, 3, 1}}, ins);
        return ins;
    }

    instruction_ref to_nchw(instruction_ref ins)
    {
Paul's avatar
Paul committed
54
        if(should_transpose(ins))
Paul's avatar
Paul committed
55
56
57
58
59
60
            return prog.add_instruction(op::transpose{{0, 3, 1, 2}}, ins);
        return ins;
    }

    instruction_ref to_kcxy(instruction_ref ins)
    {
Paul's avatar
Paul committed
61
        if(should_transpose(ins))
Paul's avatar
Paul committed
62
63
64
65
66
67
            return prog.add_instruction(op::transpose{{3, 2, 0, 1}}, ins);
        return ins;
    }

    instruction_ref make_contiguous(instruction_ref ins)
    {
Paul's avatar
Paul committed
68
        if(ins->get_shape().standard())
Paul's avatar
Paul committed
69
70
71
72
73
74
75
76
            return ins;
        else
            return prog.add_instruction(op::contiguous{}, ins);
    }

    std::vector<instruction_ref> to_nchw(const std::vector<instruction_ref>& args)
    {
        std::vector<instruction_ref> result(args.size());
Paul's avatar
Paul committed
77
        std::transform(
Paul's avatar
Paul committed
78
            args.begin(), args.end(), result.begin(), [&](auto ins) { return this->to_nchw(ins); });
Paul's avatar
Paul committed
79
80
81
        return result;
    }

kahmed10's avatar
kahmed10 committed
82
83
84
85
86
87
88
89
    std::vector<instruction_ref> to_nhwc(const std::vector<instruction_ref>& args)
    {
        std::vector<instruction_ref> result(args.size());
        std::transform(
            args.begin(), args.end(), result.begin(), [&](auto ins) { return this->to_nhwc(ins); });
        return result;
    }

Khalique's avatar
Khalique committed
90
    std::vector<size_t>
91
    parse_axes(const attribute_map& attributes, const std::string& s, const size_t num_dims) const
92
    {
93
94
95
        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
96
        if(is_nhwc)
97
        {
Khalique's avatar
Khalique committed
98
            std::transform(axes.begin(), axes.end(), axes.begin(), [&](size_t axis) {
Khalique's avatar
Khalique committed
99
                return parse_axis(axis, num_dims);
Khalique's avatar
Khalique committed
100
            });
101
102
103
104
        }
        return axes;
    }

Khalique's avatar
Khalique committed
105
    template <class T>
106
    std::vector<T> parse_axes(std::vector<T> axes, const size_t num_dims) const
Khalique's avatar
Khalique committed
107
108
109
    {
        if(is_nhwc)
        {
110
            std::vector<T> new_axes;
Khalique's avatar
Khalique committed
111
112
113
            std::transform(axes.begin(),
                           axes.end(),
                           std::back_inserter(new_axes),
Khalique's avatar
Khalique committed
114
                           [&](size_t axis) { return parse_axis(axis, num_dims); });
115
            return new_axes;
Khalique's avatar
Khalique committed
116
        }
117
        return axes;
Khalique's avatar
Khalique committed
118
119
    }

Khalique's avatar
Khalique committed
120
121
122
    // 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.
123
    template <class T>
124
    void reorder_data(std::vector<T>& prev_data) const
125
126
    {
        std::vector<T> new_data(prev_data.size());
127
        for(size_t i = 0; i < new_data.size(); i++)
128
        {
Khalique's avatar
Khalique committed
129
            auto new_idx         = parse_axis(i, new_data.size());
130
            new_data.at(new_idx) = prev_data.at(i);
131
        }
132
133
134
135
        prev_data = new_data;
    }

    template <class T>
136
    T parse_axis(const T& dim, const size_t num_dims) const
137
    {
Khalique's avatar
Khalique committed
138
        T new_dim = dim;
Khalique's avatar
Khalique committed
139
        if(is_nhwc and num_dims >= 4)
140
141
142
        {
            switch(dim)
            {
Khalique's avatar
Khalique committed
143
144
145
146
147
            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;
148
149
            }
        }
Khalique's avatar
Khalique committed
150
        return new_dim;
151
152
    }

153
154
155
156
157
158
159
    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
160
    std::vector<int64_t> get_axes_from_mask(const size_t num_axes, const uint32_t mask)
Khalique's avatar
Khalique committed
161
    {
Khalique's avatar
Khalique committed
162
        uint32_t bitwise_compare = 1;
Khalique's avatar
Khalique committed
163
164
165
166
167
168
169
170
171
172
173
174
        std::vector<int64_t> axes;
        for(size_t i = 0; i < num_axes; i++)
        {
            // the LSB corresponds to axis 0 when determining which axes to begin
            if(((mask >> i) & bitwise_compare) == 1)
                axes.push_back(1);
            else
                axes.push_back(0);
        }
        return axes;
    }

Khalique's avatar
Khalique committed
175
176
    tf_parser()
    {
Khalique's avatar
Khalique committed
177
        add_generic_op("All", op::identity{});
Khalique's avatar
Khalique committed
178
        add_generic_op("Identity", op::identity{});
Khalique's avatar
Khalique committed
179
        add_generic_op("LessEqual", op::identity{});
Khalique's avatar
Khalique committed
180
        add_generic_op("Relu", op::relu{});
Khalique's avatar
Khalique committed
181
        add_generic_op("Relu6", op::clip{6.0, 0.0});
Khalique's avatar
Khalique committed
182
        add_generic_op("Rsqrt", op::rsqrt{});
Khalique's avatar
Khalique committed
183
        add_generic_op("Tanh", op::tanh{});
Khalique's avatar
Khalique committed
184
        add_generic_op("StopGradient", op::identity{});
Khalique's avatar
Khalique committed
185

186
        add_binary_op("Add", op::add{});
Khalique's avatar
Khalique committed
187
        add_binary_op("Mul", op::mul{});
Khalique's avatar
Khalique committed
188
        add_binary_op("Pow", op::pow{});
Khalique's avatar
Khalique committed
189
        add_binary_op("SquaredDifference", op::sqdiff{});
Khalique's avatar
Khalique committed
190
        add_binary_op("Sub", op::sub{});
Khalique's avatar
Khalique committed
191

192
        add_mem_op("AvgPool", &tf_parser::parse_pooling);
Khalique's avatar
Khalique committed
193
        add_mem_op("BatchMatMul", &tf_parser::parse_matmul, false);
Khalique's avatar
Khalique committed
194
        add_mem_op("BatchMatMulV2", &tf_parser::parse_matmul, false);
195
        add_mem_op("BiasAdd", &tf_parser::parse_biasadd);
Khalique's avatar
Khalique committed
196
        add_mem_op("Cast", &tf_parser::parse_cast, false);
Paul's avatar
Paul committed
197
        add_mem_op("ConcatV2", &tf_parser::parse_concat, false);
Khalique's avatar
Khalique committed
198
        add_mem_op("Const", &tf_parser::parse_constant);
Paul's avatar
Paul committed
199
        add_mem_op("Conv2D", &tf_parser::parse_conv);
Paul's avatar
Paul committed
200
        add_mem_op("DepthwiseConv2dNative", &tf_parser::parse_depthwiseconv);
201
        add_mem_op("ExpandDims", &tf_parser::parse_expanddims, false);
Khalique's avatar
Khalique committed
202
        add_mem_op("FusedBatchNorm", &tf_parser::parse_batchnorm);
Khalique's avatar
Khalique committed
203
        add_mem_op("GatherV2", &tf_parser::parse_gather, false);
Paul's avatar
Paul committed
204
        add_mem_op("MatMul", &tf_parser::parse_matmul, false);
205
        add_mem_op("MaxPool", &tf_parser::parse_pooling);
Khalique's avatar
Khalique committed
206
        add_mem_op("Mean", &tf_parser::parse_mean, false);
Khalique's avatar
Khalique committed
207
        add_mem_op("OneHot", &tf_parser::parse_onehot, false);
Paul's avatar
Paul committed
208
        add_mem_op("Pack", &tf_parser::parse_pack, false);
Paul's avatar
Paul committed
209
        add_mem_op("Pad", &tf_parser::parse_pad);
Paul's avatar
Paul committed
210
        add_mem_op("Reshape", &tf_parser::parse_reshape, false);
Khalique's avatar
Khalique committed
211
        add_mem_op("Slice", &tf_parser::parse_slice, false);
kahmed10's avatar
kahmed10 committed
212
213
        add_mem_op("Split", &tf_parser::parse_split, false);
        add_mem_op("SplitV", &tf_parser::parse_split, false);
Khalique's avatar
Khalique committed
214
        add_mem_op("Softmax", &tf_parser::parse_softmax<op::softmax>, false);
Paul's avatar
Paul committed
215
        add_mem_op("Squeeze", &tf_parser::parse_squeeze, false);
216
        add_mem_op("StridedSlice", &tf_parser::parse_stridedslice, false);
Khalique's avatar
Khalique committed
217
        add_mem_op("Transpose", &tf_parser::parse_transpose, false);
Khalique's avatar
Khalique committed
218
219
    }

220
    template <class F>
kahmed10's avatar
kahmed10 committed
221
    void add_op(const std::string& name, F f, bool transpose = true)
222
    {
Paul's avatar
Paul committed
223
        if(transpose)
Paul's avatar
Paul committed
224
        {
kahmed10's avatar
kahmed10 committed
225
226
227
228
229
230
            ops.emplace(
                name,
                op_func{
                    [=](const attribute_map& attributes, const std::vector<instruction_ref>& args) {
                        return std::vector<instruction_ref>{to_nhwc(f(attributes, to_nchw(args)))};
                    }});
Paul's avatar
Paul committed
231
232
233
        }
        else
        {
kahmed10's avatar
kahmed10 committed
234
235
236
237
238
            ops.emplace(name,
                        op_func{[=](const attribute_map& attributes,
                                    const std::vector<instruction_ref>& args) {
                            return std::vector<instruction_ref>{f(attributes, args)};
                        }});
Paul's avatar
Paul committed
239
        }
240
241
    }

Khalique's avatar
Khalique committed
242
    template <class F>
Paul's avatar
Paul committed
243
    void add_mem_op(std::string name, F f, bool transpose = true)
Khalique's avatar
Khalique committed
244
    {
Paul's avatar
Paul committed
245
246
247
248
249
        add_op(name,
               [=](auto&&... xs) {
                   return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
               },
               transpose);
Khalique's avatar
Khalique committed
250
251
252
253
254
    }

    template <class T>
    void add_binary_op(std::string name, T x)
    {
Paul's avatar
Paul committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
        add_op(name,
               [this, x](const attribute_map&, std::vector<instruction_ref> args) {
                   if(args.size() != 2)
                       MIGRAPHX_THROW("binary operators should have 2 operands");
                   // TODO
                   // if(contains(attributes, "data_format"))
                   // {
                   //     if(is_nhwc)
                   //     {
                   //         l0 = prog.add_instruction(op::transpose{{0, 3, 1, 2}}, args[1]);
                   //     }
                   // }
                   return add_broadcastable_binary_op(args[0], args[1], x);
               },
               false);
Khalique's avatar
Khalique committed
270
271
272
273
274
    }

    template <class T>
    instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
    {
Khalique's avatar
Khalique committed
275
        if(arg0->get_shape().lens() != arg1->get_shape().lens())
Khalique's avatar
Khalique committed
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
        {
            // 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
291
292
            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
293
294
295
296
297

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

298
            std::vector<size_t> output_lens(*s1);
Khalique's avatar
Khalique committed
299
300
301
302
303
304
305
306
307
            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);
Paul's avatar
Paul committed
308
            return to_nhwc(prog.add_instruction(x, to_nchw(l0), to_nchw(l1)));
Khalique's avatar
Khalique committed
309
310
311
        }
        else
        {
Paul's avatar
Paul committed
312
            return to_nhwc(prog.add_instruction(x, {to_nchw(arg0), to_nchw(arg1)}));
Khalique's avatar
Khalique committed
313
314
315
316
        }
    }

    template <class T>
Paul's avatar
Paul committed
317
    void add_generic_op(std::string name, T x, bool transpose = true)
Khalique's avatar
Khalique committed
318
    {
Paul's avatar
Paul committed
319
320
321
322
323
        add_op(name,
               [this, x](const attribute_map&, std::vector<instruction_ref> args) {
                   return prog.add_instruction(x, args);
               },
               transpose);
Khalique's avatar
Khalique committed
324
325
326
327
328
    }

    instruction_ref
    parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
329
330
331
        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
332
333
334
335
336
337
338
339
        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));
    }

340
    instruction_ref
Khalique's avatar
Khalique committed
341
    parse_biasadd(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
342
    {
343
        uint64_t axis = 1; // assume output of previous layer is in NCHW (broadcast on channel)
Shucai Xiao's avatar
Shucai Xiao committed
344
        auto l0 = prog.add_instruction(op::broadcast{axis, args[0]->get_shape().lens()}, args[1]);
345
        return prog.add_instruction(op::add{}, args[0], l0);
346
347
    }

Khalique's avatar
Khalique committed
348
349
350
351
352
353
354
    instruction_ref
    parse_cast(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        shape::type_t type = parse_type(attributes.at("DstT").type());
        return prog.add_instruction(op::convert{type}, std::move(args));
    }

Khalique's avatar
Khalique committed
355
356
357
358
    instruction_ref
    parse_concat(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        // get index for axis within args
359
        size_t axis_idx = attributes.at("N").i();
Shucai Xiao's avatar
Shucai Xiao committed
360
        int64_t axis    = args[axis_idx]->eval().at<int64_t>();
Khalique's avatar
Khalique committed
361
        op::concat op{axis};
362
        // return only first N arguments (assuming last index is the axis value)
Paul's avatar
Paul committed
363
364
        return prog.add_instruction(
            op, std::vector<instruction_ref>(args.begin(), args.begin() + args.size() - 1));
Khalique's avatar
Khalique committed
365
366
367
368
369
370
    }

    instruction_ref parse_constant(const std::string&,
                                   attribute_map attributes,
                                   const std::vector<instruction_ref>&)
    {
Paul's avatar
Paul committed
371
        literal v = parse_tensor(attributes.at("value").tensor());
Paul's avatar
Paul committed
372
        return prog.add_literal(v);
Khalique's avatar
Khalique committed
373
374
375
376
377
378
379
380
    }

    instruction_ref
    parse_conv(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
        op::convolution op;
        if(contains(attributes, "strides"))
        {
381
            std::vector<size_t> stride;
382
            copy(attributes.at("strides").list().i(), std::back_inserter(stride));
383
            reorder_data(stride);
384
385
            if(stride.size() != 4)
            {
386
                MIGRAPHX_THROW("strides should have 4 values");
387
            }
388
389
            op.stride[0] = stride[2];
            op.stride[1] = stride[3];
Khalique's avatar
Khalique committed
390
391
392
        }
        if(contains(attributes, "dilations"))
        {
393
            std::vector<size_t> dilation;
394
            copy(attributes.at("dilations").list().i(), std::back_inserter(dilation));
395
            reorder_data(dilation);
396
397
398
399
            if(dilation.size() != 4)
            {
                MIGRAPHX_THROW("dilation should have 4 values");
            }
400
401
            op.dilation[0] = dilation[2];
            op.dilation[1] = dilation[3];
Khalique's avatar
Khalique committed
402
        }
Khalique's avatar
Khalique committed
403

Paul's avatar
Paul committed
404
        auto weights = to_kcxy(args[1]);
Paul's avatar
Paul committed
405
        auto l0      = args[0];
Khalique's avatar
Khalique committed
406
407
408
409
410
        if(contains(attributes, "padding"))
        {
            const std::string& pad_mode = attributes.at("padding").s();
            if(pad_mode.find("SAME") != std::string::npos)
            {
Khalique's avatar
Khalique committed
411
                op.padding_mode                 = op::padding_mode_t::same;
Khalique's avatar
Khalique committed
412
413
414
                std::vector<size_t> weight_dims = weights->get_shape().lens();
                size_t weight_h                 = weight_dims[2];
                size_t weight_w                 = weight_dims[3];
Khalique's avatar
Khalique committed
415
416
417

                auto input_dims = l0->get_shape().lens();
                std::vector<int64_t> pads(input_dims.size());
418
419
                calculate_padding(0, pads, input_dims[2], op.stride[0], op.dilation[0], weight_h);
                calculate_padding(1, pads, input_dims[3], op.stride[1], op.dilation[1], weight_w);
Khalique's avatar
Khalique committed
420
421
422
423
424
425
426
427

                if(pads[0] != pads[2] || pads[1] != pads[3])
                {
                    std::vector<int64_t> padding = {0, 0, pads[0], pads[1], 0, 0, pads[2], pads[3]};
                    l0 = prog.add_instruction(migraphx::op::pad{padding}, l0);
                }
                else
                {
Khalique's avatar
Khalique committed
428
429
                    op.padding[0] = pads[0];
                    op.padding[1] = pads[1];
Khalique's avatar
Khalique committed
430
                }
431
432
433
            }
            else if(pad_mode.find("VALID") != std::string::npos)
            {
434
                op.padding_mode = op::padding_mode_t::valid;
Khalique's avatar
Khalique committed
435
            }
Khalique's avatar
Khalique committed
436
            else if(pad_mode.find("EXPLICIT") != std::string::npos)
Khalique's avatar
Khalique committed
437
            {
438
                std::vector<size_t> padding;
439
                copy(attributes.at("explicit_paddings").list().i(), std::back_inserter(padding));
Khalique's avatar
Khalique committed
440
441
442
443
444
445
446
447
448
449
450
451
                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];
            }
        }
Paul's avatar
Paul committed
452
        return prog.add_instruction(op, {l0, to_kcxy(args[1])});
Khalique's avatar
Khalique committed
453
454
    }

Khalique's avatar
Khalique committed
455
456
457
    instruction_ref parse_depthwiseconv(const std::string&,
                                        attribute_map attributes,
                                        std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
458
459
460
    {
        op::convolution op;
        size_t num_channels = args[0]->get_shape().lens()[1];
Khalique's avatar
Khalique committed
461
        op.group            = num_channels;
Khalique's avatar
Khalique committed
462

Khalique's avatar
Khalique committed
463
464
        if(contains(attributes, "strides"))
        {
465
            std::vector<size_t> stride;
466
            copy(attributes.at("strides").list().i(), std::back_inserter(stride));
467
            reorder_data(stride);
468
469
            if(stride.size() != 4)
            {
470
                MIGRAPHX_THROW("strides should have 4 values");
471
            }
472
473
            op.stride[0] = stride[2];
            op.stride[1] = stride[3];
Khalique's avatar
Khalique committed
474
        }
Paul's avatar
Paul committed
475
476

        auto weights = to_kcxy(args[1]);
Khalique's avatar
Khalique committed
477
478
        if(contains(attributes, "dilations"))
        {
479
            std::vector<size_t> dilation;
480
            copy(attributes.at("dilations").list().i(), std::back_inserter(dilation));
481
            reorder_data(dilation);
482
483
484
485
            if(dilation.size() != 4)
            {
                MIGRAPHX_THROW("dilation should have 4 values");
            }
486
487
            op.dilation[0] = dilation[2];
            op.dilation[1] = dilation[3];
Khalique's avatar
Khalique committed
488
489
        }

Khalique's avatar
Khalique committed
490
        auto l0 = args[0];
Khalique's avatar
Khalique committed
491
492
493
        if(contains(attributes, "padding"))
        {
            const std::string& pad_mode = attributes.at("padding").s();
Khalique's avatar
Khalique committed
494

Khalique's avatar
Khalique committed
495
496
            if(pad_mode.find("SAME") != std::string::npos)
            {
Khalique's avatar
Khalique committed
497
                op.padding_mode                 = op::padding_mode_t::same;
Khalique's avatar
Khalique committed
498
499
500
501
502
503
                std::vector<size_t> weight_dims = weights->get_shape().lens();
                size_t weight_h                 = weight_dims[2];
                size_t weight_w                 = weight_dims[3];

                auto input_dims = l0->get_shape().lens();
                std::vector<int64_t> pads(input_dims.size());
504
505
                calculate_padding(0, pads, input_dims[2], op.stride[0], op.dilation[0], weight_h);
                calculate_padding(1, pads, input_dims[3], op.stride[1], op.dilation[1], weight_w);
Khalique's avatar
Khalique committed
506
507
508
509
510
511
512
513

                if(pads[0] != pads[2] || pads[1] != pads[3])
                {
                    std::vector<int64_t> padding = {0, 0, pads[0], pads[1], 0, 0, pads[2], pads[3]};
                    l0 = prog.add_instruction(migraphx::op::pad{padding}, l0);
                }
                else
                {
Khalique's avatar
Khalique committed
514
515
                    op.padding[0] = pads[0];
                    op.padding[1] = pads[1];
Khalique's avatar
Khalique committed
516
                }
Khalique's avatar
Khalique committed
517
            }
Khalique's avatar
Khalique committed
518
            else if(pad_mode.find("VALID") != std::string::npos)
Khalique's avatar
Khalique committed
519
            {
Khalique's avatar
Khalique committed
520
                op.padding_mode = op::padding_mode_t::valid;
Khalique's avatar
Khalique committed
521
522
            }
        }
Khalique's avatar
Khalique committed
523

Khalique's avatar
Khalique committed
524
525
        std::vector<int64_t> new_weights_shape;
        copy(weights->get_shape().lens(), std::back_inserter(new_weights_shape));
Khalique's avatar
Khalique committed
526
527
528
529

        // weight format is (out_channels, in_channels, h, w), but in depthwise_conv,
        // out_channels is equal to the multiplier. Adjust by inserting a reshape and
        // setting in_channels to 1
Khalique's avatar
Khalique committed
530
        int64_t multiplier   = new_weights_shape[0];
Khalique's avatar
Khalique committed
531
532
533
        int64_t out_channels = num_channels * multiplier;
        new_weights_shape[0] = out_channels;
        new_weights_shape[1] = 1;
Paul's avatar
Paul committed
534
        // Make sure weights are contiguous before doing reshape
Paul's avatar
Paul committed
535
536
        auto new_weights =
            prog.add_instruction(op::reshape{new_weights_shape}, make_contiguous(weights));
Khalique's avatar
Khalique committed
537

Khalique's avatar
Khalique committed
538
        return prog.add_instruction(op, {l0, new_weights});
Khalique's avatar
Khalique committed
539
540
    }

Khalique's avatar
Khalique committed
541
542
    instruction_ref
    parse_expanddims(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
543
544
    {
        std::vector<size_t> input_dims = args[0]->get_shape().lens();
Khalique's avatar
Khalique committed
545
        std::vector<int64_t> new_dims(input_dims.begin(), input_dims.end());
Khalique's avatar
Khalique committed
546
        size_t num_dims = input_dims.size();
547
        int32_t dim     = args[1]->eval().at<int32_t>();
Khalique's avatar
Khalique committed
548
549

        if(dim < 0)
Khalique's avatar
Khalique committed
550
551
552
553
554
555
556
557
558
559
        {
            new_dims.insert(new_dims.begin() + (num_dims + dim + 1), 1);
        }
        else
        {
            new_dims.insert(new_dims.begin() + dim, 1);
        }
        return prog.add_instruction(op::reshape{new_dims}, args[0]);
    }

Khalique's avatar
Khalique committed
560
561
562
563
564
565
566
567
    instruction_ref
    parse_gather(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
    {
        int axis = args[2]->eval().at<int32_t>();
        op::gather op{axis};
        return prog.add_instruction(op, {args[0], args[1]});
    }

Khalique's avatar
Khalique committed
568
569
    instruction_ref
    parse_matmul(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
570
571
572
    {
        bool transa = false;
        bool transb = false;
Khalique's avatar
Khalique committed
573

574
575
576
577
578
579
        if(contains(attributes, "transpose_a"))
        {
            transa = attributes.at("transpose_a").b();
        }
        if(contains(attributes, "transpose_b"))
        {
Khalique's avatar
Khalique committed
580
            transb = attributes.at("transpose_b").b();
581
582
        }

Khalique's avatar
Khalique committed
583
584
585
586
587
588
589
590
591
        if(contains(attributes, "adj_x"))
        {
            transa = attributes.at("adj_x").b();
        }
        if(contains(attributes, "adj_y"))
        {
            transb = attributes.at("adj_y").b();
        }

592
593
594
        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
595
        std::iter_swap(perm.end() - 1, perm.end() - 2);
596
597
598
599
600
601
602

        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
603
604
    instruction_ref
    parse_mean(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
605
    {
Khalique's avatar
Khalique committed
606
607
        bool keep_dims = attributes.at("keep_dims").b();
        auto axes      = args[1]->eval().get<int32_t>().to_vector<int64_t>();
Khalique's avatar
Khalique committed
608
609

        if(keep_dims)
Khalique's avatar
Khalique committed
610
        {
611
612
613
614
615
616
            return prog.add_instruction(op::reduce_mean{axes}, args[0]);
        }
        else
        {
            auto ins = prog.add_instruction(op::reduce_mean{axes}, args[0]);
            return prog.add_instruction(op::squeeze{axes}, ins);
Khalique's avatar
Khalique committed
617
618
619
        }
    }

Khalique's avatar
Khalique committed
620
621
622
    instruction_ref
    parse_onehot(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
    {
Khalique's avatar
Khalique committed
623
624
        size_t depth = static_cast<size_t>(args[1]->eval().at<int32_t>());

Khalique's avatar
Khalique committed
625
        int64_t axis    = -1;
Khalique's avatar
Khalique committed
626
627
        float on_value  = args[2]->eval().at<float>();
        float off_value = args[3]->eval().at<float>();
Khalique's avatar
Khalique committed
628

Khalique's avatar
Khalique committed
629
        std::vector<float> depth_input(depth * depth, off_value);
Khalique's avatar
Khalique committed
630
631
        for(int i = 0; i < depth; i++)
        {
Khalique's avatar
Khalique committed
632
            depth_input[depth * i + i] = on_value;
Khalique's avatar
Khalique committed
633
        }
Khalique's avatar
Khalique committed
634

Khalique's avatar
Khalique committed
635
        if(contains(attributes, "axis"))
Khalique's avatar
Khalique committed
636
637
638
            axis = attributes.at("axis").i();
        if(axis == -1)
        {
Khalique's avatar
Khalique committed
639
640
641
            shape s{shape::float_type, {depth, depth}};
            auto l0 = prog.add_literal({s, depth_input});
            return prog.add_instruction(op::gather{0}, {l0, args[0]});
Khalique's avatar
Khalique committed
642
643
644
645
        }
        MIGRAPHX_THROW("MIGraphX does not support axis != -1");
    }

Khalique's avatar
Khalique committed
646
647
648
    instruction_ref parse_pack(const std::string&,
                               const attribute_map& attributes,
                               std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
649
650
651
652
653
654
    {
        // reinterpret as unsqueeze with concat
        std::vector<instruction_ref> unsqueezed_args;
        int64_t axis = 0;
        if(contains(attributes, "axis"))
            axis = attributes.at("axis").i();
655
656
657
        size_t input_size = args.front()->get_shape().lens().size();
        if(axis > input_size)
        {
Khalique's avatar
Khalique committed
658
659
            MIGRAPHX_THROW("TF_PARSER: axis value of " + to_string(axis) +
                           " must be smaller than input size " + to_string(input_size));
660
661
        }

Khalique's avatar
Khalique committed
662
663
664
665
666
        std::transform(
            args.begin(),
            args.end(),
            std::back_inserter(unsqueezed_args),
            [&](instruction_ref arg) { return prog.add_instruction(op::unsqueeze{{axis}}, arg); });
Shucai Xiao's avatar
Shucai Xiao committed
667
        return to_nhwc(prog.add_instruction(op::concat{axis}, unsqueezed_args));
Khalique's avatar
Khalique committed
668
669
    }

Khalique's avatar
Khalique committed
670
671
672
673
674
    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
675
676
        // 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
677
        std::vector<std::pair<int32_t, int32_t>> pad_per_dim(ndims);
Paul's avatar
Paul committed
678
        auto tf_padding = args[1]->eval().get<int32_t>().to_vector();
Khalique's avatar
Khalique committed
679
        for(size_t i = 0; i < 2 * ndims; i += 2)
Khalique's avatar
Khalique committed
680
        {
Khalique's avatar
Khalique committed
681
682
            pad_per_dim[i / 2].first  = tf_padding[i];
            pad_per_dim[i / 2].second = tf_padding[i + 1];
Khalique's avatar
Khalique committed
683
684
685
686
        }
        reorder_data(pad_per_dim);

        op::pad op;
Khalique's avatar
Khalique committed
687
688
        std::vector<int64_t> pads(ndims * 2);
        for(size_t i = 0; i < ndims; i++)
Khalique's avatar
Khalique committed
689
        {
Khalique's avatar
Khalique committed
690
691
            pads[i]         = pad_per_dim[i].first;
            pads[i + ndims] = pad_per_dim[i].second;
Khalique's avatar
Khalique committed
692
693
        }
        op.pads = pads;
Paul's avatar
Paul committed
694
        return prog.add_instruction(op, args.front());
Khalique's avatar
Khalique committed
695
696
    }

697
698
699
700
701
    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
702

703
704
        if(contains(attributes, "strides"))
        {
705
            std::vector<size_t> stride;
706
            copy(attributes.at("strides").list().i(), std::back_inserter(stride));
707
            reorder_data(stride);
708
709
710
711
            if(stride.size() != 4)
            {
                MIGRAPHX_THROW("strides should have 4 values");
            }
712
713
            op.stride[0] = stride[2];
            op.stride[1] = stride[3];
714
715
716
        }
        if(contains(attributes, "ksize"))
        {
717
            std::vector<size_t> ksize;
718
            copy(attributes.at("ksize").list().i(), std::back_inserter(ksize));
719
            reorder_data(ksize);
720
721
722
            if(ksize.size() != 4)
            {
                MIGRAPHX_THROW("ksize should have 4 values");
Khalique's avatar
Khalique committed
723
            }
724
725
            op.lengths[0] = ksize[2];
            op.lengths[1] = ksize[3];
726
        }
Khalique's avatar
Khalique committed
727
728

        auto l0 = args[0];
Khalique's avatar
Khalique committed
729
730
731
732
733
        if(contains(attributes, "padding"))
        {
            const std::string& pad_mode = attributes.at("padding").s();
            if(pad_mode.find("SAME") != std::string::npos)
            {
Khalique's avatar
Khalique committed
734
                op.padding_mode = op::padding_mode_t::same;
Khalique's avatar
Khalique committed
735
736
                auto input_dims = l0->get_shape().lens();
                std::vector<int64_t> pads(input_dims.size());
737
738
                calculate_padding(0, pads, input_dims[2], op.stride[0], 1, op.lengths[0]);
                calculate_padding(1, pads, input_dims[3], op.stride[1], 1, op.lengths[1]);
Khalique's avatar
Khalique committed
739
740
741
742

                if(pads[0] != pads[2] || pads[1] != pads[3])
                {
                    std::vector<int64_t> padding = {0, 0, pads[0], pads[1], 0, 0, pads[2], pads[3]};
Khalique's avatar
Khalique committed
743
744
                    l0                           = prog.add_instruction(
                        migraphx::op::pad{padding, std::numeric_limits<float>::lowest()}, l0);
Khalique's avatar
Khalique committed
745
746
747
                }
                else
                {
Khalique's avatar
Khalique committed
748
749
                    op.padding[0] = pads[0];
                    op.padding[1] = pads[1];
Khalique's avatar
Khalique committed
750
                }
Khalique's avatar
Khalique committed
751
752
753
754
755
756
            }
            else if(pad_mode.find("VALID") != std::string::npos)
            {
                op.padding_mode = op::padding_mode_t::valid;
            }
        }
Khalique's avatar
Khalique committed
757
        return prog.add_instruction(op, l0);
758
    }
Khalique's avatar
Khalique committed
759

760
    instruction_ref
Khalique's avatar
Khalique committed
761
    parse_reshape(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
762
763
764
765
    {
        op::reshape op;
        if(args.size() != 2)
            MIGRAPHX_THROW("reshape needs 2 arguments (input, new_shape)");
Khalique's avatar
Khalique committed
766
        auto s = args[1]->eval();
767
        s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
768
        return prog.add_instruction(op, make_contiguous(args[0]));
769
770
    }

Khalique's avatar
Khalique committed
771
772
773
774
775
776
777
778
779
    void parse_from(std::istream& is)
    {
        tensorflow::GraphDef graph;
        if(graph.ParseFromIstream(&is))
        {
            this->parse_graph(graph);
        }
        else
        {
780
            throw std::runtime_error("Failed reading tf file");
Khalique's avatar
Khalique committed
781
782
783
        }
    }

784
    instruction_ref
Khalique's avatar
Khalique committed
785
    parse_slice(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
786
    {
Khalique's avatar
Khalique committed
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
        op::slice op;
        auto starts     = args[1]->eval().get<int32_t>().to_vector();
        auto size       = args[2]->eval().get<int32_t>().to_vector();
        auto axes       = args[0]->get_shape().lens();
        size_t num_axes = axes.size();

        op.starts = std::vector<int64_t>(starts.begin(), starts.end());
        op.ends   = std::vector<int64_t>(num_axes);
        op.axes   = std::vector<int64_t>(num_axes);
        std::iota(op.axes.begin(), op.axes.end(), 0);
        for(size_t i = 0; i < num_axes; i++)
        {
            if(size[i] == -1)
                op.ends[i] = axes[i];
            else
                op.ends[i] = starts[i] + size[i];
        }
        return prog.add_instruction(op, make_contiguous(args[0]));
    }

Khalique's avatar
Khalique committed
807
808
809
810
811
    // template to facilitate the logsoftmax later
    template <class Op>
    instruction_ref parse_softmax(const std::string&,
                                  const attribute_map& attributes,
                                  std::vector<instruction_ref> args)
812
    {
Khalique's avatar
Khalique committed
813
        int axis      = -1;
Khalique's avatar
Khalique committed
814
815
816
817
818
819
820
821
822
823
824
        auto num_dims = args[0]->get_shape().lens().size();
        if(contains(attributes, "axis"))
        {
            axis = static_cast<int>(attributes.at("axis").i());
        }
        if(axis < 0)
        {
            axis += num_dims;
        }

        return prog.add_instruction(Op{axis}, make_contiguous(args[0]));
825
826
    }

kahmed10's avatar
kahmed10 committed
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
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
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
    std::vector<instruction_ref> parse_split(const std::string&,
                                             const attribute_map& attributes,
                                             std::vector<instruction_ref> args)
    {
        bool vector_as_input = args.size() == 3;
        int num_outputs      = 1;
        auto axis_arg        = args[0];
        auto input_arg       = args[1];
        if(vector_as_input)
        {
            input_arg = args[0];
            axis_arg  = args[2];
        }

        if(contains(attributes, "num_split"))
            num_outputs = attributes.at("num_split").i();

        std::vector<int> splits(num_outputs);
        std::vector<int> slice_pos{0};
        if(vector_as_input)
        {
            splits      = args[1]->eval().get<int32_t>().to_vector();
            num_outputs = splits.size();
        }

        assert(num_outputs > 0);

        if(num_outputs == 1)
            return std::vector<instruction_ref>{prog.add_instruction(op::identity{}, input_arg)};

        auto lens     = input_arg->get_shape().lens();
        auto num_dims = lens.size();
        int axis      = axis_arg->eval().at<int32_t>();

        // ensure split is made evenly if "num_split" is used
        assert(vector_as_input or lens[axis] % num_outputs == 0);

        auto split_size = lens[axis] / num_outputs;

        // push back first end point of slice
        if(vector_as_input)
        {
            slice_pos.push_back(splits[0]);
        }
        else
        {
            slice_pos.push_back(split_size);
        }

        // calculate remaining end points for each slice
        for(auto i = 1; i < num_outputs; i++)
        {
            if(vector_as_input)
            {
                splits[i] += splits[i - 1];
                slice_pos.push_back(splits[i]);
            }
            else
            {
                slice_pos.push_back((i + 1) * split_size);
            }
        }
        std::vector<instruction_ref> result;
        for(auto i = 0; i < num_outputs; i++)
        {
            op::slice op;
            op.axes = std::vector<int64_t>(num_dims);
            std::iota(op.axes.begin(), op.axes.end(), 0);
            op.starts = std::vector<int64_t>(num_dims, 0);
            op.ends   = std::vector<int64_t>(lens.begin(), lens.end());

            op.starts[axis] = slice_pos[i];
            op.ends[axis]   = slice_pos[i + 1];
            result.push_back(prog.add_instruction(op, input_arg));
        }
        return result;
    }

Khalique's avatar
Khalique committed
905
906
907
    instruction_ref parse_squeeze(const std::string&,
                                  const attribute_map& attributes,
                                  std::vector<instruction_ref> args)
908
909
    {
        op::squeeze op;
Khalique's avatar
Khalique committed
910
        auto input_dims = args[0]->get_shape().lens();
Khalique's avatar
Khalique committed
911
        auto axes       = attributes.at("squeeze_dims").list().i();
912
        copy(axes, std::back_inserter(op.axes));
Khalique's avatar
Khalique committed
913

914
915
        if(op.axes.empty()) // no squeeze_dims provided, remove any dim that equals 1
        {
Khalique's avatar
Khalique committed
916
            for(size_t i = 0; i < input_dims.size(); i++)
917
            {
Khalique's avatar
Khalique committed
918
                if(input_dims.at(i) == 1)
919
920
921
922
                {
                    op.axes.push_back(i);
                }
            }
923
        }
Paul's avatar
Paul committed
924
        return prog.add_instruction(op, make_contiguous(args[0]));
925
926
    }

Khalique's avatar
Khalique committed
927
928
929
    instruction_ref parse_stridedslice(const std::string&,
                                       const attribute_map& attributes,
                                       std::vector<instruction_ref> args)
930
931
    {
        op::slice op;
Khalique's avatar
Khalique committed
932
933
934
935
        auto starts              = args[1]->eval().get<int32_t>().to_vector();
        auto ends                = args[2]->eval().get<int32_t>().to_vector();
        auto l0                  = args[0];
        size_t num_axes          = l0->get_shape().lens().size();
936
        std::vector<size_t> axes = l0->get_shape().lens();
937

Khalique's avatar
Khalique committed
938
939
940
941
        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);
Khalique's avatar
Khalique committed
942
943
        uint32_t begin_mask       = 0;
        uint32_t end_mask         = 0;
944
        uint32_t shrink_axis_mask = 0;
Khalique's avatar
Khalique committed
945
        uint32_t bitwise_compare  = 1;
946
947
        std::vector<int64_t> squeeze_axes;

Khalique's avatar
Khalique committed
948
949
950
951
952
953
        if(contains(attributes, "begin_mask"))
            begin_mask = static_cast<uint32_t>(attributes.at("begin_mask").i());

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

954
        if(contains(attributes, "shrink_axis_mask"))
955
            shrink_axis_mask = static_cast<uint32_t>(attributes.at("shrink_axis_mask").i());
956

Khalique's avatar
Khalique committed
957
        std::vector<int64_t> begin_axes = get_axes_from_mask(num_axes, begin_mask);
Khalique's avatar
Khalique committed
958
        std::vector<int64_t> end_axes   = get_axes_from_mask(num_axes, end_mask);
Khalique's avatar
Khalique committed
959
960
961
962
963
964
965
966
967
968
969
970
971

        for(size_t i = 0; i < num_axes; i++)
        {
            if(begin_axes.at(i) == 1)
            {
                op.starts.at(i) = 0;
            }
            if(end_axes.at(i) == 1)
            {
                op.ends.at(i) = axes.at(i);
            }
        }

972
        auto l1 = prog.add_instruction(op, l0);
Khalique's avatar
Khalique committed
973
        if(shrink_axis_mask == 0)
974
            return l1;
Khalique's avatar
Khalique committed
975

Khalique's avatar
Khalique committed
976
        for(size_t i = 0; i < num_axes; i++)
977
        {
978
            // the LSB corresponds to axis 0 when determining which axes to squeeze
Khalique's avatar
Khalique committed
979
            if(((shrink_axis_mask >> i) & bitwise_compare) == 1)
980
981
                squeeze_axes.push_back(i);
        }
Khalique's avatar
Khalique committed
982

983
        return prog.add_instruction(op::squeeze{squeeze_axes}, l1);
984
985
    }

Khalique's avatar
Khalique committed
986
987
    instruction_ref
    parse_transpose(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
Khalique's avatar
Khalique committed
988
989
990
991
992
993
994
995
    {
        auto perm = args[1]->eval().get<int32_t>().to_vector();
        op::transpose op;
        op.dims = std::vector<int64_t>(perm.begin(), perm.end());

        return prog.add_instruction(op, args.front());
    }

Khalique's avatar
Khalique committed
996
997
998
999
1000
    void parse_graph(const tensorflow::GraphDef& graph)
    {
        nodes = get_nodes(graph, input_nodes);
        for(auto&& input : input_nodes)
        {
Khalique's avatar
Khalique committed
1001
            const std::string& name   = input.name();
Khalique's avatar
Khalique committed
1002
            attribute_map input_attrs = get_attributes(input);
Khalique's avatar
Khalique committed
1003
1004
            shape::type_t shape_type  = parse_type(input_attrs.at("dtype").type());
            std::vector<size_t> dims  = parse_dims(input_attrs.at("shape").shape());
1005
            if(is_nhwc and dims.size() >= 4)
1006
            {
1007
                reorder_data(dims);
1008
            }
Khalique's avatar
Khalique committed
1009
            shape s            = shape{shape_type, dims};
Paul's avatar
Paul committed
1010
            instructions[name] = to_nhwc(prog.add_parameter(name, s));
Khalique's avatar
Khalique committed
1011
1012
1013
        }
        for(auto&& p : nodes)
        {
1014
            this->parse_node(p.first);
Khalique's avatar
Khalique committed
1015
1016
1017
1018
1019
1020
1021
1022
        }
    }

    void parse_node(const std::string& name)
    {
        if(instructions.count(name) == 0)
        {
            auto&& node = nodes.at(name);
Khalique's avatar
Khalique committed
1023
1024
1025
            // assert ops ignored
            if(node.op() == "Assert" or contains(name, "Assert"))
                return;
Khalique's avatar
Khalique committed
1026
1027
1028
1029
            std::vector<instruction_ref> args;

            for(auto&& input : node.input())
            {
Khalique's avatar
Khalique committed
1030
1031
1032
                // control dependencies (signified by ^ before the name) are ignored
                if(contains(input, "^"))
                    continue;
Khalique's avatar
Khalique committed
1033
1034
                if(nodes.count(input) > 0)
                {
kahmed10's avatar
kahmed10 committed
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
                    std::string iname;
                    // input was from a node with multiple outputs
                    if(contains(input, ':'))
                    {
                        iname = input.substr(0, input.find(':'));
                    }
                    else
                    {
                        iname = get_name(nodes.at(input));
                    }
Khalique's avatar
Khalique committed
1045
1046
                    assert(name != iname);
                    this->parse_node(iname);
kahmed10's avatar
kahmed10 committed
1047
                    args.push_back(instructions.at(input));
Khalique's avatar
Khalique committed
1048
1049
1050
1051
1052
1053
                }
                else
                {
                    args.push_back(instructions.at(input));
                }
            }
kahmed10's avatar
kahmed10 committed
1054
1055

            std::vector<instruction_ref> result;
Khalique's avatar
Khalique committed
1056
1057
            if(ops.count(node.op()) == 0)
            {
kahmed10's avatar
kahmed10 committed
1058
                result.push_back(prog.add_instruction(op::unknown{node.op()}, args));
Khalique's avatar
Khalique committed
1059
1060
1061
            }
            else
            {
kahmed10's avatar
kahmed10 committed
1062
1063
1064
1065
1066
1067
1068
1069
1070
                result = ops[node.op()](get_attributes(node), args);
            }

            assert(!result.empty());
            // First output has no ":" delimiter
            instructions[name] = result.front();
            for(size_t i = 1; i < result.size(); i++)
            {
                instructions[name + ":" + std::to_string(i)] = result.at(i);
Khalique's avatar
Khalique committed
1071
1072
1073
1074
1075
1076
1077
            }
        }
    }

    static attribute_map get_attributes(const tensorflow::NodeDef& node)
    {
        attribute_map result;
Khalique's avatar
Khalique committed
1078
        for(auto&& attr : node.attr())
Khalique's avatar
Khalique committed
1079
1080
1081
1082
1083
1084
        {
            result[attr.first] = attr.second;
        }
        return result;
    }

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

Khalique's avatar
Khalique committed
1087
1088
    static node_map get_nodes(const tensorflow::GraphDef& graph,
                              std::vector<tensorflow::NodeDef>& input_nodes)
Khalique's avatar
Khalique committed
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
    {
        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_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_INT16: shape_type = shape::int16_type; break;
        case tensorflow::DataType::DT_INT8: shape_type = shape::int8_type; break;
Paul's avatar
Paul committed
1116
1117
1118
1119
        case tensorflow::DataType::DT_INT64: shape_type = shape::int64_type; break;
        case tensorflow::DataType::DT_UINT16: shape_type = shape::uint16_type; break;
        case tensorflow::DataType::DT_HALF: shape_type = shape::half_type; break;
        case tensorflow::DataType::DT_UINT32: shape_type = shape::uint32_type; break;
Paul's avatar
Paul committed
1120
        case tensorflow::DataType::DT_UINT64: shape_type = shape::uint64_type; break;
Paul's avatar
Paul committed
1121
1122
1123

        case tensorflow::DataType::DT_INVALID:
        case tensorflow::DataType::DT_UINT8:
Khalique's avatar
Khalique committed
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
        case tensorflow::DataType::DT_STRING:
        case tensorflow::DataType::DT_COMPLEX64:
        case tensorflow::DataType::DT_BOOL:
        case tensorflow::DataType::DT_QINT8:
        case tensorflow::DataType::DT_QUINT8:
        case tensorflow::DataType::DT_QINT32:
        case tensorflow::DataType::DT_BFLOAT16:
        case tensorflow::DataType::DT_QINT16:
        case tensorflow::DataType::DT_QUINT16:
        case tensorflow::DataType::DT_COMPLEX128:
        case tensorflow::DataType::DT_RESOURCE:
        case tensorflow::DataType::DT_VARIANT:
Khalique's avatar
Khalique committed
1136
        // tf pb should not use these types
Paul's avatar
Paul committed
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
        case tensorflow::DataType::DT_FLOAT_REF:
        case tensorflow::DataType::DT_DOUBLE_REF:
        case tensorflow::DataType::DT_INT32_REF:
        case tensorflow::DataType::DT_UINT8_REF:
        case tensorflow::DataType::DT_INT16_REF:
        case tensorflow::DataType::DT_INT8_REF:
        case tensorflow::DataType::DT_STRING_REF:
        case tensorflow::DataType::DT_COMPLEX64_REF:
        case tensorflow::DataType::DT_INT64_REF:
        case tensorflow::DataType::DT_BOOL_REF:
        case tensorflow::DataType::DT_QINT8_REF:
        case tensorflow::DataType::DT_QUINT8_REF:
        case tensorflow::DataType::DT_QINT32_REF:
        case tensorflow::DataType::DT_BFLOAT16_REF:
        case tensorflow::DataType::DT_QINT16_REF:
        case tensorflow::DataType::DT_QUINT16_REF:
        case tensorflow::DataType::DT_UINT16_REF:
        case tensorflow::DataType::DT_COMPLEX128_REF:
        case tensorflow::DataType::DT_HALF_REF:
        case tensorflow::DataType::DT_RESOURCE_REF:
        case tensorflow::DataType::DT_VARIANT_REF:
        case tensorflow::DataType::DT_UINT32_REF:
        case tensorflow::DataType::DT_UINT64_REF:
Paul's avatar
Paul committed
1160
        case tensorflow::DataType::DataType_INT_MAX_SENTINEL_DO_NOT_USE_:
Khalique's avatar
Khalique committed
1161
        case tensorflow::DataType::DataType_INT_MIN_SENTINEL_DO_NOT_USE_: break;
Khalique's avatar
Khalique committed
1162
1163
1164
1165
        }
        return shape_type;
    }

Khalique's avatar
Khalique committed
1166
    static literal parse_tensor(const tensorflow::TensorProto& t)
Khalique's avatar
Khalique committed
1167
1168
    {
        std::vector<size_t> dims = parse_dims(t.tensor_shape());
1169
        size_t shape_size = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<size_t>());
Khalique's avatar
Khalique committed
1170
1171
        if(!t.tensor_content().empty()) // has raw data
        {
Khalique's avatar
Khalique committed
1172
            const std::string& s = t.tensor_content();
Khalique's avatar
Khalique committed
1173
1174
            switch(t.dtype())
            {
Khalique's avatar
Khalique committed
1175
1176
            case tensorflow::DataType::DT_FLOAT:
                return literal{{shape::float_type, dims}, s.data()};
Paul's avatar
Paul committed
1177
            case tensorflow::DataType::DT_BOOL:
1178
            case tensorflow::DataType::DT_INT8: return literal{{shape::int8_type, dims}, s.data()};
Khalique's avatar
Khalique committed
1179
1180
            case tensorflow::DataType::DT_UINT16:
            case tensorflow::DataType::DT_INT16:
1181
                return literal{{shape::int16_type, dims}, s.data()};
Khalique's avatar
Khalique committed
1182
1183
1184
1185
            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
1186
            case tensorflow::DataType::DT_HALF: return literal{{shape::half_type, dims}, s.data()};
Khalique's avatar
Khalique committed
1187
1188
            case tensorflow::DataType::DT_DOUBLE:
                return literal{{shape::double_type, dims}, s.data()};
Paul's avatar
Paul committed
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
            case tensorflow::DataType::DT_INVALID:
            case tensorflow::DataType::DT_UINT8:
            case tensorflow::DataType::DT_STRING:
            case tensorflow::DataType::DT_UINT32:
            case tensorflow::DataType::DT_UINT64:
            case tensorflow::DataType::DT_COMPLEX64:
            case tensorflow::DataType::DT_COMPLEX128:
            case tensorflow::DataType::DT_QINT8:
            case tensorflow::DataType::DT_QUINT8:
            case tensorflow::DataType::DT_QINT32:
            case tensorflow::DataType::DT_BFLOAT16:
            case tensorflow::DataType::DT_QINT16:
            case tensorflow::DataType::DT_QUINT16:
            case tensorflow::DataType::DT_RESOURCE:
            case tensorflow::DataType::DT_VARIANT:
            case tensorflow::DataType::DT_FLOAT_REF:
            case tensorflow::DataType::DT_DOUBLE_REF:
            case tensorflow::DataType::DT_INT32_REF:
            case tensorflow::DataType::DT_UINT8_REF:
            case tensorflow::DataType::DT_INT16_REF:
            case tensorflow::DataType::DT_INT8_REF:
            case tensorflow::DataType::DT_STRING_REF:
            case tensorflow::DataType::DT_COMPLEX64_REF:
            case tensorflow::DataType::DT_INT64_REF:
            case tensorflow::DataType::DT_BOOL_REF:
            case tensorflow::DataType::DT_QINT8_REF:
            case tensorflow::DataType::DT_QUINT8_REF:
            case tensorflow::DataType::DT_QINT32_REF:
            case tensorflow::DataType::DT_BFLOAT16_REF:
            case tensorflow::DataType::DT_QINT16_REF:
            case tensorflow::DataType::DT_QUINT16_REF:
            case tensorflow::DataType::DT_UINT16_REF:
            case tensorflow::DataType::DT_COMPLEX128_REF:
            case tensorflow::DataType::DT_HALF_REF:
            case tensorflow::DataType::DT_RESOURCE_REF:
            case tensorflow::DataType::DT_VARIANT_REF:
            case tensorflow::DataType::DT_UINT32_REF:
            case tensorflow::DataType::DT_UINT64_REF:
Khalique's avatar
Khalique committed
1227
1228
1229
            case tensorflow::DataType::DataType_INT_MAX_SENTINEL_DO_NOT_USE_:
            case tensorflow::DataType::DataType_INT_MIN_SENTINEL_DO_NOT_USE_:
                throw std::runtime_error("");
Khalique's avatar
Khalique committed
1230
1231
1232
1233
1234
1235
            }
            MIGRAPHX_THROW("Invalid tensor type");
        }
        switch(t.dtype())
        {
        case tensorflow::DataType::DT_FLOAT:
Khalique's avatar
Khalique committed
1236
1237
            return create_literal(
                shape::float_type, dims, get_data_vals(t.float_val(), shape_size));
Khalique's avatar
Khalique committed
1238
        case tensorflow::DataType::DT_INT8:
1239
            return create_literal(shape::int8_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
1240
        case tensorflow::DataType::DT_UINT16:
1241
            return create_literal(shape::uint16_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
1242
        case tensorflow::DataType::DT_INT16:
1243
            return create_literal(shape::int16_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
1244
        case tensorflow::DataType::DT_INT32:
1245
            return create_literal(shape::int32_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
1246
        case tensorflow::DataType::DT_INT64:
Khalique's avatar
Khalique committed
1247
1248
            return create_literal(
                shape::int64_type, dims, get_data_vals(t.int64_val(), shape_size));
Khalique's avatar
Khalique committed
1249
        case tensorflow::DataType::DT_BOOL:
1250
            return create_literal(shape::int32_type, dims, get_data_vals(t.bool_val(), shape_size));
Khalique's avatar
Khalique committed
1251
        case tensorflow::DataType::DT_HALF:
Khalique's avatar
Khalique committed
1252
        {
1253
1254
            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
1255
1256
1257
1258
1259
            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); });
1260
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
1261
        }
Khalique's avatar
Khalique committed
1262
        case tensorflow::DataType::DT_DOUBLE:
Khalique's avatar
Khalique committed
1263
            return literal{{shape::double_type, dims}, get_data_vals(t.double_val(), shape_size)};
Paul's avatar
Paul committed
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
        case tensorflow::DataType::DT_INVALID:
        case tensorflow::DataType::DT_UINT8:
        case tensorflow::DataType::DT_STRING:
        case tensorflow::DataType::DT_UINT32:
        case tensorflow::DataType::DT_UINT64:
        case tensorflow::DataType::DT_COMPLEX64:
        case tensorflow::DataType::DT_COMPLEX128:
        case tensorflow::DataType::DT_QINT8:
        case tensorflow::DataType::DT_QUINT8:
        case tensorflow::DataType::DT_QINT32:
        case tensorflow::DataType::DT_BFLOAT16:
        case tensorflow::DataType::DT_QINT16:
        case tensorflow::DataType::DT_QUINT16:
        case tensorflow::DataType::DT_RESOURCE:
        case tensorflow::DataType::DT_VARIANT:
        case tensorflow::DataType::DT_FLOAT_REF:
        case tensorflow::DataType::DT_DOUBLE_REF:
        case tensorflow::DataType::DT_INT32_REF:
        case tensorflow::DataType::DT_UINT8_REF:
        case tensorflow::DataType::DT_INT16_REF:
        case tensorflow::DataType::DT_INT8_REF:
        case tensorflow::DataType::DT_STRING_REF:
        case tensorflow::DataType::DT_COMPLEX64_REF:
        case tensorflow::DataType::DT_INT64_REF:
        case tensorflow::DataType::DT_BOOL_REF:
        case tensorflow::DataType::DT_QINT8_REF:
        case tensorflow::DataType::DT_QUINT8_REF:
        case tensorflow::DataType::DT_QINT32_REF:
        case tensorflow::DataType::DT_BFLOAT16_REF:
        case tensorflow::DataType::DT_QINT16_REF:
        case tensorflow::DataType::DT_QUINT16_REF:
        case tensorflow::DataType::DT_UINT16_REF:
        case tensorflow::DataType::DT_COMPLEX128_REF:
        case tensorflow::DataType::DT_HALF_REF:
        case tensorflow::DataType::DT_RESOURCE_REF:
        case tensorflow::DataType::DT_VARIANT_REF:
        case tensorflow::DataType::DT_UINT32_REF:
        case tensorflow::DataType::DT_UINT64_REF:
Khalique's avatar
Khalique committed
1302
1303
1304
        case tensorflow::DataType::DataType_INT_MAX_SENTINEL_DO_NOT_USE_:
        case tensorflow::DataType::DataType_INT_MIN_SENTINEL_DO_NOT_USE_:
            throw std::runtime_error("");
Khalique's avatar
Khalique committed
1305
1306
1307
1308
        }
        MIGRAPHX_THROW("Invalid tensor type");
    }

1309
    template <class T>
Khalique's avatar
Khalique committed
1310
    static std::vector<T> get_data_vals(const google::protobuf::RepeatedField<T>& data,
1311
                                        const size_t& shape_size)
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
    {
        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
1324
1325
1326
1327
    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
1328
1329
1330
        std::transform(input_dims.begin(),
                       input_dims.end(),
                       std::back_inserter(dims),
Paul's avatar
Paul committed
1331
                       [](const tensorflow::TensorShapeProto_Dim& dim) { return dim.size(); });
Khalique's avatar
Khalique committed
1332
1333
        return dims;
    }
1334
1335

    template <class T>
Khalique's avatar
Khalique committed
1336
    static literal
1337
    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, std::vector<T> data)
1338
    {
Khalique's avatar
Khalique committed
1339
        // assume if explicit value is mentioned in protobuf and dim size <= 1, treat as scalar
1340
        if(dims.empty() or (dims.size() == 1 and dims.front() == 1))
1341
            return literal{{shape_type}, data};
1342
1343
        return literal{{shape_type, dims}, data};
    }
Khalique's avatar
Khalique committed
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
};

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
Paul's avatar
Paul committed
1366
    parser.to_nchw(std::prev(parser.prog.end()));
Khalique's avatar
Khalique committed
1367
1368
1369
1370
1371
    return std::move(parser.prog);
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx