mlir.cpp 16.5 KB
Newer Older
1
#include <migraphx/gpu/mlir.hpp>
Paul's avatar
Paul committed
2

Paul's avatar
Paul committed
3
#ifdef MIGRAPHX_MLIR
Paul's avatar
Paul committed
4
5
6
7
8
9
10
#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>
Paul's avatar
Paul committed
11
#include <mlir-c/Pass.h>
Paul's avatar
Paul committed
12
#include <mlir-c/Registration.h>
Paul's avatar
Paul committed
13
#endif
Paul's avatar
Paul committed
14
15
16

#include <migraphx/manage_ptr.hpp>
#include <migraphx/module.hpp>
Paul's avatar
Paul committed
17
#include <migraphx/instruction.hpp>
Paul's avatar
Paul committed
18
#include <migraphx/config.hpp>
Paul's avatar
Paul committed
19
#include <migraphx/ranges.hpp>
Paul's avatar
Paul committed
20
#include <migraphx/gpu/code_object_op.hpp>
Paul's avatar
Paul committed
21
22
23
#include <migraphx/iterator_for.hpp>
#include <deque>
#include <variant>
Paul's avatar
Paul committed
24
25
26

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
27
namespace gpu {
Paul's avatar
Paul committed
28

Paul's avatar
Paul committed
29
#ifdef MIGRAPHX_MLIR
Paul's avatar
Paul committed
30
template <class T, class F, F f> // NOLINT
Paul's avatar
Paul committed
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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
45
        T get() const { return obj; }
Paul's avatar
Paul committed
46

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

Paul's avatar
Paul committed
49
        friend bool operator!=(ptr x, ptr y) { return !(x == y); }
Paul's avatar
Paul committed
50
51
        T obj{};
    };
Paul's avatar
Paul committed
52

Paul's avatar
Paul committed
53
54
55
56
57
58
59
60
61
62
63
64
65
    struct deleter
    {
        using pointer = ptr;

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

Paul's avatar
Paul committed
66
    mlir_handle() : handle(nullptr) {}
Paul's avatar
Paul committed
67

Paul's avatar
Paul committed
68
    mlir_handle(T p) : handle(ptr{p}) {}
Paul's avatar
Paul committed
69

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

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

Paul's avatar
Paul committed
74
    private:
Paul's avatar
Paul committed
75
76
77
    std::unique_ptr<ptr, deleter> handle;
};

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

Paul's avatar
Paul committed
80
81
82
83
84
85
86
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
87
using mlir_pass_manager      = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirPassManager, mlirPassManagerDestroy);
Paul's avatar
Paul committed
88

Paul's avatar
Paul committed
89
std::string_view to_string_view(MlirStringRef s) { return {s.data, s.length}; }
Paul's avatar
Paul committed
90
91
92
93
94
95

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

Paul's avatar
Paul committed
96
template <class F, class T, class Printer>
Paul's avatar
Paul committed
97
98
void mlir_print(F f, T x, Printer printer)
{
Paul's avatar
Paul committed
99
100
101
    f(x,
      +[](MlirStringRef s, void* data) { (*reinterpret_cast<Printer*>(data))(to_string_view(s)); },
      &printer);
Paul's avatar
Paul committed
102
103
}

Paul's avatar
Paul committed
104
template <class F, class T>
Paul's avatar
Paul committed
105
106
107
108
109
void mlir_print(F f, T x, std::ostream& os)
{
    mlir_print(f, x, [&](auto s) { os << s; });
}

Paul's avatar
Paul committed
110
111
112
113
114
115
116
117
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
118
119
struct mlir_program
{
Paul's avatar
Paul committed
120
121
122
123
    mlir_program()
        : ctx(mlirContextCreate()),
          location(mlirLocationUnknownGet(ctx.get())),
          mmodule(mlirModuleCreateEmpty(location))
Paul's avatar
Paul committed
124
    {
Paul's avatar
Paul committed
125
126
        MlirDialectHandle mixrHandle = mlirGetDialectHandle__migraphx__();
        mlirDialectHandleRegisterDialect(mixrHandle, ctx.get());
Paul's avatar
Paul committed
127
        mlirRegisterAllDialects(ctx.get());
Paul's avatar
Paul committed
128
        mlirContextSetAllowUnregisteredDialects(ctx.get(), true /*allow*/);
Paul's avatar
Paul committed
129
130
131
132
133
134
    }

    MlirType make_type(shape::type_t t) const
    {
        MlirType result;
        shape::visit(t, [&](auto as) {
Paul's avatar
Paul committed
135
            if(as.type_enum() == shape::float_type)
Paul's avatar
Paul committed
136
                result = mlirF32TypeGet(ctx.get());
Paul's avatar
Paul committed
137
            else if(as.type_enum() == shape::half_type)
Paul's avatar
Paul committed
138
                result = mlirF16TypeGet(ctx.get());
Paul's avatar
Paul committed
139
            else if(as.type_enum() == shape::double_type)
Paul's avatar
Paul committed
140
                result = mlirF64TypeGet(ctx.get());
Paul's avatar
Paul committed
141
            else if(as.is_integral())
Paul's avatar
Paul committed
142
            {
Paul's avatar
Paul committed
143
                if(as.is_signed())
Paul's avatar
Paul committed
144
145
146
147
148
149
150
151
152
153
154
155
156
157
                    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
158
159
        return mlirRankedTensorTypeGet(
            lens.size(), lens.data(), make_type(s.type()), mlirAttributeGetNull());
Paul's avatar
Paul committed
160
161
    }

Paul's avatar
Paul committed
162
    template <class Range>
Paul's avatar
Paul committed
163
164
165
166
167
168
169
170
171
172
173
    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
174
        auto in  = make_tensors(inputs);
Paul's avatar
Paul committed
175
176
177
178
        auto out = make_tensors(outputs);
        return mlirFunctionTypeGet(ctx.get(), in.size(), in.data(), out.size(), out.data());
    }

Paul's avatar
Paul committed
179
180
181
182
183
184
185
186
187
    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
188
189
    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
190
    MlirAttribute attribute(bool b) const { return mlirBoolAttrGet(ctx.get(), b ? 1 : 0); }
Paul's avatar
Paul committed
191
192
193
194
195
196
197
198
    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
199
200
    MlirAttribute attribute(std::nullptr_t) const { return {}; }
    template <class T>
Paul's avatar
Paul committed
201
202
203
204
205
206
207
208
209
210
211
212
    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
213
        v.visit_value([&](auto&& x) { attr = attribute(x); });
Paul's avatar
Paul committed
214
215
216
217
218
219
220
221
222
223
        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
224
            std::vector<MlirNamedAttribute> attributes = name_attributes(v);
Paul's avatar
Paul committed
225
226
227
228
229
230
231
232
233
234
235
236
237
            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
238
    MlirAttribute attribute(MlirType t) const { return mlirTypeAttrGet(t); }
Paul's avatar
Paul committed
239

Paul's avatar
Paul committed
240
241
    MlirAttribute attribute(MlirAttribute a) const { return a; }

Paul's avatar
Paul committed
242
    template <class T>
Paul's avatar
Paul committed
243
244
245
246
247
248
249
250
    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
251
252
253
254
255
256
257
258
259
    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
260
261
262
263
    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
264
265
        return name_attribute(na.first,
                              std::visit([&](const auto& x) { return attribute(x); }, na.second));
Paul's avatar
Paul committed
266
267
    }

Paul's avatar
Paul committed
268
269
    std::vector<MlirNamedAttribute>
    name_attributes(const std::vector<named_attribute_t>& named_attrs) const
Paul's avatar
Paul committed
270
271
272
    {
        std::vector<MlirNamedAttribute> attributes;
        attributes.reserve(named_attrs.size());
Paul's avatar
Paul committed
273
274
275
276
        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
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
        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
292
293
294
295
        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
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
336
337
338
339
340
341
342
343
344

        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
345
            for(auto& r : regions)
Paul's avatar
Paul committed
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
                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
365
        auto weak_op      = op.get();
Paul's avatar
Paul committed
366
        mlirBlockAppendOwnedOperation(body, op.release());
Paul's avatar
Paul committed
367
368
369
370
371
372
373
374
375

        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
376
377
    MlirBlock
    insert(MlirBlock body, const module& m, std::unordered_map<instruction_ref, MlirValue>& ins_map)
Paul's avatar
Paul committed
378
379
380
    {
        auto names = m.get_parameter_names();
        std::vector<shape> inputs;
Paul's avatar
Paul committed
381
382
383
384
        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
385
386
        std::vector<shape> outputs = m.get_output_shapes();

Paul's avatar
Paul committed
387
        auto body_inputs   = make_tensors(inputs);
Paul's avatar
Paul committed
388
        mlir_region region = mlirRegionCreate();
Paul's avatar
Paul committed
389
390
        mlir_block fbody   = mlirBlockCreate(body_inputs.size(), body_inputs.data());
        MlirBlock result   = fbody.get();
Paul's avatar
Paul committed
391
392
393
        mlirRegionAppendOwnedBlock(region.get(), fbody.release());

        auto ops = create_operation_state("builtin.func");
Paul's avatar
Paul committed
394
        ops.add_attributes(
Paul's avatar
Paul committed
395
            {{"type", make_function_type(inputs, outputs)}, {"sym_name", std::string("main")}});
Paul's avatar
Paul committed
396
397
398
        ops.add_region(std::move(region));
        insert(body, std::move(ops));

Paul's avatar
Paul committed
399
        for(auto i : range(names.size()))
Paul's avatar
Paul committed
400
401
402
403
            ins_map[m.get_parameter(names[i])] = mlirBlockGetArgument(result, i);
        return result;
    }

Paul's avatar
Paul committed
404
405
    static std::string get_name(instruction_ref ins)
    {
Paul's avatar
Format  
Paul committed
406
        if(ins->name() == "@return")
Paul's avatar
Paul committed
407
408
409
410
411
412
            return "std.return";
        return "migraphx." + ins->name();
    }

    static shape get_shape(instruction_ref ins)
    {
Paul's avatar
Format  
Paul committed
413
        if(ins->name() == "@return")
Paul's avatar
Paul committed
414
415
416
417
418
419
420
        {
            assert(ins->inputs().size() == 1);
            return ins->inputs().front()->get_shape();
        }
        return ins->get_shape();
    }

Paul's avatar
Paul committed
421
422
423
424
425
    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
426
        for(auto ins : iterator_for(m))
Paul's avatar
Paul committed
427
        {
Paul's avatar
Paul committed
428
            if(ins->name() == "@param")
Paul's avatar
Paul committed
429
                continue;
Paul's avatar
Paul committed
430
            auto name = get_name(ins);
Paul's avatar
Paul committed
431
            auto ops  = create_operation_state(name);
Paul's avatar
Paul committed
432
            ops.add_attribute_value(ins->get_operator().to_value());
Paul's avatar
Paul committed
433
            ops.add_results({get_shape(ins)});
Paul's avatar
Paul committed
434
435

            std::vector<MlirValue> inputs;
Paul's avatar
Paul committed
436
437
            transform(
                ins->inputs(), std::back_inserter(inputs), [&](auto i) { return ins_map.at(i); });
Paul's avatar
Paul committed
438
439
440
441
442
443
444
445
            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
446
447
448
449
450
451
    code_object_op compile()
    {
        mlir_pass_manager pm{mlirPassManagerCreate(ctx.get())};
        // 1st pipeline to call
        mlirMIGraphXAddHighLevelPipeline(pm.get());
        // 2nd pipeline to call
Paul's avatar
Format  
Paul committed
452
        const char* deviceName = "gfx908";
Paul's avatar
Paul committed
453
454
455
456
        mlirMIGraphXAddBackendPipeline(pm.get(), deviceName);
        mlirPassManagerRun(pm.get(), mmodule.get());

        code_object_op op;
Paul's avatar
Format  
Paul committed
457
        op.code_object                = get_binary();
Paul's avatar
Paul committed
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
        std::tie(op.global, op.local) = get_launch_params();
        return op;
    }

    std::pair<std::size_t, std::size_t> get_launch_params()
    {
        int attrs[2];
        // returns block and grid sizes
        mlirGetKernelAttrs(mmodule.get(), attrs);
        std::size_t local  = attrs[0];
        std::size_t global = local * attrs[1];
        return {global, local};
    }

    value::binary get_binary()
    {
        value::binary result(mlirGetBinarySize(mmodule.get()));
        if(mlirGetBinary(mmodule.get(), reinterpret_cast<char*>(result.data())))
            return result;
        MIGRAPHX_THROW("Failed to compile mlir program");
    }

Paul's avatar
Paul committed
480
    mlir_context ctx;
Paul's avatar
Paul committed
481
482
483
    MlirLocation location;
    mlir_module mmodule;
    std::deque<std::string> strings{};
Paul's avatar
Paul committed
484
485
};

Paul's avatar
Paul committed
486
487
488
489
490
491
492
493
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
494
495
496
497
498
499
500
code_object_op compile_mlir(const module& m)
{
    mlir_program mp;
    mp.parse(m);
    return mp.compile();
}

Paul's avatar
Paul committed
501
502
#else

Paul's avatar
Format  
Paul committed
503
std::string dump_mlir(const module&) { return {}; }
Paul's avatar
Paul committed
504
505
506

#endif

507
} // namespace gpu
Paul's avatar
Paul committed
508
509
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx