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

using migraphx::trim;

Paul's avatar
Paul committed
19
20
// m test_gpu_mlir && ./bin/test_gpu_mlir

Paul's avatar
Paul committed
21
22
// [  FAILED  ] test_conv_relu_half

Paul's avatar
Paul committed
23
24
25
struct mlir_gpu_target : migraphx::gpu::target
{
    std::string name() const { return "mlir"; }
Paul's avatar
Format  
Paul committed
26
27
    std::vector<migraphx::pass> get_passes(migraphx::context& gctx,
                                           const migraphx::compile_options&) const
Paul's avatar
Paul committed
28
29
    {
        auto& ctx = migraphx::any_cast<migraphx::gpu::context>(gctx);
Paul's avatar
Format  
Paul committed
30
        return {migraphx::gpu::write_literals{&ctx}};
Paul's avatar
Paul committed
31
32
33
    }
};

34
35
36
37
std::string encode(std::string s)
{
    std::stringstream ss;
    bool prespace = false;
Paul's avatar
Paul committed
38
    for(auto c : s)
39
    {
Paul's avatar
Paul committed
40
        if(std::isspace(c))
41
        {
Paul's avatar
Paul committed
42
            if(not prespace)
43
44
45
                ss << "  ";
            prespace = true;
        }
Paul's avatar
Paul committed
46
        else if(std::isprint(c))
47
48
49
50
51
52
53
54
        {
            ss << c;
            prespace = false;
        }
    }
    return migraphx::trim(ss.str());
}

Paul's avatar
Paul committed
55
56
57
migraphx::program create_program_from_mlir(const migraphx::module& mmlir)
{
    migraphx::program p;
Paul's avatar
Paul committed
58
    auto* mm   = p.get_main_module();
Paul's avatar
Paul committed
59
60
61
62
63
    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));
    });
Paul's avatar
Format  
Paul committed
64
65
66
    std::sort(inputs.begin(), inputs.end(), migraphx::by(std::less<>{}, [](auto ins) {
                  return to_string(ins->get_operator());
              }));
Paul's avatar
Paul committed
67
    inputs.push_back(mm->add_parameter("output", mmlir.get_output_shapes().front()));
Paul's avatar
Paul committed
68

Paul's avatar
Paul committed
69
70
    migraphx::gpu::context ctx;
    migraphx::gpu::insert_mlir(*mm, mm->end(), compile_mlir(ctx, mmlir), inputs);
Paul's avatar
Paul committed
71
72
73
74
75
76
77
78
79
    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())
    {
Paul's avatar
Updates  
Paul committed
80
        // m[x.first] = migraphx::fill_argument(x.second, 1);
Paul's avatar
Paul committed
81
82
83
84
85
86
87
        m[x.first] = migraphx::generate_argument(x.second, i++);
    }
    return m;
}

migraphx::argument run_gpu(migraphx::program p, const migraphx::parameter_map& inputs)
{
Paul's avatar
Paul committed
88
    mlir_gpu_target t;
Paul's avatar
Paul committed
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
    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));
}

122
123
124
TEST_CASE(conv)
{
    const std::string mlir_output = R"__migraphx__(
Paul's avatar
Paul committed
125
126
127
module {
  func @main(%arg0: tensor<2x8x3x3xf32>, %arg1: tensor<1x8x4x4xf32>) -> tensor<1x2x2x2xf32> attributes {kernel = "mixr"} {
    %0 = migraphx.convolution(%arg1, %arg0) {dilation = [1, 1], group = 1 : i64, padding = [0, 0, 0, 0], padding_mode = 0 : i64, stride = [1, 1]} : (tensor<1x8x4x4xf32>, tensor<2x8x3x3xf32>) -> tensor<1x2x2x2xf32>
128
    return %0 : tensor<1x2x2x2xf32>
129
130
131
132
  }
}
)__migraphx__";
    migraphx::module m;
Paul's avatar
Format  
Paul committed
133
134
    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
135
    auto conv = m.add_instruction(migraphx::make_op("convolution"), x, w);
Paul's avatar
Paul committed
136
    m.add_return({conv});
137
    auto s = migraphx::gpu::dump_mlir(m);
Paul's avatar
Paul committed
138
    // Skip test if MLIR is not enabled
Paul's avatar
Format  
Paul committed
139
    if(s.empty())
Paul's avatar
Paul committed
140
        return;
Paul's avatar
Paul committed
141
    CHECK(encode(s) == encode(mlir_output));
Paul's avatar
Paul committed
142
    EXPECT(verify_mlir(m));
Paul's avatar
Paul committed
143
144
145
146
147
}

TEST_CASE(conv_add_relu)
{
    const std::string mlir_output = R"__migraphx__(
Paul's avatar
Paul committed
148
149
150
151
module {
  func @main(%arg0: tensor<1x2x2x2xf32>, %arg1: tensor<2x8x3x3xf32>, %arg2: tensor<1x8x4x4xf32>) -> tensor<1x2x2x2xf32> attributes {kernel = "mixr"} {
    %0 = migraphx.convolution(%arg2, %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>
    %1 = migraphx.add(%0, %arg0) : (tensor<1x2x2x2xf32>, tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
Paul's avatar
Paul committed
152
    %2 = migraphx.relu(%1) : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
153
    return %2 : tensor<1x2x2x2xf32>
Paul's avatar
Paul committed
154
155
156
157
  }
}
)__migraphx__";
    migraphx::module m;
Paul's avatar
Format  
Paul committed
158
159
160
    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
161
    auto conv = m.add_instruction(migraphx::make_op("convolution"), x, w);
Paul's avatar
Format  
Paul committed
162
    auto add  = m.add_instruction(migraphx::make_op("add"), conv, b);
Paul's avatar
Paul committed
163
164
    auto relu = m.add_instruction(migraphx::make_op("relu"), add);
    m.add_return({relu});
Paul's avatar
Paul committed
165
166
    auto s = migraphx::gpu::dump_mlir(m);
    // Skip test if MLIR is not enabled
Paul's avatar
Format  
Paul committed
167
    if(s.empty())
Paul's avatar
Paul committed
168
        return;
Paul's avatar
Paul committed
169
    CHECK(encode(s) == encode(mlir_output));
Paul's avatar
Paul committed
170
    EXPECT(verify_mlir(m));
171
172
173
}

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