mlir.cpp 4.94 KB
Newer Older
1
#include <migraphx/gpu/mlir.hpp>
Paul's avatar
Paul committed
2
3
#include <migraphx/gpu/target.hpp>
#include <migraphx/ref/target.hpp>
4
#include <migraphx/module.hpp>
Paul's avatar
Paul committed
5
#include <migraphx/program.hpp>
6
7
8
#include <migraphx/make_op.hpp>
#include <migraphx/ranges.hpp>
#include <migraphx/stringutils.hpp>
Paul's avatar
Paul committed
9
10
#include <migraphx/generate.hpp>
#include <migraphx/verify_args.hpp>
11
12
13
14
15
16
17
18
#include <test.hpp>

using migraphx::trim;

std::string encode(std::string s)
{
    std::stringstream ss;
    bool prespace = false;
Paul's avatar
Paul committed
19
    for(auto c : s)
20
    {
Paul's avatar
Paul committed
21
        if(std::isspace(c))
22
        {
Paul's avatar
Paul committed
23
            if(not prespace)
24
25
26
                ss << "  ";
            prespace = true;
        }
Paul's avatar
Paul committed
27
        else if(std::isprint(c))
28
29
30
31
32
33
34
35
        {
            ss << c;
            prespace = false;
        }
    }
    return migraphx::trim(ss.str());
}

Paul's avatar
Paul committed
36
37
38
migraphx::program create_program_from_mlir(const migraphx::module& mmlir)
{
    migraphx::program p;
Paul's avatar
Paul committed
39
    auto* mm   = p.get_main_module();
Paul's avatar
Paul committed
40
41
42
43
44
45
46
    auto names = mmlir.get_parameter_names();
    std::vector<migraphx::instruction_ref> inputs;
    std::transform(names.begin(), names.end(), std::back_inserter(inputs), [&](const auto& name) {
        return mm->add_parameter(name, mmlir.get_parameter_shape(name));
    });
    inputs.push_back(mm->add_parameter("output", mmlir.get_output_shapes().front()));
    migraphx::gpu::insert_mlir(*mm, mm->end(), mmlir, inputs);
Paul's avatar
Paul committed
47
    std::cout << p << std::endl;
Paul's avatar
Paul committed
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
    return p;
}

migraphx::parameter_map generate_params(const migraphx::program& p)
{
    migraphx::parameter_map m;
    std::size_t i = 0;
    for(auto&& x : p.get_parameter_shapes())
    {
        m[x.first] = migraphx::generate_argument(x.second, i++);
    }
    return m;
}

migraphx::argument run_gpu(migraphx::program p, const migraphx::parameter_map& inputs)
{
    migraphx::gpu::target t;
    p.compile(t);
    migraphx::parameter_map m;
    for(auto&& input : inputs)
    {
        m[input.first] = t.copy_to(input.second);
    }
    for(auto&& x : p.get_parameter_shapes())
    {
        if(m.count(x.first) == 0)
        {
            m[x.first] = t.allocate(x.second);
        }
    }
    return t.copy_from(p.eval(m).front());
}

migraphx::argument run_ref(migraphx::program p, const migraphx::parameter_map& inputs)
{
    p.compile(migraphx::ref::target{});
    return p.eval(inputs).front();
}

bool verify_mlir(const migraphx::module& mmlir)
{
    migraphx::program ref;
    ref.get_main_module()->insert_module_instructions(ref.get_main_module()->end(), &mmlir);

    auto inputs = generate_params(ref);

    auto mlir = create_program_from_mlir(mmlir);
    return migraphx::verify_args("mlir", run_ref(ref, inputs), run_gpu(mlir, inputs));
}

98
99
100
101
102
TEST_CASE(conv)
{
    const std::string mlir_output = R"__migraphx__(
module  {
  func @main(%arg0: tensor<1x8x4x4xf32>, %arg1: tensor<2x8x3x3xf32>) -> tensor<1x2x2x2xf32> {
103
104
    %0 = migraphx.convolution(%arg0, %arg1) {dilation = [1, 1], group = 1 : i64, padding = [0, 0, 0, 0], padding_mode = 0 : i64, stride = [1, 1]} : (tensor<1x8x4x4xf32>, tensor<2x8x3x3xf32>) -> tensor<1x2x2x2xf32>
    return %0 : tensor<1x2x2x2xf32>
105
106
107
108
  }
}
)__migraphx__";
    migraphx::module m;
Paul's avatar
Format  
Paul committed
109
110
    auto x    = m.add_parameter("x", {migraphx::shape::float_type, {1, 8, 4, 4}});
    auto w    = m.add_parameter("w", {migraphx::shape::float_type, {2, 8, 3, 3}});
Paul's avatar
Paul committed
111
112
    auto conv = m.add_instruction(migraphx::make_op("convolution"), x, w);
    m.add_return({conv});
113
    auto s = migraphx::gpu::dump_mlir(m);
Paul's avatar
Paul committed
114
    // Skip test if MLIR is not enabled
Paul's avatar
Format  
Paul committed
115
    if(s.empty())
Paul's avatar
Paul committed
116
        return;
Paul's avatar
Paul committed
117
    std::cout << s << std::endl;
Paul's avatar
Paul committed
118
    EXPECT(encode(s) == encode(mlir_output));
Paul's avatar
Paul committed
119
    EXPECT(verify_mlir(m));
Paul's avatar
Paul committed
120
121
122
123
124
125
126
}

TEST_CASE(conv_add_relu)
{
    const std::string mlir_output = R"__migraphx__(
module  {
  func @main(%arg0: tensor<1x8x4x4xf32>, %arg1: tensor<2x8x3x3xf32>, %arg2: tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> {
127
    %0 = migraphx.convolution(%arg0, %arg1) {dilation = [1, 1], group = 1 : i64, padding = [0, 0, 0, 0], padding_mode = 0 : i64, stride = [1, 1]} : (tensor<1x8x4x4xf32>, tensor<2x8x3x3xf32>) -> tensor<1x2x2x2xf32>
Paul's avatar
Paul committed
128
129
    %1 = migraphx.add(%0, %arg2) : (tensor<1x2x2x2xf32>, tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
    %2 = migraphx.relu(%1) : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
130
    return %2 : tensor<1x2x2x2xf32>
Paul's avatar
Paul committed
131
132
133
134
  }
}
)__migraphx__";
    migraphx::module m;
Paul's avatar
Format  
Paul committed
135
136
137
    auto x    = m.add_parameter("x", {migraphx::shape::float_type, {1, 8, 4, 4}});
    auto w    = m.add_parameter("w", {migraphx::shape::float_type, {2, 8, 3, 3}});
    auto b    = m.add_parameter("b", {migraphx::shape::float_type, {1, 2, 2, 2}});
Paul's avatar
Paul committed
138
    auto conv = m.add_instruction(migraphx::make_op("convolution"), x, w);
Paul's avatar
Format  
Paul committed
139
    auto add  = m.add_instruction(migraphx::make_op("add"), conv, b);
Paul's avatar
Paul committed
140
141
    auto relu = m.add_instruction(migraphx::make_op("relu"), add);
    m.add_return({relu});
Paul's avatar
Paul committed
142
143
    auto s = migraphx::gpu::dump_mlir(m);
    // Skip test if MLIR is not enabled
Paul's avatar
Format  
Paul committed
144
    if(s.empty())
Paul's avatar
Paul committed
145
        return;
Paul's avatar
Paul committed
146
    std::cout << s << std::endl;
147
    EXPECT(encode(s) == encode(mlir_output));
Paul's avatar
Paul committed
148
    EXPECT(verify_mlir(m));
149
150
151
}

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