"data/train.json" did not exist on "fd31a4693ff3935e668136720345e4e7042543e4"
constant_propagate.cpp 1.21 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
#include <migraphx/constant_propagate.hpp>
#include <migraphx/program.hpp>
#include <migraphx/matcher.hpp>
#include <migraphx/literal.hpp>
Paul's avatar
Paul committed
5
#include <migraphx/functional.hpp>
6
#include <unordered_set>
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

Paul's avatar
Paul committed
11
bool skip_propogate(instruction_ref ins)
Paul's avatar
Paul committed
12
{
Paul's avatar
Paul committed
13
    if(ins->name() == "@literal")
Paul's avatar
Paul committed
14
        return true;
Paul's avatar
Paul committed
15
    if(ins->get_shape().broadcasted() and not ins->get_shape().scalar())
Paul's avatar
Paul committed
16
        return true;
Paul's avatar
Paul committed
17
    if(ins->get_shape().scalar() and ins->get_shape().elements() != 1)
Paul's avatar
Paul committed
18
19
20
        return true;
    return false;
}
Paul's avatar
Paul committed
21

Paul's avatar
Paul committed
22
void constant_propagate::apply(program& p) const
Paul's avatar
Paul committed
23
24
{
    fix([&](auto self, auto ins) {
Paul's avatar
Paul committed
25
        if(not skip_propogate(ins))
Paul's avatar
Paul committed
26
27
        {
            auto r = ins->eval();
Paul's avatar
Paul committed
28
            if(not r.empty())
Paul's avatar
Paul committed
29
            {
Paul's avatar
Paul committed
30
                assert(r.get_shape() == ins->get_shape());
Paul's avatar
Paul committed
31
32
33
34
35
                auto l = p.add_literal(r.get_shape(), r.data());
                p.replace_instruction(ins, l);
                return;
            }
        }
Paul's avatar
Paul committed
36
        std::unordered_set<instruction_ref> children(ins->inputs().begin(), ins->inputs().end());
Paul's avatar
Paul committed
37
        for(auto child : children)
Paul's avatar
Paul committed
38
39
40
            self(child);
    })(std::prev(p.end()));
}
Paul's avatar
Paul committed
41

Paul's avatar
Paul committed
42
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
43
} // namespace migraphx