mlir.cpp 14.4 KB
Newer Older
Paul's avatar
Paul committed
1
#include <migraphx/mlir.hpp>
Paul's avatar
Paul committed
2
3
4
5
6
7
8
9
10
11
12
13

#include <mlir-c/IR.h>
#include <mlir-c/BuiltinAttributes.h>
#include <mlir-c/BuiltinTypes.h>
#include <mlir-c/Diagnostics.h>
#include <mlir-c/Dialect/Standard.h>
#include <mlir-c/Dialect/MIGraphX.h>
#include <mlir-c/IntegerSet.h>
#include <mlir-c/Registration.h>

#include <migraphx/manage_ptr.hpp>
#include <migraphx/module.hpp>
Paul's avatar
Paul committed
14
#include <migraphx/instruction.hpp>
Paul's avatar
Paul committed
15
#include <migraphx/config.hpp>
Paul's avatar
Paul committed
16
17
18
19
#include <migraphx/ranges.hpp>
#include <migraphx/iterator_for.hpp>
#include <deque>
#include <variant>
Paul's avatar
Paul committed
20
21
22
23

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {

Paul's avatar
Paul committed
24
template <class T, class F, F f> // NOLINT
Paul's avatar
Paul committed
25
26
27
28
29
30
31
32
33
34
35
36
37
38
struct mlir_handle
{
    struct ptr
    {
        ptr() = default;
        ptr(std::nullptr_t) {}
        ptr(T x) : obj(x) {}

        std::intptr_t get_value() const
        {
            static_assert(sizeof(T) == sizeof(std::intptr_t), "MLIR Handle different size");
            return reinterpret_cast<const std::intptr_t&>(obj);
        }

Paul's avatar
Paul committed
39
        T get() const { return obj; }
Paul's avatar
Paul committed
40

Paul's avatar
Paul committed
41
        friend bool operator==(ptr x, ptr y) { return x.get_value() == y.get_value(); }
Paul's avatar
Paul committed
42

Paul's avatar
Paul committed
43
        friend bool operator!=(ptr x, ptr y) { return !(x == y); }
Paul's avatar
Paul committed
44
45
        T obj{};
    };
Paul's avatar
Paul committed
46

Paul's avatar
Paul committed
47
48
49
50
51
52
53
54
55
56
57
58
59
    struct deleter
    {
        using pointer = ptr;

        void operator()(pointer x) const
        {
            if(x != nullptr)
            {
                (void)f(x.obj);
            }
        }
    };

Paul's avatar
Paul committed
60
    mlir_handle() : handle(nullptr) {}
Paul's avatar
Paul committed
61

Paul's avatar
Paul committed
62
    mlir_handle(T p) : handle(ptr{p}) {}
Paul's avatar
Paul committed
63

Paul's avatar
Paul committed
64
    T get() const { return handle.get().get(); }
Paul's avatar
Paul committed
65

Paul's avatar
Paul committed
66
    T release() { return handle.release().get(); }
Paul's avatar
Paul committed
67

Paul's avatar
Paul committed
68
    private:
Paul's avatar
Paul committed
69
70
71
    std::unique_ptr<ptr, deleter> handle;
};

Paul's avatar
Paul committed
72
#define MIGRAPHX_MANAGE_MLIR_HANDLE(T, F) migraphx::mlir_handle<T, decltype(&F), &F> // NOLINT
Paul's avatar
Paul committed
73

Paul's avatar
Paul committed
74
75
76
77
78
79
80
using mlir_context           = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirContext, mlirContextDestroy);
using mlir_module            = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirModule, mlirModuleDestroy);
using mlir_operation         = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirOperation, mlirOperationDestroy);
using mlir_op_printing_flags = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirOpPrintingFlags,
                                                           mlirOpPrintingFlagsDestroy);
using mlir_region            = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirRegion, mlirRegionDestroy);
using mlir_block             = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirBlock, mlirBlockDestroy);
Paul's avatar
Paul committed
81

Paul's avatar
Paul committed
82
std::string_view to_string_view(MlirStringRef s) { return {s.data, s.length}; }
Paul's avatar
Paul committed
83
84
85
86
87
88

MlirStringRef make_mlir_string_ref(const std::string_view& s)
{
    return mlirStringRefCreate(s.data(), s.size());
}

Paul's avatar
Paul committed
89
template <class F, class T, class Printer>
Paul's avatar
Paul committed
90
91
void mlir_print(F f, T x, Printer printer)
{
Paul's avatar
Paul committed
92
93
94
    f(x,
      +[](MlirStringRef s, void* data) { (*reinterpret_cast<Printer*>(data))(to_string_view(s)); },
      &printer);
Paul's avatar
Paul committed
95
96
}

Paul's avatar
Paul committed
97
template <class F, class T>
Paul's avatar
Paul committed
98
99
100
101
102
void mlir_print(F f, T x, std::ostream& os)
{
    mlir_print(f, x, [&](auto s) { os << s; });
}

Paul's avatar
Paul committed
103
104
105
106
107
108
109
110
template <class F, class T>
std::string mlir_print(F f, T x)
{
    std::stringstream ss;
    mlir_print(f, x, [&](auto s) { ss << s; });
    return ss.str();
}

Paul's avatar
Paul committed
111
112
struct mlir_program
{
Paul's avatar
Paul committed
113
114
115
116
    mlir_program()
        : ctx(mlirContextCreate()),
          location(mlirLocationUnknownGet(ctx.get())),
          mmodule(mlirModuleCreateEmpty(location))
Paul's avatar
Paul committed
117
118
    {
        mlirRegisterAllDialects(ctx.get());
Paul's avatar
Paul committed
119
        mlirContextSetAllowUnregisteredDialects(ctx.get(), true /*allow*/);
Paul's avatar
Paul committed
120
121
122
123
124
125
    }

    MlirType make_type(shape::type_t t) const
    {
        MlirType result;
        shape::visit(t, [&](auto as) {
Paul's avatar
Paul committed
126
            if(as.type_enum() == shape::float_type)
Paul's avatar
Paul committed
127
                result = mlirF32TypeGet(ctx.get());
Paul's avatar
Paul committed
128
            else if(as.type_enum() == shape::half_type)
Paul's avatar
Paul committed
129
                result = mlirF16TypeGet(ctx.get());
Paul's avatar
Paul committed
130
            else if(as.type_enum() == shape::double_type)
Paul's avatar
Paul committed
131
                result = mlirF64TypeGet(ctx.get());
Paul's avatar
Paul committed
132
            else if(as.is_integral())
Paul's avatar
Paul committed
133
            {
Paul's avatar
Paul committed
134
                if(as.is_signed())
Paul's avatar
Paul committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
                    result = mlirIntegerTypeSignedGet(ctx.get(), as.size() * 8);
                else
                    result = mlirIntegerTypeGet(ctx.get(), as.size() * 8);
            }
            else
                MIGRAPHX_THROW("Unsupported type: " + std::to_string(as.type_enum()));
        });
        return result;
    }

    MlirType make_tensor(const shape& s) const
    {
        assert(s.standard());
        std::vector<int64_t> lens(s.lens().begin(), s.lens().end());
Paul's avatar
Format  
Paul committed
149
150
        return mlirRankedTensorTypeGet(
            lens.size(), lens.data(), make_type(s.type()), mlirAttributeGetNull());
Paul's avatar
Paul committed
151
152
    }

Paul's avatar
Paul committed
153
    template <class Range>
Paul's avatar
Paul committed
154
155
156
157
158
159
160
161
162
163
164
    std::vector<MlirType> make_tensors(const Range& r)
    {
        std::vector<MlirType> result;
        std::transform(r.begin(), r.end(), std::back_inserter(result), [&](const auto& s) {
            return make_tensor(s);
        });
        return result;
    }

    MlirType make_function_type(const std::vector<shape>& inputs, const std::vector<shape>& outputs)
    {
Paul's avatar
Paul committed
165
        auto in  = make_tensors(inputs);
Paul's avatar
Paul committed
166
167
168
169
        auto out = make_tensors(outputs);
        return mlirFunctionTypeGet(ctx.get(), in.size(), in.data(), out.size(), out.data());
    }

Paul's avatar
Paul committed
170
171
172
173
174
175
176
177
178
    MlirIdentifier id(const std::string_view& s) const
    {
        return mlirIdentifierGet(ctx.get(), make_mlir_string_ref(s));
    }

    MlirAttribute attribute(std::int64_t i) const
    {
        return mlirIntegerAttrGet(mlirIntegerTypeSignedGet(ctx.get(), 64), i);
    }
Paul's avatar
Paul committed
179
180
    MlirAttribute attribute(std::uint64_t i) const { return attribute(std::int64_t(i)); }
    MlirAttribute attribute(unsigned char i) const { return attribute(std::int64_t(i)); }
Paul's avatar
Paul committed
181
    MlirAttribute attribute(bool b) const { return mlirBoolAttrGet(ctx.get(), b ? 1 : 0); }
Paul's avatar
Paul committed
182
183
184
185
186
187
188
189
    MlirAttribute attribute(double d) const
    {
        return mlirFloatAttrDoubleGet(ctx.get(), mlirF64TypeGet(ctx.get()), d);
    }
    MlirAttribute attribute(const std::string& s) const
    {
        return mlirStringAttrGet(ctx.get(), make_mlir_string_ref(s));
    }
Paul's avatar
Paul committed
190
191
    MlirAttribute attribute(std::nullptr_t) const { return {}; }
    template <class T>
Paul's avatar
Paul committed
192
193
194
195
196
197
198
199
200
201
202
203
    MlirAttribute attribute(const std::vector<T>& v) const
    {
        std::vector<MlirAttribute> attributes;
        attributes.reserve(v.size());
        std::transform(v.begin(), v.end(), std::back_inserter(attributes), [&](auto&& x) {
            return attribute(x);
        });
        return mlirArrayAttrGet(ctx.get(), attributes.size(), attributes.data());
    }
    MlirAttribute attribute(const value& v) const
    {
        MlirAttribute attr;
Paul's avatar
Paul committed
204
        v.visit_value([&](auto&& x) { attr = attribute(x); });
Paul's avatar
Paul committed
205
206
207
208
209
210
211
212
213
214
        return attr;
    }
    MlirAttribute attribute(const std::vector<value>& v) const
    {
        if(v.empty())
        {
            return mlirArrayAttrGet(ctx.get(), 0, nullptr);
        }
        if(not v.front().get_key().empty())
        {
Paul's avatar
Paul committed
215
            std::vector<MlirNamedAttribute> attributes = name_attributes(v);
Paul's avatar
Paul committed
216
217
218
219
220
221
222
223
224
225
226
227
228
            return mlirDictionaryAttrGet(ctx.get(), attributes.size(), attributes.data());
        }
        else
        {
            std::vector<MlirAttribute> attributes;
            attributes.reserve(v.size());
            std::transform(v.begin(), v.end(), std::back_inserter(attributes), [&](auto&& x) {
                return attribute(x);
            });
            return mlirArrayAttrGet(ctx.get(), attributes.size(), attributes.data());
        }
    }

Paul's avatar
Paul committed
229
    MlirAttribute attribute(MlirType t) const { return mlirTypeAttrGet(t); }
Paul's avatar
Paul committed
230

Paul's avatar
Paul committed
231
232
    MlirAttribute attribute(MlirAttribute a) const { return a; }

Paul's avatar
Paul committed
233
    template <class T>
Paul's avatar
Paul committed
234
235
236
237
238
239
240
241
    MlirNamedAttribute name_attribute(const std::string_view& key, const T& x) const
    {
        MlirNamedAttribute attr;
        attr.name      = id(key);
        attr.attribute = attribute(x);
        return attr;
    }

Paul's avatar
Paul committed
242
243
244
245
246
247
248
249
250
    using attribute_t       = std::variant<std::nullptr_t,
                                     std::uint64_t,
                                     unsigned char,
                                     bool,
                                     double,
                                     std::string,
                                     value,
                                     std::vector<value>,
                                     MlirType>;
Paul's avatar
Paul committed
251
252
253
254
    using named_attribute_t = std::pair<std::string_view, attribute_t>;

    MlirNamedAttribute name_attribute(const named_attribute_t& na) const
    {
Paul's avatar
Paul committed
255
256
        return name_attribute(na.first,
                              std::visit([&](const auto& x) { return attribute(x); }, na.second));
Paul's avatar
Paul committed
257
258
    }

Paul's avatar
Paul committed
259
260
    std::vector<MlirNamedAttribute>
    name_attributes(const std::vector<named_attribute_t>& named_attrs) const
Paul's avatar
Paul committed
261
262
263
    {
        std::vector<MlirNamedAttribute> attributes;
        attributes.reserve(named_attrs.size());
Paul's avatar
Paul committed
264
265
266
267
        std::transform(named_attrs.begin(),
                       named_attrs.end(),
                       std::back_inserter(attributes),
                       [&](const named_attribute_t& a) { return name_attribute(a); });
Paul's avatar
Paul committed
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
        return attributes;
    }

    std::vector<MlirNamedAttribute> name_attributes(const value& v) const
    {
        std::vector<MlirNamedAttribute> attributes;
        attributes.reserve(v.size());
        std::transform(v.begin(), v.end(), std::back_inserter(attributes), [&](const value& x) {
            return name_attribute(x.get_key(), x.without_key());
        });
        return attributes;
    }

    struct mlir_operation_state
    {
Paul's avatar
Paul committed
283
284
285
286
        mlir_operation_state(mlir_program& p, const std::string_view& name)
            : prog(&p), op_state(mlirOperationStateGet(make_mlir_string_ref(name), p.location))
        {
        }
Paul's avatar
Paul committed
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335

        mlir_operation_state& add_attributes(const std::vector<named_attribute_t>& named_attrs)
        {
            auto attributes = prog->name_attributes(named_attrs);
            mlirOperationStateAddAttributes(&op_state, attributes.size(), attributes.data());
            return *this;
        }

        mlir_operation_state& add_attribute_value(const value& v)
        {
            auto attributes = prog->name_attributes(v);
            mlirOperationStateAddAttributes(&op_state, attributes.size(), attributes.data());
            return *this;
        }

        mlir_operation_state& add_regions(std::vector<mlir_region> rs)
        {
            regions = std::move(rs);
            return *this;
        }

        mlir_operation_state& add_region(mlir_region r)
        {
            regions.emplace_back(std::move(r));
            return *this;
        }

        mlir_operation_state& add_results(const std::vector<shape>& outputs)
        {
            auto x = prog->make_tensors(outputs);
            mlirOperationStateAddResults(&op_state, x.size(), x.data());
            return *this;
        }

        mlir_operation_state& add_operands(const std::vector<MlirValue>& inputs)
        {
            mlirOperationStateAddOperands(&op_state, inputs.size(), inputs.data());
            return *this;
        }

        mlir_operation create_operation()
        {
            std::vector<MlirRegion> mregions(regions.size());
            std::transform(regions.begin(), regions.end(), mregions.begin(), [](const auto& r) {
                return r.get();
            });
            mlirOperationStateAddOwnedRegions(&op_state, mregions.size(), mregions.data());
            mlir_operation op(mlirOperationCreate(&op_state));
            // Release memory since mlir_operation owns it
Paul's avatar
Paul committed
336
            for(auto& r : regions)
Paul's avatar
Paul committed
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
                r.release();
            regions.clear();
            return op;
        }

        mlir_program* prog;
        MlirOperationState op_state;
        std::vector<mlir_region> regions = {};
    };

    mlir_operation_state create_operation_state(const std::string_view& name)
    {
        return {*this, name};
    }

    std::vector<MlirValue> insert(MlirBlock body, mlir_operation_state ops)
    {
        std::vector<MlirValue> result;
        mlir_operation op = ops.create_operation();
Paul's avatar
Paul committed
356
        auto weak_op      = op.get();
Paul's avatar
Paul committed
357
358
359
360
361
362
363
364
365
366
        mlirBlockInsertOwnedOperation(body, 0, op.release());

        auto n = mlirOperationGetNumResults(weak_op);
        result.reserve(n);
        transform(range(n), std::back_inserter(result), [&](auto i) {
            return mlirOperationGetResult(weak_op, i);
        });
        return result;
    }

Paul's avatar
Paul committed
367
368
    MlirBlock
    insert(MlirBlock body, const module& m, std::unordered_map<instruction_ref, MlirValue>& ins_map)
Paul's avatar
Paul committed
369
370
371
    {
        auto names = m.get_parameter_names();
        std::vector<shape> inputs;
Paul's avatar
Paul committed
372
373
374
375
        std::transform(names.begin(),
                       names.end(),
                       std::back_inserter(inputs),
                       [&](const std::string& name) { return m.get_parameter_shape(name); });
Paul's avatar
Paul committed
376
377
        std::vector<shape> outputs = m.get_output_shapes();

Paul's avatar
Paul committed
378
        auto body_inputs   = make_tensors(inputs);
Paul's avatar
Paul committed
379
        mlir_region region = mlirRegionCreate();
Paul's avatar
Paul committed
380
381
        mlir_block fbody   = mlirBlockCreate(body_inputs.size(), body_inputs.data());
        MlirBlock result   = fbody.get();
Paul's avatar
Paul committed
382
383
384
        mlirRegionAppendOwnedBlock(region.get(), fbody.release());

        auto ops = create_operation_state("builtin.func");
Paul's avatar
Paul committed
385
        ops.add_attributes(
Paul's avatar
Paul committed
386
            {{"type", make_function_type(inputs, outputs)}, {"sym_name", std::string("main")}});
Paul's avatar
Paul committed
387
388
389
        ops.add_region(std::move(region));
        insert(body, std::move(ops));

Paul's avatar
Paul committed
390
        for(auto i : range(names.size()))
Paul's avatar
Paul committed
391
392
393
394
395
396
397
398
399
            ins_map[m.get_parameter(names[i])] = mlirBlockGetArgument(result, i);
        return result;
    }

    void parse(const module& m)
    {
        auto mbody = mlirModuleGetBody(mmodule.get());
        std::unordered_map<instruction_ref, MlirValue> ins_map;
        auto fbody = insert(mbody, m, ins_map);
Paul's avatar
Paul committed
400
        for(auto ins : iterator_for(m))
Paul's avatar
Paul committed
401
        {
Paul's avatar
Paul committed
402
            if(ins->name() == "@param")
Paul's avatar
Paul committed
403
                continue;
Paul's avatar
Paul committed
404
            auto name = "migraphx." + ins->name();
Paul's avatar
Paul committed
405
            auto ops  = create_operation_state(name);
Paul's avatar
Paul committed
406
407
408
409
            ops.add_attribute_value(ins->get_operator().to_value());
            ops.add_results({ins->get_shape()});

            std::vector<MlirValue> inputs;
Paul's avatar
Paul committed
410
411
            transform(
                ins->inputs(), std::back_inserter(inputs), [&](auto i) { return ins_map.at(i); });
Paul's avatar
Paul committed
412
413
414
415
416
417
418
419
            ops.add_operands(inputs);

            auto outputs = insert(fbody, std::move(ops));
            assert(outputs.size() == 1);
            ins_map[ins] = outputs.front();
        }
    }

Paul's avatar
Paul committed
420
    mlir_context ctx;
Paul's avatar
Paul committed
421
422
423
    MlirLocation location;
    mlir_module mmodule;
    std::deque<std::string> strings{};
Paul's avatar
Paul committed
424
425
};

Paul's avatar
Paul committed
426
427
428
429
430
431
432
433
std::string dump_mlir(const module& m)
{
    mlir_program mp;
    mp.parse(m);
    auto mod_op = mlirModuleGetOperation(mp.mmodule.get());
    return mlir_print(&mlirOperationPrint, mod_op);
}

Paul's avatar
Paul committed
434
435
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx