simplify_algebra.cpp 44.5 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
#include <migraphx/simplify_algebra.hpp>
Paul's avatar
Paul committed
25
#include <migraphx/dead_code_elimination.hpp>
Paul's avatar
Paul committed
26
#include <migraphx/program.hpp>
27
#include <migraphx/op/concat.hpp>
28
#include <migraphx/op/slice.hpp>
29
#include <migraphx/op/convolution.hpp>
Paul's avatar
Paul committed
30
#include <migraphx/op/broadcast.hpp>
31
32
#include <migraphx/op/reshape.hpp>
#include <migraphx/op/transpose.hpp>
Paul's avatar
Paul committed
33
34
#include <migraphx/matcher.hpp>
#include <migraphx/literal.hpp>
35
36
37
#include <migraphx/make_op.hpp>
#include <migraphx/serialize.hpp>

38
#include <migraphx/algorithm.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
39
#include <unordered_set>
Paul's avatar
Paul committed
40

Paul's avatar
Paul committed
41
namespace migraphx {
Paul's avatar
Paul committed
42
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
43

Paul's avatar
Paul committed
44
auto lit_broadcast() { return match::any_of(match::is_constant(), match::name("broadcast")); }
Paul's avatar
Paul committed
45
auto not_lit_broadcast() { return match::none_of(match::is_constant(), match::name("broadcast")); }
Paul's avatar
Paul committed
46
47
auto op_lit_broadcast(std::string op, std::string x, std::string y)
{
Paul's avatar
Paul committed
48
49
    return match::name(std::move(op))(match::either_arg(0, 1)(
        lit_broadcast().bind(std::move(x)), not_lit_broadcast().bind(std::move(y))));
Paul's avatar
Paul committed
50
51
}

Paul's avatar
Paul committed
52
53
auto conv_const_weights()
{
Paul's avatar
Paul committed
54
    return match::name("convolution")(match::used_once(),
Paul's avatar
Paul committed
55
                                      match::args(match::any(), match::is_constant().bind("w")));
Paul's avatar
Paul committed
56
57
}

Shucai Xiao's avatar
Shucai Xiao committed
58
59
auto reduction() { return match::name_contains("reduce"); }

60
// conv(x, w) * a => conv(x, a * w)
Paul's avatar
Paul committed
61
62
63
struct find_mul_conv
{
    auto matcher() const
Paul's avatar
Paul committed
64
    {
65
66
67
        return match::name("mul")(
            match::either_arg(0, 1)(conv_const_weights().bind("conv"),
                                    match::name("broadcast", "multibroadcast").bind("a")));
Paul's avatar
Paul committed
68
    }
Paul's avatar
Paul committed
69

70
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
71
    {
Paul's avatar
Paul committed
72
        auto ins      = r.result;
Paul's avatar
Paul committed
73
        auto conv_ins = r.instructions["conv"];
Paul's avatar
Paul committed
74
75
76
        auto a_ins    = r.instructions["a"];
        auto w_ins    = r.instructions["w"];

77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
        const auto& a_input_lens = a_ins->inputs().front()->get_shape().lens();

        std::size_t num_not_one_dims = std::count_if(
            a_input_lens.cbegin(), a_input_lens.cend(), [](auto dim) { return dim != 1; });
        if(num_not_one_dims > 1)
            return;

        // check broadcasted along channels
        const auto& a_lens    = a_ins->get_shape().lens();
        const auto& a_strides = a_ins->get_shape().strides();

        auto is_broadcasted_axis = [](auto len, auto stride) { return len == 1 or stride == 0; };

        if(a_strides.at(1) != 1)
            return;

        if(not is_broadcasted_axis(a_lens.front(), a_strides.front()))
            return;

        if(not std::equal(a_lens.begin() + 2,
                          a_lens.end(),
                          a_strides.begin() + 2,
                          a_strides.end(),
                          is_broadcasted_axis))
Paul's avatar
Paul committed
101
102
            return;

103
        auto sq    = m.insert_instruction(ins, make_op("squeeze"), a_ins->inputs().front());
104
        auto new_a = m.insert_instruction(
105
            ins, make_op("broadcast", {{"axis", 0}, {"out_lens", w_ins->get_shape().lens()}}), sq);
106
107
        auto new_mul  = m.insert_instruction(ins, make_op("mul"), new_a, w_ins);
        auto new_conv = m.insert_instruction(
Paul's avatar
Paul committed
108
            ins, conv_ins->get_operator(), conv_ins->inputs().front(), new_mul);
109
        m.replace_instruction(ins, new_conv);
Paul's avatar
Paul committed
110
    }
Paul's avatar
Paul committed
111
112
};

113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
struct find_mul_slice_conv
{
    static auto conv()
    {
        return match::name("convolution")(
            match::all_of[match::outputs()](match::name("slice")),
            match::args(match::any(), match::is_constant().bind("w")));
    }
    auto matcher() const
    {
        return match::name("mul")(match::either_arg(0, 1)(
            match::name("slice")(match::used_once(), match::arg(0)(conv().bind("conv")))
                .bind("slice"),
            match::name("broadcast")(match::is_constant()).bind("a")));
    }

129
    void apply(module& m, const match::matcher_result& r) const
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
158
159
160
161
162
163
164
    {
        auto ins       = r.result;
        auto slice_ins = r.instructions["slice"];
        auto conv_ins  = r.instructions["conv"];
        auto a_ins     = r.instructions["a"];
        auto w_ins     = r.instructions["w"];

        auto broadcast_op = any_cast<op::broadcast>(a_ins->get_operator());
        if(broadcast_op.axis != 1)
            return;

        auto slice_op = any_cast<op::slice>(slice_ins->get_operator());
        if(slice_op.axes.size() != 1)
            return;
        if(slice_op.axes.front() != 1)
            return;

        auto slice_idx = std::distance(conv_ins, slice_ins);
        if(std::any_of(conv_ins->outputs().begin(), conv_ins->outputs().end(), [&](auto i) {
               if(i == slice_ins)
                   return false;
               if(std::distance(conv_ins, i) < slice_idx)
                   return true;
               auto sop = any_cast<op::slice>(i->get_operator());
               if(sop.axes != slice_op.axes)
                   return true;
               if(std::max(sop.starts.front(), slice_op.starts.front()) <
                  std::min(sop.ends.front(), slice_op.ends.front()))
                   return true;
               return false;
           }))
            return;

        auto w_slice_op  = slice_op;
        w_slice_op.axes  = {0};
165
        auto slice_w_ins = m.insert_instruction(ins, w_slice_op, w_ins);
166

167
        auto new_a = m.insert_instruction(
168
            ins,
169
            make_op("broadcast", {{"axis", 0}, {"out_lens", slice_w_ins->get_shape().lens()}}),
170
            a_ins->inputs().front());
171
        auto new_mul = m.insert_instruction(ins, make_op("mul"), new_a, slice_w_ins);
172
173
174

        std::vector<instruction_ref> sliced_weights;
        if(slice_op.starts.front() != 0)
175
            sliced_weights.push_back(m.insert_instruction(
176
177
178
                ins,
                make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", slice_op.starts}}),
                w_ins));
179
180
181
        sliced_weights.push_back(new_mul);
        int64_t end_axis = w_ins->get_shape().lens().at(0);
        if(slice_op.ends.front() != end_axis)
182
            sliced_weights.push_back(m.insert_instruction(
183
184
185
                ins,
                make_op("slice", {{"axes", {0}}, {"starts", slice_op.ends}, {"ends", {end_axis}}}),
                w_ins));
186

187
        auto new_weights =
188
            m.insert_instruction(ins, make_op("concat", {{"axis", 0}}), sliced_weights);
189

190
        auto new_conv = m.insert_instruction(
191
192
193
            ins, conv_ins->get_operator(), conv_ins->inputs().front(), new_weights);
        assert(conv_ins->get_shape() == new_conv->get_shape());

194
        auto slice1 = m.insert_instruction(ins, slice_op, new_conv);
195
        assert(ins->get_shape().lens() == slice1->get_shape().lens());
196
        m.replace_instruction(ins, slice1);
197
        // TODO: Check each slice doesn't overlap and that it occurs after slice_ins
198
199
        auto outputs = conv_ins->outputs();
        for(auto output : outputs)
200
201
202
203
204
            if(output != slice_ins)
                instruction::replace_argument(output, conv_ins, new_conv);
    }
};

Paul's avatar
Paul committed
205
// a * (x + b) => a * x + a * b
Paul's avatar
Paul committed
206
207
208
209
210
struct find_mul_add
{
    auto matcher() const
    {
        return match::name("mul")(match::either_arg(0, 1)(
Paul's avatar
Paul committed
211
212
213
            match::name("add")(
                match::either_arg(0, 1)(
                    match::any().bind("x"),
Paul's avatar
Paul committed
214
                    match::any_of(conv_const_weights(), match::is_constant()).bind("b")),
Paul's avatar
Paul committed
215
                match::none_of(match::args(match::is_constant(), match::is_constant())),
Paul's avatar
Paul committed
216
                match::used_once()),
Paul's avatar
Paul committed
217
            match::is_constant().bind("a")));
Paul's avatar
Paul committed
218
219
    }

220
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
221
    {
Paul's avatar
Paul committed
222
        auto ins   = r.result;
Paul's avatar
Paul committed
223
        auto a_ins = r.instructions["a"];
Paul's avatar
Paul committed
224
        auto b_ins = r.instructions["b"];
Paul's avatar
Paul committed
225
        auto x_ins = r.instructions["x"];
Paul's avatar
Paul committed
226
        assert(x_ins != b_ins);
Paul's avatar
Paul committed
227

228
229
230
        auto ax_ins = m.insert_instruction(ins, make_op("mul"), a_ins, x_ins);
        auto ab_ins = m.insert_instruction(ins, make_op("mul"), a_ins, b_ins);
        m.replace_instruction(ins, make_op("add"), ax_ins, ab_ins);
Paul's avatar
Paul committed
231
232
233
    }
};

234
235
236
237
238
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
struct find_dot_add
{
    auto matcher() const
    {
        return match::name("dot")(match::either_arg(0, 1)(
            match::name("add")(
                match::either_arg(0, 1)(match::any().bind("x"),
                                        match::any_of(match::is_constant()).bind("b")),
                match::none_of(match::args(match::is_constant(), match::is_constant())),
                match::used_once()),
            match::is_constant().bind("a")));
    }

    void apply(module& m, const match::matcher_result& r) const
    {
        auto ins   = r.result;
        auto a_ins = r.instructions["a"];
        auto b_ins = r.instructions["b"];
        auto x_ins = r.instructions["x"];
        assert(x_ins != b_ins);

        const bool flipped = a_ins == ins->inputs().back();

        auto insert_dot = [&](auto x, auto y) {
            if(flipped)
                return m.insert_instruction(ins, make_op("dot"), y, x);
            else
                return m.insert_instruction(ins, make_op("dot"), x, y);
        };

        auto ax_ins = insert_dot(a_ins, x_ins);
        auto ab_ins = insert_dot(a_ins, b_ins);
        m.replace_instruction(ins, make_op("add"), ax_ins, ab_ins);
    }
};

Paul's avatar
Paul committed
270
struct find_add_lit_broadcast
Paul's avatar
Paul committed
271
272
273
274
275
276
277
{
    auto matcher() const
    {
        return match::name("add")(
            match::either_arg(0, 1)(op_lit_broadcast("add", "a", "x"), lit_broadcast().bind("b")));
    }

278
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
279
280
281
282
283
284
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto a_ins = r.instructions["a"];
        auto b_ins = r.instructions["b"];

285
286
        auto sumab = m.insert_instruction(ins, make_op("add"), a_ins, b_ins);
        m.replace_instruction(ins, make_op("add"), x_ins, sumab);
Paul's avatar
Paul committed
287
288
289
290
    }
};

struct find_double_add_lit_broadcast
Paul's avatar
Paul committed
291
{
Paul's avatar
Paul committed
292
293
    auto matcher() const
    {
Paul's avatar
Paul committed
294
        return match::name("add")(
Paul's avatar
Paul committed
295
            match::args(op_lit_broadcast("add", "a", "x"), op_lit_broadcast("add", "b", "y")));
Paul's avatar
Paul committed
296
297
    }

298
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
299
    {
Paul's avatar
Paul committed
300
301
302
303
304
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto y_ins = r.instructions["y"];
        auto a_ins = r.instructions["a"];
        auto b_ins = r.instructions["b"];
Paul's avatar
Paul committed
305
306
307

        instruction_ref sumab;

Paul's avatar
Paul committed
308
        if(a_ins->name() == "broadcast" and b_ins->name() == "broadcast")
Paul's avatar
Paul committed
309
310
311
        {
            if(a_ins->inputs().at(0)->get_shape() != b_ins->inputs().at(0)->get_shape())
                return;
312
            auto op     = a_ins->get_operator();
313
            auto presum = m.insert_instruction(
314
                ins, make_op("add"), a_ins->inputs().at(0), b_ins->inputs().at(0));
315
            sumab = m.insert_instruction(ins, op, presum);
Paul's avatar
Paul committed
316
317
318
        }
        else
        {
319
            sumab = m.insert_instruction(ins, make_op("add"), a_ins, b_ins);
Paul's avatar
Paul committed
320
321
        }

322
323
        auto sumxy = m.insert_instruction(ins, make_op("add"), x_ins, y_ins);
        m.replace_instruction(ins, make_op("add"), sumxy, sumab);
Paul's avatar
Paul committed
324
325
326
    }
};

Paul's avatar
Paul committed
327
328
struct find_inner_broadcast
{
329
    auto matcher() const { return pointwise(match::all_of[match::inputs()](match::broadcast())); }
Paul's avatar
Paul committed
330

331
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
332
    {
333
334
335
336
337
338
339
340
341
342
343
344
        auto ins        = r.result;
        auto broadcasts = ins->inputs();
        if(broadcasts.empty())
            return;
        std::vector<instruction_ref> inputs;
        std::transform(broadcasts.begin(),
                       broadcasts.end(),
                       std::back_inserter(inputs),
                       [](auto i) { return i->inputs().front(); });
        if(std::any_of(inputs.begin(), inputs.end(), [&](auto i) {
               return i->get_shape() != inputs.front()->get_shape();
           }))
Paul's avatar
Paul committed
345
346
            return;

347
348
        auto op = m.insert_instruction(ins, ins->get_operator(), inputs);
        m.replace_instruction(ins, broadcasts.front()->get_operator(), op);
Paul's avatar
Paul committed
349
350
351
    }
};

352
struct find_concat_op
353
354
355
{
    auto matcher() const
    {
356
        return match::name("concat")(match::any_of[match::inputs()](
357
            match::any_of(match::pointwise(), match::name("broadcast")), match::used_once()));
358
359
    }

360
361
    template <class Iterator>
    static std::vector<std::size_t> get_output_lens(Iterator start, Iterator last, std::size_t axis)
362
    {
363
364
365
        assert(start != last);
        std::size_t dim = 0;
        for(auto ins : range(start, last))
366
        {
367
            dim += ins->get_shape().lens().at(axis);
368
        }
369
370
371
        auto lens  = (*start)->get_shape().lens();
        lens[axis] = dim;
        return lens;
372
373
    }

374
375
376
377
378
    static bool is_valid_op(const operation& op)
    {
        return op.name() == "broadcast" or op.attributes().contains("pointwise");
    }

379
    void apply(module& m, const match::matcher_result& r) const
380
    {
381
382
        auto ins  = r.result;
        auto axis = any_cast<op::concat>(ins->get_operator()).axis;
383

384
385
386
387
388
389
        auto each = [&](auto start, auto last) -> std::vector<instruction_ref> {
            if(std::distance(start, last) < 2)
                return {start, last};
            auto x = *start;
            if(x->inputs().size() > 2 or x->inputs().empty() or x->outputs().size() > 1)
                return {start, last};
390
391
            auto op = x->get_operator();
            if(not is_valid_op(op))
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
                return {start, last};
            auto iaxis = axis;
            // Adjust broadcast lens
            if(op.name() == "broadcast")
            {
                auto b = any_cast<op::broadcast>(op);
                if(b.axis != iaxis)
                    return {start, last};
                b.broadcast_lens = get_output_lens(start, last, iaxis);
                op               = b;
                iaxis            = 0;
            }

            std::vector<instruction_ref> concats;
            for(std::size_t i = 0; i < x->inputs().size(); i++)
            {
                std::vector<instruction_ref> inputs;
                std::transform(start, last, std::back_inserter(inputs), [&](auto j) {
                    return j->inputs().at(i);
                });
412
                auto concat =
413
                    m.insert_instruction(ins, make_op("concat", {{"axis", iaxis}}), inputs);
414
415
                concats.push_back(concat);
            }
416
            auto y = m.insert_instruction(ins, op, concats);
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
            return {y};
        };

        std::vector<instruction_ref> args;
        auto update_args = [&](auto start, auto last) {
            auto x = each(start, last);
            args.insert(args.end(), x.begin(), x.end());
        };
        auto pred = [](auto i, auto j) {
            return i->get_operator() == j->get_operator() and
                   i->inputs().size() == i->inputs().size() and
                   i->outputs().size() == i->outputs().size();
        };
        group_unique(ins->inputs().begin(), ins->inputs().end(), update_args, pred);
        if(args.size() == 1)
432
            m.replace_instruction(ins, args.front());
433
        else
434
            m.replace_instruction(ins, make_op("concat", {{"axis", axis}}), args);
435
436
437
    }
};

438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
std::vector<instruction_ref> get_splits(instruction_ref ins)
{
    std::vector<instruction_ref> result;
    std::copy_if(ins->outputs().begin(),
                 ins->outputs().end(),
                 std::back_inserter(result),
                 [&](auto i) { return i->name() == "slice"; });
    if(result.size() < 2)
        return {};
    auto get_slice = [](auto& i) -> auto& { return any_cast<op::slice>(i->get_operator()); };
    auto&& axes    = get_slice(result.front()).axes;
    if(std::any_of(result.begin(), result.end(), [&](auto i) { return get_slice(i).axes != axes; }))
        return {};
    auto get_start = [&](auto& i) -> auto& { return get_slice(i).starts; };
    auto get_end   = [&](auto& i) -> auto& { return get_slice(i).ends; };
    std::sort(
        result.begin(), result.end(), [&](auto x, auto y) { return get_start(x) < get_start(y); });
    if(std::any_of(get_start(result.front()).begin(), get_start(result.front()).end(), [&](auto i) {
           return i != 0;
       }))
        return {};
    auto it = std::adjacent_find(
        result.begin(), result.end(), [&](auto x, auto y) { return get_end(x) != get_start(y); });
    if(it != result.end())
        return {};
    for(std::size_t i = 0; i < axes.size(); i++)
    {
        auto axis = axes[i];
        if(ins->get_shape().lens()[axis] != get_slice(result.back()).ends[i])
            return {};
    }
    return result;
}

struct find_splits
{
    auto matcher() const
    {
476
477
478
        return match::any(
            match::any_of[match::outputs()](match::name("slice")(match::any_of[match::outputs()](
                match::pointwise(match::any_of(match::nargs(1), match::nargs(2))), reduction()))));
479
480
    }

Shucai Xiao's avatar
Shucai Xiao committed
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
    static bool is_dependent(const module& m, instruction_ref ins1, instruction_ref ins2)
    {

        std::unordered_set<instruction_ref> traversed;
        return fix<bool>([&](auto self, auto ins) -> bool {
            if(ins == ins2)
                return true;

            if(contains(traversed, ins))
                return false;

            traversed.insert(ins);
            const auto& inputs = ins->inputs();
            return std::any_of(inputs.begin(), inputs.end(), [&](auto in) {
                return m.has_instruction(in) and self(in);
            });
        })(ins1);
    }

500
    static std::vector<std::vector<instruction_ref>>
Shucai Xiao's avatar
Shucai Xiao committed
501
    get_split_groups(const module& m, const std::vector<instruction_ref>& splits)
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
    {
        std::vector<std::vector<instruction_ref>> groups;
        for(auto out : splits.front()->outputs())
        {
            if(out->name() == "slice")
                continue;
            std::vector<instruction_ref> group;
            for(auto split : splits)
            {
                auto it =
                    std::find_if(split->outputs().begin(), split->outputs().end(), [&](auto i) {
                        return i->get_operator() == out->get_operator();
                    });
                if(it == split->outputs().end())
                    break;
                assert((*it)->name() != "slice");
Shucai Xiao's avatar
Shucai Xiao committed
518

519
                // If there is a duplicate bail
Shucai Xiao's avatar
Shucai Xiao committed
520
521
522
523
524
                // there are should be no dependency between instructions in the group
                if(std::any_of(group.begin(), group.end(), [&](auto i) {
                       return is_dependent(m, *it, i) or is_dependent(m, i, *it);
                   }))
                {
525
                    return {};
Shucai Xiao's avatar
Shucai Xiao committed
526
527
                }

528
529
530
531
532
533
534
535
536
                group.push_back(*it);
            }
            if(group.size() != splits.size())
                continue;
            groups.push_back(group);
        }
        return groups;
    }

Shucai Xiao's avatar
Shucai Xiao committed
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
    bool is_fusable(instruction_ref start, instruction_ref split_front) const
    {
        auto op = start->get_operator();
        if(contains(op.name(), "reduce"))
        {
            auto slc         = any_cast<op::slice>(split_front->get_operator());
            auto slc_axes    = slc.axes;
            auto reduce_axes = start->get_operator().to_value()["axes"].to_vector<int64_t>();
            // axes of slice and reduce op cannot have overlap
            if(std::any_of(slc_axes.begin(), slc_axes.end(), [&](auto axis) {
                   return (std::find(reduce_axes.begin(), reduce_axes.end(), axis) !=
                           reduce_axes.end());
               }))
            {
                return false;
            }
        }
        else if(not op.attributes().contains("pointwise"))
        {
            return false;
        }

        return true;
    }

Paul's avatar
Format  
Paul committed
562
563
    static std::vector<instruction_ref> split_nary(const std::vector<instruction_ref>& group,
                                                   bool commutative)
Paul's avatar
Paul committed
564
565
    {
        // All inputs have the same slices
Paul's avatar
Format  
Paul committed
566
567
568
569
570
571
572
573
574
575
        if(not std::all_of(group.begin(), group.end(), [](auto ins) {
               if(ins->inputs().empty())
                   return false;
               auto first = ins->inputs().front();
               if(first->name() != "slice")
                   return false;
               return std::all_of(ins->inputs().begin() + 1, ins->inputs().end(), [&](auto input) {
                   return input->get_operator() == first->get_operator();
               });
           }))
Paul's avatar
Paul committed
576
            return {};
Paul's avatar
Format  
Paul committed
577
        auto start      = group.front();
Paul's avatar
Paul committed
578
579
580
581
582
583
584
585
586
        auto get_inputs = [](auto ins) {
            std::vector<instruction_ref> result;
            std::transform(ins->inputs().begin(),
                           ins->inputs().end(),
                           std::back_inserter(result),
                           [](auto slice) { return slice->inputs().front(); });
            return result;
        };
        auto inputs = get_inputs(start);
Paul's avatar
Format  
Paul committed
587
        if(commutative and inputs.size() > 1)
Paul's avatar
Paul committed
588
589
590
        {
            std::sort(inputs.begin(), inputs.end(), compare_instruction_ref{});
            if(not std::all_of(group.begin(), group.end(), [&](auto ins) {
Paul's avatar
Format  
Paul committed
591
592
593
594
                   auto inputs2 = get_inputs(ins);
                   std::sort(inputs2.begin(), inputs2.end(), compare_instruction_ref{});
                   return inputs == inputs2;
               }))
Paul's avatar
Paul committed
595
596
597
598
599
600
601
602
603
604
605
606
607
                return {};
        }
        else
        {
            if(not std::all_of(group.begin(), group.end(), [&](auto ins) {
                   return std::equal(
                       ins->inputs().begin(),
                       ins->inputs().end(),
                       inputs.begin(),
                       inputs.end(),
                       [](auto slice, auto input) { return slice->inputs().front() == input; });
               }))
                return {};
Paul's avatar
Format  
Paul committed
608
        }
Paul's avatar
Paul committed
609
610
611
        return inputs;
    }

Paul's avatar
Format  
Paul committed
612
    template <class Range>
Paul's avatar
Paul committed
613
614
615
616
617
618
619
    static instruction_ref find_last_instruction(const module& m, const Range& r)
    {
        auto rm = reverse(m);
        auto it = std::find_first_of(rm.begin(), rm.end(), r.begin(), r.end());
        return std::prev(it.base());
    }

620
    void apply(module& m, const match::matcher_result& r) const
621
    {
Shucai Xiao's avatar
Shucai Xiao committed
622
        auto ins    = r.result;
623
624
625
        auto splits = get_splits(ins);
        if(splits.empty())
            return;
Shucai Xiao's avatar
Shucai Xiao committed
626

627
        for(const auto& group : get_split_groups(m, splits))
628
        {
Shucai Xiao's avatar
Shucai Xiao committed
629
630
631
632
633
            auto start       = group.front();
            auto split_front = splits.front();
            auto op          = start->get_operator();
            if(not is_fusable(start, split_front))
            {
634
                continue;
Shucai Xiao's avatar
Shucai Xiao committed
635
            }
636
637
638
639
640
641

            // Make sure there is no duplicates
            assert(std::none_of(
                std::next(group.begin()), group.end(), [&](auto i) { return i == start; }));

            auto split_idx    = 0;
642
            instruction_ref c = m.end();
643
644
            if(start->inputs().size() == 1)
            {
645
                c = m.insert_instruction(std::next(ins), op, ins);
646
647
648
649
650
651
            }
            else if(start->inputs().size() == 2)
            {
                assert(not std::none_of(start->inputs().begin(), start->inputs().end(), [](auto i) {
                    return i->name() == "slice";
                }) && "one argument must be a split");
Paul's avatar
Paul committed
652

Paul's avatar
Paul committed
653
                auto split_inputs = split_nary(group, op.attributes().get("commutative", false));
Paul's avatar
Format  
Paul committed
654
                if(not split_inputs.empty())
655
                {
Paul's avatar
Paul committed
656
                    auto last = find_last_instruction(m, split_inputs);
Paul's avatar
Format  
Paul committed
657
                    c         = m.insert_instruction(std::next(last), op, split_inputs);
Paul's avatar
Paul committed
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
                }
                else
                {
                    auto data_idx = 1;
                    if(start->inputs().back()->name() == "slice")
                    {
                        split_idx = 1;
                        data_idx  = 0;
                    }

                    std::vector<instruction_ref> data_args;
                    std::transform(group.begin(),
                                   group.end(),
                                   std::back_inserter(data_args),
                                   [&](auto i) { return i->inputs()[data_idx]; });

                    // Data arguments must be a constant
                    if(std::any_of(data_args.begin(), data_args.end(), [](auto i) {
                           return not i->can_eval();
                       }))
                        return;

                    for(auto data : data_args)
                        m.move_instructions(data, ins);

                    auto slice_op = any_cast<op::slice>(splits.front()->get_operator());
                    assert(not slice_op.axes.empty());
                    if(slice_op.axes.size() > 1)
                        return;
                    auto concat_axis = slice_op.axes.front();
                    // TODO: Check if axises match
                    auto concat = m.insert_instruction(
                        ins, make_op("concat", {{"axis", concat_axis}}), data_args);

                    std::vector<instruction_ref> args;
                    args.resize(2);
                    args[split_idx] = ins;
                    args[data_idx]  = concat;
                    c               = m.insert_instruction(std::next(ins), op, args);
697
698
                }
            }
699
            if(c != m.end())
700
701
702
703
704
705
            {
                for(auto i : group)
                {
                    auto split = i->inputs()[split_idx];
                    assert(split->name() == "slice");
                    // Insert contiguous for reshapes
706
707
                    auto outputs = i->outputs();
                    for(auto output : outputs)
708
                    {
709
                        if(output->name() != "reshape")
710
                            continue;
711
                        auto x = m.insert_instruction(output, make_op("contiguous"), i);
712
                        m.replace_instruction(output, output->get_operator(), x);
713
714
                    }

715
                    m.replace_instruction(i, split->get_operator(), c);
716
717
718
719
720
721
722
723
724
725
726
727
728
729
                }
            }
        }
    }
};

struct find_split_concat
{
    auto matcher() const
    {
        return match::any(match::any_of[match::outputs()](
            match::name("slice")(match::all_of[match::outputs()](match::name("concat")))));
    }

730
    void apply(module& m, const match::matcher_result& r) const
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
    {
        auto ins = r.result;

        auto splits = get_splits(ins);
        if(splits.empty())
            return;
        if(std::any_of(
               splits.begin(), splits.end(), [](auto i) { return i->outputs().size() != 1; }))
            return;
        // Check for concat operator
        auto concat = splits.front()->outputs().front();
        if(std::any_of(splits.begin(), splits.end(), [&](auto i) {
               return i->outputs().front() != concat;
           }))
            return;
        // Check axis match
        auto concat_op = any_cast<op::concat>(concat->get_operator());
        auto split_op  = any_cast<op::slice>(splits.front()->get_operator());
        if(split_op.axes.size() != 1)
            return;
        if(split_op.axes.front() != concat_op.axis)
            return;
        // Replace args
        auto args = concat->inputs();
        auto it =
            std::find_if(args.begin(), args.end(), [&](auto i) { return i == splits.front(); });
        if(std::distance(it, args.end()) < splits.size())
            return;
759
760
761
762
763
764
765
        // If the slices are not in order then stop
        if(not std::is_sorted(it, it + splits.size(), [](instruction_ref x, instruction_ref y) {
               auto xop = any_cast<op::slice>(x->get_operator());
               auto yop = any_cast<op::slice>(y->get_operator());
               return std::tie(xop.starts, xop.ends) < std::tie(yop.starts, yop.ends);
           }))
            return;
766
767
768
769
        *it = splits.front()->inputs().front();
        args.erase(std::next(it), it + splits.size());

        if(args.size() == 1)
770
            m.replace_instruction(concat, args.front());
771
        else
772
            m.replace_instruction(concat, concat->get_operator(), args);
773
774
775
    }
};

776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
bool axis_equal(const std::vector<std::size_t>& x,
                const std::vector<std::size_t>& y,
                std::size_t axis)
{
    return x.size() == y.size() and x.size() > axis and
           std::equal(x.begin(), x.begin() + axis, y.begin()) and
           std::equal(x.begin() + axis + 1, x.end(), y.begin() + axis + 1);
}

bool axis_shape_equal(const shape& x, const shape& y, std::size_t axis)
{
    // TODO: Check strides
    return axis_equal(x.lens(), y.lens(), axis);
}

struct find_add_convs
{
    auto matcher() const
    {
        return match::name("add")(
            match::args(conv_const_weights().bind("a"), conv_const_weights().bind("b")));
    }

    static bool symmetrical_strides(const op::convolution& op)
    {
        return op.stride[0] == op.stride[1];
    }

    static std::size_t compute_stride_factor(const op::convolution& x, const op::convolution& y)
    {
        if(not symmetrical_strides(x))
            return 0;
        if(not symmetrical_strides(y))
            return 0;
        if((x.stride[0] % y.stride[0]) != 0)
            return 0;
        return x.stride[0] / y.stride[0];
    }

815
    void apply(module& m, const match::matcher_result& r) const
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
    {
        auto ins       = r.result;
        auto a_conv    = r.instructions["a"];
        auto a_input   = a_conv->inputs().at(0);
        auto a_weights = a_conv->inputs().at(1);
        auto b_conv    = r.instructions["b"];
        auto b_input   = b_conv->inputs().at(0);
        auto b_weights = b_conv->inputs().at(1);

        if(not axis_shape_equal(a_weights->get_shape(), b_weights->get_shape(), 1))
            return;

        auto a_op   = any_cast<op::convolution>(a_conv->get_operator());
        auto b_op   = any_cast<op::convolution>(b_conv->get_operator());
        auto new_op = a_op;

        if(a_op != b_op)
        {
            if(std::tie(a_op.padding, a_op.dilation, a_op.group) ==
                   std::tie(b_op.padding, b_op.dilation, b_op.group) and
               a_weights->get_shape().lens()[2] == 1 and a_weights->get_shape().lens()[3] == 1)
            {
                if(a_op.stride < b_op.stride)
                {
                    auto n = compute_stride_factor(b_op, a_op);
                    if(n == 0)
                        return;
                    new_op  = a_op;
844
                    b_input = m.insert_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
845
                        ins, make_op("step", {{"axes", {2, 3}}, {"steps", {n, n}}}), b_input);
846
847
848
849
850
851
852
                }
                else if(b_op.stride < a_op.stride)
                {
                    auto n = compute_stride_factor(a_op, b_op);
                    if(n == 0)
                        return;
                    new_op  = b_op;
853
                    a_input = m.insert_instruction(
Shucai Xiao's avatar
Shucai Xiao committed
854
                        ins, make_op("step", {{"axes", {2, 3}}, {"steps", {n, n}}}), a_input);
855
856
857
858
859
860
861
862
                }
                else
                    return;
            }
            else
                return;
        }

863
        auto concat_input =
864
            m.insert_instruction(ins, make_op("concat", {{"axis", 1}}), a_input, b_input);
865
        auto concat_weights =
866
867
            m.insert_instruction(ins, make_op("concat", {{"axis", 1}}), a_weights, b_weights);
        m.replace_instruction(ins, new_op, concat_input, concat_weights);
868
869
870
    }
};

Paul's avatar
Paul committed
871
872
873
874
struct find_add_dots
{
    auto matcher() const
    {
Paul's avatar
Format  
Paul committed
875
876
877
878
879
880
881
        auto dot_const_weights =
            match::name("dot")(match::used_once(), match::arg(1)(match::is_constant()));
        auto dot_const_inputs =
            match::name("dot")(match::used_once(), match::arg(0)(match::is_constant()));
        return match::name("add")(
            match::any_of(match::args(dot_const_weights.bind("a"), dot_const_weights.bind("b")),
                          match::args(dot_const_inputs.bind("a"), dot_const_inputs.bind("b"))));
Paul's avatar
Paul committed
882
883
884
885
886
    }

    void apply(module& m, const match::matcher_result& r) const
    {
        auto ins = r.result;
Paul's avatar
Format  
Paul committed
887
888
        auto a   = r.instructions["a"];
        auto b   = r.instructions["b"];
Paul's avatar
Paul committed
889
890
891

        auto n = ins->get_shape().lens().size();

Paul's avatar
Format  
Paul committed
892
893
894
895
        auto x = m.insert_instruction(
            ins, make_op("concat", {{"axis", (n - 1)}}), a->inputs()[0], b->inputs()[0]);
        auto w = m.insert_instruction(
            ins, make_op("concat", {{"axis", (n - 2)}}), a->inputs()[1], b->inputs()[1]);
Paul's avatar
Paul committed
896
897
898
899
        m.replace_instruction(ins, make_op("dot"), x, w);
    }
};

900
901
902
903
904
905
906
907
908
909
MIGRAPHX_PRED_MATCHER(horiz_conv_dot, instruction_ref ins)
{
    auto pred = [&](auto name) {
        return [=](auto i) {
            return i->name() == name and i->inputs().front() == ins and
                   i->inputs().at(1)->can_eval();
        };
    };
    auto dots  = std::count_if(ins->outputs().begin(), ins->outputs().end(), pred("dot"));
    auto convs = std::count_if(ins->outputs().begin(), ins->outputs().end(), pred("convolution"));
910
    return not(dots < 2 and convs < 2);
911
912
913
914
915
916
}

struct find_conv_dot_horiz_fusion
{
    auto matcher() const { return horiz_conv_dot(); }

917
    void apply(module& m, const match::matcher_result& r) const
918
919
920
921
922
923
924
925
926
927
928
929
    {
        auto ins = r.result;

        auto pred = [](auto i, auto j) {
            if(i->get_operator() != j->get_operator())
                return false;
            if(not contains({"dot", "convolution"}, i->name()))
                return true;
            auto x = i->inputs()[1]->get_shape().lens();
            auto y = j->inputs()[1]->get_shape().lens();
            if(x.size() != y.size())
                return false;
930
            // Check that non-axes match
931
932
933
934
935
936
937
938
939
940
941
942
943
944
            int axis = 1;
            if(i->name() == "dot")
            {
                axis = x.size() - 1;
            }
            return axis_equal(x, y, axis);
        };

        auto each = [&](auto start, auto last) {
            if(std::distance(start, last) < 2)
                return;
            auto&& name = (*start)->name();
            if(not contains({"dot", "convolution"}, name))
                return;
945
946
947
948
949
950
951
            auto op   = (*start)->get_operator();
            int group = 1;
            if(name == "convolution")
                group = any_cast<op::convolution>(op).group;
            // Skip group convolution
            if(group != 1)
                return;
952
953
954
955
956
957
958
959
960
961
962
963
964
            auto input = (*start)->inputs().front();
            std::vector<instruction_ref> args;
            std::transform(
                start, last, std::back_inserter(args), [&](auto x) { return x->inputs().at(1); });
            int axis        = 1;
            int concat_axis = 0;
            if(name == "dot")
            {
                axis        = int(args.front()->get_shape().lens().size() - 1);
                concat_axis = axis;
            }

            for(auto arg : args)
965
                m.move_instructions(arg, input);
966
            // TODO: Check if axes match
967
            auto concat =
968
969
                m.insert_instruction(input, make_op("concat", {{"axis", concat_axis}}), args);
            auto fused     = m.insert_instruction(std::next(input), op, input, concat);
970
971
972
            int64_t offset = 0;
            for(auto arg : range(start, last))
            {
973
974
975
976
977
978
979
980
981
                auto outputs = arg->outputs();
                for(auto output : outputs)
                {
                    if(output->name() != "reshape")
                        continue;
                    auto x = m.insert_instruction(output, make_op("contiguous"), arg);
                    m.replace_instruction(output, output->get_operator(), x);
                }

982
                int64_t len = arg->get_shape().lens()[axis];
983
                m.replace_instruction(
984
985
986
987
                    arg,
                    make_op("slice",
                            {{"axes", {axis}}, {"starts", {offset}}, {"ends", {offset + len}}}),
                    fused);
988
989
990
991
992
993
994
995
996
                offset += len;
            }
        };

        auto outputs = ins->outputs();
        group_by(outputs.begin(), outputs.end(), each, pred);
    }
};

997
998
999
1000
1001
1002
1003
struct find_div_const
{
    auto matcher() const
    {
        return match::name("div")(match::arg(1)(match::is_constant().bind("c")));
    }

1004
    void apply(module& m, const match::matcher_result& r) const
1005
1006
1007
1008
    {
        auto ins   = r.result;
        auto c_ins = r.instructions["c"];

1009
        auto recip = m.insert_instruction(std::next(c_ins), make_op("recip"), c_ins);
1010
1011
1012

        auto args = ins->inputs();

1013
        m.replace_instruction(ins, make_op("mul"), args.front(), recip);
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
    }
};

struct find_sub_const
{
    auto matcher() const
    {
        return match::name("sub")(match::arg(1)(match::is_constant().bind("c")));
    }

1024
    void apply(module& m, const match::matcher_result& r) const
1025
1026
1027
1028
    {
        auto ins   = r.result;
        auto c_ins = r.instructions["c"];

1029
        auto neg = m.insert_instruction(std::next(c_ins), make_op("neg"), c_ins);
1030
1031
1032

        auto args = ins->inputs();

1033
        m.replace_instruction(ins, make_op("add"), args.front(), neg);
1034
1035
1036
    }
};

kahmed10's avatar
kahmed10 committed
1037
1038
1039
1040
1041
1042
1043
1044
struct find_rsqrt
{
    auto matcher() const
    {
        return match::name("recip")(match::args(
            match::name("sqrt")(match::used_once(), match::args(match::any().bind("x")))));
    }

1045
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
1046
1047
1048
1049
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];

1050
        m.replace_instruction(ins, make_op("rsqrt"), x_ins);
kahmed10's avatar
kahmed10 committed
1051
1052
1053
    }
};

1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
static bool same_ops(const std::vector<instruction_ref>& vec_ins)
{
    return std::all_of(vec_ins.begin(), vec_ins.end(), [&](auto i) {
        return i->get_operator() == vec_ins.front()->get_operator();
    });
}

struct find_split_reshape
{
    auto matcher() const
    {
        return match::name("reshape")(match::arg(0)(match::name("contiguous")(
                                          match::arg(0)(match::name("slice").bind("slice")))))
            .bind("reshape");
    }

1070
    void apply(module& m, const match::matcher_result& r) const
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
    {
        auto slc = r.instructions["slice"];
        auto rsp = r.instructions["reshape"];

        auto input         = slc->inputs().front();
        auto split_outputs = get_splits(input);
        if(split_outputs.empty())
        {
            return;
        }

        std::vector<instruction_ref> vec_rsp(split_outputs.size());
        std::transform(split_outputs.begin(), split_outputs.end(), vec_rsp.begin(), [](auto i) {
            assert(i->outputs().size() == 1);
            auto cont = i->outputs().front();
            assert(cont->outputs().size() == 1);
            return cont->outputs().front();
        });

        // all outputs are reshape and of the same shape
        auto dims = any_cast<op::reshape>(rsp->get_operator()).dims;
1092
        if(not same_ops(vec_rsp))
1093
1094
1095
1096
1097
        {
            return;
        }

        // ensure reshape happens after the axis dimension
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
        auto axis         = any_cast<op::slice>(slc->get_operator()).axes[0];
        auto slc_lens     = slc->get_shape().lens();
        auto slc_dim_size = std::accumulate(
            slc_lens.begin() + axis, slc_lens.end(), 1, std::multiplies<std::size_t>());

        // search the reshape output (standard shape) to decide which axis are
        // in its output corresponding to the slc_dim_size
        auto rsp_lens    = rsp->get_shape().lens();
        auto rsp_strides = rsp->get_shape().strides();
        rsp_strides.insert(rsp_strides.begin(), rsp_strides[0] * rsp_lens[0]);
1108
1109
1110

        auto ait     = std::find(rsp_strides.begin(), rsp_strides.end(), slc_dim_size);
        int rsp_axis = -1;
1111
        if(ait == rsp_strides.end())
1112
1113
1114
        {
            return;
        }
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
        else if(ait == rsp_strides.end() - 1)
        {
            // edge case
            // slice_dim == 1, in that case it could match with last stride of 1.
            // it should accumulate lengths from last dim in that case. discount 1 to avoid going
            // out of bounds.
            assert(slc_dim_size == 1);
            rsp_axis = std::distance(rsp_strides.begin(), ait) - 1;
        }
        else
        {
            rsp_axis = std::distance(rsp_strides.begin(), ait);
        }
1128
        // calculate reshape output shape
1129
        std::vector<int64_t> vec_dims(vec_rsp.size());
1130

1131
1132
1133
1134
1135
        std::transform(vec_rsp.begin(), vec_rsp.end(), vec_dims.begin(), [&](auto is) {
            return is->get_shape().lens()[rsp_axis];
        });

        std::vector<int64_t> rsp_out_lens(rsp_lens.begin(), rsp_lens.end());
1136

1137
        rsp_out_lens[rsp_axis] = std::accumulate(vec_dims.begin(), vec_dims.end(), std::int64_t{0});
1138

1139
1140
1141
1142
1143
        // insert the reshape instruction and add contiguous if needed
        if(not input->get_shape().standard())
        {
            input = m.insert_instruction(std::next(input), make_op("contiguous"), input);
        }
1144
        auto rsp_ins = m.insert_instruction(
1145
            std::next(input), make_op("reshape", {{"dims", rsp_out_lens}}), input);
1146
1147

        // replace the original reshape with slice
1148
1149
        int64_t start = 0;
        for(std::size_t i = 0; i < vec_rsp.size(); ++i)
1150
        {
1151
            m.replace_instruction(
1152
1153
1154
1155
1156
                vec_rsp[i],
                make_op(
                    "slice",
                    {{"axes", {rsp_axis}}, {"starts", {start}}, {"ends", {start + vec_dims[i]}}}),
                rsp_ins);
1157
            start += vec_dims[i];
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
        }
    }
};

struct find_split_transpose
{
    auto matcher() const
    {
        return match::name("transpose")(match::arg(0)(match::name("slice").bind("slice")))
            .bind("trans");
    }

1170
    void apply(module& m, const match::matcher_result& r) const
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
    {
        auto slc   = r.instructions["slice"];
        auto trans = r.instructions["trans"];

        auto input         = slc->inputs().front();
        auto split_outputs = get_splits(input);
        if(split_outputs.empty())
        {
            return;
        }

        std::vector<instruction_ref> vec_trans(split_outputs.size());
        std::transform(split_outputs.begin(), split_outputs.end(), vec_trans.begin(), [](auto i) {
            assert(i->outputs().size() == 1);
            return i->outputs().front();
        });

        // all transpose are the same
        auto perm = any_cast<op::transpose>(trans->get_operator()).dims;
1190
        if(not same_ops(vec_trans))
1191
1192
1193
1194
1195
        {
            return;
        }

        // insert an transpose instruction
1196
        auto tr = m.insert_instruction(
1197
            std::next(input), make_op("transpose", {{"permutation", perm}}), input);
1198
1199
1200
1201
1202

        // compute the axis in the slice
        auto axis = any_cast<op::slice>(slc->get_operator()).axes.front();
        auto it   = std::find(perm.begin(), perm.end(), axis);
        assert(it != perm.end());
Paul Fultz II's avatar
Paul Fultz II committed
1203
        int64_t axis_new = std::distance(perm.begin(), it);
1204
1205
1206
1207
1208
1209
1210

        for(auto in : split_outputs)
        {
            auto oper    = any_cast<op::slice>(in->get_operator());
            auto starts  = oper.starts;
            auto ends    = oper.ends;
            auto tr_orig = in->outputs().front();
1211
            m.replace_instruction(
1212
1213
1214
                tr_orig,
                make_op("slice", {{"axes", {axis_new}}, {"starts", starts}, {"ends", ends}}),
                tr);
1215
1216
1217
1218
        }
    }
};

1219
void simplify_algebra::apply(module& m) const
Paul's avatar
Paul committed
1220
{
Paul's avatar
Paul committed
1221
    // Run simplifications multiple times
1222
    for(int i = 0; i < 8; i++)
Paul's avatar
Paul committed
1223
    {
1224
        match::find_matches(m,
Paul's avatar
Paul committed
1225
                            find_inner_broadcast{},
Paul's avatar
Paul committed
1226
1227
                            find_double_add_lit_broadcast{},
                            find_add_lit_broadcast{},
1228
                            find_add_convs{},
Paul's avatar
Paul committed
1229
                            find_add_dots{},
1230
                            find_conv_dot_horiz_fusion{},
Paul's avatar
Paul committed
1231
                            find_mul_conv{},
1232
                            find_mul_slice_conv{},
1233
                            find_mul_add{},
1234
                            find_dot_add{},
1235
1236
                            find_div_const{},
                            find_sub_const{},
kahmed10's avatar
kahmed10 committed
1237
                            find_rsqrt{},
1238
                            find_concat_op{},
1239
                            find_split_concat{},
1240
1241
1242
                            find_splits{},
                            find_split_reshape{},
                            find_split_transpose{});
1243
        dead_code_elimination{}.apply(m);
Paul's avatar
Paul committed
1244
    }
Paul's avatar
Paul committed
1245
}
Paul's avatar
Paul committed
1246

Paul's avatar
Paul committed
1247
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
1248
} // namespace migraphx