decompose_test.cpp 6.04 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
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
#include <migraphx/decompose.hpp>
#include <migraphx/pass_manager.hpp>
#include <basic_ops.hpp>
#include <migraphx/op/add.hpp>
#include <migraphx/op/identity.hpp>
#include <migraphx/op/dot.hpp>
#include <migraphx/op/mul.hpp>
#include <migraphx/op/multibroadcast.hpp>
#include <test.hpp>

void run_pass(migraphx::program& p) { migraphx::run_passes(p, {migraphx::decompose{}}); }

TEST_CASE(dot_add)
{
    migraphx::program p1;
    {
        auto x   = p1.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto y   = p1.add_parameter("y", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto z   = p1.add_parameter("z", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto dot = p1.add_instruction(migraphx::op::dot{}, x, y, z);
        p1.add_instruction(migraphx::op::identity{}, dot);
    }
    run_pass(p1);
    migraphx::program p2;
    {
        auto x   = p2.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto y   = p2.add_parameter("y", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto z   = p2.add_parameter("z", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto dot = p2.add_instruction(migraphx::op::dot{1, 0}, x, y);
        auto add = p2.add_instruction(migraphx::op::add{}, dot, z);
        p2.add_instruction(migraphx::op::identity{}, add);
    }
    EXPECT(p1 == p2);
}

TEST_CASE(dot_add_beta_float)
{
    migraphx::program p1;
    {
        auto x   = p1.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto y   = p1.add_parameter("y", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto z   = p1.add_parameter("z", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto dot = p1.add_instruction(migraphx::op::dot{1.0, 0.5}, x, y, z);
        p1.add_instruction(migraphx::op::identity{}, dot);
    }
    run_pass(p1);
    migraphx::program p2;
    {
        auto x   = p2.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto y   = p2.add_parameter("y", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto z   = p2.add_parameter("z", migraphx::shape{migraphx::shape::float_type, {2, 2}});
        auto dot = p2.add_instruction(migraphx::op::dot{1, 0}, x, y);
        auto beta =
            p2.add_literal(migraphx::literal{migraphx::shape{migraphx::shape::float_type}, {0.5}});
        auto beta_broadcast = p2.add_instruction(migraphx::op::multibroadcast{{2, 2}}, beta);
        auto mul            = p2.add_instruction(migraphx::op::mul{}, z, beta_broadcast);
        auto add            = p2.add_instruction(migraphx::op::add{}, dot, mul);
        p2.add_instruction(migraphx::op::identity{}, add);
    }
    EXPECT(p1 == p2);
}

TEST_CASE(dot_add_beta_half)
{
    migraphx::program p1;
    {
        auto x   = p1.add_parameter("x", migraphx::shape{migraphx::shape::half_type, {2, 2}});
        auto y   = p1.add_parameter("y", migraphx::shape{migraphx::shape::half_type, {2, 2}});
        auto z   = p1.add_parameter("z", migraphx::shape{migraphx::shape::half_type, {2, 2}});
        auto dot = p1.add_instruction(migraphx::op::dot{1.0, 0.5}, x, y, z);
        p1.add_instruction(migraphx::op::identity{}, dot);
    }
    run_pass(p1);
    migraphx::program p2;
    {
        auto x   = p2.add_parameter("x", migraphx::shape{migraphx::shape::half_type, {2, 2}});
        auto y   = p2.add_parameter("y", migraphx::shape{migraphx::shape::half_type, {2, 2}});
        auto z   = p2.add_parameter("z", migraphx::shape{migraphx::shape::half_type, {2, 2}});
        auto dot = p2.add_instruction(migraphx::op::dot{1, 0}, x, y);
        auto beta =
            p2.add_literal(migraphx::literal{migraphx::shape{migraphx::shape::half_type}, {0.5}});
        auto beta_broadcast = p2.add_instruction(migraphx::op::multibroadcast{{2, 2}}, beta);
        auto mul            = p2.add_instruction(migraphx::op::mul{}, z, beta_broadcast);
        auto add            = p2.add_instruction(migraphx::op::add{}, dot, mul);
        p2.add_instruction(migraphx::op::identity{}, add);
    }
    EXPECT(p1 == p2);
}

TEST_CASE(dot_add_beta_double)
{
    migraphx::program p1;
    {
        auto x   = p1.add_parameter("x", migraphx::shape{migraphx::shape::double_type, {2, 2}});
        auto y   = p1.add_parameter("y", migraphx::shape{migraphx::shape::double_type, {2, 2}});
        auto z   = p1.add_parameter("z", migraphx::shape{migraphx::shape::double_type, {2, 2}});
        auto dot = p1.add_instruction(migraphx::op::dot{1.0, 0.5}, x, y, z);
        p1.add_instruction(migraphx::op::identity{}, dot);
    }
    run_pass(p1);
    migraphx::program p2;
    {
        auto x   = p2.add_parameter("x", migraphx::shape{migraphx::shape::double_type, {2, 2}});
        auto y   = p2.add_parameter("y", migraphx::shape{migraphx::shape::double_type, {2, 2}});
        auto z   = p2.add_parameter("z", migraphx::shape{migraphx::shape::double_type, {2, 2}});
        auto dot = p2.add_instruction(migraphx::op::dot{1, 0}, x, y);
        auto beta =
            p2.add_literal(migraphx::literal{migraphx::shape{migraphx::shape::double_type}, {0.5}});
        auto beta_broadcast = p2.add_instruction(migraphx::op::multibroadcast{{2, 2}}, beta);
        auto mul            = p2.add_instruction(migraphx::op::mul{}, z, beta_broadcast);
        auto add            = p2.add_instruction(migraphx::op::add{}, dot, mul);
        p2.add_instruction(migraphx::op::identity{}, add);
    }
    EXPECT(p1 == p2);
}

TEST_CASE(dot_add_beta_int)
{
    migraphx::program p1;
    {
        auto x   = p1.add_parameter("x", migraphx::shape{migraphx::shape::int32_type, {2, 2}});
        auto y   = p1.add_parameter("y", migraphx::shape{migraphx::shape::int32_type, {2, 2}});
        auto z   = p1.add_parameter("z", migraphx::shape{migraphx::shape::int32_type, {2, 2}});
        auto dot = p1.add_instruction(migraphx::op::dot{1.0, 0.5}, x, y, z);
        p1.add_instruction(migraphx::op::identity{}, dot);
    }
    migraphx::program p2 = p1;
    run_pass(p1);
    EXPECT(p1 == p2);
}

int main(int argc, const char* argv[]) { test::run(argc, argv); }