auto_contiguous.cpp 1.48 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
#include <migraphx/auto_contiguous.hpp>
#include <migraphx/program.hpp>
#include <migraphx/instruction.hpp>
4
5
#include <migraphx/make_op.hpp>

Paul's avatar
Paul committed
6
#include <migraphx/iterator_for.hpp>
Paul's avatar
Paul committed
7

Paul's avatar
Paul committed
8
namespace migraphx {
Paul's avatar
Paul committed
9
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
10

11
void auto_contiguous::apply(module& p) const
Paul's avatar
Paul committed
12
{
Shucai Xiao's avatar
Shucai Xiao committed
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
    std::string key = "require_std_shape";
    for(auto ins : reverse_iterator_for(p))
    {
        auto&& attr = ins->get_operator().attributes();
        if((attr.get(key, false)))
        {
            auto args     = ins->inputs();
            auto new_args = args;
            std::transform(args.begin(), args.end(), new_args.begin(), [&](auto in) {
                if(in->name() == "contiguous")
                {
                    return in;
                }
                return p.insert_instruction(ins, make_op("contiguous"), in);
            });

            if(new_args != args)
            {
                p.replace_instruction(ins, ins->get_operator(), new_args);
            }
        }
    }

    auto last = std::prev(p.end());
Paul's avatar
Paul committed
37
38
    for(auto ins : iterator_for(p))
    {
Shucai Xiao's avatar
Shucai Xiao committed
39
40
41
        // for last instruction that is NOT a return
        if(ins->outputs().empty() and ins != last)
            continue;
Paul's avatar
Paul committed
42
        shape s = ins->get_shape();
43
        if(not s.standard() and s.elements() != 0)
Paul's avatar
Paul committed
44
        {
45
            auto c = p.insert_instruction(std::next(ins), make_op("contiguous"), ins);
Paul's avatar
Paul committed
46
            p.replace_instruction(ins, c);
Paul's avatar
Paul committed
47
48
49
50
        }
    }
}

Paul's avatar
Paul committed
51
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
52
} // namespace migraphx