tf_test.cpp 21.7 KB
Newer Older
1
2
3
#include <iostream>
#include <vector>
#include <migraphx/literal.hpp>
Paul's avatar
Paul committed
4
5
6
7
#include <migraphx/pass_manager.hpp>
#include <migraphx/simplify_reshapes.hpp>
#include <migraphx/dead_code_elimination.hpp>
#include <migraphx/eliminate_identity.hpp>
8
9
10
11
12
13
#include <migraphx/operators.hpp>
#include <migraphx/program.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/tf.hpp>
#include "test.hpp"

Paul's avatar
Paul committed
14
15
16
migraphx::program optimize_tf(const std::string& name, bool is_nhwc)
{
    auto prog = migraphx::parse_tf(name, is_nhwc);
Paul's avatar
Paul committed
17
18
19
20
21
    if(is_nhwc)
        migraphx::run_passes(prog,
                             {migraphx::simplify_reshapes{},
                              migraphx::dead_code_elimination{},
                              migraphx::eliminate_identity{}});
Paul's avatar
Paul committed
22
23
24
    return prog;
}

25
26
27
28
29
30
TEST_CASE(add_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    p.add_instruction(migraphx::op::add{}, l0, l1);
Paul's avatar
Paul committed
31
    auto prog = optimize_tf("add_test.pb", false);
32
33
34
35
36
37

    EXPECT(p == prog);
}

TEST_CASE(add_bcast_test)
{
Khalique's avatar
Khalique committed
38

39
40
41
42
43
44
45
    migraphx::program p;
    migraphx::shape s0{migraphx::shape::float_type, {2, 3}};
    auto l0 = p.add_parameter("0", s0);
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {2, 1}});
    auto l2 = p.add_instruction(migraphx::op::multibroadcast{s0.lens()}, l0);
    auto l3 = p.add_instruction(migraphx::op::multibroadcast{s0.lens()}, l1);
    p.add_instruction(migraphx::op::add{}, l2, l3);
Paul's avatar
Paul committed
46
    auto prog = optimize_tf("add_bcast_test.pb", false);
47
48
49
50

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
TEST_CASE(assert_less_equal_test)
{
    migraphx::program p;
    migraphx::shape s0{migraphx::shape::float_type, {2, 3}};
    auto l0 = p.add_parameter("0", s0);
    auto l1 = p.add_parameter("1", s0);
    migraphx::literal l{migraphx::shape{migraphx::shape::int32_type, {2}}, {0, 1}};
    auto l2 = p.add_literal(l);
    p.add_instruction(migraphx::op::add{}, l0, l1);
    auto l3 = p.add_instruction(migraphx::op::identity{}, l0, l1);
    p.add_instruction(migraphx::op::identity{}, l3, l2);
    auto prog = optimize_tf("assert_less_equal_test.pb", false);

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
TEST_CASE(batchmatmul_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 8, 4}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 4, 8}});

    auto trans_l0 = p.add_instruction(migraphx::op::transpose{{0, 1, 3, 2}}, l0);
    auto trans_l1 = p.add_instruction(migraphx::op::transpose{{0, 1, 3, 2}}, l1);

    p.add_instruction(migraphx::op::dot{}, trans_l0, trans_l1);
    auto prog = optimize_tf("batchmatmul_test.pb", false);

    EXPECT(p == prog);
}

82
83
TEST_CASE(batchnorm_test)
{
Khalique's avatar
Khalique committed
84
85
    float epsilon  = 1.001e-5f;
    float momentum = 0.9f;
86
87

    migraphx::program p;
Khalique's avatar
Khalique committed
88
89
    migraphx::op::batch_norm_inference op{
        epsilon, momentum, migraphx::op::batch_norm_inference::spatial};
90
    migraphx::shape s0{migraphx::shape::float_type, {32}};
91
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 32, 16, 16}});
92
93
    std::vector<float> const_vals(32);
    std::fill(const_vals.begin(), const_vals.end(), 1.0f);
Khalique's avatar
Khalique committed
94

95
96
97
98
99
    auto l2 = p.add_parameter("2", s0);
    auto l3 = p.add_parameter("3", s0);
    auto l4 = p.add_parameter("4", s0);
    auto l1 = p.add_literal(migraphx::literal{s0, const_vals});
    p.add_instruction(op, l0, l1, l2, l3, l4);
Paul's avatar
Paul committed
100
    auto prog = optimize_tf("batchnorm_test.pb", true);
101
102
103
104

    EXPECT(p == prog);
}

105
106
107
TEST_CASE(biasadd_test)
{
    migraphx::program p;
108
    migraphx::shape s0{migraphx::shape::float_type, {1, 500, 1, 1}};
109
    uint64_t axis = 1;
Khalique's avatar
Khalique committed
110
    auto l0       = p.add_parameter("0", s0);
111
    auto l1       = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {500}});
112
    auto l2       = p.add_instruction(migraphx::op::broadcast{axis, l0->get_shape().lens()}, l1);
113
    p.add_instruction(migraphx::op::add{}, l0, l2);
Paul's avatar
Paul committed
114
    auto prog = optimize_tf("biasadd_test.pb", true);
115
116
117
118

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
119
120
121
122
123
124
125
126
127
128
TEST_CASE(cast_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    p.add_instruction(migraphx::op::convert{migraphx::shape::int32_type}, l0);
    auto prog = optimize_tf("cast_test.pb", false);

    EXPECT(p == prog);
}

129
130
131
TEST_CASE(concat_test)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
132

133
134
135
136
137
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {4, 7, 3}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {4, 2, 3}});

    int axis = 1;
    // tf uses axis as the third input, and it is in int32 format
Khalique's avatar
Khalique committed
138
    // add the literal using a vector in order to set stride to 1 (like in tf parser)
139
    p.add_literal(migraphx::shape{migraphx::shape::int32_type}, std::vector<int>{axis});
140
141

    p.add_instruction(migraphx::op::concat{static_cast<std::size_t>(axis)}, l0, l1);
Paul's avatar
Paul committed
142
    auto prog = optimize_tf("concat_test.pb", false);
143
144
145
146
147
148
149

    EXPECT(p == prog);
}

TEST_CASE(const_test)
{
    migraphx::program p;
150
    p.add_literal(migraphx::shape{migraphx::shape::float_type}, std::vector<float>{1.0f});
Paul's avatar
Paul committed
151
    auto prog = optimize_tf("constant_test.pb", false);
152
153
154
155
156
157
158

    EXPECT(p == prog);
}

TEST_CASE(conv_test)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
159

160
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
Khalique's avatar
Khalique committed
161
    std::vector<float> weight_data(3 * 3 * 3 * 32);
162
    std::fill(weight_data.begin(), weight_data.end(), 1.0f);
Khalique's avatar
Khalique committed
163
164
    auto l1 =
        p.add_literal(migraphx::shape{migraphx::shape::float_type, {3, 3, 3, 32}}, weight_data);
165
166
167

    migraphx::op::convolution op;
    op.padding_mode = migraphx::op::padding_mode_t::same;
Khalique's avatar
Khalique committed
168
    op.padding      = {1, 1};
Khalique's avatar
Khalique committed
169
170
    op.stride       = {1, 1};
    op.dilation     = {1, 1};
Paul's avatar
Paul committed
171
172
173
    auto l2         = p.add_instruction(migraphx::op::transpose{{3, 2, 0, 1}}, l1);
    p.add_instruction(op, l0, l2);
    auto prog = optimize_tf("conv_test.pb", true);
174
175
176
177

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
178
179
180
181
182
183
184
185
186
187
188
189
TEST_CASE(depthwiseconv_test)
{
    migraphx::program p;

    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    std::vector<float> weight_data(3 * 3 * 3 * 1);
    std::fill(weight_data.begin(), weight_data.end(), 1.0f);
    auto l1 =
        p.add_literal(migraphx::shape{migraphx::shape::float_type, {3, 3, 3, 1}}, weight_data);

    migraphx::op::convolution op;
    op.padding_mode = migraphx::op::padding_mode_t::same;
Khalique's avatar
Khalique committed
190
    op.padding      = {1, 1};
Khalique's avatar
Khalique committed
191
192
193
    op.stride       = {1, 1};
    op.dilation     = {1, 1};
    op.group        = 3;
Paul's avatar
Paul committed
194
    auto l3         = p.add_instruction(migraphx::op::transpose{{3, 2, 0, 1}}, l1);
Paul's avatar
Paul committed
195
196
197
    auto l4         = p.add_instruction(migraphx::op::contiguous{}, l3);
    auto l5         = p.add_instruction(migraphx::op::reshape{{3, 1, 3, 3}}, l4);
    p.add_instruction(op, l0, l5);
Paul's avatar
Paul committed
198
    auto prog = optimize_tf("depthwise_conv_test.pb", true);
Khalique's avatar
Khalique committed
199
200
201
202

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
203
204
205
TEST_CASE(expanddims_test)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
206
207

    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {2, 3, 4}});
208
    p.add_literal(0);
Khalique's avatar
Khalique committed
209
    p.add_instruction(migraphx::op::reshape{{1, 2, 3, 4}}, l0);
210
    auto prog = optimize_tf("expanddims_test.pb", false);
Khalique's avatar
Khalique committed
211
212
213
214
215
216
217
218
219
220

    EXPECT(p == prog);
}

TEST_CASE(expanddims_test_neg_dims)
{
    // this check makes sure the pb parses negative dim value correctly
    migraphx::program p;

    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {2, 3, 4}});
221
    p.add_literal(-1);
Khalique's avatar
Khalique committed
222
    p.add_instruction(migraphx::op::reshape{{2, 3, 4, 1}}, l0);
223
    auto prog = optimize_tf("expanddims_neg_test.pb", false);
Khalique's avatar
Khalique committed
224
225
226
227

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
228
229
230
231
TEST_CASE(gather_test)
{
    migraphx::program p;

Khalique's avatar
Khalique committed
232
233
234
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {2, 4}});
    auto l1 =
        p.add_literal(migraphx::literal{migraphx::shape{migraphx::shape::int32_type, {2}}, {1, 1}});
Khalique's avatar
Khalique committed
235
236
237
238
239
240
    p.add_literal(1);

    int axis = 1;
    p.add_instruction(migraphx::op::gather{axis}, l0, l1);
    auto prog = optimize_tf("gather_test.pb", false);

Khalique's avatar
Khalique committed
241
242
243
    EXPECT(p == prog);
}

244
245
246
247
248
TEST_CASE(identity_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    p.add_instruction(migraphx::op::identity{}, l0);
Paul's avatar
Paul committed
249
    auto prog = optimize_tf("identity_test.pb", false);
250
251
252
253

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
254
255
256
257
258
TEST_CASE(matmul_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {8, 4}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {4, 8}});
Khalique's avatar
Khalique committed
259
260
261

    auto trans_l0 = p.add_instruction(migraphx::op::transpose{{1, 0}}, l0);
    auto trans_l1 = p.add_instruction(migraphx::op::transpose{{1, 0}}, l1);
Khalique's avatar
Khalique committed
262
263

    p.add_instruction(migraphx::op::dot{}, trans_l0, trans_l1);
Paul's avatar
Paul committed
264
    auto prog = optimize_tf("matmul_test.pb", false);
Khalique's avatar
Khalique committed
265
266
267
268

    EXPECT(p == prog);
}

269
270
271
TEST_CASE(mean_test)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
272
    migraphx::literal l{migraphx::shape{migraphx::shape::int32_type, {2}}, {2, 3}};
273
274
275
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    p.add_literal(l);
    p.add_literal(l);
Paul's avatar
Paul committed
276
    migraphx::op::reduce_mean op{{2, 3}};
Paul's avatar
Paul committed
277
    p.add_instruction(op, l0);
Paul's avatar
Paul committed
278
    auto l3 = p.add_instruction(op, l0);
Khalique's avatar
Khalique committed
279
    p.add_instruction(migraphx::op::squeeze{{2, 3}}, l3);
Paul's avatar
Paul committed
280
    auto prog = optimize_tf("mean_test.pb", false);
281
282
283
284
285
286
287

    EXPECT(p == prog);
}

TEST_CASE(mean_test_nhwc)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
288
    migraphx::literal l{migraphx::shape{migraphx::shape::int32_type, {2}}, {1, 2}};
289
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
Khalique's avatar
Khalique committed
290
    auto l1 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
Khalique's avatar
Khalique committed
291
292
    migraphx::op::reduce_mean op{{1, 2}};
    auto l2 = p.add_instruction(op, l1);
293
    p.add_instruction(migraphx::op::squeeze{{1, 2}}, l2);
Paul's avatar
Paul committed
294
    auto prog = optimize_tf("mean_test_nhwc.pb", true);
295
296
297
298

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
299
300
301
TEST_CASE(mul_test)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
302
303
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 1, 1, 16}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 1, 1, 16}});
Khalique's avatar
Khalique committed
304
305

    p.add_instruction(migraphx::op::mul{}, l0, l1);
Paul's avatar
Paul committed
306
    auto prog = optimize_tf("mul_test.pb", false);
Khalique's avatar
Khalique committed
307
308
309
310

    EXPECT(p == prog);
}

311
312
313
TEST_CASE(onehot_test)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
314
315
    auto l0 = p.add_literal(
        migraphx::literal{migraphx::shape{migraphx::shape::int32_type, {5}}, {1, 1, 1, 1, 1}});
316
317
318
    p.add_literal(2);
    p.add_literal(1.0f);
    p.add_literal(0.0f);
Khalique's avatar
Khalique committed
319
320
    auto l1 = p.add_literal(
        migraphx::literal{migraphx::shape{migraphx::shape::float_type, {2, 2}}, {1, 0, 0, 1}});
321
322
323
324
325
326
327
    int axis = 0;
    p.add_instruction(migraphx::op::gather{axis}, l1, l0);
    auto prog = optimize_tf("onehot_test.pb", false);

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
328
329
330
331
332
333
334
335
336
337
TEST_CASE(pack_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {2}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {2}});
    auto l2 = p.add_parameter("2", migraphx::shape{migraphx::shape::float_type, {2}});
    std::vector<migraphx::instruction_ref> args{l0, l1, l2};
    std::vector<migraphx::instruction_ref> unsqueezed_args;
    int64_t axis = 1;

Khalique's avatar
Khalique committed
338
339
340
341
342
343
    std::transform(args.begin(),
                   args.end(),
                   std::back_inserter(unsqueezed_args),
                   [&](migraphx::instruction_ref arg) {
                       return p.add_instruction(migraphx::op::unsqueeze{{axis}}, arg);
                   });
Khalique's avatar
Khalique committed
344
    p.add_instruction(migraphx::op::concat{static_cast<size_t>(axis)}, unsqueezed_args);
Paul's avatar
Paul committed
345
    auto prog = optimize_tf("pack_test.pb", false);
Khalique's avatar
Khalique committed
346
347
348
349

    EXPECT(p == prog);
}

350
351
352
TEST_CASE(pack_test_nhwc)
{
    migraphx::program p;
Paul's avatar
Paul committed
353
    auto l0  = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
Paul's avatar
Paul committed
354
    auto lt0 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
Paul's avatar
Paul committed
355
    auto l1  = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
Paul's avatar
Paul committed
356
    auto lt1 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l1);
Paul's avatar
Paul committed
357
    auto l2  = p.add_parameter("2", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
Paul's avatar
Paul committed
358
359
    auto lt2 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l2);
    std::vector<migraphx::instruction_ref> args{lt0, lt1, lt2};
360
    std::vector<migraphx::instruction_ref> unsqueezed_args;
Paul's avatar
Paul committed
361
    int64_t nchw_axis = 3;
362
363
364
365
366
367
368
369

    std::transform(args.begin(),
                   args.end(),
                   std::back_inserter(unsqueezed_args),
                   [&](migraphx::instruction_ref arg) {
                       return p.add_instruction(migraphx::op::unsqueeze{{nchw_axis}}, arg);
                   });
    p.add_instruction(migraphx::op::concat{static_cast<size_t>(nchw_axis)}, unsqueezed_args);
Paul's avatar
Paul committed
370
    auto prog = optimize_tf("pack_test_nhwc.pb", true);
371
372
373
374

    EXPECT(p == prog);
}

375
376
377
TEST_CASE(pooling_test)
{
    migraphx::program p;
378
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
379
380
381
382
    migraphx::op::pooling avg_pool_op{"average"};
    migraphx::op::pooling max_pool_op{"max"};
    avg_pool_op.padding_mode = migraphx::op::padding_mode_t::valid;
    max_pool_op.padding_mode = migraphx::op::padding_mode_t::valid;
Khalique's avatar
Khalique committed
383
384
385
386
    avg_pool_op.stride       = {2, 2};
    max_pool_op.stride       = {2, 2};
    avg_pool_op.lengths      = {2, 2};
    max_pool_op.lengths      = {2, 2};
387
    p.add_instruction(max_pool_op, l0);
Paul's avatar
Paul committed
388
    auto prog = optimize_tf("pooling_test.pb", true);
389
390
391
392

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
393
394
395
396
397
398
399
400
401
402
403
TEST_CASE(pow_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    p.add_instruction(migraphx::op::pow{}, l0, l1);
    auto prog = optimize_tf("pow_test.pb", false);

    EXPECT(p == prog);
}

404
405
406
407
408
TEST_CASE(relu_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    p.add_instruction(migraphx::op::relu{}, l0);
Paul's avatar
Paul committed
409
    auto prog = optimize_tf("relu_test.pb", false);
410
411
412
413

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
414
415
416
417
418
TEST_CASE(relu6_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    p.add_instruction(migraphx::op::clip{6.0, 0.0}, l0);
Paul's avatar
Paul committed
419
    auto prog = optimize_tf("relu6_test.pb", false);
Khalique's avatar
Khalique committed
420
421
422
423

    EXPECT(p == prog);
}

424
425
426
427
428
429
TEST_CASE(reshape_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {16}});
    migraphx::shape s0{migraphx::shape::int32_type, {4}};
    // in tf, the second arg is a literal that contains new dimensions
Khalique's avatar
Khalique committed
430
431
    p.add_literal(migraphx::literal{s0, {1, 1, 1, 16}});
    p.add_instruction(migraphx::op::reshape{{1, 1, 1, 16}}, l0);
Paul's avatar
Paul committed
432
    auto prog = optimize_tf("reshape_test.pb", false);
433
434
435
436

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
437
438
439
440
441
442
443
444
445
446
TEST_CASE(rsqrt_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    p.add_instruction(migraphx::op::rsqrt{}, l0);
    auto prog = optimize_tf("rsqrt_test.pb", false);

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
447
448
449
450
TEST_CASE(slice_test)
{
    migraphx::program p;
    std::size_t num_axes = 2;
Khalique's avatar
Khalique committed
451
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {5, 10}});
Khalique's avatar
Khalique committed
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
    migraphx::shape s0{migraphx::shape::int32_type, {num_axes}};
    p.add_literal(migraphx::literal{s0, {1, 0}});
    p.add_literal(migraphx::literal{s0, {2, -1}});

    migraphx::op::slice op;
    op.starts = {1, 0};
    op.ends   = {3, 10};
    op.axes   = std::vector<int64_t>(num_axes);
    std::iota(op.axes.begin(), op.axes.end(), 0);
    p.add_instruction(op, l0);
    auto prog = optimize_tf("slice_test.pb", false);

    EXPECT(p == prog);
}

467
468
469
TEST_CASE(softmax_test)
{
    migraphx::program p;
Shucai Xiao's avatar
Shucai Xiao committed
470
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3}});
471
    p.add_instruction(migraphx::op::softmax{1}, l0);
Paul's avatar
Paul committed
472
    auto prog = optimize_tf("softmax_test.pb", false);
473
474
475
476

    EXPECT(p == prog);
}

kahmed10's avatar
kahmed10 committed
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
TEST_CASE(split_test)
{
    migraphx::program p;
    std::vector<int64_t> axes{0, 1};
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {5, 30}});
    p.add_literal(3); // num_splits
    p.add_literal(1); // split axis
    p.add_literal(1); // concat axis
    p.add_literal(1); // concat axis
    auto l1 = p.add_instruction(migraphx::op::slice{axes, {0, 0}, {5, 10}}, l0);
    auto l2 = p.add_instruction(migraphx::op::slice{axes, {0, 10}, {5, 20}}, l0);
    auto l3 = p.add_instruction(migraphx::op::slice{axes, {0, 20}, {5, 30}}, l0);
    p.add_instruction(migraphx::op::concat{1}, l1, l2);
    p.add_instruction(migraphx::op::concat{1}, l2, l3);

    auto prog = migraphx::parse_tf("split_test.pb", false);

    EXPECT(p == prog);
}

TEST_CASE(split_test_one_output)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {5, 30}});
    p.add_literal(1); // num_splits
    p.add_literal(1); // split axis
    p.add_instruction(migraphx::op::identity{}, l0);

    auto prog = migraphx::parse_tf("split_test_one_output.pb", false);

    EXPECT(p == prog);
}

TEST_CASE(split_test_vector_as_input)
{
    migraphx::program p;
    std::vector<int64_t> axes{0, 1};
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {5, 30}});
    // split sizes
    p.add_literal(
        migraphx::literal{migraphx::shape{migraphx::shape::int32_type, {3}}, {4, 15, 11}});
    p.add_literal(1); // split axis
    p.add_literal(1); // concat axis
    p.add_literal(1); // concat axis
    auto l1 = p.add_instruction(migraphx::op::slice{axes, {0, 0}, {5, 4}}, l0);
    auto l2 = p.add_instruction(migraphx::op::slice{axes, {0, 4}, {5, 19}}, l0);
    auto l3 = p.add_instruction(migraphx::op::slice{axes, {0, 19}, {5, 30}}, l0);
    p.add_instruction(migraphx::op::concat{1}, l1, l2);
    p.add_instruction(migraphx::op::concat{1}, l2, l3);

    auto prog = migraphx::parse_tf("split_test_vector_as_input.pb", false);

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
532
533
534
535
536
537
538
539
540
541
542
TEST_CASE(sqdiff_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    p.add_instruction(migraphx::op::sqdiff{}, l0, l1);
    auto prog = optimize_tf("sqdiff_test.pb", false);

    EXPECT(p == prog);
}

543
544
545
TEST_CASE(squeeze_test)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
546
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 3, 1}});
547
    p.add_instruction(migraphx::op::squeeze{{0, 3}}, l0);
Paul's avatar
Paul committed
548
    auto prog = optimize_tf("squeeze_test.pb", false);
549
550
551

    EXPECT(p == prog);
}
Khalique's avatar
Khalique committed
552

Khalique's avatar
Khalique committed
553
554
555
556
557
558
559
TEST_CASE(stopgradient_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    p.add_instruction(migraphx::op::identity{}, l0);
    auto prog = optimize_tf("stopgradient_test.pb", false);

Khalique's avatar
Khalique committed
560
    EXPECT(p == prog);
Khalique's avatar
Khalique committed
561
562
}

Khalique's avatar
Khalique committed
563
564
565
566
TEST_CASE(stridedslice_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 10, 1, 1}});
Khalique's avatar
Khalique committed
567
    auto l1 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
Khalique's avatar
Khalique committed
568
569
    std::size_t num_axes = 4;
    migraphx::op::slice op;
Khalique's avatar
Khalique committed
570
    op.starts = {0, 0, 0, 0};
Paul's avatar
Paul committed
571
    op.ends   = {1, 1, 1, 5};
Khalique's avatar
Khalique committed
572
573
    op.axes   = std::vector<int64_t>(num_axes);
    std::iota(op.axes.begin(), op.axes.end(), 0);
574
    auto l2          = p.add_instruction(op, l1);
Paul's avatar
Paul committed
575
    auto shrink_axis = 1;
576
    p.add_instruction(migraphx::op::squeeze{{shrink_axis}}, l2);
Paul's avatar
Paul committed
577
    auto prog = optimize_tf("stridedslice_test.pb", true);
Khalique's avatar
Khalique committed
578
579
580
581

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
582
583
584
585
586
587
TEST_CASE(stridedslice_masks_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 10, 3, 3}});
    std::size_t num_axes = 4;
    migraphx::op::slice op;
588
589
    op.starts = {0, 1, 1, 0};
    op.ends   = {1, 3, 3, 10};
Khalique's avatar
Khalique committed
590
591
592
593
594
595
596
    op.axes   = std::vector<int64_t>(num_axes);
    std::iota(op.axes.begin(), op.axes.end(), 0);
    // add literals for starts, ends, and strides in tf (NHWC format)
    p.add_literal(migraphx::shape{migraphx::shape::int32_type, {4}}, std::vector<int>{0, 1, 1, 0});
    p.add_literal(migraphx::shape{migraphx::shape::int32_type, {4}}, std::vector<int>{0, 0, 0, 0});
    p.add_literal(migraphx::shape{migraphx::shape::int32_type, {4}}, std::vector<int>{1, 1, 1, 1});

Khalique's avatar
Khalique committed
597
    auto l1 = p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
598
    auto l2 = p.add_instruction(op, l1);
Khalique's avatar
Khalique committed
599
    p.add_instruction(migraphx::op::transpose{{0, 3, 1, 2}}, l2);
Khalique's avatar
Khalique committed
600
601
602
603
604
    auto prog = migraphx::parse_tf("stridedslice_masks_test.pb", true);

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
605
606
607
608
609
610
611
612
613
614
615
TEST_CASE(sub_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    p.add_instruction(migraphx::op::sub{}, l0, l1);
    auto prog = migraphx::parse_tf("sub_test.pb", false);

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
616
617
618
TEST_CASE(tanh_test)
{
    migraphx::program p;
Khalique's avatar
Khalique committed
619
620
621
622
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 2, 3}});
    p.add_instruction(migraphx::op::sub{}, l0, l1);
    auto prog = migraphx::parse_tf("sub_test.pb", false);
Khalique's avatar
Khalique committed
623
624
625
626

    EXPECT(p == prog);
}

Khalique's avatar
Khalique committed
627
628
629
630
631
TEST_CASE(transpose_test)
{
    migraphx::program p;
    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
    migraphx::shape s0{migraphx::shape::int32_type, {4}};
Khalique's avatar
Khalique committed
632
633
    p.add_literal(migraphx::literal{s0, {0, 2, 3, 1}});
    p.add_instruction(migraphx::op::transpose{{0, 2, 3, 1}}, l0);
Khalique's avatar
Khalique committed
634
635
636
637
638
    auto prog = optimize_tf("transpose_test.pb", false);

    EXPECT(p == prog);
}

639
int main(int argc, const char* argv[]) { test::run(argc, argv); }