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();
Paul's avatar
Paul committed
360
        size_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); });
Paul's avatar
Paul committed
667
668
        return to_nhwc(
            prog.add_instruction(op::concat{static_cast<size_t>(axis)}, unsqueezed_args));
Khalique's avatar
Khalique committed
669
670
    }

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

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

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

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

        auto l0 = args[0];
Khalique's avatar
Khalique committed
730
731
732
733
734
        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
735
                op.padding_mode = op::padding_mode_t::same;
Khalique's avatar
Khalique committed
736
737
                auto input_dims = l0->get_shape().lens();
                std::vector<int64_t> pads(input_dims.size());
738
739
                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
740
741
742
743

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

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

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

785
    instruction_ref
Khalique's avatar
Khalique committed
786
    parse_slice(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
787
    {
Khalique's avatar
Khalique committed
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
        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
808
809
810
811
812
    // 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)
813
    {
Khalique's avatar
Khalique committed
814
        int axis      = -1;
Khalique's avatar
Khalique committed
815
816
817
818
819
820
821
822
823
824
825
        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]));
826
827
    }

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

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

Khalique's avatar
Khalique committed
928
929
930
    instruction_ref parse_stridedslice(const std::string&,
                                       const attribute_map& attributes,
                                       std::vector<instruction_ref> args)
931
932
    {
        op::slice op;
Khalique's avatar
Khalique committed
933
934
935
936
        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();
937
        std::vector<size_t> axes = l0->get_shape().lens();
938

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

Khalique's avatar
Khalique committed
949
950
951
952
953
954
        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());

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

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

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

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

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

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

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

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

            for(auto&& input : node.input())
            {
Khalique's avatar
Khalique committed
1031
1032
1033
                // control dependencies (signified by ^ before the name) are ignored
                if(contains(input, "^"))
                    continue;
Khalique's avatar
Khalique committed
1034
1035
                if(nodes.count(input) > 0)
                {
kahmed10's avatar
kahmed10 committed
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
                    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
1046
1047
                    assert(name != iname);
                    this->parse_node(iname);
kahmed10's avatar
kahmed10 committed
1048
                    args.push_back(instructions.at(input));
Khalique's avatar
Khalique committed
1049
1050
1051
1052
1053
1054
                }
                else
                {
                    args.push_back(instructions.at(input));
                }
            }
kahmed10's avatar
kahmed10 committed
1055
1056

            std::vector<instruction_ref> result;
Khalique's avatar
Khalique committed
1057
1058
            if(ops.count(node.op()) == 0)
            {
kahmed10's avatar
kahmed10 committed
1059
                result.push_back(prog.add_instruction(op::unknown{node.op()}, args));
Khalique's avatar
Khalique committed
1060
1061
1062
            }
            else
            {
kahmed10's avatar
kahmed10 committed
1063
1064
1065
1066
1067
1068
1069
1070
1071
                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
1072
1073
1074
1075
1076
1077
1078
            }
        }
    }

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

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

Khalique's avatar
Khalique committed
1088
1089
    static node_map get_nodes(const tensorflow::GraphDef& graph,
                              std::vector<tensorflow::NodeDef>& input_nodes)
Khalique's avatar
Khalique committed
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
1116
    {
        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
1117
1118
1119
1120
        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
1121
        case tensorflow::DataType::DT_UINT64: shape_type = shape::uint64_type; break;
Paul's avatar
Paul committed
1122
1123
1124

        case tensorflow::DataType::DT_INVALID:
        case tensorflow::DataType::DT_UINT8:
Khalique's avatar
Khalique committed
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
        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
1137
        // tf pb should not use these types
Paul's avatar
Paul committed
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
        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
1161
        case tensorflow::DataType::DataType_INT_MAX_SENTINEL_DO_NOT_USE_:
Khalique's avatar
Khalique committed
1162
        case tensorflow::DataType::DataType_INT_MIN_SENTINEL_DO_NOT_USE_: break;
Khalique's avatar
Khalique committed
1163
1164
1165
1166
        }
        return shape_type;
    }

Khalique's avatar
Khalique committed
1167
    static literal parse_tensor(const tensorflow::TensorProto& t)
Khalique's avatar
Khalique committed
1168
1169
    {
        std::vector<size_t> dims = parse_dims(t.tensor_shape());
1170
        size_t shape_size = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<size_t>());
Khalique's avatar
Khalique committed
1171
1172
        if(!t.tensor_content().empty()) // has raw data
        {
Khalique's avatar
Khalique committed
1173
            const std::string& s = t.tensor_content();
Khalique's avatar
Khalique committed
1174
1175
            switch(t.dtype())
            {
Khalique's avatar
Khalique committed
1176
1177
            case tensorflow::DataType::DT_FLOAT:
                return literal{{shape::float_type, dims}, s.data()};
Paul's avatar
Paul committed
1178
            case tensorflow::DataType::DT_BOOL:
1179
            case tensorflow::DataType::DT_INT8: return literal{{shape::int8_type, dims}, s.data()};
Khalique's avatar
Khalique committed
1180
1181
            case tensorflow::DataType::DT_UINT16:
            case tensorflow::DataType::DT_INT16:
1182
                return literal{{shape::int16_type, dims}, s.data()};
Khalique's avatar
Khalique committed
1183
1184
1185
1186
            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
1187
            case tensorflow::DataType::DT_HALF: return literal{{shape::half_type, dims}, s.data()};
Khalique's avatar
Khalique committed
1188
1189
            case tensorflow::DataType::DT_DOUBLE:
                return literal{{shape::double_type, dims}, s.data()};
Paul's avatar
Paul committed
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
1227
            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
1228
1229
1230
            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
1231
1232
1233
1234
1235
1236
            }
            MIGRAPHX_THROW("Invalid tensor type");
        }
        switch(t.dtype())
        {
        case tensorflow::DataType::DT_FLOAT:
Khalique's avatar
Khalique committed
1237
1238
            return create_literal(
                shape::float_type, dims, get_data_vals(t.float_val(), shape_size));
Khalique's avatar
Khalique committed
1239
        case tensorflow::DataType::DT_INT8:
1240
            return create_literal(shape::int8_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
1241
        case tensorflow::DataType::DT_UINT16:
1242
            return create_literal(shape::uint16_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
1243
        case tensorflow::DataType::DT_INT16:
1244
            return create_literal(shape::int16_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
1245
        case tensorflow::DataType::DT_INT32:
1246
            return create_literal(shape::int32_type, dims, get_data_vals(t.int_val(), shape_size));
Khalique's avatar
Khalique committed
1247
        case tensorflow::DataType::DT_INT64:
Khalique's avatar
Khalique committed
1248
1249
            return create_literal(
                shape::int64_type, dims, get_data_vals(t.int64_val(), shape_size));
Khalique's avatar
Khalique committed
1250
        case tensorflow::DataType::DT_BOOL:
1251
            return create_literal(shape::int32_type, dims, get_data_vals(t.bool_val(), shape_size));
Khalique's avatar
Khalique committed
1252
        case tensorflow::DataType::DT_HALF:
Khalique's avatar
Khalique committed
1253
        {
1254
1255
            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
1256
1257
1258
1259
1260
            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); });
1261
            return create_literal(shape::half_type, dims, data_half);
Khalique's avatar
Khalique committed
1262
        }
Khalique's avatar
Khalique committed
1263
        case tensorflow::DataType::DT_DOUBLE:
Khalique's avatar
Khalique committed
1264
            return literal{{shape::double_type, dims}, get_data_vals(t.double_val(), shape_size)};
Paul's avatar
Paul committed
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
1302
        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
1303
1304
1305
        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
1306
1307
1308
1309
        }
        MIGRAPHX_THROW("Invalid tensor type");
    }

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

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

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

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx