"vscode:/vscode.git/clone" did not exist on "d1305d0c165002c040f76b37bddca463ecf0f522"
simplify_algebra_test.cpp 113 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
Paul's avatar
Paul committed
24
25
#include <migraphx/simplify_algebra.hpp>
#include <migraphx/dead_code_elimination.hpp>
26
#include <migraphx/pass_manager.hpp>
Paul's avatar
Paul committed
27
#include <migraphx/operators.hpp>
Paul's avatar
Paul committed
28
29
30
#include <migraphx/generate.hpp>
#include <migraphx/ranges.hpp>
#include <migraphx/instruction.hpp>
Paul's avatar
Paul committed
31
#include <basic_ops.hpp>
32
#include <migraphx/make_op.hpp>
Paul's avatar
Paul committed
33
34
#include <test.hpp>

Paul Fultz II's avatar
Paul Fultz II committed
35
void run_pass(migraphx::module& m)
Paul's avatar
Paul committed
36
{
Paul Fultz II's avatar
Paul Fultz II committed
37
    migraphx::run_passes(m, {migraphx::simplify_algebra{}, migraphx::dead_code_elimination{}});
38
}
Paul's avatar
Paul committed
39

Paul's avatar
Paul committed
40
TEST_CASE(simplify_add1)
Paul's avatar
Paul committed
41
{
Paul Fultz II's avatar
Paul Fultz II committed
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
    migraphx::module m1;
    {
        auto x    = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto y    = m1.add_parameter("y", {migraphx::shape::int32_type, {1}});
        auto one  = m1.add_literal(1);
        auto two  = m1.add_literal(2);
        auto sum1 = m1.add_instruction(migraphx::make_op("add"), x, one);
        auto sum2 = m1.add_instruction(migraphx::make_op("add"), y, two);
        auto sum3 = m1.add_instruction(migraphx::make_op("add"), sum1, sum2);
        m1.add_instruction(pass_op{}, sum3);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x    = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto y    = m2.add_parameter("y", {migraphx::shape::int32_type, {1}});
        auto one  = m2.add_literal(1);
        auto two  = m2.add_literal(2);
        auto sum1 = m2.add_instruction(migraphx::make_op("add"), one, two);
        auto sum2 = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto sum3 = m2.add_instruction(migraphx::make_op("add"), sum2, sum1);
        m2.add_instruction(pass_op{}, sum3);
    }
    EXPECT(m1 == m2);
Paul's avatar
Paul committed
67
68
}

Paul's avatar
Paul committed
69
TEST_CASE(simplify_add2)
Paul's avatar
Paul committed
70
{
Paul Fultz II's avatar
Paul Fultz II committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    migraphx::module m1;
    {
        auto x    = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto y    = m1.add_parameter("y", {migraphx::shape::int32_type, {1}});
        auto one  = m1.add_literal(1);
        auto two  = m1.add_literal(2);
        auto sum1 = m1.add_instruction(migraphx::make_op("add"), one, x);
        auto sum2 = m1.add_instruction(migraphx::make_op("add"), two, y);
        auto sum3 = m1.add_instruction(migraphx::make_op("add"), sum1, sum2);
        m1.add_instruction(pass_op{}, sum3);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x    = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto y    = m2.add_parameter("y", {migraphx::shape::int32_type, {1}});
        auto one  = m2.add_literal(1);
        auto two  = m2.add_literal(2);
        auto sum1 = m2.add_instruction(migraphx::make_op("add"), one, two);
        auto sum2 = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto sum3 = m2.add_instruction(migraphx::make_op("add"), sum2, sum1);
        m2.add_instruction(pass_op{}, sum3);
    }
    EXPECT(m1 == m2);
Paul's avatar
Paul committed
96
97
}

Paul's avatar
Paul committed
98
TEST_CASE(simplify_add3)
Paul's avatar
Paul committed
99
{
Paul Fultz II's avatar
Paul Fultz II committed
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
    migraphx::module m1;
    {
        auto x    = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto one  = m1.add_literal(1);
        auto two  = m1.add_literal(2);
        auto sum1 = m1.add_instruction(migraphx::make_op("add"), one, x);
        auto sum2 = m1.add_instruction(migraphx::make_op("add"), one, two);
        auto sum3 = m1.add_instruction(migraphx::make_op("add"), sum1, sum2);
        m1.add_instruction(pass_op{}, sum3);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x    = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto one  = m2.add_literal(1);
        auto two  = m2.add_literal(2);
        auto sum1 = m2.add_instruction(migraphx::make_op("add"), one, two);
        auto sum2 = m2.add_instruction(migraphx::make_op("add"), one, sum1);
        auto sum3 = m2.add_instruction(migraphx::make_op("add"), x, sum2);
        m2.add_instruction(pass_op{}, sum3);
    }
    EXPECT(m1 == m2);
Paul's avatar
Paul committed
123
124
}

125
126
127
128
129
TEST_CASE(simplify_add_broadcast1)
{
    migraphx::shape inner{migraphx::shape::int32_type, {2}};
    migraphx::shape outer{migraphx::shape::int32_type, {1, 2, 3, 3}};
    migraphx::op::broadcast b{1, {1, 2, 3, 3}};
Paul Fultz II's avatar
Paul Fultz II committed
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
    migraphx::module m1;
    {
        auto x    = m1.add_parameter("x", outer);
        auto y    = m1.add_parameter("y", outer);
        auto one  = m1.add_literal({inner, {1, 1}});
        auto oneb = m1.add_instruction(b, one);
        auto two  = m1.add_literal({inner, {2, 2}});
        auto twob = m1.add_instruction(b, two);
        auto sum1 = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto sum2 = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto sum3 = m1.add_instruction(migraphx::make_op("add"), sum1, sum2);
        m1.add_instruction(pass_op{}, sum3);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x     = m2.add_parameter("x", outer);
        auto y     = m2.add_parameter("y", outer);
        auto one   = m2.add_literal({inner, {1, 1}});
        auto two   = m2.add_literal({inner, {2, 2}});
        auto sum1  = m2.add_instruction(migraphx::make_op("add"), one, two);
        auto sum1b = m2.add_instruction(b, sum1);
        auto sum2  = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto sum3  = m2.add_instruction(migraphx::make_op("add"), sum2, sum1b);
        m2.add_instruction(pass_op{}, sum3);
    }
    EXPECT(m1 == m2);
158
159
160
161
162
163
164
165
}

TEST_CASE(simplify_add_broadcast2)
{
    migraphx::shape inner{migraphx::shape::int32_type, {2}};
    migraphx::shape outer{migraphx::shape::int32_type, {1, 2, 3, 3}};
    migraphx::op::broadcast b{1, {1, 2, 3, 3}};
    auto create_program = [&] {
Paul Fultz II's avatar
Paul Fultz II committed
166
167
168
169
170
171
172
173
174
175
176
        migraphx::module m;
        auto x    = m.add_parameter("x", outer);
        auto y    = m.add_parameter("y", outer);
        auto one  = m.add_literal({inner, {1, 1}});
        auto oneb = m.add_instruction(b, one);
        auto two  = m.add_literal({outer, {2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2}});
        auto sum1 = m.add_instruction(migraphx::make_op("add"), x, y);
        auto sum2 = m.add_instruction(migraphx::make_op("add"), oneb, two);
        auto sum3 = m.add_instruction(migraphx::make_op("add"), sum2, sum1);
        m.add_instruction(pass_op{}, sum3);
        return m;
177
    };
Paul Fultz II's avatar
Paul Fultz II committed
178
179
    migraphx::module m1 = create_program();
    run_pass(m1);
180

Paul Fultz II's avatar
Paul Fultz II committed
181
182
    migraphx::module m2 = create_program();
    EXPECT(m1 == m2);
183
184
}

Paul's avatar
Paul committed
185
// TODO: Add test case
186
// TEST_CASE(simplify_add4)
Paul's avatar
Paul committed
187
188
void simplify_add4()
{
Paul Fultz II's avatar
Paul Fultz II committed
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
    migraphx::module m1;
    {
        auto x    = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto y    = m1.add_parameter("y", {migraphx::shape::int32_type, {1}});
        auto one  = m1.add_literal(1);
        auto two  = m1.add_literal(2);
        auto sum1 = m1.add_instruction(migraphx::make_op("add"), one, x);
        auto sum2 = m1.add_instruction(migraphx::make_op("add"), sum1, y);
        auto sum3 = m1.add_instruction(migraphx::make_op("add"), sum2, two);
        m1.add_instruction(pass_op{}, sum3);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x    = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto y    = m2.add_parameter("y", {migraphx::shape::int32_type, {1}});
        auto one  = m2.add_literal(1);
        auto two  = m2.add_literal(2);
        auto sum1 = m2.add_instruction(migraphx::make_op("add"), one, two);
        auto sum2 = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto sum3 = m2.add_instruction(migraphx::make_op("add"), sum2, sum1);
        m2.add_instruction(pass_op{}, sum3);
    }
    EXPECT(m1 == m2);
Paul's avatar
Paul committed
214
215
}

Paul's avatar
Paul committed
216
217
TEST_CASE(simplify_mul_conv1)
{
Paul Fultz II's avatar
Paul Fultz II committed
218
219
220
221
222
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::int32_type, {1, 128, 28, 28}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {256, 128, 3, 3}}));
    auto conv = m.add_instruction(
223
224
225
226
        migraphx::make_op("convolution",
                          {{"padding", {1, 1}}, {"stride", {2, 2}}, {"dilation", {1, 1}}}),
        x,
        w);
Paul Fultz II's avatar
Paul Fultz II committed
227
228
    auto a = m.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {256}}));
    auto b = m.add_instruction(
229
        migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 256, 14, 14}}}), a);
Paul Fultz II's avatar
Paul Fultz II committed
230
231
    auto mul = m.add_instruction(migraphx::make_op("mul"), conv, b);
    m.add_instruction(pass_op{}, mul);
Paul's avatar
Paul committed
232
    EXPECT(conv->outputs().front()->name() == "mul");
Paul Fultz II's avatar
Paul Fultz II committed
233
234
235
    run_pass(m);
    auto new_conv =
        std::find_if(m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; });
Paul's avatar
Paul committed
236
237
238
    EXPECT(new_conv->outputs().front()->name() != "mul");
}

239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
TEST_CASE(simplify_mul_conv2)
{
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::int32_type, {1, 128, 28, 28}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {256, 128, 3, 3}}));
    auto conv = m.add_instruction(
        migraphx::make_op("convolution",
                          {{"padding", {1, 1}}, {"stride", {2, 2}}, {"dilation", {1, 1}}}),
        x,
        w);
    auto a      = m.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {256}}));
    auto unsq_a = m.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1, 2}}}), a);
    auto b      = m.add_instruction(
        migraphx::make_op("multibroadcast", {{"out_lens", {1, 256, 14, 14}}}), unsq_a);
    auto mul = m.add_instruction(migraphx::make_op("mul"), conv, b);
    m.add_instruction(pass_op{}, mul);
    EXPECT(conv->outputs().front()->name() == "mul");
    run_pass(m);
    auto new_conv =
        std::find_if(m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; });
    EXPECT(new_conv->outputs().front()->name() != "mul");
}

// len = 1 case
TEST_CASE(simplify_mul_conv3)
{
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::int32_type, {1, 128, 28, 28}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {256, 128, 3, 3}}));
    auto conv = m.add_instruction(
        migraphx::make_op("convolution",
                          {{"padding", {1, 1}}, {"stride", {2, 2}}, {"dilation", {1, 1}}}),
        x,
        w);
    auto a = m.add_literal(
        migraphx::generate_literal({migraphx::shape::int32_type, {256, 1, 1}, {1, 18, 1}}));
    auto b =
        m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 256, 14, 14}}}), a);
    auto mul = m.add_instruction(migraphx::make_op("mul"), conv, b);
    m.add_instruction(pass_op{}, mul);
    EXPECT(conv->outputs().front()->name() == "mul");
    run_pass(m);
    auto new_conv =
        std::find_if(m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; });
    EXPECT(new_conv->outputs().front()->name() != "mul");
}

// Previously broadcasted literal case, should skip
TEST_CASE(simplify_mul_conv_skip1)
{
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::int32_type, {1, 128, 28, 28}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {256, 128, 3, 3}}));
    auto conv = m.add_instruction(
        migraphx::make_op("convolution",
                          {{"padding", {1, 1}}, {"stride", {2, 2}}, {"dilation", {1, 1}}}),
        x,
        w);
    auto a = m.add_literal(
        migraphx::generate_literal({migraphx::shape::int32_type, {256, 14, 14}, {1, 0, 0}}));
    auto b = m.add_instruction(
        migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 256, 14, 14}}}), a);
    auto mul = m.add_instruction(migraphx::make_op("mul"), conv, b);
    m.add_instruction(pass_op{}, mul);
    EXPECT(conv->outputs().front()->name() == "mul");
    run_pass(m);
    auto new_conv =
        std::find_if(m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; });
    EXPECT(new_conv->outputs().front()->name() == "mul");
}

// Another previously broadcasted literal case, should skip
TEST_CASE(simplify_mul_conv_skip2)
{
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::int32_type, {1, 128, 28, 28}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {256, 128, 3, 3}}));
    auto conv = m.add_instruction(
        migraphx::make_op("convolution",
                          {{"padding", {1, 1}}, {"stride", {2, 2}}, {"dilation", {1, 1}}}),
        x,
        w);
    auto a = m.add_literal(
        migraphx::generate_literal({migraphx::shape::int32_type, {256, 14, 14}, {1, 0, 0}}));
    auto b =
        m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 256, 14, 14}}}), a);
    auto mul = m.add_instruction(migraphx::make_op("mul"), conv, b);
    m.add_instruction(pass_op{}, mul);
    EXPECT(conv->outputs().front()->name() == "mul");
    run_pass(m);
    auto new_conv =
        std::find_if(m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; });
    EXPECT(new_conv->outputs().front()->name() == "mul");
}

338
339
TEST_CASE(simplify_mul_slice_conv1)
{
Paul Fultz II's avatar
Paul Fultz II committed
340
    migraphx::module m1;
341
    {
Paul Fultz II's avatar
Paul Fultz II committed
342
343
        auto x = m1.add_parameter("x", {migraphx::shape::int32_type, {1, 1024, 17, 17}});
        auto w = m1.add_literal(
344
            migraphx::generate_literal({migraphx::shape::int32_type, {768, 1024, 1, 1}}));
Paul Fultz II's avatar
Paul Fultz II committed
345
346
        auto conv   = m1.add_instruction(migraphx::make_op("convolution"), x, w);
        auto slice1 = m1.add_instruction(
347
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {384}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
348
349
        auto a = m1.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}));
        auto b = m1.add_instruction(
350
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 384, 17, 17}}}), a);
Paul Fultz II's avatar
Paul Fultz II committed
351
352
        auto mul    = m1.add_instruction(migraphx::make_op("mul"), slice1, b);
        auto slice2 = m1.add_instruction(
353
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {384}}, {"ends", {768}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
354
355
        auto add = m1.add_instruction(migraphx::make_op("add"), mul, slice2);
        m1.add_instruction(pass_op{}, add);
356
    }
Paul Fultz II's avatar
Paul Fultz II committed
357
    run_pass(m1);
358

Paul Fultz II's avatar
Paul Fultz II committed
359
    migraphx::module m2;
360
    {
Paul Fultz II's avatar
Paul Fultz II committed
361
362
        auto x = m2.add_parameter("x", {migraphx::shape::int32_type, {1, 1024, 17, 17}});
        auto w = m2.add_literal(
363
            migraphx::generate_literal({migraphx::shape::int32_type, {768, 1024, 1, 1}}));
Paul Fultz II's avatar
Paul Fultz II committed
364
        auto wslice1 = m2.add_instruction(
365
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {384}}}), w);
Paul Fultz II's avatar
Paul Fultz II committed
366
367
        auto a = m2.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}));
        auto b = m2.add_instruction(
368
            migraphx::make_op("broadcast", {{"axis", 0}, {"out_lens", {384, 1024, 1, 1}}}), a);
Paul Fultz II's avatar
Paul Fultz II committed
369
370
        auto mul     = m2.add_instruction(migraphx::make_op("mul"), b, wslice1);
        auto wslice2 = m2.add_instruction(
371
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {384}}, {"ends", {768}}}), w);
Paul Fultz II's avatar
Paul Fultz II committed
372
373
374
        auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), mul, wslice2);
        auto conv   = m2.add_instruction(migraphx::make_op("convolution"), x, concat);
        auto slice1 = m2.add_instruction(
375
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {384}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
376
        auto slice2 = m2.add_instruction(
377
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {384}}, {"ends", {768}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
378
379
        auto add = m2.add_instruction(migraphx::make_op("add"), slice1, slice2);
        m2.add_instruction(pass_op{}, add);
380
    }
Paul Fultz II's avatar
Paul Fultz II committed
381
    EXPECT(m1 == m2);
382
383
384
385
}

TEST_CASE(simplify_mul_slice_conv_overlapping_slice)
{
Paul Fultz II's avatar
Paul Fultz II committed
386
    migraphx::module m1;
387
    {
Paul Fultz II's avatar
Paul Fultz II committed
388
389
        auto x = m1.add_parameter("x", {migraphx::shape::int32_type, {1, 1024, 17, 17}});
        auto w = m1.add_literal(
390
            migraphx::generate_literal({migraphx::shape::int32_type, {768, 1024, 1, 1}}));
Paul Fultz II's avatar
Paul Fultz II committed
391
392
        auto conv   = m1.add_instruction(migraphx::make_op("convolution"), x, w);
        auto slice1 = m1.add_instruction(
393
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {384}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
394
395
        auto a = m1.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}));
        auto b = m1.add_instruction(
396
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 384, 17, 17}}}), a);
Paul Fultz II's avatar
Paul Fultz II committed
397
398
        auto mul    = m1.add_instruction(migraphx::make_op("mul"), slice1, b);
        auto slice2 = m1.add_instruction(
399
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {383}}, {"ends", {767}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
400
401
        auto add = m1.add_instruction(migraphx::make_op("add"), mul, slice2);
        m1.add_instruction(pass_op{}, add);
402
    }
Paul Fultz II's avatar
Paul Fultz II committed
403
404
405
    migraphx::module m2 = m1;
    run_pass(m1);
    EXPECT(m1 == m2);
406
407
408
409
}

TEST_CASE(simplify_mul_slice_conv_not_all_slice)
{
Paul Fultz II's avatar
Paul Fultz II committed
410
    migraphx::module m1;
411
    {
Paul Fultz II's avatar
Paul Fultz II committed
412
413
        auto x = m1.add_parameter("x", {migraphx::shape::int32_type, {1, 1024, 17, 17}});
        auto w = m1.add_literal(
414
            migraphx::generate_literal({migraphx::shape::int32_type, {768, 1024, 1, 1}}));
Paul Fultz II's avatar
Paul Fultz II committed
415
416
        auto conv   = m1.add_instruction(migraphx::make_op("convolution"), x, w);
        auto slice1 = m1.add_instruction(
417
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {384}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
418
419
        auto a = m1.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}));
        auto b = m1.add_instruction(
420
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 384, 17, 17}}}), a);
Paul Fultz II's avatar
Paul Fultz II committed
421
422
        auto mul = m1.add_instruction(migraphx::make_op("mul"), slice1, b);
        auto c   = m1.add_literal(
423
            migraphx::generate_literal({migraphx::shape::int32_type, {1, 768, 17, 17}}));
Paul Fultz II's avatar
Paul Fultz II committed
424
425
426
        auto add    = m1.add_instruction(migraphx::make_op("add"), conv, c);
        auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), mul, add);
        m1.add_instruction(pass_op{}, concat);
427
    }
Paul Fultz II's avatar
Paul Fultz II committed
428
429
430
    migraphx::module m2 = m1;
    run_pass(m1);
    EXPECT(m1 == m2);
431
432
}

Paul's avatar
Paul committed
433
434
TEST_CASE(simplify_mul_add)
{
Paul Fultz II's avatar
Paul Fultz II committed
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
    migraphx::module m1;
    {
        auto x   = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto one = m1.add_literal(1);
        auto two = m1.add_literal(2);
        auto sum = m1.add_instruction(migraphx::make_op("add"), one, x);
        auto mul = m1.add_instruction(migraphx::make_op("mul"), sum, two);
        m1.add_instruction(pass_op{}, mul);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x    = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto one  = m2.add_literal(1);
        auto two  = m2.add_literal(2);
        auto mul1 = m2.add_instruction(migraphx::make_op("mul"), two, x);
        auto mul2 = m2.add_instruction(migraphx::make_op("mul"), two, one);
        auto sum  = m2.add_instruction(migraphx::make_op("add"), mul1, mul2);
        m2.add_instruction(pass_op{}, sum);
    }
    EXPECT(m1 == m2);
Paul's avatar
Paul committed
457
458
}

459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
TEST_CASE(simplify_dot_add)
{
    migraphx::module m1;
    {
        auto x   = m1.add_parameter("x", {migraphx::shape::float_type, {2, 2}});
        auto one = m1.add_literal(get_2x2());
        auto two = m1.add_literal(get_2x2(1));
        auto sum = m1.add_instruction(migraphx::make_op("add"), one, x);
        auto dot = m1.add_instruction(migraphx::make_op("dot"), sum, two);
        m1.add_instruction(pass_op{}, dot);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x    = m2.add_parameter("x", {migraphx::shape::float_type, {2, 2}});
        auto one  = m2.add_literal(get_2x2());
        auto two  = m2.add_literal(get_2x2(1));
        auto dot1 = m2.add_instruction(migraphx::make_op("dot"), x, two);
        auto dot2 = m2.add_instruction(migraphx::make_op("dot"), one, two);
        auto sum  = m2.add_instruction(migraphx::make_op("add"), dot1, dot2);
        m2.add_instruction(pass_op{}, sum);
    }
    EXPECT(m1 == m2);
}

TEST_CASE(simplify_inner_broadcast1)
Paul's avatar
Paul committed
486
487
{
    auto b = migraphx::op::broadcast{1, {2, 1, 4, 5}};
Paul Fultz II's avatar
Paul Fultz II committed
488
    migraphx::module m1;
Paul's avatar
Paul committed
489
    {
Paul Fultz II's avatar
Paul Fultz II committed
490
491
492
493
494
495
        auto x   = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto y   = m1.add_parameter("y", {migraphx::shape::int32_type, {1}});
        auto xb  = m1.add_instruction(b, x);
        auto yb  = m1.add_instruction(b, y);
        auto sum = m1.add_instruction(migraphx::make_op("add"), xb, yb);
        m1.add_instruction(pass_op{}, sum);
Paul's avatar
Paul committed
496
    }
Paul Fultz II's avatar
Paul Fultz II committed
497
    run_pass(m1);
Paul's avatar
Paul committed
498

Paul Fultz II's avatar
Paul Fultz II committed
499
    migraphx::module m2;
Paul's avatar
Paul committed
500
    {
Paul Fultz II's avatar
Paul Fultz II committed
501
502
503
504
505
        auto x    = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto y    = m2.add_parameter("y", {migraphx::shape::int32_type, {1}});
        auto sum  = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto sumb = m2.add_instruction(b, sum);
        m2.add_instruction(pass_op{}, sumb);
Paul's avatar
Paul committed
506
    }
Paul Fultz II's avatar
Paul Fultz II committed
507
    EXPECT(m1 == m2);
Paul's avatar
Paul committed
508
509
}

510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
TEST_CASE(simplify_inner_broadcast2)
{
    auto b = migraphx::op::multibroadcast{{2, 1, 4, 5}};
    migraphx::module m1;
    {
        auto x   = m1.add_parameter("x", {migraphx::shape::int32_type, {1, 1, 1, 1}});
        auto y   = m1.add_parameter("y", {migraphx::shape::int32_type, {1, 1, 1, 1}});
        auto xb  = m1.add_instruction(b, x);
        auto yb  = m1.add_instruction(b, y);
        auto sum = m1.add_instruction(migraphx::make_op("add"), xb, yb);
        m1.add_instruction(pass_op{}, sum);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x    = m2.add_parameter("x", {migraphx::shape::int32_type, {1, 1, 1, 1}});
        auto y    = m2.add_parameter("y", {migraphx::shape::int32_type, {1, 1, 1, 1}});
        auto sum  = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto sumb = m2.add_instruction(b, sum);
        m2.add_instruction(pass_op{}, sumb);
    }
    EXPECT(m1 == m2);
}

535
536
TEST_CASE(simplify_add_conv1)
{
Paul Fultz II's avatar
Paul Fultz II committed
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::float_type, {1, 128, 28, 28}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 3, 3}}));
    auto y = m.add_parameter("y", {migraphx::shape::float_type, {1, 128, 28, 28}});
    auto v =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 3, 3}}));
    auto conv1 = m.add_instruction(migraphx::make_op("convolution"), x, w);
    auto conv2 = m.add_instruction(migraphx::make_op("convolution"), y, v);
    auto sum   = m.add_instruction(migraphx::make_op("add"), conv1, conv2);
    m.add_instruction(pass_op{}, sum);
    auto s = m.get_output_shapes().back();
    run_pass(m);
    EXPECT(s == m.get_output_shapes().back());
    EXPECT(std::count_if(
               m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; }) == 1);
553
554
555
556
}

TEST_CASE(simplify_add_conv_no_fusion_7x7_diff_strides)
{
Paul Fultz II's avatar
Paul Fultz II committed
557
558
559
560
561
562
563
564
565
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::float_type, {1, 128, 14, 14}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 7, 7}}));
    auto y = m.add_parameter("y", {migraphx::shape::float_type, {1, 128, 28, 28}});
    auto v =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 7, 7}}));
    auto conv1 = m.add_instruction(migraphx::make_op("convolution"), x, w);
    auto conv2 = m.add_instruction(
566
        migraphx::make_op("convolution", {{"padding", {0, 0}}, {"stride", {3, 3}}}), y, v);
Paul Fultz II's avatar
Paul Fultz II committed
567
568
569
570
571
    auto sum = m.add_instruction(migraphx::make_op("add"), conv1, conv2);
    m.add_instruction(pass_op{}, sum);
    auto s = m.get_output_shapes().back();
    run_pass(m);
    EXPECT(s == m.get_output_shapes().back());
572
    // No fusion
Paul Fultz II's avatar
Paul Fultz II committed
573
574
    EXPECT(std::count_if(
               m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; }) == 2);
575
576
577
578
}

TEST_CASE(simplify_add_conv_1x1_diff_strides1)
{
Paul Fultz II's avatar
Paul Fultz II committed
579
580
581
582
583
584
585
586
587
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::float_type, {1, 128, 14, 14}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
    auto y = m.add_parameter("y", {migraphx::shape::float_type, {1, 128, 28, 28}});
    auto v =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
    auto conv1 = m.add_instruction(migraphx::make_op("convolution"), x, w);
    auto conv2 = m.add_instruction(
588
        migraphx::make_op("convolution", {{"padding", {0, 0}}, {"stride", {2, 2}}}), y, v);
Paul Fultz II's avatar
Paul Fultz II committed
589
590
591
592
593
594
595
    auto sum = m.add_instruction(migraphx::make_op("add"), conv1, conv2);
    m.add_instruction(pass_op{}, sum);
    auto s = m.get_output_shapes().back();
    run_pass(m);
    EXPECT(s == m.get_output_shapes().back());
    EXPECT(std::count_if(
               m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; }) == 1);
596
597
598
599
}

TEST_CASE(simplify_add_conv_1x1_diff_strides2)
{
Paul Fultz II's avatar
Paul Fultz II committed
600
601
602
603
604
605
606
607
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::float_type, {1, 128, 28, 28}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
    auto y = m.add_parameter("y", {migraphx::shape::float_type, {1, 128, 14, 14}});
    auto v =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
    auto conv1 = m.add_instruction(
608
        migraphx::make_op("convolution", {{"padding", {0, 0}}, {"stride", {2, 2}}}), x, w);
Paul Fultz II's avatar
Paul Fultz II committed
609
610
611
612
613
614
615
616
    auto conv2 = m.add_instruction(migraphx::make_op("convolution"), y, v);
    auto sum   = m.add_instruction(migraphx::make_op("add"), conv1, conv2);
    m.add_instruction(pass_op{}, sum);
    auto s = m.get_output_shapes().back();
    run_pass(m);
    EXPECT(s == m.get_output_shapes().back());
    EXPECT(std::count_if(
               m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; }) == 1);
617
618
619
620
}

TEST_CASE(simplify_add_conv_1x1_diff_strides_odd)
{
Paul Fultz II's avatar
Paul Fultz II committed
621
622
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::float_type, {1, 54, 83, 83}});
623
    auto w =
Paul Fultz II's avatar
Paul Fultz II committed
624
625
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {54, 54, 1, 1}}));
    auto y = m.add_parameter("y", {migraphx::shape::float_type, {1, 54, 165, 165}});
626
    auto v =
Paul Fultz II's avatar
Paul Fultz II committed
627
628
629
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {54, 54, 1, 1}}));
    auto conv1 = m.add_instruction(migraphx::make_op("convolution"), x, w);
    auto conv2 = m.add_instruction(
630
        migraphx::make_op("convolution", {{"padding", {0, 0}}, {"stride", {2, 2}}}), y, v);
Paul Fultz II's avatar
Paul Fultz II committed
631
632
633
634
635
636
637
    auto sum = m.add_instruction(migraphx::make_op("add"), conv1, conv2);
    m.add_instruction(pass_op{}, sum);
    auto s = m.get_output_shapes().back();
    run_pass(m);
    EXPECT(s == m.get_output_shapes().back());
    EXPECT(std::count_if(
               m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; }) == 1);
638
639
640
641
}

TEST_CASE(simplify_add_conv_no_fusion_asymetrical_strides1)
{
Paul Fultz II's avatar
Paul Fultz II committed
642
643
644
645
646
647
648
649
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::float_type, {1, 128, 28, 14}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
    auto y = m.add_parameter("y", {migraphx::shape::float_type, {1, 128, 14, 14}});
    auto v =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
    auto conv1 = m.add_instruction(
650
        migraphx::make_op("convolution", {{"padding", {0, 0}}, {"stride", {2, 1}}}), x, w);
Paul Fultz II's avatar
Paul Fultz II committed
651
652
653
654
655
656
    auto conv2 = m.add_instruction(migraphx::make_op("convolution"), y, v);
    auto sum   = m.add_instruction(migraphx::make_op("add"), conv1, conv2);
    m.add_instruction(pass_op{}, sum);
    auto s = m.get_output_shapes().back();
    run_pass(m);
    EXPECT(s == m.get_output_shapes().back());
657
    // No fusion
Paul Fultz II's avatar
Paul Fultz II committed
658
659
    EXPECT(std::count_if(
               m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; }) == 2);
660
661
662
663
}

TEST_CASE(simplify_add_conv_no_fusion_asymetrical_strides2)
{
Paul Fultz II's avatar
Paul Fultz II committed
664
665
666
667
668
669
670
671
672
    migraphx::module m;
    auto x = m.add_parameter("x", {migraphx::shape::float_type, {1, 128, 14, 14}});
    auto w =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
    auto y = m.add_parameter("y", {migraphx::shape::float_type, {1, 128, 28, 14}});
    auto v =
        m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {256, 128, 1, 1}}));
    auto conv1 = m.add_instruction(migraphx::make_op("convolution"), x, w);
    auto conv2 = m.add_instruction(
673
        migraphx::make_op("convolution", {{"padding", {0, 0}}, {"stride", {2, 1}}}), y, v);
Paul Fultz II's avatar
Paul Fultz II committed
674
675
676
677
678
    auto sum = m.add_instruction(migraphx::make_op("add"), conv1, conv2);
    m.add_instruction(pass_op{}, sum);
    auto s = m.get_output_shapes().back();
    run_pass(m);
    EXPECT(s == m.get_output_shapes().back());
679
    // No fusion
Paul Fultz II's avatar
Paul Fultz II committed
680
681
    EXPECT(std::count_if(
               m.begin(), m.end(), [](auto&& ins) { return ins.name() == "convolution"; }) == 2);
682
683
}

684
685
686
TEST_CASE(simplify_concat_add_relu)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {1}};
Paul Fultz II's avatar
Paul Fultz II committed
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
    migraphx::module m1;
    {
        auto x      = m1.add_parameter("x", s);
        auto y      = m1.add_parameter("y", s);
        auto one    = m1.add_literal({s, {1}});
        auto two    = m1.add_literal({s, {2}});
        auto sum1   = m1.add_instruction(migraphx::make_op("add"), x, one);
        auto relu1  = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2   = m1.add_instruction(migraphx::make_op("add"), y, two);
        auto relu2  = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), relu1, relu2);
        m1.add_instruction(pass_op{}, concat);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x       = m2.add_parameter("x", s);
        auto y       = m2.add_parameter("y", s);
        auto one     = m2.add_literal({s, {1}});
        auto two     = m2.add_literal({s, {2}});
        auto concat1 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), x, y);
        auto concat2 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto sum     = m2.add_instruction(migraphx::make_op("add"), concat1, concat2);
        auto relu    = m2.add_instruction(migraphx::make_op("relu"), sum);
        m2.add_instruction(pass_op{}, relu);
    }
    EXPECT(m1 == m2);
715
716
}

717
718
719
TEST_CASE(simplify_concat_add_relu_partial)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {1}};
Paul Fultz II's avatar
Paul Fultz II committed
720
721
722
723
724
725
726
727
728
729
730
    migraphx::module m1;
    {
        auto x     = m1.add_parameter("x", s);
        auto y     = m1.add_parameter("y", s);
        auto one   = m1.add_literal({s, {1}});
        auto two   = m1.add_literal({s, {2}});
        auto sum1  = m1.add_instruction(migraphx::make_op("add"), x, one);
        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2  = m1.add_instruction(migraphx::make_op("add"), y, two);
        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto sum3  = m1.add_instruction(migraphx::make_op("add"), x, y);
731
        auto concat =
Paul Fultz II's avatar
Paul Fultz II committed
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
            m1.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), sum3, relu1, relu2);
        m1.add_instruction(pass_op{}, concat);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto x       = m2.add_parameter("x", s);
        auto y       = m2.add_parameter("y", s);
        auto one     = m2.add_literal({s, {1}});
        auto two     = m2.add_literal({s, {2}});
        auto concat1 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), x, y);
        auto concat2 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto sum1    = m2.add_instruction(migraphx::make_op("add"), concat1, concat2);
        auto relu    = m2.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2    = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto concat  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), sum2, relu);
        m2.add_instruction(pass_op{}, concat);
    }
    EXPECT(m1.sort() == m2.sort());
752
753
754
755
756
}

TEST_CASE(simplify_concat_add_relu_partial_broadcast)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {2, 1, 4, 5}};
Paul Fultz II's avatar
Paul Fultz II committed
757
    migraphx::module m1;
758
    {
759
        auto b    = migraphx::op::broadcast{1, {2, 1, 4, 5}};
Paul Fultz II's avatar
Paul Fultz II committed
760
761
762
763
764
765
766
        auto x    = m1.add_parameter("x", s);
        auto y    = m1.add_parameter("y", s);
        auto one  = m1.add_literal(1);
        auto oneb = m1.add_instruction(b, one);
        auto two  = m1.add_literal(2);
        auto twob = m1.add_instruction(b, two);
        auto sum  = m1.add_instruction(migraphx::make_op("add"), x, y);
767
        auto concat =
Paul Fultz II's avatar
Paul Fultz II committed
768
769
            m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), sum, oneb, twob);
        m1.add_instruction(pass_op{}, concat);
770
    }
Paul Fultz II's avatar
Paul Fultz II committed
771
    run_pass(m1);
772

Paul Fultz II's avatar
Paul Fultz II committed
773
    migraphx::module m2;
774
775
    {
        auto b       = migraphx::op::broadcast{1, {2, 2, 4, 5}};
Paul Fultz II's avatar
Paul Fultz II committed
776
777
778
779
780
781
782
783
784
785
786
        auto x       = m2.add_parameter("x", s);
        auto y       = m2.add_parameter("y", s);
        auto one     = m2.add_literal(1);
        auto two     = m2.add_literal(2);
        auto concat1 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto concatb = m2.add_instruction(b, concat1);
        auto sum     = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto concat2 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), sum, concatb);
        m2.add_instruction(pass_op{}, concat2);
    }
    EXPECT(m1.sort() == m2.sort());
787
788
}

789
790
791
TEST_CASE(simplify_concat_add_relu_broadcast_different_axis)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {2, 1, 4, 5}};
Paul Fultz II's avatar
Paul Fultz II committed
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
    migraphx::module m1;
    {
        auto b      = migraphx::op::broadcast{1, {2, 1, 4, 5}};
        auto x      = m1.add_parameter("x", s);
        auto y      = m1.add_parameter("y", s);
        auto one    = m1.add_literal(1);
        auto oneb   = m1.add_instruction(b, one);
        auto two    = m1.add_literal(2);
        auto twob   = m1.add_instruction(b, two);
        auto sum1   = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1  = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2   = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2  = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), relu1, relu2);
        m1.add_instruction(pass_op{}, concat);
    }
    run_pass(m1);

    migraphx::module m2;
811
812
    {
        auto b        = migraphx::op::broadcast{1, {2, 2, 4, 5}};
Paul Fultz II's avatar
Paul Fultz II committed
813
814
815
816
817
818
819
820
821
822
823
824
        auto x        = m2.add_parameter("x", s);
        auto y        = m2.add_parameter("y", s);
        auto one      = m2.add_literal(1);
        auto two      = m2.add_literal(2);
        auto concat1  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), x, y);
        auto concat2  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto concat2b = m2.add_instruction(b, concat2);
        auto sum      = m2.add_instruction(migraphx::make_op("add"), concat1, concat2b);
        auto relu     = m2.add_instruction(migraphx::make_op("relu"), sum);
        m2.add_instruction(pass_op{}, relu);
    }
    EXPECT(m1 == m2);
825
826
827
828
829
}

TEST_CASE(simplify_concat_add_relu_broadcast_same_axis)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {2, 1, 4, 5}};
Paul Fultz II's avatar
Paul Fultz II committed
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
    migraphx::module m1;
    {
        auto b      = migraphx::op::broadcast{1, {2, 1, 4, 5}};
        auto x      = m1.add_parameter("x", s);
        auto y      = m1.add_parameter("y", s);
        auto one    = m1.add_literal(1);
        auto oneb   = m1.add_instruction(b, one);
        auto two    = m1.add_literal(2);
        auto twob   = m1.add_instruction(b, two);
        auto sum1   = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1  = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2   = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2  = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), relu1, relu2);
        m1.add_instruction(pass_op{}, concat);
    }
    run_pass(m1);

    migraphx::module m2;
849
850
    {
        auto b       = migraphx::op::broadcast{1, {2, 1, 4, 5}};
Paul Fultz II's avatar
Paul Fultz II committed
851
852
853
854
855
856
857
858
859
860
861
862
863
        auto x       = m2.add_parameter("x", s);
        auto y       = m2.add_parameter("y", s);
        auto one     = m2.add_literal(1);
        auto oneb    = m2.add_instruction(b, one);
        auto two     = m2.add_literal(2);
        auto twob    = m2.add_instruction(b, two);
        auto concat1 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), x, y);
        auto concat2 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), oneb, twob);
        auto sum     = m2.add_instruction(migraphx::make_op("add"), concat1, concat2);
        auto relu    = m2.add_instruction(migraphx::make_op("relu"), sum);
        m2.add_instruction(pass_op{}, relu);
    }
    EXPECT(m1 == m2);
864
865
}

866
867
TEST_CASE(simplify_div_const)
{
Paul Fultz II's avatar
Paul Fultz II committed
868
    migraphx::module m1;
869
    {
Paul Fultz II's avatar
Paul Fultz II committed
870
871
872
        auto x   = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto two = m1.add_literal(2);
        m1.add_instruction(migraphx::make_op("div"), x, two);
873
    }
Paul Fultz II's avatar
Paul Fultz II committed
874
    run_pass(m1);
875

Paul Fultz II's avatar
Paul Fultz II committed
876
    migraphx::module m2;
877
    {
Paul Fultz II's avatar
Paul Fultz II committed
878
879
880
881
        auto x     = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto two   = m2.add_literal(2);
        auto recip = m2.insert_instruction(std::next(two), migraphx::make_op("recip"), two);
        m2.add_instruction(migraphx::make_op("mul"), x, recip);
882
    }
Paul Fultz II's avatar
Paul Fultz II committed
883
    EXPECT(m1 == m2);
884
885
886
887
}

TEST_CASE(simplify_sub_const)
{
Paul Fultz II's avatar
Paul Fultz II committed
888
    migraphx::module m1;
889
    {
Paul Fultz II's avatar
Paul Fultz II committed
890
891
892
        auto x   = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto two = m1.add_literal(2);
        m1.add_instruction(migraphx::make_op("sub"), x, two);
893
    }
Paul Fultz II's avatar
Paul Fultz II committed
894
    run_pass(m1);
895

Paul Fultz II's avatar
Paul Fultz II committed
896
    migraphx::module m2;
897
    {
Paul Fultz II's avatar
Paul Fultz II committed
898
899
900
901
        auto x   = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto two = m2.add_literal(2);
        auto neg = m2.insert_instruction(std::next(two), migraphx::make_op("neg"), two);
        m2.add_instruction(migraphx::make_op("add"), x, neg);
902
    }
Paul Fultz II's avatar
Paul Fultz II committed
903
    EXPECT(m1 == m2);
904
905
}

kahmed10's avatar
kahmed10 committed
906
907
TEST_CASE(simplify_rsqrt)
{
Paul Fultz II's avatar
Paul Fultz II committed
908
    migraphx::module m1;
kahmed10's avatar
kahmed10 committed
909
    {
Paul Fultz II's avatar
Paul Fultz II committed
910
911
912
        auto x    = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto sqrt = m1.add_instruction(migraphx::make_op("sqrt"), x);
        m1.add_instruction(migraphx::make_op("recip"), sqrt);
kahmed10's avatar
kahmed10 committed
913
    }
Paul Fultz II's avatar
Paul Fultz II committed
914
    run_pass(m1);
kahmed10's avatar
kahmed10 committed
915

Paul Fultz II's avatar
Paul Fultz II committed
916
    migraphx::module m2;
kahmed10's avatar
kahmed10 committed
917
    {
Paul Fultz II's avatar
Paul Fultz II committed
918
919
        auto x = m2.add_parameter("x", {migraphx::shape::int32_type, {1}});
        m2.add_instruction(migraphx::make_op("rsqrt"), x);
kahmed10's avatar
kahmed10 committed
920
    }
Paul Fultz II's avatar
Paul Fultz II committed
921
    EXPECT(m1 == m2);
kahmed10's avatar
kahmed10 committed
922
923
924
925
}

TEST_CASE(simplify_rsqrt_multi_use)
{
Paul Fultz II's avatar
Paul Fultz II committed
926
    migraphx::module m1;
kahmed10's avatar
kahmed10 committed
927
    {
Paul Fultz II's avatar
Paul Fultz II committed
928
929
930
931
932
        auto x     = m1.add_parameter("x", {migraphx::shape::int32_type, {1}});
        auto sqrt  = m1.add_instruction(migraphx::make_op("sqrt"), x);
        auto add   = m1.add_instruction(migraphx::make_op("add"), sqrt, sqrt);
        auto rsqrt = m1.add_instruction(migraphx::make_op("recip"), sqrt);
        m1.add_instruction(migraphx::make_op("add"), rsqrt, add);
kahmed10's avatar
kahmed10 committed
933
    }
Paul Fultz II's avatar
Paul Fultz II committed
934
    migraphx::module m2{m1};
kahmed10's avatar
kahmed10 committed
935

Paul Fultz II's avatar
Paul Fultz II committed
936
937
    run_pass(m1);
    EXPECT(m1 == m2);
kahmed10's avatar
kahmed10 committed
938
939
}

940
941
942
943
TEST_CASE(simplify_slice_concat)
{
    auto s = migraphx::shape{migraphx::shape::float_type, {256}};

Paul Fultz II's avatar
Paul Fultz II committed
944
    migraphx::module m1;
945
    {
Paul Fultz II's avatar
Paul Fultz II committed
946
947
948
        auto x       = m1.add_parameter("x", s);
        auto y       = m1.add_parameter("y", s);
        auto xslice1 = m1.add_instruction(
949
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {128}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
950
        auto xslice2 = m1.add_instruction(
951
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {128}}, {"ends", {256}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
952
        auto yslice1 = m1.add_instruction(
953
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {128}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
954
        auto yslice2 = m1.add_instruction(
955
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {128}}, {"ends", {256}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
956
        auto concat = m1.add_instruction(
957
            migraphx::make_op("concat", {{"axis", 0}}), xslice1, xslice2, yslice1, yslice2);
Paul Fultz II's avatar
Paul Fultz II committed
958
        m1.add_instruction(pass_op{}, concat);
959
    }
Paul Fultz II's avatar
Paul Fultz II committed
960
    run_pass(m1);
961

Paul Fultz II's avatar
Paul Fultz II committed
962
    migraphx::module m2;
963
    {
Paul Fultz II's avatar
Paul Fultz II committed
964
965
966
967
        auto x      = m2.add_parameter("x", s);
        auto y      = m2.add_parameter("y", s);
        auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), x, y);
        m2.add_instruction(pass_op{}, concat);
968
    }
Paul Fultz II's avatar
Paul Fultz II committed
969
    EXPECT(m1 == m2);
970
971
972
973
974
975
}

TEST_CASE(simplify_slice_concat_non_uniform)
{
    auto s = migraphx::shape{migraphx::shape::float_type, {256}};

Paul Fultz II's avatar
Paul Fultz II committed
976
    migraphx::module m1;
977
    {
Paul Fultz II's avatar
Paul Fultz II committed
978
979
980
        auto x       = m1.add_parameter("x", s);
        auto y       = m1.add_parameter("y", s);
        auto xslice1 = m1.add_instruction(
981
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {64}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
982
        auto xslice2 = m1.add_instruction(
983
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {64}}, {"ends", {192}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
984
        auto xslice3 = m1.add_instruction(
985
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {192}}, {"ends", {256}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
986
        auto yslice1 = m1.add_instruction(
987
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {64}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
988
        auto yslice2 = m1.add_instruction(
989
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {64}}, {"ends", {192}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
990
        auto yslice3 = m1.add_instruction(
991
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {192}}, {"ends", {256}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
992
993
994
995
996
997
998
999
        auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 0}}),
                                         xslice1,
                                         xslice2,
                                         xslice3,
                                         yslice1,
                                         yslice2,
                                         yslice3);
        m1.add_instruction(pass_op{}, concat);
1000
    }
Paul Fultz II's avatar
Paul Fultz II committed
1001
    run_pass(m1);
1002

Paul Fultz II's avatar
Paul Fultz II committed
1003
    migraphx::module m2;
1004
    {
Paul Fultz II's avatar
Paul Fultz II committed
1005
1006
1007
1008
        auto x      = m2.add_parameter("x", s);
        auto y      = m2.add_parameter("y", s);
        auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), x, y);
        m2.add_instruction(pass_op{}, concat);
1009
1010
    }

Paul Fultz II's avatar
Paul Fultz II committed
1011
    EXPECT(m1 == m2);
1012
1013
1014
1015
1016
1017
}

TEST_CASE(simplify_slice_concat_flipped)
{
    auto s = migraphx::shape{migraphx::shape::float_type, {256}};

Paul Fultz II's avatar
Paul Fultz II committed
1018
    migraphx::module m1;
1019
    {
Paul Fultz II's avatar
Paul Fultz II committed
1020
1021
1022
        auto x       = m1.add_parameter("x", s);
        auto y       = m1.add_parameter("y", s);
        auto xslice1 = m1.add_instruction(
1023
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {64}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
1024
        auto xslice2 = m1.add_instruction(
1025
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {192}}, {"ends", {256}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
1026
        auto xslice3 = m1.add_instruction(
1027
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {64}}, {"ends", {192}}}), x);
Paul Fultz II's avatar
Paul Fultz II committed
1028
        auto yslice1 = m1.add_instruction(
1029
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {64}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
1030
        auto yslice2 = m1.add_instruction(
1031
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {192}}, {"ends", {256}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
1032
        auto yslice3 = m1.add_instruction(
1033
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {64}}, {"ends", {192}}}), y);
Paul Fultz II's avatar
Paul Fultz II committed
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
        auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 0}}),
                                         xslice1,
                                         xslice2,
                                         xslice3,
                                         yslice1,
                                         yslice2,
                                         yslice3);
        m1.add_instruction(pass_op{}, concat);
    }
    migraphx::module m2 = m1;
    run_pass(m1);

    EXPECT(m1 == m2);
1047
1048
}

1049
1050
1051
TEST_CASE(simplify_split_add_relu)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1052
    migraphx::module m1;
1053
1054
    {
        auto b     = migraphx::op::broadcast{1, {3, 1, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1055
1056
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1057
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1058
        auto y = m1.add_instruction(
1059
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
        auto one   = m1.add_literal(1);
        auto oneb  = m1.add_instruction(b, one);
        auto two   = m1.add_literal(2);
        auto twob  = m1.add_instruction(b, two);
        auto sum1  = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2  = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto add   = m1.add_instruction(migraphx::make_op("add"), relu1, relu2);
        m1.add_instruction(pass_op{}, add);
    }
    run_pass(m1);

    migraphx::module m2;
    {
1075
        auto b       = migraphx::op::broadcast{1, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1076
1077
1078
1079
1080
1081
1082
1083
        auto input   = m2.add_parameter("input", s);
        auto one     = m2.add_literal(1);
        auto two     = m2.add_literal(2);
        auto concat  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto concatb = m2.add_instruction(b, concat);
        auto sum     = m2.add_instruction(migraphx::make_op("add"), input, concatb);
        auto relu    = m2.add_instruction(migraphx::make_op("relu"), sum);
        auto x       = m2.add_instruction(
1084
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), relu);
Paul Fultz II's avatar
Paul Fultz II committed
1085
        auto y = m2.add_instruction(
1086
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), relu);
Paul Fultz II's avatar
Paul Fultz II committed
1087
1088
        auto add = m2.add_instruction(migraphx::make_op("add"), x, y);
        m2.add_instruction(pass_op{}, add);
1089
    }
Paul Fultz II's avatar
Paul Fultz II committed
1090
    EXPECT(m1.sort() == m2.sort());
1091
1092
}

Shucai Xiao's avatar
Shucai Xiao committed
1093
1094
1095
1096
1097
1098
1099
1100
TEST_CASE(simplify_split_reduce0)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
kahmed10's avatar
kahmed10 committed
1101
        auto y = m1.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1102
1103
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);

kahmed10's avatar
kahmed10 committed
1104
1105
        auto one = m1.add_literal(1);
        auto two = m1.add_literal(2);
Shucai Xiao's avatar
Shucai Xiao committed
1106

kahmed10's avatar
kahmed10 committed
1107
1108
        auto arx   = m1.add_instruction(migraphx::make_op("contiguous"), x);
        auto ary   = m1.add_instruction(migraphx::make_op("contiguous"), y);
Shucai Xiao's avatar
Shucai Xiao committed
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
        auto rmax0 = m1.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {0, 1}}}), x);
        auto rmin0 = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 1}}}), x);
        auto rmax1 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 1}}), arx, one);
        auto rmin1 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 1}}), ary, two);
        auto rmax2 = m1.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {0, 1}}}), y);
        auto rmin2 = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 1}}}), y);
        m1.add_return({rmax0, rmin0, rmax1, rmin1, rmax2, rmin2});
    }

    migraphx::module m2 = m1;
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
}

TEST_CASE(simplify_split_reduce1)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
        auto y = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);

        auto rmax0 = m1.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {0, 2}}}), x);
        auto rmin0 = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 2}}}), x);
        auto rmax2 = m1.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {0, 2}}}), y);
        auto rmin2 = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 2}}}), y);
        m1.add_return({rmax0, rmin0, rmax2, rmin2});
    }

    migraphx::module m2;
    {
        auto input = m2.add_parameter("input", s);

kahmed10's avatar
kahmed10 committed
1145
1146
        auto rmn  = m2.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 2}}}), input);
        auto slc0 = m2.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1147
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), rmn);
kahmed10's avatar
kahmed10 committed
1148
1149
        auto rmx  = m2.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {0, 2}}}), input);
        auto slc1 = m2.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1150
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), rmx);
kahmed10's avatar
kahmed10 committed
1151
        auto slc2 = m2.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1152
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), rmn);
kahmed10's avatar
kahmed10 committed
1153
        auto slc3 = m2.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), rmx);
        m2.add_return({slc3, slc2, slc1, slc0});
    }

    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
}

TEST_CASE(simplify_split_reduce2)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
        auto y = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
        auto rmax0 = m1.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {0, 2}}}), x);
        auto rmin0 = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 1}}}), x);
        auto rmax2 = m1.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {0, 2}}}), y);
        auto rmin2 = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 1}}}), y);
        m1.add_return({rmax0, rmin0, rmax2, rmin2});
    }

    migraphx::module m2;
    {
        auto input = m2.add_parameter("input", s);
        auto x     = m2.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
        auto rmn1 = m2.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 1}}}), x);
kahmed10's avatar
kahmed10 committed
1185
        auto y    = m2.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1186
1187
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
        auto rmn2 = m2.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0, 1}}}), y);
kahmed10's avatar
kahmed10 committed
1188
1189
        auto rms  = m2.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {0, 2}}}), input);
        auto slc0 = m2.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1190
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), rms);
kahmed10's avatar
kahmed10 committed
1191
        auto slc1 = m2.add_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
1192
1193
1194
1195
1196
1197
1198
1199
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), rms);
        m2.add_return({slc1, rmn2, slc0, rmn1});
    }

    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
}

1200
1201
1202
TEST_CASE(simplify_split_add_relu_reshape)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1203
    migraphx::module m1;
1204
    {
1205
1206
        auto b     = migraphx::op::broadcast{1, {3, 1, 4}};
        auto r     = migraphx::op::reshape{{3, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1207
1208
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1209
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1210
        auto y = m1.add_instruction(
1211
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
        auto one      = m1.add_literal(1);
        auto oneb     = m1.add_instruction(b, one);
        auto two      = m1.add_literal(2);
        auto twob     = m1.add_instruction(b, two);
        auto sum1     = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1    = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto reshape1 = m1.add_instruction(r, relu1);
        auto sum2     = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2    = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto reshape2 = m1.add_instruction(r, relu2);
        auto add      = m1.add_instruction(migraphx::make_op("add"), reshape1, reshape2);
        m1.add_instruction(pass_op{}, add);
    }
    run_pass(m1);

    migraphx::module m2;
    {
1229
        auto b       = migraphx::op::broadcast{1, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1230
1231
1232
1233
1234
1235
1236
1237
1238
        auto input   = m2.add_parameter("input", s);
        auto one     = m2.add_literal(1);
        auto two     = m2.add_literal(2);
        auto concat  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto concatb = m2.add_instruction(b, concat);
        auto sum     = m2.add_instruction(migraphx::make_op("add"), input, concatb);
        auto relu    = m2.add_instruction(migraphx::make_op("relu"), sum);
        auto rsp     = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {3, 8}}}), relu);
        auto slc1    = m2.add_instruction(
1239
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {4}}}), rsp);
Paul Fultz II's avatar
Paul Fultz II committed
1240
        auto slc2 = m2.add_instruction(
1241
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {4}}, {"ends", {8}}}), rsp);
Paul Fultz II's avatar
Paul Fultz II committed
1242
1243
        auto add = m2.add_instruction(migraphx::make_op("add"), slc1, slc2);
        m2.add_instruction(pass_op{}, add);
1244
    }
Paul Fultz II's avatar
Paul Fultz II committed
1245
    EXPECT(m1.sort() == m2.sort());
1246
1247
1248
1249
1250
}

TEST_CASE(simplify_slice_different_axis)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4, 2}};
Paul Fultz II's avatar
Paul Fultz II committed
1251
    migraphx::module m1;
1252
    {
1253
        auto r     = migraphx::op::reshape{{3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1254
1255
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1256
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1257
        auto y = m1.add_instruction(
1258
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1259
1260
        auto one  = m1.add_literal(1);
        auto oneb = m1.add_instruction(
1261
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {3, 1, 4, 2}}}), one);
Paul Fultz II's avatar
Paul Fultz II committed
1262
1263
        auto two  = m1.add_literal(2);
        auto twob = m1.add_instruction(
1264
            migraphx::make_op("broadcast", {{"axis", 3}, {"out_lens", {3, 2, 4, 1}}}), two);
Paul Fultz II's avatar
Paul Fultz II committed
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
        auto sum1     = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1    = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto reshape1 = m1.add_instruction(r, relu1);
        auto sum2     = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2    = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto reshape2 = m1.add_instruction(r, relu2);
        auto add      = m1.add_instruction(migraphx::make_op("add"), reshape1, reshape2);
        m1.add_instruction(pass_op{}, add);
    }
    migraphx::module m2 = m1;
    run_pass(m1);

    EXPECT(m1.sort() == m2.sort());
1278
1279
1280
1281
1282
}

TEST_CASE(simplify_slice_missing_begining_slice)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 3, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1283
    migraphx::module m1;
1284
1285
    {
        auto b     = migraphx::op::broadcast{1, {3, 1, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1286
1287
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1288
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {2}}, {"ends", {3}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1289
        auto y = m1.add_instruction(
1290
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
        auto one   = m1.add_literal(1);
        auto oneb  = m1.add_instruction(b, one);
        auto two   = m1.add_literal(2);
        auto twob  = m1.add_instruction(b, two);
        auto sum1  = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2  = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto add   = m1.add_instruction(migraphx::make_op("add"), relu1, relu2);
        m1.add_instruction(pass_op{}, add);
    }
    migraphx::module m2 = m1;
    run_pass(m1);

    EXPECT(m1.sort() == m2.sort());
1306
1307
1308
1309
1310
}

TEST_CASE(simplify_slice_missing_middle_slice)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 3, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1311
    migraphx::module m1;
1312
1313
    {
        auto b     = migraphx::op::broadcast{1, {3, 1, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1314
1315
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1316
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {2}}, {"ends", {3}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1317
        auto y = m1.add_instruction(
1318
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
        auto one   = m1.add_literal(1);
        auto oneb  = m1.add_instruction(b, one);
        auto two   = m1.add_literal(2);
        auto twob  = m1.add_instruction(b, two);
        auto sum1  = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2  = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto add   = m1.add_instruction(migraphx::make_op("add"), relu1, relu2);
        m1.add_instruction(pass_op{}, add);
    }
    migraphx::module m2 = m1;
    run_pass(m1);

    EXPECT(m1.sort() == m2.sort());
1334
1335
1336
1337
1338
}

TEST_CASE(simplify_slice_missing_end_slice)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 3, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1339
    migraphx::module m1;
1340
1341
    {
        auto b     = migraphx::op::broadcast{1, {3, 1, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1342
1343
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1344
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1345
        auto y = m1.add_instruction(
1346
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
        auto one   = m1.add_literal(1);
        auto oneb  = m1.add_instruction(b, one);
        auto two   = m1.add_literal(2);
        auto twob  = m1.add_instruction(b, two);
        auto sum1  = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2  = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto add   = m1.add_instruction(migraphx::make_op("add"), relu1, relu2);
        m1.add_instruction(pass_op{}, add);
    }
    migraphx::module m2 = m1;
    run_pass(m1);

    EXPECT(m1.sort() == m2.sort());
1362
1363
1364
1365
1366
}

TEST_CASE(simplify_split_add_relu_concat_same_axis)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1367
    migraphx::module m1;
1368
    {
1369
        auto b     = migraphx::op::broadcast{1, {3, 1, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1370
1371
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1372
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1373
        auto y = m1.add_instruction(
1374
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
        auto one    = m1.add_literal(1);
        auto oneb   = m1.add_instruction(b, one);
        auto two    = m1.add_literal(2);
        auto twob   = m1.add_instruction(b, two);
        auto sum1   = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1  = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2   = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2  = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), relu1, relu2);
        m1.add_instruction(pass_op{}, concat);
1385
    }
Paul Fultz II's avatar
Paul Fultz II committed
1386
    run_pass(m1);
1387

Paul Fultz II's avatar
Paul Fultz II committed
1388
    migraphx::module m2;
1389
1390
    {
        auto b       = migraphx::op::broadcast{1, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
        auto input   = m2.add_parameter("input", s);
        auto one     = m2.add_literal(1);
        auto two     = m2.add_literal(2);
        auto concat  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto concatb = m2.add_instruction(b, concat);
        auto sum     = m2.add_instruction(migraphx::make_op("add"), input, concatb);
        auto relu    = m2.add_instruction(migraphx::make_op("relu"), sum);
        m2.add_instruction(pass_op{}, relu);
    }
    EXPECT(m1.sort() == m2.sort());
1401
1402
1403
1404
1405
}

TEST_CASE(simplify_split_add_relu_multi_axes)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4, 6}};
Paul Fultz II's avatar
Paul Fultz II committed
1406
    migraphx::module m1;
1407
1408
    {
        auto b     = migraphx::op::broadcast{1, {3, 1, 4, 3}};
Paul Fultz II's avatar
Paul Fultz II committed
1409
1410
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1411
1412
            migraphx::make_op("slice", {{"axes", {1, 3}}, {"starts", {0, 0}}, {"ends", {1, 3}}}),
            input);
Paul Fultz II's avatar
Paul Fultz II committed
1413
        auto y = m1.add_instruction(
1414
1415
            migraphx::make_op("slice", {{"axes", {1, 3}}, {"starts", {1, 3}}, {"ends", {2, 6}}}),
            input);
Paul Fultz II's avatar
Paul Fultz II committed
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
        auto one   = m1.add_literal(1);
        auto oneb  = m1.add_instruction(b, one);
        auto two   = m1.add_literal(2);
        auto twob  = m1.add_instruction(b, two);
        auto sum1  = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2  = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto add   = m1.add_instruction(migraphx::make_op("add"), relu1, relu2);
        m1.add_instruction(pass_op{}, add);
    }
    migraphx::module m2 = m1;
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
1430
1431
1432
1433
1434
}

TEST_CASE(simplify_split_add_relu_used_multiple_split1)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1435
    migraphx::module m1;
1436
1437
    {
        auto b     = migraphx::op::broadcast{1, {3, 1, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1438
1439
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1440
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1441
        auto y = m1.add_instruction(
1442
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
        auto one   = m1.add_literal(1);
        auto oneb  = m1.add_instruction(b, one);
        auto two   = m1.add_literal(2);
        auto twob  = m1.add_instruction(b, two);
        auto sum1  = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2  = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto add1  = m1.add_instruction(migraphx::make_op("add"), relu1, relu2);
        auto add2  = m1.add_instruction(migraphx::make_op("add"), x, add1);
        m1.add_instruction(pass_op{}, add2);
    }
    run_pass(m1);

    migraphx::module m2;
    {
1459
        auto b     = migraphx::op::broadcast{1, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1460
1461
        auto input = m2.add_parameter("input", s);
        auto slice = m2.add_instruction(
1462
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1463
1464
1465
1466
1467
1468
1469
        auto one     = m2.add_literal(1);
        auto two     = m2.add_literal(2);
        auto concat  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto concatb = m2.add_instruction(b, concat);
        auto sum     = m2.add_instruction(migraphx::make_op("add"), input, concatb);
        auto relu    = m2.add_instruction(migraphx::make_op("relu"), sum);
        auto x       = m2.add_instruction(
1470
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), relu);
Paul Fultz II's avatar
Paul Fultz II committed
1471
        auto y = m2.add_instruction(
1472
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), relu);
Paul Fultz II's avatar
Paul Fultz II committed
1473
1474
1475
        auto add1 = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto add2 = m2.add_instruction(migraphx::make_op("add"), slice, add1);
        m2.add_instruction(pass_op{}, add2);
1476
    }
Paul Fultz II's avatar
Paul Fultz II committed
1477
    EXPECT(m1.sort() == m2.sort());
1478
1479
1480
1481
1482
}

TEST_CASE(simplify_split_add_relu_used_multiple_split2)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1483
    migraphx::module m1;
1484
1485
    {
        auto b     = migraphx::op::broadcast{1, {3, 1, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1486
1487
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1488
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1489
        auto y = m1.add_instruction(
1490
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
        auto z     = m1.add_instruction(migraphx::make_op("relu"), x);
        auto one   = m1.add_literal(1);
        auto oneb  = m1.add_instruction(b, one);
        auto two   = m1.add_literal(2);
        auto twob  = m1.add_instruction(b, two);
        auto sum1  = m1.add_instruction(migraphx::make_op("add"), x, oneb);
        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), sum1);
        auto sum2  = m1.add_instruction(migraphx::make_op("add"), y, twob);
        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), sum2);
        auto add1  = m1.add_instruction(migraphx::make_op("add"), relu1, relu2);
        auto add2  = m1.add_instruction(migraphx::make_op("add"), z, add1);
        m1.add_instruction(pass_op{}, add2);
    }
    run_pass(m1);

    migraphx::module m2;
    {
1508
        auto b     = migraphx::op::broadcast{1, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1509
1510
        auto input = m2.add_parameter("input", s);
        auto slice = m2.add_instruction(
1511
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1512
1513
1514
1515
1516
1517
1518
1519
        auto z       = m2.add_instruction(migraphx::make_op("relu"), slice);
        auto one     = m2.add_literal(1);
        auto two     = m2.add_literal(2);
        auto concat  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), one, two);
        auto concatb = m2.add_instruction(b, concat);
        auto sum     = m2.add_instruction(migraphx::make_op("add"), input, concatb);
        auto relu    = m2.add_instruction(migraphx::make_op("relu"), sum);
        auto x       = m2.add_instruction(
1520
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), relu);
Paul Fultz II's avatar
Paul Fultz II committed
1521
        auto y = m2.add_instruction(
1522
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), relu);
Paul Fultz II's avatar
Paul Fultz II committed
1523
1524
1525
        auto add1 = m2.add_instruction(migraphx::make_op("add"), x, y);
        auto add2 = m2.add_instruction(migraphx::make_op("add"), z, add1);
        m2.add_instruction(pass_op{}, add2);
1526
    }
Paul Fultz II's avatar
Paul Fultz II committed
1527
    EXPECT(m1.sort() == m2.sort());
1528
1529
1530
1531
1532
}

TEST_CASE(simplify_split_between_add)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 4}};
Paul Fultz II's avatar
Paul Fultz II committed
1533
    migraphx::module m1;
1534
    {
Paul Fultz II's avatar
Paul Fultz II committed
1535
1536
        auto input = m1.add_parameter("input", s);
        auto x     = m1.add_instruction(
1537
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1538
        auto y = m1.add_instruction(
1539
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
1540
1541
        auto sum = m1.add_instruction(migraphx::make_op("add"), x, y);
        m1.add_instruction(pass_op{}, sum);
1542
    }
Paul Fultz II's avatar
Paul Fultz II committed
1543
1544
1545
    migraphx::module m2 = m1;
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
1546
1547
}

Paul's avatar
Paul committed
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
TEST_CASE(simplify_split_binary_dot)
{

    migraphx::module m1;
    {
        auto x1 = m1.add_parameter("x1", {migraphx::shape::float_type, {1, 160, 4}});
        auto w1 = m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 64}}, 1));
        auto w1b = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 64}}}), w1);
        auto dot1 = m1.add_instruction(migraphx::make_op("dot"), x1, w1b);

        auto x2 = m1.add_parameter("x2", {migraphx::shape::float_type, {1, 160, 4}});
        auto w2 = m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 64}}, 2));
        auto w2b = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 64}}}), w2);
        auto dot2 = m1.add_instruction(migraphx::make_op("dot"), x2, w2b);

        auto slice11  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {32}}}), dot1);
        auto slice12  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {32}}, {"ends", {64}}}), dot1);

        auto slice21  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {32}}}), dot2);
        auto slice22  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {32}}, {"ends", {64}}}), dot2);
        auto sum1 = m1.add_instruction(migraphx::make_op("add"), slice11, slice21);
        auto sum2 = m1.add_instruction(migraphx::make_op("add"), slice12, slice22);
        auto ret = m1.add_instruction(migraphx::make_op("mul"), sum1, sum2);
        m1.add_return({ret});
    };

    migraphx::module m2;
    {
        // TODO: Fuse these dot operators
        auto x1 = m2.add_parameter("x1", {migraphx::shape::float_type, {1, 160, 4}});
        auto w1 = m2.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 64}}, 1));
        auto w1b = m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 64}}}), w1);
        auto dot1 = m2.add_instruction(migraphx::make_op("dot"), x1, w1b);

        auto x2 = m2.add_parameter("x2", {migraphx::shape::float_type, {1, 160, 4}});
        auto w2 = m2.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 64}}, 2));
        auto w2b = m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 64}}}), w2);
        auto dot2 = m2.add_instruction(migraphx::make_op("dot"), x2, w2b);

        auto sum = m2.add_instruction(migraphx::make_op("add"), dot1, dot2);
        auto slice1  = m2.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {32}}}), sum);
        auto slice2  = m2.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {32}}, {"ends", {64}}}), sum);
        auto ret = m2.add_instruction(migraphx::make_op("mul"), slice1, slice2);

        m2.add_return({ret});
    };

    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
}

TEST_CASE(simplify_split_binary_same_input)
{
    migraphx::module m1;
    {
        auto x1 = m1.add_parameter("x1", {migraphx::shape::float_type, {1, 160, 4}});
        auto w1 = m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 64}}, 1));
        auto w1b = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 64}}}), w1);
        auto dot1 = m1.add_instruction(migraphx::make_op("dot"), x1, w1b);

        auto x2 = m1.add_parameter("x2", {migraphx::shape::float_type, {1, 160, 4}});
        auto w2 = m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 64}}, 2));
        auto w2b = m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 64}}}), w2);
        auto dot2 = m1.add_instruction(migraphx::make_op("dot"), x2, w2b);

        auto slice11  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {32}}}), dot1);
        auto slice12  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {32}}, {"ends", {64}}}), dot1);
        auto sum1 = m1.add_instruction(migraphx::make_op("add"), slice11, slice12);

        auto slice21  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {32}}}), dot2);
        auto slice22  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {32}}, {"ends", {64}}}), dot2);
        auto sum2 = m1.add_instruction(migraphx::make_op("add"), slice21, slice22);
        auto ret = m1.add_instruction(migraphx::make_op("mul"), sum1, sum2);
        m1.add_return({ret});
    };

    migraphx::module m2 = m1;
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
}

1639
1640
1641
TEST_CASE(simplify_dot_horiz)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 2}};
Paul Fultz II's avatar
Paul Fultz II committed
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto a     = m1.add_literal(migraphx::generate_literal(s, 0));
        auto b     = m1.add_literal(migraphx::generate_literal(s, 1));
        auto x     = m1.add_instruction(migraphx::make_op("dot"), input, a);
        auto y     = m1.add_instruction(migraphx::make_op("dot"), input, b);
        auto sum   = m1.add_instruction(migraphx::make_op("add"), x, y);
        m1.add_instruction(pass_op{}, sum);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto input  = m2.add_parameter("input", s);
        auto a      = m2.add_literal(migraphx::generate_literal(s, 0));
        auto b      = m2.add_literal(migraphx::generate_literal(s, 1));
        auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 2}}), a, b);
        auto dot    = m2.add_instruction(migraphx::make_op("dot"), input, concat);
        auto x      = m2.add_instruction(
1662
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {2}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1663
        auto y = m2.add_instruction(
1664
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {2}}, {"ends", {4}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1665
1666
        auto sum = m2.add_instruction(migraphx::make_op("add"), x, y);
        m2.add_instruction(pass_op{}, sum);
1667
    }
Paul Fultz II's avatar
Paul Fultz II committed
1668
    EXPECT(m1.sort() == m2.sort());
1669
1670
1671
1672
1673
}

TEST_CASE(simplify_dot_horiz_same_constant)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 2}};
Paul Fultz II's avatar
Paul Fultz II committed
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto a     = m1.add_literal(migraphx::generate_literal(s, 0));
        auto x     = m1.add_instruction(migraphx::make_op("dot"), input, a);
        auto y     = m1.add_instruction(migraphx::make_op("dot"), input, a);
        auto sum   = m1.add_instruction(migraphx::make_op("add"), x, y);
        m1.add_instruction(pass_op{}, sum);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto input  = m2.add_parameter("input", s);
        auto a      = m2.add_literal(migraphx::generate_literal(s, 0));
        auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 2}}), a, a);
        auto dot    = m2.add_instruction(migraphx::make_op("dot"), input, concat);
        auto x      = m2.add_instruction(
1692
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {2}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1693
        auto y = m2.add_instruction(
1694
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {2}}, {"ends", {4}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1695
1696
        auto sum = m2.add_instruction(migraphx::make_op("add"), x, y);
        m2.add_instruction(pass_op{}, sum);
1697
    }
Paul Fultz II's avatar
Paul Fultz II committed
1698
    EXPECT(m1.sort() == m2.sort());
1699
1700
1701
1702
1703
}

TEST_CASE(simplify_dot_horiz_flipped)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 2, 2}};
Paul Fultz II's avatar
Paul Fultz II committed
1704
    migraphx::module m1;
1705
    {
Paul Fultz II's avatar
Paul Fultz II committed
1706
1707
1708
1709
1710
1711
1712
        auto input = m1.add_parameter("input", s);
        auto a     = m1.add_literal(migraphx::generate_literal(s, 0));
        auto b     = m1.add_literal(migraphx::generate_literal(s, 1));
        auto x     = m1.add_instruction(migraphx::make_op("dot"), input, a);
        auto y     = m1.add_instruction(migraphx::make_op("dot"), b, input);
        auto sum   = m1.add_instruction(migraphx::make_op("add"), x, y);
        m1.add_instruction(pass_op{}, sum);
1713
1714
    }

Paul Fultz II's avatar
Paul Fultz II committed
1715
1716
1717
    migraphx::module m2 = m1;
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
1718
1719
}

1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
// test if contiguous is added as necessary for reshapes
TEST_CASE(simplify_dot_horiz_reshape)
{
    auto s = migraphx::shape{migraphx::shape::int32_type, {3, 4, 4}};
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto a     = m1.add_literal(migraphx::generate_literal(s, 0));
        auto b     = m1.add_literal(migraphx::generate_literal(s, 1));
        auto x     = m1.add_instruction(migraphx::make_op("dot"), input, a);
        auto y     = m1.add_instruction(migraphx::make_op("dot"), input, b);
        auto x_rsp = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {3, 4, 2, 2}}}), x);
        auto y_rsp =
            m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {2}}, {"steps", {2}}}), y);
        auto sum = m1.add_instruction(migraphx::make_op("add"), {x_rsp, y_rsp});
        m1.add_instruction(pass_op{}, sum);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto input  = m2.add_parameter("input", s);
        auto a      = m2.add_literal(migraphx::generate_literal(s, 0));
        auto b      = m2.add_literal(migraphx::generate_literal(s, 1));
        auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 2}}), a, b);
        auto dot    = m2.add_instruction(migraphx::make_op("dot"), input, concat);
        auto x      = m2.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {4}}}), dot);
        auto y = m2.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {4}}, {"ends", {8}}}), dot);
        auto x_cont = m2.add_instruction(migraphx::make_op("contiguous"), x);
        auto x_rsp =
            m2.add_instruction(migraphx::make_op("reshape", {{"dims", {3, 4, 2, 2}}}), x_cont);
        auto y_rsp =
            m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {2}}, {"steps", {2}}}), y);
        auto sum = m2.add_instruction(migraphx::make_op("add"), {x_rsp, y_rsp});
        m2.add_instruction(pass_op{}, sum);
    }

    EXPECT(m1.sort() == m2.sort());
}

1762
1763
1764
1765
TEST_CASE(simplify_conv_horiz)
{
    auto s  = migraphx::shape{migraphx::shape::int32_type, {8, 3, 64, 64}};
    auto ws = migraphx::shape{migraphx::shape::int32_type, {12, 3, 3, 3}};
Paul Fultz II's avatar
Paul Fultz II committed
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto a     = m1.add_literal(migraphx::generate_literal(ws, 0));
        auto b     = m1.add_literal(migraphx::generate_literal(ws, 1));
        auto x     = m1.add_instruction(migraphx::make_op("convolution"), input, a);
        auto y     = m1.add_instruction(migraphx::make_op("convolution"), input, b);
        auto sum   = m1.add_instruction(migraphx::make_op("add"), x, y);
        m1.add_instruction(pass_op{}, sum);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto input  = m2.add_parameter("input", s);
        auto a      = m2.add_literal(migraphx::generate_literal(ws, 0));
        auto b      = m2.add_literal(migraphx::generate_literal(ws, 1));
        auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), a, b);
        auto conv   = m2.add_instruction(migraphx::make_op("convolution"), input, concat);
        auto x      = m2.add_instruction(
1786
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {12}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
1787
        auto y = m2.add_instruction(
1788
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {12}}, {"ends", {24}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
1789
1790
        auto sum = m2.add_instruction(migraphx::make_op("add"), x, y);
        m2.add_instruction(pass_op{}, sum);
1791
    }
Paul Fultz II's avatar
Paul Fultz II committed
1792
    EXPECT(m1.sort() == m2.sort());
1793
1794
}

1795
1796
1797
1798
TEST_CASE(simplify_group_conv_horiz)
{
    auto s  = migraphx::shape{migraphx::shape::int32_type, {1, 32, 111, 111}};
    auto ws = migraphx::shape{migraphx::shape::int32_type, {32, 1, 7, 7}};
Paul Fultz II's avatar
Paul Fultz II committed
1799
    migraphx::module m1;
1800
    {
Paul Fultz II's avatar
Paul Fultz II committed
1801
1802
1803
1804
        auto x     = m1.add_parameter("x", s);
        auto w1    = m1.add_literal(migraphx::generate_literal(ws, 1));
        auto w2    = m1.add_literal(migraphx::generate_literal(ws, 2));
        auto conv1 = m1.add_instruction(
1805
1806
1807
1808
1809
            migraphx::make_op(
                "convolution",
                {{"padding", {3, 3}}, {"stride", {2, 2}}, {"dilation", {1, 1}}, {"group", 32}}),
            x,
            w1);
Paul Fultz II's avatar
Paul Fultz II committed
1810
        auto conv2 = m1.add_instruction(
1811
1812
1813
1814
1815
            migraphx::make_op(
                "convolution",
                {{"padding", {3, 3}}, {"stride", {2, 2}}, {"dilation", {1, 1}}, {"group", 32}}),
            x,
            w2);
Paul Fultz II's avatar
Paul Fultz II committed
1816
        m1.add_instruction(pass_op{}, conv1, conv2);
1817
    }
Paul Fultz II's avatar
Paul Fultz II committed
1818
1819
    migraphx::module m2 = m1;
    run_pass(m1);
1820

Paul Fultz II's avatar
Paul Fultz II committed
1821
    EXPECT(m1.sort() == m2.sort());
1822
1823
1824
}

TEST_CASE(simplify_conv_horiz_grouped)
1825
1826
1827
1828
{
    auto s   = migraphx::shape{migraphx::shape::int32_type, {8, 6, 64, 64}};
    auto ws1 = migraphx::shape{migraphx::shape::int32_type, {6, 6, 3, 3}};
    auto ws2 = migraphx::shape{migraphx::shape::int32_type, {8, 6, 64, 64}};
Paul Fultz II's avatar
Paul Fultz II committed
1829
1830
1831
1832
1833
1834
1835
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto a     = m1.add_literal(migraphx::generate_literal(ws1, 0));
        auto b     = m1.add_literal(migraphx::generate_literal(ws1, 1));
        auto c     = m1.add_literal(migraphx::generate_literal(ws2, 2));
        auto d     = m1.add_literal(migraphx::generate_literal(ws2, 3));
1836
        auto convx =
Paul Fultz II's avatar
Paul Fultz II committed
1837
            m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, a);
1838
        auto convy =
Paul Fultz II's avatar
Paul Fultz II committed
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
            m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, b);
        auto dotx = m1.add_instruction(migraphx::make_op("dot"), input, c);
        auto doty = m1.add_instruction(migraphx::make_op("dot"), input, d);
        auto sum1 = m1.add_instruction(migraphx::make_op("add"), convx, convy);
        auto sum2 = m1.add_instruction(migraphx::make_op("add"), dotx, doty);
        auto sum3 = m1.add_instruction(migraphx::make_op("add"), sum1, sum2);

        m1.add_instruction(pass_op{}, sum3);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto input   = m2.add_parameter("input", s);
        auto a       = m2.add_literal(migraphx::generate_literal(ws1, 0));
        auto b       = m2.add_literal(migraphx::generate_literal(ws1, 1));
        auto c       = m2.add_literal(migraphx::generate_literal(ws2, 2));
        auto d       = m2.add_literal(migraphx::generate_literal(ws2, 3));
        auto concat1 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), a, b);
        auto concat2 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 3}}), c, d);
        auto conv    = m2.add_instruction(
1860
            migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, concat1);
Paul Fultz II's avatar
Paul Fultz II committed
1861
        auto convx = m2.add_instruction(
1862
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {6}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
1863
        auto convy = m2.add_instruction(
1864
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {6}}, {"ends", {12}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
1865
1866
1867
        auto sum1 = m2.add_instruction(migraphx::make_op("add"), convx, convy);
        auto dot  = m2.add_instruction(migraphx::make_op("dot"), input, concat2);
        auto dotx = m2.add_instruction(
1868
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {0}}, {"ends", {64}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1869
        auto doty = m2.add_instruction(
1870
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {64}}, {"ends", {128}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1871
1872
1873
        auto sum2 = m2.add_instruction(migraphx::make_op("add"), dotx, doty);
        auto sum3 = m2.add_instruction(migraphx::make_op("add"), sum1, sum2);
        m2.add_instruction(pass_op{}, sum3);
1874
    }
Paul Fultz II's avatar
Paul Fultz II committed
1875
    EXPECT(m1.sort() == m2.sort());
1876
1877
}

1878
TEST_CASE(simplify_conv_horiz_grouped_extra1)
1879
1880
1881
1882
{
    auto s   = migraphx::shape{migraphx::shape::int32_type, {8, 6, 64, 64}};
    auto ws1 = migraphx::shape{migraphx::shape::int32_type, {6, 6, 3, 3}};
    auto ws2 = migraphx::shape{migraphx::shape::int32_type, {8, 6, 64, 64}};
Paul Fultz II's avatar
Paul Fultz II committed
1883
1884
1885
1886
1887
1888
1889
1890
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto a     = m1.add_literal(migraphx::generate_literal(ws1, 0));
        auto b     = m1.add_literal(migraphx::generate_literal(ws1, 1));
        auto c     = m1.add_literal(migraphx::generate_literal(ws2, 2));
        auto d     = m1.add_literal(migraphx::generate_literal(ws2, 3));
        auto e     = m1.add_literal(migraphx::generate_literal(s, 4));
1891
        auto convx =
Paul Fultz II's avatar
Paul Fultz II committed
1892
            m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, a);
1893
        auto convy =
Paul Fultz II's avatar
Paul Fultz II committed
1894
1895
1896
1897
1898
1899
            m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, b);
        auto dotx    = m1.add_instruction(migraphx::make_op("dot"), input, c);
        auto doty    = m1.add_instruction(migraphx::make_op("dot"), input, d);
        auto sqdiffx = m1.add_instruction(migraphx::make_op("sqdiff"), input, e);
        auto sum1    = m1.add_instruction(migraphx::make_op("add"), convx, convy);
        auto sum2    = m1.add_instruction(migraphx::make_op("add"), dotx, doty);
1900
        auto sum3    = sqdiffx;
Paul Fultz II's avatar
Paul Fultz II committed
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
        auto sum4    = m1.add_instruction(migraphx::make_op("add"), sum1, sum2);
        auto sum5    = m1.add_instruction(migraphx::make_op("add"), sum4, sum3);
        m1.add_instruction(pass_op{}, sum5);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto input   = m2.add_parameter("input", s);
        auto a       = m2.add_literal(migraphx::generate_literal(ws1, 0));
        auto b       = m2.add_literal(migraphx::generate_literal(ws1, 1));
        auto c       = m2.add_literal(migraphx::generate_literal(ws2, 2));
        auto d       = m2.add_literal(migraphx::generate_literal(ws2, 3));
        auto e       = m2.add_literal(migraphx::generate_literal(s, 4));
        auto concat1 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), a, b);
        auto concat2 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 3}}), c, d);
        auto conv    = m2.add_instruction(
1918
            migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, concat1);
Paul Fultz II's avatar
Paul Fultz II committed
1919
        auto convx = m2.add_instruction(
1920
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {6}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
1921
        auto convy = m2.add_instruction(
1922
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {6}}, {"ends", {12}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
1923
1924
1925
        auto sum1 = m2.add_instruction(migraphx::make_op("add"), convx, convy);
        auto dot  = m2.add_instruction(migraphx::make_op("dot"), input, concat2);
        auto dotx = m2.add_instruction(
1926
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {0}}, {"ends", {64}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1927
        auto doty = m2.add_instruction(
1928
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {64}}, {"ends", {128}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1929
1930
        auto sum2    = m2.add_instruction(migraphx::make_op("add"), dotx, doty);
        auto sqdiffx = m2.add_instruction(migraphx::make_op("sqdiff"), input, e);
1931
        auto sum3    = sqdiffx;
Paul Fultz II's avatar
Paul Fultz II committed
1932
1933
1934
        auto sum4    = m2.add_instruction(migraphx::make_op("add"), sum1, sum2);
        auto sum5    = m2.add_instruction(migraphx::make_op("add"), sum4, sum3);
        m2.add_instruction(pass_op{}, sum5);
1935
    }
Paul Fultz II's avatar
Paul Fultz II committed
1936
    EXPECT(m1.sort() == m2.sort());
1937
1938
}

1939
TEST_CASE(simplify_conv_horiz_grouped_extra2)
1940
1941
1942
1943
{
    auto s   = migraphx::shape{migraphx::shape::int32_type, {8, 6, 64, 64}};
    auto ws1 = migraphx::shape{migraphx::shape::int32_type, {6, 6, 3, 3}};
    auto ws2 = migraphx::shape{migraphx::shape::int32_type, {8, 6, 64, 64}};
Paul Fultz II's avatar
Paul Fultz II committed
1944
1945
1946
1947
1948
1949
1950
1951
1952
    migraphx::module m1;
    {
        auto input = m1.add_parameter("input", s);
        auto a     = m1.add_literal(migraphx::generate_literal(ws1, 0));
        auto b     = m1.add_literal(migraphx::generate_literal(ws1, 1));
        auto c     = m1.add_literal(migraphx::generate_literal(ws2, 2));
        auto d     = m1.add_literal(migraphx::generate_literal(ws2, 3));
        auto e     = m1.add_literal(migraphx::generate_literal(s, 4));
        auto f     = m1.add_literal(migraphx::generate_literal(s, 5));
1953
        auto convx =
Paul Fultz II's avatar
Paul Fultz II committed
1954
            m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, a);
1955
        auto convy =
Paul Fultz II's avatar
Paul Fultz II committed
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
            m1.add_instruction(migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, b);
        auto dotx    = m1.add_instruction(migraphx::make_op("dot"), input, c);
        auto doty    = m1.add_instruction(migraphx::make_op("dot"), input, d);
        auto sqdiffx = m1.add_instruction(migraphx::make_op("sqdiff"), input, e);
        auto sqdiffy = m1.add_instruction(migraphx::make_op("sqdiff"), input, f);
        auto sum1    = m1.add_instruction(migraphx::make_op("add"), convx, convy);
        auto sum2    = m1.add_instruction(migraphx::make_op("add"), dotx, doty);
        auto sum3    = m1.add_instruction(migraphx::make_op("add"), sqdiffx, sqdiffy);
        auto sum4    = m1.add_instruction(migraphx::make_op("add"), sum1, sum2);
        auto sum5    = m1.add_instruction(migraphx::make_op("add"), sum4, sum3);
        m1.add_instruction(pass_op{}, sum5);
    }
    run_pass(m1);

    migraphx::module m2;
    {
        auto input   = m2.add_parameter("input", s);
        auto a       = m2.add_literal(migraphx::generate_literal(ws1, 0));
        auto b       = m2.add_literal(migraphx::generate_literal(ws1, 1));
        auto c       = m2.add_literal(migraphx::generate_literal(ws2, 2));
        auto d       = m2.add_literal(migraphx::generate_literal(ws2, 3));
        auto e       = m2.add_literal(migraphx::generate_literal(s, 4));
        auto f       = m2.add_literal(migraphx::generate_literal(s, 5));
        auto concat1 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), a, b);
        auto concat2 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 3}}), c, d);
        auto conv    = m2.add_instruction(
1982
            migraphx::make_op("convolution", {{"padding", {1, 1}}}), input, concat1);
Paul Fultz II's avatar
Paul Fultz II committed
1983
        auto convx = m2.add_instruction(
1984
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {6}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
1985
        auto convy = m2.add_instruction(
1986
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {6}}, {"ends", {12}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
1987
1988
1989
        auto sum1 = m2.add_instruction(migraphx::make_op("add"), convx, convy);
        auto dot  = m2.add_instruction(migraphx::make_op("dot"), input, concat2);
        auto dotx = m2.add_instruction(
1990
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {0}}, {"ends", {64}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1991
        auto doty = m2.add_instruction(
1992
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {64}}, {"ends", {128}}}), dot);
Paul Fultz II's avatar
Paul Fultz II committed
1993
1994
1995
1996
1997
1998
1999
2000
2001
        auto sum2    = m2.add_instruction(migraphx::make_op("add"), dotx, doty);
        auto sqdiffx = m2.add_instruction(migraphx::make_op("sqdiff"), input, e);
        auto sqdiffy = m2.add_instruction(migraphx::make_op("sqdiff"), input, f);
        auto sum3    = m2.add_instruction(migraphx::make_op("add"), sqdiffx, sqdiffy);
        auto sum4    = m2.add_instruction(migraphx::make_op("add"), sum1, sum2);
        auto sum5    = m2.add_instruction(migraphx::make_op("add"), sum4, sum3);
        m2.add_instruction(pass_op{}, sum5);
    }
    EXPECT(m1.sort() == m2.sort());
2002
2003
}

2004
2005
TEST_CASE(simplify_mul_slice_conv_horiz_fusion)
{
Paul Fultz II's avatar
Paul Fultz II committed
2006
    migraphx::module m1;
2007
    {
Paul Fultz II's avatar
Paul Fultz II committed
2008
2009
        auto x = m1.add_parameter("x", {migraphx::shape::int32_type, {1, 1024, 17, 17}});
        auto w = m1.add_literal(
2010
            migraphx::generate_literal({migraphx::shape::int32_type, {768, 1024, 1, 1}}));
Paul Fultz II's avatar
Paul Fultz II committed
2011
2012
        auto conv   = m1.add_instruction(migraphx::make_op("convolution"), x, w);
        auto slice1 = m1.add_instruction(
2013
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {384}}}), conv);
2014
        auto a1 =
Paul Fultz II's avatar
Paul Fultz II committed
2015
2016
            m1.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}, 1));
        auto b1 = m1.add_instruction(
2017
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 384, 17, 17}}}), a1);
Paul Fultz II's avatar
Paul Fultz II committed
2018
        auto mul = m1.add_instruction(migraphx::make_op("mul"), slice1, b1);
2019
        auto a2 =
Paul Fultz II's avatar
Paul Fultz II committed
2020
2021
            m1.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}, 2));
        auto b2 = m1.add_instruction(
2022
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 384, 17, 17}}}), a2);
Paul Fultz II's avatar
Paul Fultz II committed
2023
        auto add1 = m1.add_instruction(migraphx::make_op("add"), mul, b2);
2024
        auto a3 =
Paul Fultz II's avatar
Paul Fultz II committed
2025
2026
            m1.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}, 3));
        auto b3 = m1.add_instruction(
2027
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 384, 17, 17}}}), a3);
Paul Fultz II's avatar
Paul Fultz II committed
2028
        auto slice2 = m1.add_instruction(
2029
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {384}}, {"ends", {768}}}), conv);
Paul Fultz II's avatar
Paul Fultz II committed
2030
2031
        auto add2 = m1.add_instruction(migraphx::make_op("add"), slice2, b3);
        m1.add_instruction(pass_op{}, add1, add2);
2032
    }
Paul Fultz II's avatar
Paul Fultz II committed
2033
    run_pass(m1);
2034

Paul Fultz II's avatar
Paul Fultz II committed
2035
    migraphx::module m2;
2036
    {
Paul Fultz II's avatar
Paul Fultz II committed
2037
2038
        auto x = m2.add_parameter("x", {migraphx::shape::int32_type, {1, 1024, 17, 17}});
        auto w = m2.add_literal(
2039
            migraphx::generate_literal({migraphx::shape::int32_type, {768, 1024, 1, 1}}));
Paul Fultz II's avatar
Paul Fultz II committed
2040
        auto wslice1 = m2.add_instruction(
2041
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {384}}}), w);
2042
        auto a1 =
Paul Fultz II's avatar
Paul Fultz II committed
2043
2044
            m2.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}, 1));
        auto b1 = m2.add_instruction(
2045
            migraphx::make_op("broadcast", {{"axis", 0}, {"out_lens", {384, 1024, 1, 1}}}), a1);
Paul Fultz II's avatar
Paul Fultz II committed
2046
2047
        auto mul     = m2.add_instruction(migraphx::make_op("mul"), b1, wslice1);
        auto wslice2 = m2.add_instruction(
2048
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {384}}, {"ends", {768}}}), w);
Paul Fultz II's avatar
Paul Fultz II committed
2049
2050
        auto concat1 = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), mul, wslice2);
        auto conv    = m2.add_instruction(migraphx::make_op("convolution"), x, concat1);
2051
        auto a2 =
Paul Fultz II's avatar
Paul Fultz II committed
2052
            m2.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}, 2));
2053
        auto a3 =
Paul Fultz II's avatar
Paul Fultz II committed
2054
2055
2056
            m2.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {384}}, 3));
        auto concat2 = m2.add_instruction(migraphx::make_op("concat"), a2, a3);
        auto b4      = m2.add_instruction(
2057
            migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {1, 768, 17, 17}}}), concat2);
Paul Fultz II's avatar
Paul Fultz II committed
2058
2059
        auto add    = m2.add_instruction(migraphx::make_op("add"), conv, b4);
        auto slice1 = m2.add_instruction(
2060
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {384}}}), add);
Paul Fultz II's avatar
Paul Fultz II committed
2061
        auto slice2 = m2.add_instruction(
2062
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {384}}, {"ends", {768}}}), add);
Paul Fultz II's avatar
Paul Fultz II committed
2063
        m2.add_instruction(pass_op{}, slice1, slice2);
2064
    }
Paul Fultz II's avatar
Paul Fultz II committed
2065
    EXPECT(m1.sort() == m2.sort());
2066
}
2067
2068
2069

template <std::size_t BS, bool TransposeInput>
void reorder_reshape_slice()
2070
2071
2072
{
    std::vector<int64_t> perm0 = {0, 2, 1, 3};
    std::vector<int64_t> perm1 = {0, 2, 3, 1};
2073
2074
2075
2076
2077
2078
2079
    migraphx::module m1;
    {
        auto s = migraphx::shape{migraphx::shape::float_type, {BS, 128, 1920}};
        if(TransposeInput)
        {
            s = migraphx::shape{migraphx::shape::float_type, {BS, 128, 1920}, {165120, 1, 128}};
        }
Paul Fultz II's avatar
Paul Fultz II committed
2080
2081
        auto input = m1.add_parameter("input", s);
        auto slc0  = m1.add_instruction(
2082
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {640}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2083
        auto slc1 = m1.add_instruction(
2084
2085
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {640}}, {"ends", {1280}}}),
            input);
Paul Fultz II's avatar
Paul Fultz II committed
2086
        auto slc2 = m1.add_instruction(
2087
2088
2089
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {1280}}, {"ends", {1920}}}),
            input);

Paul Fultz II's avatar
Paul Fultz II committed
2090
2091
2092
        auto c0 = m1.add_instruction(migraphx::make_op("contiguous"), slc0);
        auto c1 = m1.add_instruction(migraphx::make_op("contiguous"), slc1);
        auto c2 = m1.add_instruction(migraphx::make_op("contiguous"), slc2);
2093

2094
        std::vector<int64_t> lens = {static_cast<int64_t>(BS), 128, 10, 64};
Paul Fultz II's avatar
Paul Fultz II committed
2095
2096
2097
        auto r0 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c0);
        auto r1 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c1);
        auto r2 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c2);
2098

2099
2100
2101
        auto t0 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), r0);
        auto t1 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), r1);
        auto t2 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm1}}), r2);
2102

Paul Fultz II's avatar
Paul Fultz II committed
2103
2104
2105
        auto sum = m1.add_instruction(migraphx::make_op("add"), t0, t1);
        auto ret = m1.add_instruction(migraphx::make_op("dot"), sum, t2);
        m1.add_return({ret});
2106
2107
    };

2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
    migraphx::module m2;
    {
        auto s = migraphx::shape{migraphx::shape::float_type, {BS, 128, 1920}};
        if(TransposeInput)
        {
            s = migraphx::shape{migraphx::shape::float_type, {BS, 128, 1920}, {165120, 1, 128}};
        }
        auto input     = m2.add_parameter("input", s);
        auto rsp_input = input;
        if(TransposeInput)
        {
            rsp_input = m2.add_instruction(migraphx::make_op("contiguous"), {input});
        }
        std::vector<int64_t> lens = {static_cast<int64_t>(BS), 128, 30, 64};
        auto r = m2.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), rsp_input);
2123

Paul Fultz II's avatar
Paul Fultz II committed
2124
        auto slc0 = m2.add_instruction(
2125
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {10}}}), r);
Paul Fultz II's avatar
Paul Fultz II committed
2126
        auto slc1 = m2.add_instruction(
2127
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {10}}, {"ends", {20}}}), r);
Paul Fultz II's avatar
Paul Fultz II committed
2128
        auto slc2 = m2.add_instruction(
2129
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {20}}, {"ends", {30}}}), r);
2130

2131
2132
2133
2134
2135
2136
        auto t0 =
            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), slc0);
        auto t1 =
            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), slc1);
        auto t2 =
            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", perm1}}), slc2);
2137

Paul Fultz II's avatar
Paul Fultz II committed
2138
2139
2140
        auto sum = m2.add_instruction(migraphx::make_op("add"), t0, t1);
        auto ret = m2.add_instruction(migraphx::make_op("dot"), sum, t2);
        m2.add_return({ret});
2141
    };
2142
2143
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
2144
2145
}

2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
TEST_CASE_REGISTER(reorder_reshape_slice<1, true>); // test if contiguous is added as necessary if
                                                    // input is transposed
TEST_CASE_REGISTER(reorder_reshape_slice<4, true>);
TEST_CASE_REGISTER(reorder_reshape_slice<8, true>);
TEST_CASE_REGISTER(reorder_reshape_slice<1, false>);
TEST_CASE_REGISTER(reorder_reshape_slice<4, false>);
TEST_CASE_REGISTER(reorder_reshape_slice<8, false>);

template <std::size_t BS>
void reorder_reshape_slice_move_axis1()
2156
{
2157
2158
2159
    migraphx::module m1;
    {
        auto s                     = migraphx::shape{migraphx::shape::float_type, {BS, 256, 96}};
2160
2161
        std::vector<int64_t> perm0 = {0, 2, 1, 3};
        std::vector<int64_t> perm1 = {0, 2, 3, 1};
Paul Fultz II's avatar
Paul Fultz II committed
2162
2163
        auto input                 = m1.add_parameter("input", s);
        auto slc0                  = m1.add_instruction(
2164
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {32}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2165
        auto slc1 = m1.add_instruction(
2166
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {32}}, {"ends", {64}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2167
        auto slc2 = m1.add_instruction(
2168
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {64}}, {"ends", {96}}}), input);
2169

Paul Fultz II's avatar
Paul Fultz II committed
2170
2171
2172
        auto c0 = m1.add_instruction(migraphx::make_op("contiguous"), slc0);
        auto c1 = m1.add_instruction(migraphx::make_op("contiguous"), slc1);
        auto c2 = m1.add_instruction(migraphx::make_op("contiguous"), slc2);
2173

2174
        std::vector<int64_t> lens = {static_cast<int64_t>(BS), 64, 4, 32};
Paul Fultz II's avatar
Paul Fultz II committed
2175
2176
2177
        auto r0 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c0);
        auto r1 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c1);
        auto r2 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c2);
2178

2179
2180
2181
        auto t0 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), r0);
        auto t1 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), r1);
        auto t2 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm1}}), r2);
2182

Paul Fultz II's avatar
Paul Fultz II committed
2183
2184
2185
        auto sum = m1.add_instruction(migraphx::make_op("add"), t0, t1);
        auto ret = m1.add_instruction(migraphx::make_op("dot"), sum, t2);
        m1.add_return({ret});
2186
2187
    };

2188
2189
2190
    migraphx::module m2;
    {
        auto s                     = migraphx::shape{migraphx::shape::float_type, {BS, 256, 96}};
2191
2192
        std::vector<int64_t> perm0 = {0, 2, 1, 3};
        std::vector<int64_t> perm1 = {0, 2, 3, 1};
2193
2194
2195
2196
        auto input                 = m2.add_parameter("input", s);
        std::vector<int64_t> lens  = {static_cast<int64_t>(BS), 64, 4, 96};
        auto rsp  = m2.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), input);
        auto slc0 = m2.add_instruction(
2197
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {0}}, {"ends", {32}}}), rsp);
2198
2199
2200
        auto t0 =
            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), slc0);
        auto slc1 = m2.add_instruction(
2201
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {32}}, {"ends", {64}}}), rsp);
2202
2203
2204
        auto t1 =
            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), slc1);
        auto slc2 = m2.add_instruction(
2205
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {64}}, {"ends", {96}}}), rsp);
2206
2207
        auto t2 =
            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", perm1}}), slc2);
2208

2209
2210
2211
        auto sum = m2.add_instruction(migraphx::make_op("add"), t0, t1);
        auto ret = m2.add_instruction(migraphx::make_op("dot"), sum, t2);
        m2.add_return({ret});
2212
2213
    };

2214
2215
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
2216
2217
}

2218
2219
2220
TEST_CASE_REGISTER(reorder_reshape_slice_move_axis1<4>);
TEST_CASE_REGISTER(reorder_reshape_slice_move_axis1<8>);

2221
2222
TEST_CASE(reorder_reshape_slice_move_axis2)
{
2223
2224
    migraphx::module m1;
    {
2225
        migraphx::shape s{migraphx::shape::float_type, {128, 96}};
Paul Fultz II's avatar
Paul Fultz II committed
2226
2227
        auto input = m1.add_parameter("input", s);
        auto slc0  = m1.add_instruction(
2228
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {32}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2229
        auto slc1 = m1.add_instruction(
2230
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {32}}, {"ends", {64}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2231
        auto slc2 = m1.add_instruction(
2232
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {64}}, {"ends", {96}}}), input);
2233

Paul Fultz II's avatar
Paul Fultz II committed
2234
2235
2236
        auto c0 = m1.add_instruction(migraphx::make_op("contiguous"), slc0);
        auto c1 = m1.add_instruction(migraphx::make_op("contiguous"), slc1);
        auto c2 = m1.add_instruction(migraphx::make_op("contiguous"), slc2);
2237
2238

        std::vector<int64_t> lens = {1, 16, 8, 32};
Paul Fultz II's avatar
Paul Fultz II committed
2239
2240
2241
        auto r0 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c0);
        auto r1 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c1);
        auto r2 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c2);
2242

Paul Fultz II's avatar
Paul Fultz II committed
2243
2244
2245
        auto sum = m1.add_instruction(migraphx::make_op("add"), r0, r1);
        auto ret = m1.add_instruction(migraphx::make_op("mul"), sum, r2);
        m1.add_return({ret});
2246
2247
    };

2248
2249
    migraphx::module m2;
    {
2250
        auto s                    = migraphx::shape{migraphx::shape::float_type, {128, 96}};
2251
        auto input                = m2.add_parameter("input", s);
2252
        std::vector<int64_t> lens = {1, 16, 8, 96};
2253
2254
        auto rsp  = m2.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), input);
        auto slc0 = m2.add_instruction(
2255
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {0}}, {"ends", {32}}}), rsp);
2256
        auto slc1 = m2.add_instruction(
2257
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {32}}, {"ends", {64}}}), rsp);
2258
        auto slc2 = m2.add_instruction(
2259
2260
            migraphx::make_op("slice", {{"axes", {3}}, {"starts", {64}}, {"ends", {96}}}), rsp);

2261
2262
        auto sum = m2.add_instruction(migraphx::make_op("add"), slc0, slc1);
        auto ret = m2.add_instruction(migraphx::make_op("mul"), sum, slc2);
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
        m2.add_return({ret});
    };

    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
}

TEST_CASE(reorder_reshape_slice_len_1)
{
    migraphx::module m1;
    {
        migraphx::shape s{migraphx::shape::float_type, {1, 128, 3}};
        auto input = m1.add_parameter("input", s);
        auto slc0  = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {1}}}), input);
        auto slc1 = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {2}}}), input);
        auto slc2 = m1.add_instruction(
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {2}}, {"ends", {3}}}), input);

        auto c0 = m1.add_instruction(migraphx::make_op("contiguous"), slc0);
        auto c1 = m1.add_instruction(migraphx::make_op("contiguous"), slc1);
        auto c2 = m1.add_instruction(migraphx::make_op("contiguous"), slc2);

        std::vector<int64_t> lens = {1, 128};
        auto r0 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c0);
        auto r1 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c1);
        auto r2 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c2);

        auto sum = m1.add_instruction(migraphx::make_op("add"), r0, r1);
        auto ret = m1.add_instruction(migraphx::make_op("mul"), sum, r2);
        m1.add_return({ret});
    };

    migraphx::module m2;
    {
        auto s                    = migraphx::shape{migraphx::shape::float_type, {1, 128, 3}};
        auto input                = m2.add_parameter("input", s);
        std::vector<int64_t> lens = {1, 384};
        auto rsp  = m2.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), input);
        auto slc0 = m2.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {128}}}), rsp);
        auto slc1 = m2.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {128}}, {"ends", {256}}}), rsp);
        auto slc2 = m2.add_instruction(
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {256}}, {"ends", {384}}}), rsp);

        auto sum = m2.add_instruction(migraphx::make_op("add"), slc0, slc1);
        auto ret = m2.add_instruction(migraphx::make_op("mul"), sum, slc2);
2312
        m2.add_return({ret});
2313
2314
    };

Paul Fultz II's avatar
Paul Fultz II committed
2315
2316
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
2317
2318
2319
2320
2321
}

TEST_CASE(reorder_reshape_slice_not_apply)
{
    auto create_p = [] {
Paul Fultz II's avatar
Paul Fultz II committed
2322
        migraphx::module m;
2323
        migraphx::shape s{migraphx::shape::float_type, {128, 96}};
Paul Fultz II's avatar
Paul Fultz II committed
2324
2325
        auto input = m.add_parameter("input", s);
        auto slc0  = m.add_instruction(
2326
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {32}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2327
        auto slc1 = m.add_instruction(
2328
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {32}}, {"ends", {64}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2329
        auto slc2 = m.add_instruction(
2330
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {64}}, {"ends", {96}}}), input);
2331

Paul Fultz II's avatar
Paul Fultz II committed
2332
2333
2334
        auto c0 = m.add_instruction(migraphx::make_op("contiguous"), slc0);
        auto c1 = m.add_instruction(migraphx::make_op("contiguous"), slc1);
        auto c2 = m.add_instruction(migraphx::make_op("contiguous"), slc2);
2335
2336

        std::vector<int64_t> lens = {1, 16, 16, 16};
Paul Fultz II's avatar
Paul Fultz II committed
2337
2338
2339
        auto r0 = m.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c0);
        auto r1 = m.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c1);
        auto r2 = m.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c2);
2340

Paul Fultz II's avatar
Paul Fultz II committed
2341
2342
2343
        auto sum = m.add_instruction(migraphx::make_op("add"), r0, r1);
        auto ret = m.add_instruction(migraphx::make_op("mul"), sum, r2);
        m.add_return({ret});
2344

Paul Fultz II's avatar
Paul Fultz II committed
2345
        return m;
2346
2347
    };

Paul Fultz II's avatar
Paul Fultz II committed
2348
2349
2350
2351
    auto m1 = create_p();
    auto m2 = m1;
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
2352
2353
}

2354
2355
template <std::size_t BS>
void reorder_reshape_slice_diff_dims()
2356
{
2357
2358
2359
2360
2361
    migraphx::module m1;
    {
        auto s     = migraphx::shape{migraphx::shape::float_type, {BS, 96, 96}};
        auto input = m1.add_parameter("input", s);
        auto slc0  = m1.add_instruction(
2362
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {32}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2363
        auto slc1 = m1.add_instruction(
2364
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {32}}, {"ends", {64}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2365
        auto slc2 = m1.add_instruction(
2366
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {64}}, {"ends", {96}}}), input);
2367

Paul Fultz II's avatar
Paul Fultz II committed
2368
2369
2370
        auto c0 = m1.add_instruction(migraphx::make_op("contiguous"), slc0);
        auto c1 = m1.add_instruction(migraphx::make_op("contiguous"), slc1);
        auto c2 = m1.add_instruction(migraphx::make_op("contiguous"), slc2);
2371

2372
2373
        std::vector<int64_t> lens  = {static_cast<int64_t>(BS), 32, 3, 32};
        std::vector<int64_t> lens1 = {static_cast<int64_t>(BS), 48, 2, 32};
Paul Fultz II's avatar
Paul Fultz II committed
2374
2375
2376
        auto r0 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c0);
        auto r1 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens}}), c1);
        auto r2 = m1.add_instruction(migraphx::make_op("reshape", {{"dims", lens1}}), c2);
2377

Paul Fultz II's avatar
Paul Fultz II committed
2378
        m1.add_return({r0, r1, r2});
2379
2380
    };

2381
2382
2383
    auto m2 = m1;
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
2384
2385
}

2386
2387
2388
2389
2390
TEST_CASE_REGISTER(reorder_reshape_slice_diff_dims<4>);
TEST_CASE_REGISTER(reorder_reshape_slice_diff_dims<8>);

template <std::size_t BS>
void reorder_slice_trans()
2391
2392
{
    std::vector<int64_t> perm = {0, 2, 1};
2393
2394
2395
2396

    migraphx::module m1;
    {
        auto s     = migraphx::shape{migraphx::shape::float_type, {BS, 128, 1920}};
Paul Fultz II's avatar
Paul Fultz II committed
2397
2398
        auto input = m1.add_parameter("input", s);
        auto slc0  = m1.add_instruction(
2399
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {640}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2400
        auto slc1 = m1.add_instruction(
2401
2402
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {640}}, {"ends", {1280}}}),
            input);
Paul Fultz II's avatar
Paul Fultz II committed
2403
        auto slc2 = m1.add_instruction(
2404
2405
2406
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {1280}}, {"ends", {1920}}}),
            input);

2407
2408
2409
        auto t0 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm}}), slc0);
        auto t1 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm}}), slc1);
        auto t2 = m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm}}), slc2);
2410

Paul Fultz II's avatar
Paul Fultz II committed
2411
2412
2413
        auto sum = m1.add_instruction(migraphx::make_op("add"), t0, t1);
        auto ret = m1.add_instruction(migraphx::make_op("mul"), sum, t2);
        m1.add_return({ret});
2414
2415
    };

2416
2417
2418
    migraphx::module m2;
    {
        auto s     = migraphx::shape{migraphx::shape::float_type, {BS, 128, 1920}};
Paul Fultz II's avatar
Paul Fultz II committed
2419
        auto input = m2.add_parameter("input", s);
2420
        auto r = m2.add_instruction(migraphx::make_op("transpose", {{"permutation", perm}}), input);
2421

Paul Fultz II's avatar
Paul Fultz II committed
2422
        auto slc0 = m2.add_instruction(
2423
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {640}}}), r);
Paul Fultz II's avatar
Paul Fultz II committed
2424
        auto slc1 = m2.add_instruction(
2425
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {640}}, {"ends", {1280}}}), r);
Paul Fultz II's avatar
Paul Fultz II committed
2426
        auto slc2 = m2.add_instruction(
2427
            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1280}}, {"ends", {1920}}}), r);
2428

Paul Fultz II's avatar
Paul Fultz II committed
2429
2430
2431
        auto sum = m2.add_instruction(migraphx::make_op("add"), slc0, slc1);
        auto ret = m2.add_instruction(migraphx::make_op("mul"), sum, slc2);
        m2.add_return({ret});
2432
2433
    };

2434
2435
    run_pass(m1);
    EXPECT(m1.sort() == m2.sort());
2436
2437
}

2438
2439
2440
2441
2442
TEST_CASE_REGISTER(reorder_slice_trans<1>);
TEST_CASE_REGISTER(reorder_slice_trans<8>);

template <std::size_t BS>
void reorder_slice_trans_diff_perm()
2443
{
2444
2445
2446
    migraphx::module m1;
    {
        auto s                     = migraphx::shape{migraphx::shape::float_type, {BS, 128, 1920}};
2447
2448
        std::vector<int64_t> perm0 = {0, 2, 1};
        std::vector<int64_t> perm1 = {0, 1, 2};
Paul Fultz II's avatar
Paul Fultz II committed
2449
2450
        auto input                 = m1.add_parameter("input", s);
        auto slc0                  = m1.add_instruction(
2451
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {640}}}), input);
Paul Fultz II's avatar
Paul Fultz II committed
2452
        auto slc1 = m1.add_instruction(
2453
2454
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {640}}, {"ends", {1280}}}),
            input);
Paul Fultz II's avatar
Paul Fultz II committed
2455
        auto slc2 = m1.add_instruction(
2456
2457
2458
            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {1280}}, {"ends", {1920}}}),
            input);

2459
2460
2461
2462
2463
2464
        auto t0 =
            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), slc0);
        auto t1 =
            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm0}}), slc1);
        auto t2 =
            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", perm1}}), slc2);
2465

Paul Fultz II's avatar
Paul Fultz II committed
2466
2467
2468
        auto sum = m1.add_instruction(migraphx::make_op("add"), t0, t1);
        auto ret = m1.add_instruction(migraphx::make_op("dot"), sum, t2);
        m1.add_return({ret});
2469
2470
    };

2471
2472
2473
    run_pass(m1);
    auto m2 = m1;
    EXPECT(m1.sort() == m2.sort());
2474
2475
}

2476
2477
2478
TEST_CASE_REGISTER(reorder_slice_trans_diff_perm<1>);
TEST_CASE_REGISTER(reorder_slice_trans_diff_perm<4>);

Shucai Xiao's avatar
Shucai Xiao committed
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
TEST_CASE(reorder_slice_ins_deps)
{
    auto create_module = [] {
        migraphx::module m;
        migraphx::shape sx{migraphx::shape::float_type, {4, 2}};
        migraphx::shape sy{migraphx::shape::float_type, {2, 2}};
        std::vector<float> datax = {0, 1, 2, 3, 4, 5, 6, 7};
        std::vector<float> datay = {0, 1, 2, 3};
        auto inx                 = m.add_literal(migraphx::literal(sx, datax));
        auto iny                 = m.add_literal(migraphx::literal(sy, datay));
        auto slc0                = m.add_instruction(
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}), inx);
        auto slc1 = m.add_instruction(
            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {2}}, {"ends", {4}}}), inx);
        auto n0 = m.add_instruction(migraphx::make_op("neg"), slc0);
        auto a0 = m.add_instruction(migraphx::make_op("add"), n0, slc1);
        auto m0 = m.add_instruction(migraphx::make_op("mul"), a0, iny);
        auto r  = m.add_instruction(migraphx::make_op("add"), m0, slc0);
        m.add_return({r});

        return m;
    };

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

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