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

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

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

12
void auto_contiguous::apply(module& p) const
Paul's avatar
Paul committed
13
{
14
15
    std::string key = "standard_input_shape";
    for(auto ins : reverse_iterator_for(p))
Paul's avatar
Paul committed
16
    {
17
18
        auto&& attr = ins->get_operator().attributes();
        if((attr.contains(key) and attr.at(key).to<bool>()))
Paul's avatar
Paul committed
19
        {
Shucai Xiao's avatar
Shucai Xiao committed
20
            auto args     = ins->inputs();
21
22
            auto new_args = args;
            std::transform(args.begin(), args.end(), new_args.begin(), [&](auto in) {
Shucai Xiao's avatar
Shucai Xiao committed
23
                return p.insert_instruction(ins, make_op("contiguous"), in);
24
25
26
27
28
29
            });

            if(new_args != args)
            {
                p.replace_instruction(ins, ins->get_operator(), new_args);
            }
Paul's avatar
Paul committed
30
31
        }
    }
Shucai Xiao's avatar
Shucai Xiao committed
32
33
34
35
36
37
38
39
40
41
42
43

    auto last = std::prev(p.end());
    for(auto ins : iterator_for(p))
    {
        if (ins->outputs().empty() and ins != last) continue;
        shape s = ins->get_shape();
        if(not s.standard() and s.elements() != 0)
        {
            auto c = p.insert_instruction(std::next(ins), make_op("contiguous"), ins);
            p.replace_instruction(ins, c);
        }
    }
Paul's avatar
Paul committed
44
45
}

Paul's avatar
Paul committed
46
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
47
} // namespace migraphx