migraphx_py.cpp 15.6 KB
Newer Older
Paul's avatar
Paul committed
1
2
3

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
Shucai Xiao's avatar
Shucai Xiao committed
4
#include <pybind11/numpy.h>
Paul's avatar
Paul committed
5
#include <migraphx/program.hpp>
6
7
#include <migraphx/instruction_ref.hpp>
#include <migraphx/operation.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
8
#include <migraphx/quantization.hpp>
Paul's avatar
Paul committed
9
#include <migraphx/generate.hpp>
10
#include <migraphx/ref/target.hpp>
Paul's avatar
Paul committed
11
#include <migraphx/stringutils.hpp>
12
13
#include <migraphx/tf.hpp>
#include <migraphx/onnx.hpp>
14
#include <migraphx/type_name.hpp>
15
16
#include <migraphx/load_save.hpp>
#include <migraphx/register_target.hpp>
17
18
#include <migraphx/json.hpp>
#include <migraphx/make_op.hpp>
19

Paul's avatar
Paul committed
20
21
22
#ifdef HAVE_GPU
#include <migraphx/gpu/hip.hpp>
#endif
Paul's avatar
Paul committed
23

Shucai Xiao's avatar
Shucai Xiao committed
24
using half   = half_float::half;
Paul's avatar
Paul committed
25
26
namespace py = pybind11;

27
28
29
30
31
32
33
34
35
36
37
38
39
40
#ifdef __clang__
#define MIGRAPHX_PUSH_UNUSED_WARNING \
    _Pragma("clang diagnostic push") \
        _Pragma("clang diagnostic ignored \"-Wused-but-marked-unused\"")
#define MIGRAPHX_POP_WARNING _Pragma("clang diagnostic pop")
#else
#define MIGRAPHX_PUSH_UNUSED_WARNING
#define MIGRAPHX_POP_WARNING
#endif
#define MIGRAPHX_PYBIND11_MODULE(...) \
    MIGRAPHX_PUSH_UNUSED_WARNING      \
    PYBIND11_MODULE(__VA_ARGS__)      \
    MIGRAPHX_POP_WARNING

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
namespace migraphx {

migraphx::value to_value(py::kwargs kwargs);
migraphx::value to_value(py::list lst);

template <class T, class F>
void visit_py(T x, F f)
{
    if(py::isinstance<py::kwargs>(x))
    {
        f(to_value(x.template cast<py::kwargs>()));
    }
    else if(py::isinstance<py::list>(x))
    {
        f(to_value(x.template cast<py::list>()));
    }
    else if(py::isinstance<py::bool_>(x))
    {
        f(x.template cast<bool>());
    }
    else if(py::isinstance<py::int_>(x))
    {
        f(x.template cast<int>());
    }
    else if(py::isinstance<py::float_>(x))
    {
        f(x.template cast<float>());
    }
    else if(py::isinstance<py::str>(x))
    {
        f(x.template cast<std::string>());
    }
    else
    {
        MIGRAPHX_THROW("VISIT_PY: Unsupported data type!");
    }
}

migraphx::value to_value(py::list lst)
{
    migraphx::value v = migraphx::value::array{};
    for(auto val : lst)
    {
        visit_py(val, [&](auto py_val) { v.push_back(py_val); });
    }

    return v;
}

migraphx::value to_value(py::kwargs kwargs)
{
    migraphx::value v = migraphx::value::object{};

    for(auto arg : kwargs)
    {
        auto&& key = py::str(arg.first);
        auto&& val = arg.second;
        visit_py(val, [&](auto py_val) { v[key] = py_val; });
    }
    return v;
}
} // namespace migraphx

Shucai Xiao's avatar
Shucai Xiao committed
104
105
namespace pybind11 {
namespace detail {
Paul's avatar
Paul committed
106

Shucai Xiao's avatar
Shucai Xiao committed
107
108
template <>
struct npy_format_descriptor<half>
Paul's avatar
Paul committed
109
{
Shucai Xiao's avatar
Shucai Xiao committed
110
    static std::string format()
Paul's avatar
Paul committed
111
    {
Shucai Xiao's avatar
Shucai Xiao committed
112
113
        // following: https://docs.python.org/3/library/struct.html#format-characters
        return "e";
Paul's avatar
Paul committed
114
    }
Shucai Xiao's avatar
Shucai Xiao committed
115
    static constexpr auto name() { return _("half"); }
Paul's avatar
Paul committed
116
117
};

Shucai Xiao's avatar
Shucai Xiao committed
118
119
120
} // namespace detail
} // namespace pybind11

Paul's avatar
Paul committed
121
template <class F>
Paul's avatar
Paul committed
122
123
void visit_type(const migraphx::shape& s, F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
124
    s.visit_type(f);
Paul's avatar
Paul committed
125
126
}

Paul's avatar
Paul committed
127
128
129
template <class T, class F>
void visit(const migraphx::raw_data<T>& x, F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
130
    x.visit(f);
Paul's avatar
Paul committed
131
132
}

Paul's avatar
Paul committed
133
134
135
template <class F>
void visit_types(F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
136
    migraphx::shape::visit_types(f);
Paul's avatar
Paul committed
137
138
}

Paul's avatar
Paul committed
139
template <class T>
Paul's avatar
Paul committed
140
141
142
py::buffer_info to_buffer_info(T& x)
{
    migraphx::shape s = x.get_shape();
Paul's avatar
Paul committed
143
144
145
    auto strides      = s.strides();
    std::transform(
        strides.begin(), strides.end(), strides.begin(), [&](auto i) { return i * s.type_size(); });
Paul's avatar
Paul committed
146
147
    py::buffer_info b;
    visit_type(s, [&](auto as) {
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
        // migraphx use int8_t data to store bool type, we need to
        // explicitly specify the data type as bool for python
        if(s.type() == migraphx::shape::bool_type)
        {
            b = py::buffer_info(x.data(),
                                as.size(),
                                py::format_descriptor<bool>::format(),
                                s.lens().size(),
                                s.lens(),
                                strides);
        }
        else
        {
            b = py::buffer_info(x.data(),
                                as.size(),
                                py::format_descriptor<decltype(as())>::format(),
                                s.lens().size(),
                                s.lens(),
                                strides);
        }
Paul's avatar
Paul committed
168
169
170
171
    });
    return b;
}

Paul's avatar
Paul committed
172
173
174
migraphx::shape to_shape(const py::buffer_info& info)
{
    migraphx::shape::type_t t;
175
    std::size_t n = 0;
Paul's avatar
Paul committed
176
    visit_types([&](auto as) {
Shucai Xiao's avatar
Shucai Xiao committed
177
178
179
        if(info.format == py::format_descriptor<decltype(as())>::format() or
           (info.format == "l" and py::format_descriptor<decltype(as())>::format() == "q") or
           (info.format == "L" and py::format_descriptor<decltype(as())>::format() == "Q"))
Paul's avatar
Paul committed
180
        {
Paul's avatar
Paul committed
181
            t = as.type_enum();
182
183
            n = sizeof(as());
        }
Shucai Xiao's avatar
Shucai Xiao committed
184
185
186
187
188
        else if(info.format == "?" and py::format_descriptor<decltype(as())>::format() == "b")
        {
            t = migraphx::shape::bool_type;
            n = sizeof(bool);
        }
189
    });
190

Shucai Xiao's avatar
Shucai Xiao committed
191
    if(n == 0)
192
    {
Shucai Xiao's avatar
Shucai Xiao committed
193
        MIGRAPHX_THROW("MIGRAPHX PYTHON: Unsupported data type " + info.format);
194
195
    }

196
197
    auto strides = info.strides;
    std::transform(strides.begin(), strides.end(), strides.begin(), [&](auto i) -> std::size_t {
Paul's avatar
Paul committed
198
        return n > 0 ? i / n : 0;
Paul's avatar
Paul committed
199
    });
200
201
202
203
204
205
206
207
208
209

    // scalar support
    if(info.shape.empty())
    {
        return migraphx::shape{t};
    }
    else
    {
        return migraphx::shape{t, info.shape, strides};
    }
Paul's avatar
Paul committed
210
211
}

212
MIGRAPHX_PYBIND11_MODULE(migraphx, m)
Paul's avatar
Paul committed
213
{
Paul's avatar
Paul committed
214
    py::class_<migraphx::shape>(m, "shape")
215
216
        .def(py::init([](py::kwargs kwargs) {
            auto v    = migraphx::to_value(kwargs);
217
            auto t    = migraphx::shape::parse_type(v.get("type", "float"));
218
219
220
221
222
223
            auto lens = v.get<std::size_t>("lens", {1});
            if(v.contains("strides"))
                return migraphx::shape(t, lens, v.at("strides").to_vector<std::size_t>());
            else
                return migraphx::shape(t, lens);
        }))
Paul's avatar
Paul committed
224
225
226
227
228
        .def("type", &migraphx::shape::type)
        .def("lens", &migraphx::shape::lens)
        .def("strides", &migraphx::shape::strides)
        .def("elements", &migraphx::shape::elements)
        .def("bytes", &migraphx::shape::bytes)
229
        .def("type_string", &migraphx::shape::type_string)
Paul's avatar
Paul committed
230
231
232
233
234
        .def("type_size", &migraphx::shape::type_size)
        .def("packed", &migraphx::shape::packed)
        .def("transposed", &migraphx::shape::transposed)
        .def("broadcasted", &migraphx::shape::broadcasted)
        .def("standard", &migraphx::shape::standard)
Paul's avatar
Paul committed
235
        .def("scalar", &migraphx::shape::scalar)
Paul's avatar
Paul committed
236
237
        .def("__eq__", std::equal_to<migraphx::shape>{})
        .def("__ne__", std::not_equal_to<migraphx::shape>{})
Paul's avatar
Paul committed
238
        .def("__repr__", [](const migraphx::shape& s) { return migraphx::to_string(s); });
Paul's avatar
Paul committed
239
240

    py::class_<migraphx::argument>(m, "argument", py::buffer_protocol())
Paul's avatar
Paul committed
241
        .def_buffer([](migraphx::argument& x) -> py::buffer_info { return to_buffer_info(x); })
Paul's avatar
Paul committed
242
243
244
245
246
        .def("__init__",
             [](migraphx::argument& x, py::buffer b) {
                 py::buffer_info info = b.request();
                 new(&x) migraphx::argument(to_shape(info), info.ptr);
             })
Paul's avatar
Paul committed
247
        .def("get_shape", &migraphx::argument::get_shape)
Paul's avatar
Paul committed
248
249
250
251
252
253
        .def("tolist",
             [](migraphx::argument& x) {
                 py::list l{x.get_shape().elements()};
                 visit(x, [&](auto data) { l = py::cast(data.to_vector()); });
                 return l;
             })
Paul's avatar
Paul committed
254
255
256
        .def("__eq__", std::equal_to<migraphx::argument>{})
        .def("__ne__", std::not_equal_to<migraphx::argument>{})
        .def("__repr__", [](const migraphx::argument& x) { return migraphx::to_string(x); });
Paul's avatar
Paul committed
257

Paul's avatar
Paul committed
258
259
    py::class_<migraphx::target>(m, "target");

260
261
262
    py::class_<migraphx::instruction_ref>(m, "instruction_ref");

    py::class_<migraphx::module, std::unique_ptr<migraphx::module, py::nodelete>>(m, "module")
Shucai Xiao's avatar
Shucai Xiao committed
263
        .def("print", [](const migraphx::module& mm) { std::cout << mm << std::endl; })
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
        .def(
            "add_instruction",
            [](migraphx::module& mm,
               const migraphx::operation& op,
               std::vector<migraphx::instruction_ref>& args,
               std::vector<migraphx::module*>& mod_args) {
                return mm.add_instruction(op, args, mod_args);
            },
            py::arg("op"),
            py::arg("args"),
            py::arg("mod_args") = std::vector<migraphx::module*>{})
        .def(
            "add_parameter",
            [](migraphx::module& mm, const std::string& name, const migraphx::shape shape) {
                return mm.add_parameter(name, shape);
            },
            py::arg("name"),
            py::arg("shape"))
        .def(
            "add_return",
            [](migraphx::module& mm, std::vector<migraphx::instruction_ref>& args) {
                return mm.add_return(args);
            },
            py::arg("args"))
Shucai Xiao's avatar
Shucai Xiao committed
288
        .def("__repr__", [](const migraphx::module& mm) { return migraphx::to_string(mm); });
Shucai Xiao's avatar
Shucai Xiao committed
289

Paul's avatar
Paul committed
290
    py::class_<migraphx::program>(m, "program")
291
        .def(py::init([]() { return migraphx::program(); }))
292
        .def("get_parameter_names", &migraphx::program::get_parameter_names)
Paul's avatar
Paul committed
293
        .def("get_parameter_shapes", &migraphx::program::get_parameter_shapes)
294
        .def("get_output_shapes", &migraphx::program::get_output_shapes)
kahmed10's avatar
kahmed10 committed
295
296
297
298
299
300
301
302
303
304
305
        .def(
            "compile",
            [](migraphx::program& p, const migraphx::target& t, bool offload_copy, bool fast_math) {
                migraphx::compile_options options;
                options.offload_copy = offload_copy;
                options.fast_math    = fast_math;
                p.compile(t, options);
            },
            py::arg("t"),
            py::arg("offload_copy") = true,
            py::arg("fast_math")    = true)
306
307
308
309
310
        .def("get_main_module", [](const migraphx::program& p) { return p.get_main_module(); })
        .def(
            "create_module",
            [](migraphx::program& p, const std::string& name) { return p.create_module(name); },
            py::arg("name"))
311
312
        .def("run",
             [](migraphx::program& p, py::dict params) {
313
                 migraphx::parameter_map pm;
314
315
316
317
318
319
320
321
322
                 for(auto x : params)
                 {
                     std::string key      = x.first.cast<std::string>();
                     py::buffer b         = x.second.cast<py::buffer>();
                     py::buffer_info info = b.request();
                     pm[key]              = migraphx::argument(to_shape(info), info.ptr);
                 }
                 return p.eval(pm);
             })
323
        .def("sort", &migraphx::program::sort)
324
        .def("print", [](const migraphx::program& p) { std::cout << p << std::endl; })
Paul's avatar
Paul committed
325
326
        .def("__eq__", std::equal_to<migraphx::program>{})
        .def("__ne__", std::not_equal_to<migraphx::program>{})
Paul's avatar
Paul committed
327
        .def("__repr__", [](const migraphx::program& p) { return migraphx::to_string(p); });
Paul's avatar
Paul committed
328

329
330
331
332
333
334
335
336
337
338
339
340
    py::class_<migraphx::operation>(m, "op")
        .def(py::init([](const std::string& name, py::kwargs kwargs) {
            migraphx::value v = migraphx::value::object{};
            if(kwargs)
            {
                v = migraphx::to_value(kwargs);
            }
            return migraphx::make_op(name, v);
        }))

        .def("name", &migraphx::operation::name);

bpickrel's avatar
bpickrel committed
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
    m.def(
        "parse_tf",
        [](const std::string& filename,
           bool is_nhwc,
           unsigned int batch_size,
           std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims,
           std::vector<std::string> output_names) {
            return migraphx::parse_tf(
                filename, migraphx::tf_options{is_nhwc, batch_size, map_input_dims, output_names});
        },
        "Parse tf protobuf (default format is nhwc)",
        py::arg("filename"),
        py::arg("is_nhwc")        = true,
        py::arg("batch_size")     = 1,
        py::arg("map_input_dims") = std::unordered_map<std::string, std::vector<std::size_t>>(),
        py::arg("output_names")   = std::vector<std::string>());

    m.def(
        "parse_onnx",
        [](const std::string& filename,
           unsigned int default_dim_value,
           std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims,
           bool skip_unknown_operators,
           bool print_program_on_error,
           int64_t max_loop_iterations) {
            migraphx::onnx_options options;
            options.default_dim_value      = default_dim_value;
            options.map_input_dims         = map_input_dims;
            options.skip_unknown_operators = skip_unknown_operators;
            options.print_program_on_error = print_program_on_error;
            options.max_loop_iterations    = max_loop_iterations;
            return migraphx::parse_onnx(filename, options);
        },
        "Parse onnx file",
        py::arg("filename"),
        py::arg("default_dim_value") = 1,
        py::arg("map_input_dims")    = std::unordered_map<std::string, std::vector<std::size_t>>(),
        py::arg("skip_unknown_operators") = false,
        py::arg("print_program_on_error") = false,
        py::arg("max_loop_iterations")    = 10);

    m.def(
        "parse_onnx_buffer",
        [](const std::string& onnx_buffer,
           unsigned int default_dim_value,
           std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims,
           bool skip_unknown_operators,
           bool print_program_on_error) {
            migraphx::onnx_options options;
            options.default_dim_value      = default_dim_value;
            options.map_input_dims         = map_input_dims;
            options.skip_unknown_operators = skip_unknown_operators;
            options.print_program_on_error = print_program_on_error;
            return migraphx::parse_onnx_buffer(onnx_buffer, options);
        },
        "Parse onnx file",
        py::arg("filename"),
        py::arg("default_dim_value") = 1,
        py::arg("map_input_dims")    = std::unordered_map<std::string, std::vector<std::size_t>>(),
        py::arg("skip_unknown_operators") = false,
        py::arg("print_program_on_error") = false);

    m.def(
        "load",
        [](const std::string& name, const std::string& format) {
            migraphx::file_options options;
            options.format = format;
            return migraphx::load(name, options);
        },
        "Load MIGraphX program",
        py::arg("filename"),
        py::arg("format") = "msgpack");

    m.def(
        "save",
        [](const migraphx::program& p, const std::string& name, const std::string& format) {
            migraphx::file_options options;
            options.format = format;
            return migraphx::save(p, name, options);
        },
        "Save MIGraphX program",
        py::arg("p"),
        py::arg("filename"),
        py::arg("format") = "msgpack");
Paul's avatar
Paul committed
425

426
    m.def("get_target", &migraphx::make_target);
Paul's avatar
Paul committed
427
    m.def("generate_argument", &migraphx::generate_argument, py::arg("s"), py::arg("seed") = 0);
428
    m.def("fill_argument", &migraphx::fill_argument, py::arg("s"), py::arg("value"));
Shucai Xiao's avatar
Shucai Xiao committed
429
430
431
432
433
434
435
436
    m.def("quantize_fp16",
          &migraphx::quantize_fp16,
          py::arg("prog"),
          py::arg("ins_names") = std::vector<std::string>{"all"});
    m.def("quantize_int8",
          &migraphx::quantize_int8,
          py::arg("prog"),
          py::arg("t"),
437
          py::arg("calibration") = std::vector<migraphx::parameter_map>{},
Shucai Xiao's avatar
Shucai Xiao committed
438
          py::arg("ins_names")   = std::vector<std::string>{"dot", "convolution"});
Shucai Xiao's avatar
Shucai Xiao committed
439

Paul's avatar
Paul committed
440
441
442
443
#ifdef HAVE_GPU
    m.def("allocate_gpu", &migraphx::gpu::allocate_gpu, py::arg("s"), py::arg("host") = false);
    m.def("to_gpu", &migraphx::gpu::to_gpu, py::arg("arg"), py::arg("host") = false);
    m.def("from_gpu", &migraphx::gpu::from_gpu);
444
    m.def("gpu_sync", [] { migraphx::gpu::gpu_sync(); });
Paul's avatar
Paul committed
445
446
#endif

Paul's avatar
Paul committed
447
448
449
450
451
452
#ifdef VERSION_INFO
    m.attr("__version__") = VERSION_INFO;
#else
    m.attr("__version__") = "dev";
#endif
}