decompose.cpp 2.94 KB
Newer Older
1
2
3
4
5
6
7
8
9
#include <migraphx/decompose.hpp>
#include <migraphx/program.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/iterator_for.hpp>
#include <migraphx/functional.hpp>
#include <migraphx/ranges.hpp>
#include <migraphx/float_equal.hpp>
#include <migraphx/matcher.hpp>
#include <migraphx/op/dot.hpp>
10
11
#include <migraphx/make_op.hpp>

12
13
14
namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
namespace {
15
16
17
18
19
20
21
22
23
24
25
26
27

struct alpha_beta
{
    float alpha = 0.0;
    float beta  = 0.0;
};

alpha_beta get_alpha_beta(const operation& op)
{
    auto v = op.to_value();
    return {v.at("alpha").to<float>(), v.at("beta").to<float>()};
}

28
29
struct find_dot_add
{
30
    auto matcher() const { return match::name("dot", "quant_dot")(match::nargs(3)); }
31

32
    void apply(module& p, const match::matcher_result& r) const
33
    {
34
35
        auto ins   = r.result;
        auto dot   = get_alpha_beta(ins->get_operator());
36
37
38
39
        auto a_ins = ins->inputs()[0];
        auto b_ins = ins->inputs()[1];
        if(not float_equal(dot.alpha, 1))
        {
40
            auto alpha = p.add_literal(literal{shape{a_ins->get_shape().type()}, {dot.alpha}});
41
            auto alpha_broadcast = p.insert_instruction(
42
                ins, make_op("multibroadcast", {{"out_lens", a_ins->get_shape().lens()}}), alpha);
43
44
            a_ins = p.insert_instruction(ins, make_op("mul"), a_ins, alpha_broadcast);
        }
45
        auto dot_ins = p.insert_instruction(ins, make_op(ins->name(), {{"beta", 0}}), a_ins, b_ins);
46
47

        auto c_ins = ins->inputs()[2];
48
49
        if(not float_equal(dot.beta, 1))
        {
50
            auto beta = p.add_literal(literal{shape{c_ins->get_shape().type()}, {dot.beta}});
51
            auto beta_broadcast = p.insert_instruction(
52
                ins, make_op("multibroadcast", {{"out_lens", ins->get_shape().lens()}}), beta);
53
            c_ins = p.insert_instruction(ins, make_op("mul"), c_ins, beta_broadcast);
54
        }
55
        p.replace_instruction(ins, make_op("add"), dot_ins, c_ins);
56
57
58
    }
};

59
60
struct find_dot_alpha
{
61
    auto matcher() const { return match::name("dot", "quant_dot")(match::nargs(2)); }
62

Shucai Xiao's avatar
Shucai Xiao committed
63
    void apply(module& p, const match::matcher_result& r) const
64
65
    {
        auto ins   = r.result;
66
        auto dot   = get_alpha_beta(ins->get_operator());
67
68
69
70
        auto a_ins = ins->inputs()[0];
        auto b_ins = ins->inputs()[1];
        if(not float_equal(dot.alpha, 1))
        {
71
            auto alpha = p.add_literal(literal{shape{a_ins->get_shape().type()}, {dot.alpha}});
72
            auto alpha_broadcast = p.insert_instruction(
73
                ins, make_op("multibroadcast", {{"out_lens", a_ins->get_shape().lens()}}), alpha);
74
75
            a_ins = p.insert_instruction(ins, make_op("mul"), a_ins, alpha_broadcast);
        }
76
        p.replace_instruction(ins, make_op(ins->name(), {{"beta", 0}}), a_ins, b_ins);
77
78
79
    }
};

80
81
} // namespace

82
void decompose::apply(module& p) const { match::find_matches(p, find_dot_add{}, find_dot_alpha{}); }
83
84
85

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx