simplify_reshapes_test.cpp 44.1 KB
Newer Older
Paul's avatar
Paul committed
1
2
#include <migraphx/simplify_reshapes.hpp>
#include <migraphx/dead_code_elimination.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
3
#include <migraphx/eliminate_contiguous.hpp>
4
#include <migraphx/pass_manager.hpp>
Paul's avatar
Paul committed
5
#include <migraphx/operators.hpp>
Paul's avatar
Paul committed
6
#include <migraphx/instruction.hpp>
Paul Fultz II's avatar
Paul Fultz II committed
7
#include <migraphx/generate.hpp>
Paul's avatar
Paul committed
8
#include <basic_ops.hpp>
9
10
11
12
#include <migraphx/make_op.hpp>

#include <migraphx/serialize.hpp>

Paul's avatar
Paul committed
13
14
#include <test.hpp>

Paul Fultz II's avatar
Paul Fultz II committed
15
void run_pass(migraphx::module& m)
Paul's avatar
Paul committed
16
{
Shucai Xiao's avatar
Shucai Xiao committed
17
18
19
    migraphx::run_passes(m,
                         {migraphx::simplify_reshapes{},
                          migraphx::dead_code_elimination{}});
20
}
Paul's avatar
Paul committed
21

Paul's avatar
Paul committed
22
TEST_CASE(double_contig)
Paul's avatar
Paul committed
23
{
Paul's avatar
Paul committed
24
    migraphx::program p;
25
    auto* mm = p.get_main_module();
Paul Fultz II's avatar
Paul Fultz II committed
26
27

    auto l  = mm->add_literal(get_2x2());
28
    auto t1 = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), l);
Paul Fultz II's avatar
Paul Fultz II committed
29
30
    auto c1 = mm->add_instruction(migraphx::make_op("contiguous"), t1);
    auto c2 = mm->add_instruction(migraphx::make_op("contiguous"), c1);
31
    mm->add_return({c2});
Paul Fultz II's avatar
Paul Fultz II committed
32
33
34
35
36
    EXPECT(mm->get_output_shapes().back().standard());
    EXPECT(not mm->get_output_shapes().back().transposed());
    run_pass(*mm);
    EXPECT(mm->get_output_shapes().back().standard());
    EXPECT(not mm->get_output_shapes().back().transposed());
Shucai Xiao's avatar
Shucai Xiao committed
37
    EXPECT(std::distance(mm->begin(), mm->end()) == 4);
38
    auto result = p.eval({}).back();
Paul's avatar
Paul committed
39
    EXPECT(result != get_2x2());
Paul's avatar
Paul committed
40
41
}

Paul's avatar
Paul committed
42
TEST_CASE(double_transpose)
Paul's avatar
Paul committed
43
{
Paul's avatar
Paul committed
44
    migraphx::program p;
45
    auto* mm = p.get_main_module();
Paul Fultz II's avatar
Paul Fultz II committed
46
47

    auto l  = mm->add_literal(get_2x2());
48
49
    auto t1 = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), l);
    auto t2 = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), t1);
50
    mm->add_return({t2});
Paul Fultz II's avatar
Paul Fultz II committed
51
52
53
54
55
    EXPECT(mm->get_output_shapes().back().standard());
    EXPECT(not mm->get_output_shapes().back().transposed());
    run_pass(*mm);
    EXPECT(mm->get_output_shapes().back().standard());
    EXPECT(not mm->get_output_shapes().back().transposed());
Shucai Xiao's avatar
Shucai Xiao committed
56
    EXPECT(std::distance(mm->begin(), mm->end()) == 2);
57
    auto result = p.eval({}).back();
Paul's avatar
Paul committed
58
59
60
    EXPECT(result == get_2x2());
}

Paul's avatar
Paul committed
61
TEST_CASE(double_transpose_contig)
Paul's avatar
Paul committed
62
{
Paul's avatar
Paul committed
63
    migraphx::program p;
64
    auto* mm = p.get_main_module();
Paul Fultz II's avatar
Paul Fultz II committed
65
66

    auto l  = mm->add_literal(get_2x2());
67
    auto t1 = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), l);
Paul Fultz II's avatar
Paul Fultz II committed
68
    auto c1 = mm->add_instruction(migraphx::make_op("contiguous"), t1);
69
    auto t2 = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), c1);
Paul Fultz II's avatar
Paul Fultz II committed
70
    auto c2 = mm->add_instruction(migraphx::make_op("contiguous"), t2);
71
    mm->add_return({c2});
Paul Fultz II's avatar
Paul Fultz II committed
72
73
74
75
76
    EXPECT(mm->get_output_shapes().back().standard());
    EXPECT(not mm->get_output_shapes().back().transposed());
    run_pass(*mm);
    EXPECT(mm->get_output_shapes().back().standard());
    EXPECT(not mm->get_output_shapes().back().transposed());
Shucai Xiao's avatar
Shucai Xiao committed
77
    EXPECT(std::distance(mm->begin(), mm->end()) == 2);
78
    auto result = p.eval({}).back();
Paul's avatar
Paul committed
79
80
81
    EXPECT(result == get_2x2());
}

Paul's avatar
Paul committed
82
TEST_CASE(single_transpose)
Paul's avatar
Paul committed
83
{
Paul's avatar
Paul committed
84
    migraphx::program p;
85
    auto* mm = p.get_main_module();
Paul Fultz II's avatar
Paul Fultz II committed
86
87

    auto l  = mm->add_literal(get_2x2());
88
    auto t1 = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), l);
89
    mm->add_return({t1});
Paul Fultz II's avatar
Paul Fultz II committed
90
91
92
93
94
    EXPECT(not mm->get_output_shapes().back().standard());
    EXPECT(mm->get_output_shapes().back().transposed());
    run_pass(*mm);
    EXPECT(not mm->get_output_shapes().back().standard());
    EXPECT(mm->get_output_shapes().back().transposed());
Shucai Xiao's avatar
Shucai Xiao committed
95
    EXPECT(std::distance(mm->begin(), mm->end()) == 3);
96
    auto result = p.eval({}).back();
Paul's avatar
Paul committed
97
98
99
    EXPECT(result != get_2x2());
}

Paul's avatar
Paul committed
100
TEST_CASE(double_transpose_sin_pass)
Paul's avatar
Paul committed
101
{
Paul's avatar
Paul committed
102
    migraphx::program p;
103
    auto* mm = p.get_main_module();
Paul Fultz II's avatar
Paul Fultz II committed
104
105

    auto l  = mm->add_literal(get_2x2());
106
107
    auto t1 = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), l);
    mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), t1);
Paul Fultz II's avatar
Paul Fultz II committed
108
109
110
111
112
    EXPECT(mm->get_output_shapes().back().standard());
    EXPECT(not mm->get_output_shapes().back().transposed());
    run_pass(*mm);
    EXPECT(mm->get_output_shapes().back().standard());
    EXPECT(not mm->get_output_shapes().back().transposed());
Paul's avatar
Paul committed
113
    // TODO: Fix this
Shucai Xiao's avatar
Shucai Xiao committed
114
    // EXPECT(std::distance(mm->begin(), mm->end()) == 1);
115
    auto result = p.eval({}).back();
Paul's avatar
Paul committed
116
117
118
    EXPECT(result == get_2x2());
}

Paul's avatar
Paul committed
119
TEST_CASE(single_transpose_sin_pass)
Paul's avatar
Paul committed
120
{
Paul's avatar
Paul committed
121
    migraphx::program p;
122
    auto* mm = p.get_main_module();
Paul Fultz II's avatar
Paul Fultz II committed
123
124

    auto l = mm->add_literal(get_2x2());
125
    mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), l);
Paul Fultz II's avatar
Paul Fultz II committed
126
127
128
129
130
    EXPECT(not mm->get_output_shapes().back().standard());
    EXPECT(mm->get_output_shapes().back().transposed());
    run_pass(*mm);
    EXPECT(not mm->get_output_shapes().back().standard());
    EXPECT(mm->get_output_shapes().back().transposed());
Shucai Xiao's avatar
Shucai Xiao committed
131
    EXPECT(std::distance(mm->begin(), mm->end()) == 2);
132
    auto result = p.eval({}).back();
Paul's avatar
Paul committed
133
134
135
    EXPECT(result != get_2x2());
}

Paul's avatar
Paul committed
136
137
TEST_CASE(reshape_transpose)
{
Paul Fultz II's avatar
Paul Fultz II committed
138
139
140
141
142
    migraphx::module m;

    auto s  = migraphx::shape{migraphx::shape::float_type, {1, 112, 56, 56}};
    auto x  = m.add_parameter("x", s);
    auto r1 = m.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 4, 28, 56, 56}}}), x);
143
144
    auto t =
        m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 1, 3, 4}}}), r1);
Paul Fultz II's avatar
Paul Fultz II committed
145
146
147
148
149
150
151
152
    auto ct = m.add_instruction(migraphx::make_op("contiguous"), t);
    auto r2 = m.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 112, 56, 56}}}), ct);
    m.add_return({r2});
    EXPECT(m.get_output_shapes().back() == s);
    auto n = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == s);
    EXPECT(std::distance(m.begin(), m.end()) == n);
Paul's avatar
Paul committed
153
154
}

Paul's avatar
Paul committed
155
156
TEST_CASE(transpose_contiguous)
{
Paul Fultz II's avatar
Paul Fultz II committed
157
158
159
160
    migraphx::module m;

    auto s  = migraphx::shape{migraphx::shape::float_type, {4, 4}};
    auto x  = m.add_parameter("x", s);
161
    auto t  = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
162
163
164
165
166
167
168
    auto c1 = m.add_instruction(migraphx::make_op("contiguous"), t);
    m.add_return({c1});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n);
Paul's avatar
Paul committed
169
170
171
172
}

TEST_CASE(transpose_double_contiguous)
{
Paul Fultz II's avatar
Paul Fultz II committed
173
174
175
176
    migraphx::module m;

    auto s  = migraphx::shape{migraphx::shape::float_type, {4, 4}};
    auto x  = m.add_parameter("x", s);
177
    auto t  = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
178
179
180
181
182
183
184
185
186
    auto c1 = m.add_instruction(migraphx::make_op("contiguous"), t);
    auto c2 = m.add_instruction(migraphx::make_op("contiguous"), c1);
    m.add_return({c2});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n - 1);
    EXPECT(m.has_instruction(t));
Paul's avatar
Paul committed
187
188
}

189
190
TEST_CASE(transpose_partial1)
{
Paul Fultz II's avatar
Paul Fultz II committed
191
192
193
194
    migraphx::module m;

    auto s  = migraphx::shape{migraphx::shape::float_type, {1, 2, 3}};
    auto x  = m.add_parameter("x", s);
195
196
    auto t1 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}}), x);
    auto t2 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 2, 0}}}), t1);
Paul Fultz II's avatar
Paul Fultz II committed
197
198
199
200
201
202
    m.add_return({t2});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n - 1);
203
204
205
206
}

TEST_CASE(transpose_partial2)
{
Paul Fultz II's avatar
Paul Fultz II committed
207
208
209
210
    migraphx::module m;

    auto s  = migraphx::shape{migraphx::shape::float_type, {1, 2, 3}};
    auto x  = m.add_parameter("x", s);
211
212
213
    auto t1 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}}), x);
    auto t2 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 2, 0}}}), t1);
    auto t3 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}}), t2);
Paul Fultz II's avatar
Paul Fultz II committed
214
215
216
217
218
219
    m.add_return({t3});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n - 2);
220
221
222
223
}

TEST_CASE(transpose_partial3)
{
Paul Fultz II's avatar
Paul Fultz II committed
224
225
226
227
    migraphx::module m;

    auto s  = migraphx::shape{migraphx::shape::float_type, {1, 2, 3}};
    auto x  = m.add_parameter("x", s);
228
229
230
231
    auto t1 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}}), x);
    auto t2 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 2, 0}}}), t1);
    auto t3 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}}), t2);
    auto t4 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}}), t3);
Paul Fultz II's avatar
Paul Fultz II committed
232
233
234
235
236
237
    m.add_return({t4});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n - 3);
238
239
}

Paul's avatar
Paul committed
240
241
TEST_CASE(nop_transpose1)
{
Paul Fultz II's avatar
Paul Fultz II committed
242
243
244
245
    migraphx::module m;

    auto s = migraphx::shape{migraphx::shape::float_type, {1, 2, 3}};
    auto x = m.add_parameter("x", s);
246
    auto t = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 2}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
247
248
249
250
251
252
    m.add_return({t});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n - 1);
Paul's avatar
Paul committed
253
254
255
256
}

TEST_CASE(nop_transpose2)
{
Paul Fultz II's avatar
Paul Fultz II committed
257
258
259
260
    migraphx::module m;

    auto s  = migraphx::shape{migraphx::shape::float_type, {1, 2, 3}};
    auto x  = m.add_parameter("x", s);
261
262
263
264
    auto t1 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 2}}}), x);
    auto t2 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 2}}}), t1);
    auto t3 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 2}}}), t2);
    auto t4 = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 2}}}), t3);
Paul Fultz II's avatar
Paul Fultz II committed
265
266
267
268
269
270
    m.add_instruction(pass_op{}, t4);
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n - 4);
Paul's avatar
Paul committed
271
272
273
274
}

TEST_CASE(nop_transpose3)
{
Paul Fultz II's avatar
Paul Fultz II committed
275
    migraphx::module m;
276

Paul's avatar
Paul committed
277
    auto s      = migraphx::shape{migraphx::shape::float_type, {1, 2, 3, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
278
279
280
    auto x      = m.add_parameter("x", s);
    auto y      = m.add_parameter("y", s);
    auto concat = m.add_instruction(migraphx::make_op("concat", {{"axis", 3}}), x, y);
281
282
283
284
    auto t1 =
        m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 2, 3}}}), concat);
    auto t2 =
        m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2}}}), t1);
Paul Fultz II's avatar
Paul Fultz II committed
285
286
287
288
289
290
    m.add_return({t2});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n - 1);
Paul's avatar
Paul committed
291
}
Paul Fultz II's avatar
Paul Fultz II committed
292
293
294

TEST_CASE(nop_convert)
{
Paul Fultz II's avatar
Paul Fultz II committed
295
    migraphx::module m;
296

Paul Fultz II's avatar
Paul Fultz II committed
297
298
299
    auto s = migraphx::shape{migraphx::shape::float_type, {1, 2, 3}};
    auto x = m.add_parameter("x", s);
    auto t = m.add_instruction(
300
301
302
        migraphx::make_op("convert",
                          {{"target_type", migraphx::to_value(migraphx::shape::float_type)}}),
        x);
Paul Fultz II's avatar
Paul Fultz II committed
303
304
305
306
307
308
    m.add_return({t});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back() == out_shape);
    EXPECT(std::distance(m.begin(), m.end()) == n - 1);
Paul Fultz II's avatar
Paul Fultz II committed
309
}
Paul's avatar
Paul committed
310
311
312

TEST_CASE(concat_transpose1)
{
Paul Fultz II's avatar
Paul Fultz II committed
313
    migraphx::module m;
314

315
316
317
318
319
    auto s  = migraphx::shape{migraphx::shape::float_type, {1, 2, 3, 4}};
    auto x  = m.add_parameter("x", s);
    auto y  = m.add_parameter("y", s);
    auto xt = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2}}}), x);
    auto yt = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
320
    auto concat = m.add_instruction(migraphx::make_op("concat", {{"axis", 2}}), xt, yt);
321
322
    auto t =
        m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2}}}), concat);
Paul Fultz II's avatar
Paul Fultz II committed
323
324
325
326
327
328
    m.add_return({t});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back().lens() == out_shape.lens());
    EXPECT(std::distance(m.begin(), m.end()) == n - 3);
Paul's avatar
Paul committed
329
    auto new_concat =
Paul Fultz II's avatar
Paul Fultz II committed
330
331
        std::find_if(m.begin(), m.end(), [](auto ins) { return ins.name() == "concat"; });
    EXPECT(bool{new_concat != m.end()});
Paul's avatar
Paul committed
332
333
334
    EXPECT(migraphx::any_cast<migraphx::op::concat>(new_concat->get_operator()).axis == 3);
}

Paul's avatar
Paul committed
335
336
TEST_CASE(concat_transpose2)
{
Paul Fultz II's avatar
Paul Fultz II committed
337
    migraphx::module m;
338

339
340
341
342
343
    auto s  = migraphx::shape{migraphx::shape::float_type, {1, 2, 3, 4}};
    auto x  = m.add_parameter("x", s);
    auto y  = m.add_parameter("y", s);
    auto xt = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), x);
    auto yt = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
344
    auto concat = m.add_instruction(migraphx::make_op("concat", {{"axis", -1}}), xt, yt);
345
346
    auto t =
        m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), concat);
Paul Fultz II's avatar
Paul Fultz II committed
347
348
349
350
351
352
    m.add_return({t});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back().lens() == out_shape.lens());
    EXPECT(std::distance(m.begin(), m.end()) == n - 2);
Paul's avatar
Paul committed
353
    auto new_concat =
Paul Fultz II's avatar
Paul Fultz II committed
354
355
        std::find_if(m.begin(), m.end(), [](auto ins) { return ins.name() == "concat"; });
    EXPECT(bool{new_concat != m.end()});
Paul's avatar
Paul committed
356
357
358
    EXPECT(migraphx::any_cast<migraphx::op::concat>(new_concat->get_operator()).axis == 1);
}

359
360
TEST_CASE(concat_transpose3)
{
Paul Fultz II's avatar
Paul Fultz II committed
361
    migraphx::module m;
362

363
364
365
366
367
    auto s  = migraphx::shape{migraphx::shape::float_type, {1, 2, 3, 4}};
    auto x  = m.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 2, 3, 4}});
    auto y  = m.add_parameter("y", migraphx::shape{migraphx::shape::float_type, {1, 5, 3, 4}});
    auto xt = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), x);
    auto yt = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
368
    auto concat = m.add_instruction(migraphx::make_op("concat", {{"axis", 3}}), xt, yt);
369
370
    auto t =
        m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), concat);
Paul Fultz II's avatar
Paul Fultz II committed
371
372
373
374
375
376
    m.add_return({t});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back().lens() == out_shape.lens());
    EXPECT(std::distance(m.begin(), m.end()) == n - 2);
377
    auto new_concat =
Paul Fultz II's avatar
Paul Fultz II committed
378
379
        std::find_if(m.begin(), m.end(), [](auto ins) { return ins.name() == "concat"; });
    EXPECT(bool{new_concat != m.end()});
380
381
382
    EXPECT(migraphx::any_cast<migraphx::op::concat>(new_concat->get_operator()).axis == 1);
}

Shucai Xiao's avatar
Shucai Xiao committed
383
384
TEST_CASE(concat_transpose4)
{
Paul Fultz II's avatar
Paul Fultz II committed
385
    migraphx::module m;
386
387
388
389
390
391
    auto sx = migraphx::shape{migraphx::shape::float_type, {1, 1, 12, 64}};
    auto sy = migraphx::shape{migraphx::shape::float_type, {1, 12, 1, 64}};
    auto x  = m.add_parameter("x", sx);
    auto y  = m.add_parameter("y", sy);
    auto xt = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), x);
    auto yt = m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
392
    auto concat = m.add_instruction(migraphx::make_op("concat", {{"axis", 3}}), xt, yt);
393
394
    auto t =
        m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), concat);
Paul Fultz II's avatar
Paul Fultz II committed
395
396
397
398
399
400
    m.add_return({t});

    migraphx::module m1 = m;
    run_pass(m);

    EXPECT(m1 == m);
Shucai Xiao's avatar
Shucai Xiao committed
401
402
}

Paul Fultz II's avatar
Paul Fultz II committed
403
404
TEST_CASE(nested_concat)
{
Paul Fultz II's avatar
Paul Fultz II committed
405
    migraphx::module m;
406

Paul Fultz II's avatar
Paul Fultz II committed
407
    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 2, 3, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
408
409
410
411
412
413
414
415
416
417
418
419
    auto x       = m.add_parameter("x", s);
    auto y       = m.add_parameter("y", s);
    auto concat1 = m.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), x, y);
    auto concat2 = m.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), y, x);
    auto concat3 = m.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), concat1, concat2);
    m.add_return({concat3});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back().lens() == out_shape.lens());
    EXPECT(std::distance(m.begin(), m.end()) == n - 2);
    EXPECT(std::count_if(m.begin(), m.end(), [](auto ins) { return ins.name() == "concat"; }) == 1);
Paul Fultz II's avatar
Paul Fultz II committed
420
421
422
423
}

TEST_CASE(nested_concat_partial)
{
Paul Fultz II's avatar
Paul Fultz II committed
424
    migraphx::module m;
425

Paul Fultz II's avatar
Paul Fultz II committed
426
427
428
429
    auto s = migraphx::shape{migraphx::shape::float_type, {1, 2, 3, 4}};
    auto x = m.add_parameter("x", s);
    auto y = m.add_parameter("y", s);
    auto l = m.add_literal(
Paul Fultz II's avatar
Paul Fultz II committed
430
        migraphx::generate_literal(migraphx::shape{migraphx::shape::float_type, {1, 4, 3, 4}}));
Paul Fultz II's avatar
Paul Fultz II committed
431
432
    auto concat1 = m.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), x, y);
    auto concat2 = m.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), y, x);
433
    auto concat3 =
Paul Fultz II's avatar
Paul Fultz II committed
434
435
436
437
438
439
440
441
        m.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), concat1, concat2, l);
    m.add_return({concat3});
    auto out_shape = m.get_output_shapes().back();
    auto n         = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(m.get_output_shapes().back().lens() == out_shape.lens());
    EXPECT(std::distance(m.begin(), m.end()) == n - 2);
    EXPECT(std::count_if(m.begin(), m.end(), [](auto ins) { return ins.name() == "concat"; }) == 1);
Paul Fultz II's avatar
Paul Fultz II committed
442
443
}

444
445
TEST_CASE(multibroadcast_simplify)
{
Paul Fultz II's avatar
Paul Fultz II committed
446
    migraphx::module m;
447

448
449
    std::vector<size_t> s_lens{1, 2, 3, 4};
    auto s = migraphx::shape{migraphx::shape::float_type, s_lens};
Paul Fultz II's avatar
Paul Fultz II committed
450
    auto x = m.add_parameter("x", s);
451
    auto y = m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s_lens}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
452
453
454
455
    m.add_instruction(migraphx::make_op("mul"), y, y);
    auto n = std::distance(m.begin(), m.end());
    run_pass(m);
    EXPECT(std::distance(m.begin(), m.end()) == n - 1);
456
457
}

458
459
TEST_CASE(double_slice1)
{
Paul Fultz II's avatar
Paul Fultz II committed
460
    migraphx::module m1;
461
    {
Paul Fultz II's avatar
Paul Fultz II committed
462
463
        auto x      = m1.add_parameter("x", {migraphx::shape::int32_type, {256}});
        auto slice1 = m1.add_instruction(
464
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {32}}, {"ends", {256}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
465
        auto slice2 = m1.add_instruction(
466
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {32}}, {"ends", {64}}}), slice1);
Paul Fultz II's avatar
Paul Fultz II committed
467
        m1.add_return({slice2});
468
    }
Paul Fultz II's avatar
Paul Fultz II committed
469
    run_pass(m1);
470

Paul Fultz II's avatar
Paul Fultz II committed
471
    migraphx::module m2;
472
    {
Paul Fultz II's avatar
Paul Fultz II committed
473
474
        auto x     = m2.add_parameter("x", {migraphx::shape::int32_type, {256}});
        auto slice = m2.add_instruction(
475
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {64}}, {"ends", {96}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
476
        m2.add_return({slice});
477
    }
Paul Fultz II's avatar
Paul Fultz II committed
478
    EXPECT(m1 == m2);
479
480
481
482
}

TEST_CASE(double_slice2)
{
Paul Fultz II's avatar
Paul Fultz II committed
483
    migraphx::module m1;
484
    {
Paul Fultz II's avatar
Paul Fultz II committed
485
486
        auto x      = m1.add_parameter("x", {migraphx::shape::int32_type, {256}});
        auto slice1 = m1.add_instruction(
487
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {32}}, {"ends", {128}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
488
        auto slice2 = m1.add_instruction(
489
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {32}}}), slice1);
Paul Fultz II's avatar
Paul Fultz II committed
490
        m1.add_return({slice2});
491
    }
Paul Fultz II's avatar
Paul Fultz II committed
492
    run_pass(m1);
493

Paul Fultz II's avatar
Paul Fultz II committed
494
    migraphx::module m2;
495
    {
Paul Fultz II's avatar
Paul Fultz II committed
496
497
        auto x     = m2.add_parameter("x", {migraphx::shape::int32_type, {256}});
        auto slice = m2.add_instruction(
498
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {32}}, {"ends", {64}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
499
        m2.add_return({slice});
500
    }
Paul Fultz II's avatar
Paul Fultz II committed
501
    EXPECT(m1 == m2);
502
503
504
505
}

TEST_CASE(double_slice_multi_axes)
{
Paul Fultz II's avatar
Paul Fultz II committed
506
    migraphx::module m1;
507
    {
Paul Fultz II's avatar
Paul Fultz II committed
508
509
        auto x      = m1.add_parameter("x", {migraphx::shape::int32_type, {256, 128}});
        auto slice1 = m1.add_instruction(
510
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {32}}, {"ends", {128}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
511
        auto slice2 = m1.add_instruction(
512
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {32}}}), slice1);
Paul Fultz II's avatar
Paul Fultz II committed
513
        m1.add_return({slice2});
514
    }
Paul Fultz II's avatar
Paul Fultz II committed
515
    run_pass(m1);
516

Paul Fultz II's avatar
Paul Fultz II committed
517
    migraphx::module m2;
518

519
    {
Paul Fultz II's avatar
Paul Fultz II committed
520
521
        auto x     = m2.add_parameter("x", {migraphx::shape::int32_type, {256, 128}});
        auto slice = m2.add_instruction(
522
523
524
            migraphx::make_op("slice",
                              {{"axes", {0, 1}}, {"starts", {32, 0}}, {"ends", {128, 32}}}),
            x);
Paul Fultz II's avatar
Paul Fultz II committed
525
        m2.add_return({slice});
526
    }
Paul Fultz II's avatar
Paul Fultz II committed
527
    EXPECT(m1 == m2);
528
529
}

530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
TEST_CASE(optimize_resize)
{
    migraphx::shape sx{migraphx::shape::float_type, {1, 1, 2, 2}};
    auto create_resize_module = [&] {
        migraphx::module m;
        auto inx = m.add_parameter("X", sx);

        migraphx::shape si{migraphx::shape::int32_type, {1, 2, 4, 6}};
        std::vector<int> ind = {0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3,
                                3, 3, 2, 2, 2, 3, 3, 3, 0, 0, 0, 1, 1, 1, 0, 0,
                                0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 2, 2, 2, 3, 3, 3};
        auto li              = m.add_literal(migraphx::literal(si, ind));

        auto lrsp = m.add_instruction(migraphx::make_op("reshape", {{"dims", {4}}}), inx);
        auto gr   = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), lrsp, li);
        auto r    = m.add_instruction(migraphx::make_op("softmax", {{"axis", 1}}), gr);
        m.add_return({r});

        return m;
    };

    auto m1 = create_resize_module();
    run_pass(m1);

    auto create_optimized_module = [&] {
        migraphx::module m;
        auto inx                  = m.add_parameter("X", sx);
        std::vector<int64_t> dims = {1, 1, 2, 1, 2, 1};
        auto rspx = m.add_instruction(migraphx::make_op("reshape", {{"dims", dims}}), inx);
        std::vector<int64_t> mb_dims = {1, 2, 2, 2, 2, 3};
560
561
        auto mbx =
            m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", mb_dims}}), rspx);
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
        auto std_mb                    = m.add_instruction(migraphx::make_op("contiguous"), mbx);
        std::vector<int64_t> orig_dims = {1, 2, 4, 6};
        auto rmb = m.add_instruction(migraphx::make_op("reshape", {{"dims", orig_dims}}), std_mb);
        auto r   = m.add_instruction(migraphx::make_op("softmax", {{"axis", 1}}), rmb);
        m.add_return({r});

        return m;
    };

    EXPECT(m1 == create_optimized_module());
}

TEST_CASE(optimize_resize_ind_not_apply)
{
    migraphx::shape sx{migraphx::shape::float_type, {1, 1, 2, 2}};
    auto create_resize_module = [&] {
        migraphx::module m;
        auto inx = m.add_parameter("X", sx);

        migraphx::shape si{migraphx::shape::int32_type, {1, 2, 4, 6}};
        std::vector<int> ind = {0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 2, 2, 2, 3,
                                3, 3, 2, 2, 2, 3, 3, 3, 0, 0, 0, 1, 1, 1, 0, 0,
                                0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 2, 2, 2, 3, 3, 3};
        auto li              = m.add_literal(migraphx::literal(si, ind));

        auto lrsp = m.add_instruction(migraphx::make_op("reshape", {{"dims", {4}}}), inx);
        auto gr   = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), lrsp, li);
        auto r    = m.add_instruction(migraphx::make_op("softmax", {{"axis", 1}}), gr);
        m.add_return({r});

        return m;
    };

    auto m1 = create_resize_module();
    run_pass(m1);
    EXPECT(m1 == create_resize_module());
}

TEST_CASE(optimize_resize_rsp_dim_1)
{
    migraphx::shape sx{migraphx::shape::float_type, {1, 1, 2, 2}};
    auto create_resize_module = [&] {
        migraphx::module m;
        auto inx = m.add_parameter("X", sx);

        migraphx::shape si{migraphx::shape::int32_type, {1, 1, 4, 3, 2}};
        std::vector<int> ind = {0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1,
                                2, 2, 2, 3, 3, 3, 2, 2, 2, 3, 3, 3};
        auto li              = m.add_literal(migraphx::literal(si, ind));

        auto lrsp = m.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2}}}), inx);
        auto r    = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), lrsp, li);
        m.add_return({r});

        return m;
    };

    auto m = create_resize_module();
    run_pass(m);
    EXPECT(m == create_resize_module());
}

TEST_CASE(optimize_resize_ndims_unequal)
{
    migraphx::shape sx{migraphx::shape::float_type, {1, 1, 2, 2}};
    migraphx::shape sy{migraphx::shape::float_type, {1, 1, 4, 3, 2}};
    auto create_resize_module = [&] {
        migraphx::module m;
        auto inx = m.add_parameter("X", sx);
        auto iny = m.add_parameter("Y", sy);

        migraphx::shape si{migraphx::shape::int32_type, {1, 1, 4, 3, 2}};
        std::vector<int> ind = {0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1,
                                2, 2, 2, 3, 3, 3, 2, 2, 2, 3, 3, 3};
        auto li              = m.add_literal(migraphx::literal(si, ind));

        auto lrsp = m.add_instruction(migraphx::make_op("reshape", {{"dims", {4}}}), inx);
        auto gr   = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), lrsp, li);
        auto r    = m.add_instruction(migraphx::make_op("sub"), iny, gr);
        m.add_return({r});

        return m;
    };

    auto m = create_resize_module();
    run_pass(m);
    EXPECT(m == create_resize_module());
}

TEST_CASE(optimize_resize_ind_non_brcst)
{
    migraphx::shape sx{migraphx::shape::float_type, {1, 1, 3, 2}};
    migraphx::shape sy{migraphx::shape::float_type, {1, 1, 4, 6}};
    auto create_resize_module = [&] {
        migraphx::module m;
        auto inx = m.add_parameter("X", sx);
        auto iny = m.add_parameter("Y", sy);

        migraphx::shape si{migraphx::shape::int32_type, {1, 1, 4, 6}};
        std::vector<int> ind = {0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1,
                                2, 2, 2, 3, 3, 3, 2, 2, 2, 3, 3, 3};
        auto li              = m.add_literal(migraphx::literal(si, ind));

        auto lrsp = m.add_instruction(migraphx::make_op("reshape", {{"dims", {6}}}), inx);
        auto gr   = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), lrsp, li);
        auto r    = m.add_instruction(migraphx::make_op("sub"), iny, gr);
        m.add_return({r});

        return m;
    };

    auto m = create_resize_module();
    run_pass(m);
    EXPECT(m == create_resize_module());
}

TEST_CASE(optimize_resize_ind_non_const)
{
    migraphx::shape sx{migraphx::shape::float_type, {1, 1, 3, 2}};
    migraphx::shape sy{migraphx::shape::float_type, {1, 1, 4, 6}};
    auto create_resize_module = [&] {
        migraphx::module m;
        auto inx = m.add_parameter("X", sx);
        auto iny = m.add_parameter("Y", sy);

        migraphx::shape si{migraphx::shape::int32_type, {1, 1, 4, 6}};
        auto li   = m.add_parameter("ind", si);
        auto lrsp = m.add_instruction(migraphx::make_op("reshape", {{"dims", {6}}}), inx);
        auto gr   = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), lrsp, li);
        auto r    = m.add_instruction(migraphx::make_op("sub"), iny, gr);
        m.add_return({r});

        return m;
    };

    auto m = create_resize_module();
    run_pass(m);
    EXPECT(m == create_resize_module());
}

TEST_CASE(optimize_where_true)
{
    migraphx::shape s{migraphx::shape::float_type, {1, 1, 3, 2}};
    auto create_where_module = [&](bool cond) {
        migraphx::module m;
        auto inx = m.add_parameter("X", s);
        auto iny = m.add_parameter("Y", s);

        migraphx::shape si{migraphx::shape::bool_type, {1, 1, 3, 2}};
        std::vector<char> idata(si.elements(), static_cast<char>(cond));
        auto li     = m.add_literal(migraphx::literal(si, idata));
        auto data   = m.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), inx, iny);
        auto data_1 = m.add_instruction(migraphx::make_op("reshape", {{"dims", {12}}}), data);
        auto r      = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), data_1, li);
        m.add_return({r});
        return m;
    };

720
    auto return_xy = [&](bool cond) {
721
        migraphx::module m;
722
723
724
        auto x = m.add_parameter("X", s);
        auto y = m.add_parameter("Y", s);
        cond ? m.add_return({x}) : m.add_return({y});
725
726
727
728
729
        return m;
    };

    auto m = create_where_module(true);
    run_pass(m);
730
    EXPECT(m == return_xy(true));
731
732
733

    auto m1 = create_where_module(false);
    run_pass(m1);
734
    EXPECT(m1 == return_xy(false));
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
}

TEST_CASE(where_different_cond_values)
{
    auto create_where_module = [] {
        migraphx::module m;
        migraphx::shape s{migraphx::shape::float_type, {1, 1, 3, 2}};
        auto inx = m.add_parameter("X", s);
        auto iny = m.add_parameter("Y", s);

        migraphx::shape si{migraphx::shape::bool_type, {1, 1, 3, 2}};
        std::vector<char> idata = {1, 1, 0, 1, 0, 1};
        auto li                 = m.add_literal(migraphx::literal(si, idata));
        auto data   = m.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), inx, iny);
        auto data_1 = m.add_instruction(migraphx::make_op("reshape", {{"dims", {12}}}), data);
        auto r      = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), data_1, li);
        m.add_return({r});
        return m;
    };

    auto m = create_where_module();
    run_pass(m);
    EXPECT(m == create_where_module());
}

TEST_CASE(where_axis_nonzero)
{
    auto create_where_module = [] {
        migraphx::module m;
        migraphx::shape s{migraphx::shape::float_type, {1, 1, 3, 2}};
        auto inx = m.add_parameter("X", s);
        auto iny = m.add_parameter("Y", s);

        migraphx::shape si{migraphx::shape::bool_type, {1, 1, 3, 2}};
        std::vector<char> idata(6, 1);
        auto li     = m.add_literal(migraphx::literal(si, idata));
        auto data   = m.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), inx, iny);
        auto data_1 = m.add_instruction(migraphx::make_op("reshape", {{"dims", {12}}}), data);
        auto r      = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), data_1, li);
        m.add_return({r});
        return m;
    };

    auto m = create_where_module();
    run_pass(m);
    EXPECT(m == create_where_module());
}

TEST_CASE(where_three_concat_inputs)
{
    auto create_where_module = [] {
        migraphx::module m;
        migraphx::shape s{migraphx::shape::float_type, {1, 1, 3, 2}};
        auto inx = m.add_parameter("X", s);
        auto iny = m.add_parameter("Y", s);

        migraphx::shape si{migraphx::shape::bool_type, {1, 1, 3, 2}};
        std::vector<char> idata(6, 1);
        auto li     = m.add_literal(migraphx::literal(si, idata));
        auto data   = m.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), inx, iny, inx);
        auto data_1 = m.add_instruction(migraphx::make_op("reshape", {{"dims", {18}}}), data);
        auto r      = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), data_1, li);
        m.add_return({r});
        return m;
    };

    auto m = create_where_module();
    run_pass(m);
    EXPECT(m == create_where_module());
}

TEST_CASE(where_three_inputs_diff_shapes)
{
    auto create_where_module = [] {
        migraphx::module m;
        migraphx::shape sx{migraphx::shape::float_type, {1, 1, 3, 2}};
        migraphx::shape sy{migraphx::shape::float_type, {2, 1, 3, 2}};
        auto inx = m.add_parameter("X", sx);
        auto iny = m.add_parameter("Y", sy);

        migraphx::shape si{migraphx::shape::bool_type, {1, 1, 3, 2}};
        std::vector<char> idata(6, 1);
        auto li     = m.add_literal(migraphx::literal(si, idata));
        auto data   = m.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), inx, iny);
        auto data_1 = m.add_instruction(migraphx::make_op("reshape", {{"dims", {18}}}), data);
        auto r      = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), data_1, li);
        m.add_return({r});
        return m;
    };

    auto m = create_where_module();
    run_pass(m);
    EXPECT(m == create_where_module());
}

TEST_CASE(where_three_lens_diff)
{
    auto create_where_module = [] {
        migraphx::module m;
        migraphx::shape sx{migraphx::shape::float_type, {1, 1, 3, 2}};
        migraphx::shape sy{migraphx::shape::float_type, {1, 1, 3, 2}};
        auto inx = m.add_parameter("X", sx);
        auto iny = m.add_parameter("Y", sy);

        migraphx::shape si{migraphx::shape::bool_type, {1, 1, 6}};
        std::vector<char> idata(6, 1);
        auto li     = m.add_literal(migraphx::literal(si, idata));
        auto data   = m.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), inx, iny);
        auto data_1 = m.add_instruction(migraphx::make_op("reshape", {{"dims", {12}}}), data);
        auto r      = m.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), data_1, li);
        m.add_return({r});
        return m;
    };

    auto m = create_where_module();
    run_pass(m);
    EXPECT(m == create_where_module());
}

TEST_CASE(reshape_cont)
{
    auto create_module = [] {
        migraphx::module m;
        migraphx::shape sx{migraphx::shape::float_type, {1, 4, 1}};
        migraphx::shape sy{migraphx::shape::float_type, {2, 2, 2, 6}};

861
862
863
864
        auto inx = m.add_parameter("x", sx);
        auto iny = m.add_parameter("y", sy);
        auto mb_inx =
            m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 6}}}), inx);
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
        auto std_inx = m.add_instruction(migraphx::make_op("contiguous"), mb_inx);
        auto rsp =
            m.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 2, 6}}}), std_inx);
        auto r = m.add_instruction(migraphx::make_op("add"), rsp, iny);
        m.add_return({r});

        return m;
    };

    auto m1 = create_module();
    run_pass(m1);

    auto create_opt_module = [] {
        migraphx::module m;
        migraphx::shape sx{migraphx::shape::float_type, {1, 4, 1}};
        migraphx::shape sy{migraphx::shape::float_type, {2, 2, 2, 6}};

882
883
884
885
        auto inx = m.add_parameter("x", sx);
        auto iny = m.add_parameter("y", sy);
        auto mb_inx =
            m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 6}}}), inx);
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
        auto rsp_iny = m.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 4, 6}}}), iny);
        auto sum     = m.add_instruction(migraphx::make_op("add"), mb_inx, rsp_iny);
        auto r = m.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 2, 6}}}), sum);
        m.add_return({r});

        return m;
    };

    EXPECT(m1 == create_opt_module());
}

TEST_CASE(reshape_input_non_std)
{
    auto create_module = [] {
        migraphx::module m;
        migraphx::shape sx{migraphx::shape::float_type, {1, 4, 1}};
        migraphx::shape sy{migraphx::shape::float_type, {2, 6, 2, 2}};

904
905
906
907
        auto inx = m.add_parameter("x", sx);
        auto iny = m.add_parameter("y", sy);
        auto mb_inx =
            m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 6}}}), inx);
908
909
910
        auto std_inx = m.add_instruction(migraphx::make_op("contiguous"), mb_inx);
        auto rsp =
            m.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 2, 6}}}), std_inx);
911
912
913
        auto ty =
            m.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), iny);
        auto r = m.add_instruction(migraphx::make_op("add"), rsp, ty);
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
        m.add_return({r});

        return m;
    };

    auto m1 = create_module();
    run_pass(m1);

    EXPECT(m1 == create_module());
}

TEST_CASE(reshape_cont_nonpw)
{
    auto create_module = [] {
        migraphx::module m;
        migraphx::shape sx{migraphx::shape::float_type, {1, 4, 1}};
        migraphx::shape sy{migraphx::shape::float_type, {2, 2, 2, 6}};

932
933
934
935
        auto inx = m.add_parameter("x", sx);
        auto iny = m.add_parameter("y", sy);
        auto mb_inx =
            m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 6}}}), inx);
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
        auto std_inx = m.add_instruction(migraphx::make_op("contiguous"), mb_inx);
        auto rsp =
            m.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 2, 6}}}), std_inx);
        auto r = m.add_instruction(migraphx::make_op("convolution"), rsp, iny);
        m.add_return({r});

        return m;
    };

    auto m1 = create_module();
    run_pass(m1);

    EXPECT(m1 == create_module());
}

951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
TEST_CASE(transpose_contiguous_reshape_unary)
{
    migraphx::module m1;
    {
        auto x = m1.add_parameter("x", {migraphx::shape::float_type, {2, 8, 5, 5}});
        auto reshape_ins1 =
            m1.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 2, 2, 5, 5}}}), x);
        auto transpose_ins = m1.add_instruction(
            migraphx::make_op("transpose", {{"permutation", {0, 3, 4, 1, 5, 2}}}), reshape_ins1);
        auto cont_ins = m1.add_instruction(migraphx::make_op("contiguous"), transpose_ins);
        auto reshape_ins2 =
            m1.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 10, 10}}}), cont_ins);
        auto relu = m1.add_instruction(migraphx::make_op("relu"), reshape_ins2);
        m1.add_instruction(pass_op{}, relu);
    }
    run_pass(m1);
    migraphx::module m2;
    {
        auto x = m2.add_parameter("x", {migraphx::shape::float_type, {2, 8, 5, 5}});
        auto reshape_ins1 =
            m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 2, 2, 5, 5}}}), x);
        auto transpose_ins = m2.add_instruction(
            migraphx::make_op("transpose", {{"permutation", {0, 3, 4, 1, 5, 2}}}), reshape_ins1);
        auto relu     = m2.add_instruction(migraphx::make_op("relu"), transpose_ins);
        auto cont_ins = m2.add_instruction(migraphx::make_op("contiguous"), relu);
        auto reshape_ins2 =
            m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 10, 10}}}), cont_ins);
        m2.add_instruction(pass_op{}, reshape_ins2);
    }
    EXPECT(m1 == m2);
}

TEST_CASE(transpose_contiguous_squeeze_unary)
{
    migraphx::module m1;
    {
        auto x = m1.add_parameter("x", {migraphx::shape::float_type, {2, 8, 1, 5}});
        auto transpose_ins =
            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), x);
        auto cont_ins = m1.add_instruction(migraphx::make_op("contiguous"), transpose_ins);
        auto sq_ins   = m1.add_instruction(migraphx::make_op("squeeze", {{"axes", {1}}}), cont_ins);
        auto rsqrt    = m1.add_instruction(migraphx::make_op("rsqrt"), sq_ins);
        m1.add_instruction(pass_op{}, rsqrt);
    }
    run_pass(m1);
    migraphx::module m2;
    {
        auto x = m2.add_parameter("x", {migraphx::shape::float_type, {2, 8, 1, 5}});
        auto transpose_ins =
            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), x);
        auto rsqrt    = m2.add_instruction(migraphx::make_op("rsqrt"), transpose_ins);
        auto cont_ins = m2.add_instruction(migraphx::make_op("contiguous"), rsqrt);
        auto sq_ins   = m2.add_instruction(migraphx::make_op("squeeze", {{"axes", {1}}}), cont_ins);
        m2.add_instruction(pass_op{}, sq_ins);
    }
    EXPECT(m1 == m2);
}

TEST_CASE(transpose_contiguous_unsqueeze_unary)
{
    migraphx::module m1;
    {
        auto x = m1.add_parameter("x", {migraphx::shape::float_type, {2, 8, 5, 5}});
        auto transpose_ins =
            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), x);
        auto cont_ins = m1.add_instruction(migraphx::make_op("contiguous"), transpose_ins);
        auto unsq_ins =
            m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {2}}}), cont_ins);
        auto round = m1.add_instruction(migraphx::make_op("round"), unsq_ins);
        m1.add_instruction(pass_op{}, round);
    }
    run_pass(m1);
    migraphx::module m2;
    {
        auto x = m2.add_parameter("x", {migraphx::shape::float_type, {2, 8, 5, 5}});
        auto transpose_ins =
            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), x);
        auto round    = m2.add_instruction(migraphx::make_op("round"), transpose_ins);
        auto cont_ins = m2.add_instruction(migraphx::make_op("contiguous"), round);
        auto unsq_ins =
            m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {2}}}), cont_ins);
        m2.add_instruction(pass_op{}, unsq_ins);
    }
    EXPECT(m1 == m2);
}

1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
TEST_CASE(transpose_contiguous_reshape_binary_packed)
{
    migraphx::module m1;
    {
        auto x  = m1.add_parameter("x", {migraphx::shape::float_type, {2, 128, 28, 28}});
        auto w1 = m1.add_literal(
            migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
        auto conv1 = m1.add_instruction(
            migraphx::make_op("convolution",
                              {{"padding", {0, 0}}, {"stride", {1, 1}}, {"dilation", {1, 1}}}),
            x,
            w1); // (2, 256, 28, 28)
        auto w2 = m1.add_literal(
            migraphx::generate_literal({migraphx::shape::float_type, {512, 256, 1, 1}}));
        auto conv2 = m1.add_instruction(
            migraphx::make_op("convolution",
                              {{"padding", {0, 0}}, {"stride", {2, 2}}, {"dilation", {1, 1}}}),
            conv1,
            w2); // (2, 512, 14, 14)

        auto conv2_rsp1 = m1.add_instruction(
            migraphx::make_op("reshape", {{"dims", {2, 2, 2, 128, 14, 14}}}), conv2);
        auto conv2_trans = m1.add_instruction(
            migraphx::make_op("transpose", {{"permutation", {0, 3, 4, 1, 5, 2}}}), conv2_rsp1);
        auto conv2_cont = m1.add_instruction(migraphx::make_op("contiguous"), conv2_trans);
        auto conv2_rsp2 = m1.add_instruction(
            migraphx::make_op("reshape", {{"dims", {2, 128, 28, 28}}}), conv2_cont);
        auto add_ins = m1.add_instruction(migraphx::make_op("add"), conv2_rsp2, x);
        m1.add_instruction(pass_op{}, add_ins);
    }
    run_pass(m1);
    migraphx::module m2;
    {
        auto x  = m2.add_parameter("x", {migraphx::shape::float_type, {2, 128, 28, 28}});
        auto w1 = m2.add_literal(
            migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
        auto conv1 = m2.add_instruction(
            migraphx::make_op("convolution",
                              {{"padding", {0, 0}}, {"stride", {1, 1}}, {"dilation", {1, 1}}}),
            x,
            w1); // (2, 256, 28, 28)
        auto w2 = m2.add_literal(
            migraphx::generate_literal({migraphx::shape::float_type, {512, 256, 1, 1}}));
        auto conv2 = m2.add_instruction(
            migraphx::make_op("convolution",
                              {{"padding", {0, 0}}, {"stride", {2, 2}}, {"dilation", {1, 1}}}),
            conv1,
            w2); // (2, 512, 14, 14)

        auto conv2_rsp = m2.add_instruction(
            migraphx::make_op("reshape", {{"dims", {2, 2, 2, 128, 14, 14}}}), conv2);
        auto conv2_trans = m2.add_instruction(
            migraphx::make_op("transpose", {{"permutation", {0, 3, 4, 1, 5, 2}}}), conv2_rsp);
        auto x_rsp =
            m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 128, 14, 2, 14, 2}}}), x);
        auto add_ins = m2.add_instruction(migraphx::make_op("add"), conv2_trans, x_rsp);
        auto add_rsp =
            m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 128, 28, 28}}}), add_ins);
        m2.add_instruction(pass_op{}, add_rsp);
    }
    EXPECT(m1 == m2);
}

TEST_CASE(transpose_contiguous_reshape_binary_broadcast)
{
    migraphx::module m1;
    {
        migraphx::shape sx{migraphx::shape::float_type, {4}};
        migraphx::shape sy{migraphx::shape::float_type, {2, 6, 2, 2}};

        auto x       = m1.add_parameter("x", sx);
        auto y       = m1.add_parameter("y", sy);
        auto x_brcst = m1.add_instruction(
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {2, 4, 6}}}), x);
        auto y_trans =
            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), y);
        auto y_cont = m1.add_instruction(migraphx::make_op("contiguous"), y_trans);
        auto y_rsp =
            m1.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 4, 6}}}), y_cont);
        auto r = m1.add_instruction(migraphx::make_op("add"), y_rsp, x_brcst);
        m1.add_return({r});
    }
    migraphx::module m2 = m1;
    run_pass(m1);
    EXPECT(m1 == m2);
}

Paul's avatar
Paul committed
1124
int main(int argc, const char* argv[]) { test::run(argc, argv); }