migraphx_py.cpp 19 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
Paul's avatar
Paul committed
24
25
26

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
Shucai Xiao's avatar
Shucai Xiao committed
27
#include <pybind11/numpy.h>
Paul's avatar
Paul committed
28
#include <migraphx/program.hpp>
29
30
#include <migraphx/instruction_ref.hpp>
#include <migraphx/operation.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
31
#include <migraphx/quantization.hpp>
Paul's avatar
Paul committed
32
#include <migraphx/generate.hpp>
33
#include <migraphx/instruction.hpp>
34
#include <migraphx/ref/target.hpp>
Paul's avatar
Paul committed
35
#include <migraphx/stringutils.hpp>
36
37
#include <migraphx/tf.hpp>
#include <migraphx/onnx.hpp>
38
39
#include <migraphx/load_save.hpp>
#include <migraphx/register_target.hpp>
40
41
#include <migraphx/json.hpp>
#include <migraphx/make_op.hpp>
42
#include <migraphx/op/common.hpp>
43

Paul's avatar
Paul committed
44
45
46
#ifdef HAVE_GPU
#include <migraphx/gpu/hip.hpp>
#endif
Paul's avatar
Paul committed
47

Shucai Xiao's avatar
Shucai Xiao committed
48
using half   = half_float::half;
Paul's avatar
Paul committed
49
50
namespace py = pybind11;

51
52
53
54
55
56
57
58
59
60
61
62
63
64
#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

65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
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>());
    }
85
    else if(py::isinstance<py::int_>(x) or py::hasattr(x, "__index__"))
86
87
88
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
122
123
124
125
126
127
    {
        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
128
129
namespace pybind11 {
namespace detail {
Paul's avatar
Paul committed
130

Shucai Xiao's avatar
Shucai Xiao committed
131
132
template <>
struct npy_format_descriptor<half>
Paul's avatar
Paul committed
133
{
Shucai Xiao's avatar
Shucai Xiao committed
134
    static std::string format()
Paul's avatar
Paul committed
135
    {
Shucai Xiao's avatar
Shucai Xiao committed
136
137
        // following: https://docs.python.org/3/library/struct.html#format-characters
        return "e";
Paul's avatar
Paul committed
138
    }
Shucai Xiao's avatar
Shucai Xiao committed
139
    static constexpr auto name() { return _("half"); }
Paul's avatar
Paul committed
140
141
};

Shucai Xiao's avatar
Shucai Xiao committed
142
143
144
} // namespace detail
} // namespace pybind11

Paul's avatar
Paul committed
145
template <class F>
Paul's avatar
Paul committed
146
147
void visit_type(const migraphx::shape& s, F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
148
    s.visit_type(f);
Paul's avatar
Paul committed
149
150
}

Paul's avatar
Paul committed
151
152
153
template <class T, class F>
void visit(const migraphx::raw_data<T>& x, F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
154
    x.visit(f);
Paul's avatar
Paul committed
155
156
}

Paul's avatar
Paul committed
157
158
159
template <class F>
void visit_types(F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
160
    migraphx::shape::visit_types(f);
Paul's avatar
Paul committed
161
162
}

Paul's avatar
Paul committed
163
template <class T>
Paul's avatar
Paul committed
164
165
166
py::buffer_info to_buffer_info(T& x)
{
    migraphx::shape s = x.get_shape();
Paul's avatar
Paul committed
167
168
169
    auto strides      = s.strides();
    std::transform(
        strides.begin(), strides.end(), strides.begin(), [&](auto i) { return i * s.type_size(); });
Paul's avatar
Paul committed
170
171
    py::buffer_info b;
    visit_type(s, [&](auto as) {
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
        // 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
192
193
194
195
    });
    return b;
}

Paul's avatar
Paul committed
196
197
198
migraphx::shape to_shape(const py::buffer_info& info)
{
    migraphx::shape::type_t t;
199
    std::size_t n = 0;
Paul's avatar
Paul committed
200
    visit_types([&](auto as) {
Shucai Xiao's avatar
Shucai Xiao committed
201
202
203
        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
204
        {
Paul's avatar
Paul committed
205
            t = as.type_enum();
206
207
            n = sizeof(as());
        }
Shucai Xiao's avatar
Shucai Xiao committed
208
209
210
211
212
        else if(info.format == "?" and py::format_descriptor<decltype(as())>::format() == "b")
        {
            t = migraphx::shape::bool_type;
            n = sizeof(bool);
        }
213
    });
214

Shucai Xiao's avatar
Shucai Xiao committed
215
    if(n == 0)
216
    {
Shucai Xiao's avatar
Shucai Xiao committed
217
        MIGRAPHX_THROW("MIGRAPHX PYTHON: Unsupported data type " + info.format);
218
219
    }

220
221
    auto strides = info.strides;
    std::transform(strides.begin(), strides.end(), strides.begin(), [&](auto i) -> std::size_t {
Paul's avatar
Paul committed
222
        return n > 0 ? i / n : 0;
Paul's avatar
Paul committed
223
    });
224
225
226
227
228
229
230
231
232
233

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

236
MIGRAPHX_PYBIND11_MODULE(migraphx, m)
Paul's avatar
Paul committed
237
{
Paul's avatar
Paul committed
238
    py::class_<migraphx::shape>(m, "shape")
239
240
        .def(py::init([](py::kwargs kwargs) {
            auto v    = migraphx::to_value(kwargs);
241
            auto t    = migraphx::shape::parse_type(v.get("type", "float"));
242
243
244
245
246
247
            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
248
249
250
251
252
        .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)
253
        .def("type_string", &migraphx::shape::type_string)
Paul's avatar
Paul committed
254
255
256
257
258
        .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
259
        .def("scalar", &migraphx::shape::scalar)
Paul's avatar
Paul committed
260
261
        .def("__eq__", std::equal_to<migraphx::shape>{})
        .def("__ne__", std::not_equal_to<migraphx::shape>{})
Paul's avatar
Paul committed
262
        .def("__repr__", [](const migraphx::shape& s) { return migraphx::to_string(s); });
Paul's avatar
Paul committed
263
264

    py::class_<migraphx::argument>(m, "argument", py::buffer_protocol())
Paul's avatar
Paul committed
265
        .def_buffer([](migraphx::argument& x) -> py::buffer_info { return to_buffer_info(x); })
266
267
268
269
        .def(py::init([](py::buffer b) {
            py::buffer_info info = b.request();
            return migraphx::argument(to_shape(info), info.ptr);
        }))
Paul's avatar
Paul committed
270
        .def("get_shape", &migraphx::argument::get_shape)
271
272
        .def("data_ptr",
             [](migraphx::argument& x) { return reinterpret_cast<std::uintptr_t>(x.data()); })
Paul's avatar
Paul committed
273
274
275
276
277
278
        .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
279
280
281
        .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
282

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

285
286
287
    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
288
        .def("print", [](const migraphx::module& mm) { std::cout << mm << std::endl; })
289
290
291
292
293
294
295
296
297
298
299
        .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*>{})
300
301
302
303
304
305
306
307
        .def(
            "add_literal",
            [](migraphx::module& mm, py::buffer data) {
                py::buffer_info info = data.request();
                auto literal_shape   = to_shape(info);
                return mm.add_literal(literal_shape, reinterpret_cast<char*>(info.ptr));
            },
            py::arg("data"))
308
309
310
311
312
313
314
315
316
317
318
319
320
        .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
321
        .def("__repr__", [](const migraphx::module& mm) { return migraphx::to_string(mm); });
Shucai Xiao's avatar
Shucai Xiao committed
322

Paul's avatar
Paul committed
323
    py::class_<migraphx::program>(m, "program")
324
        .def(py::init([]() { return migraphx::program(); }))
325
        .def("get_parameter_names", &migraphx::program::get_parameter_names)
Paul's avatar
Paul committed
326
        .def("get_parameter_shapes", &migraphx::program::get_parameter_shapes)
327
        .def("get_output_shapes", &migraphx::program::get_output_shapes)
328
        .def("is_compiled", &migraphx::program::is_compiled)
kahmed10's avatar
kahmed10 committed
329
330
        .def(
            "compile",
331
332
333
334
335
            [](migraphx::program& p,
               const migraphx::target& t,
               bool offload_copy,
               bool fast_math,
               bool exhaustive_tune) {
kahmed10's avatar
kahmed10 committed
336
                migraphx::compile_options options;
337
338
339
                options.offload_copy    = offload_copy;
                options.fast_math       = fast_math;
                options.exhaustive_tune = exhaustive_tune;
kahmed10's avatar
kahmed10 committed
340
341
342
                p.compile(t, options);
            },
            py::arg("t"),
343
344
345
            py::arg("offload_copy")    = true,
            py::arg("fast_math")       = true,
            py::arg("exhaustive_tune") = false)
346
347
348
349
350
        .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"))
351
352
        .def("run",
             [](migraphx::program& p, py::dict params) {
353
                 migraphx::parameter_map pm;
354
355
356
357
358
359
360
361
362
                 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);
             })
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
        .def("run_async",
             [](migraphx::program& p,
                py::dict params,
                std::uintptr_t stream,
                std::string stream_name) {
                 migraphx::parameter_map pm;
                 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);
                 }
                 migraphx::execution_environment exec_env{
                     migraphx::any_ptr(reinterpret_cast<void*>(stream), stream_name), true};
                 return p.eval(pm, exec_env);
             })
380
        .def("sort", &migraphx::program::sort)
381
        .def("print", [](const migraphx::program& p) { std::cout << p << std::endl; })
Paul's avatar
Paul committed
382
383
        .def("__eq__", std::equal_to<migraphx::program>{})
        .def("__ne__", std::not_equal_to<migraphx::program>{})
Paul's avatar
Paul committed
384
        .def("__repr__", [](const migraphx::program& p) { return migraphx::to_string(p); });
Paul's avatar
Paul committed
385

386
387
388
389
390
391
392
393
394
    py::class_<migraphx::operation> op(m, "op");
    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);
      }))
395
396
        .def("name", &migraphx::operation::name);

397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
    py::enum_<migraphx::op::pooling_mode>(op, "pooling_mode")
        .value("average", migraphx::op::pooling_mode::average)
        .value("max", migraphx::op::pooling_mode::max)
        .value("lpnorm", migraphx::op::pooling_mode::lpnorm);

    py::enum_<migraphx::op::rnn_direction>(op, "rnn_direction")
        .value("forward", migraphx::op::rnn_direction::forward)
        .value("reverse", migraphx::op::rnn_direction::reverse)
        .value("bidirectional", migraphx::op::rnn_direction::bidirectional);

    m.def(
        "argument_from_pointer",
        [](const migraphx::shape shape, const int64_t address) {
            return migraphx::argument(shape, reinterpret_cast<void*>(address));
        },
        py::arg("shape"),
        py::arg("address"));

bpickrel's avatar
bpickrel committed
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
    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
499

500
    m.def("get_target", &migraphx::make_target);
Paul's avatar
Paul committed
501
    m.def("generate_argument", &migraphx::generate_argument, py::arg("s"), py::arg("seed") = 0);
502
    m.def("fill_argument", &migraphx::fill_argument, py::arg("s"), py::arg("value"));
Shucai Xiao's avatar
Shucai Xiao committed
503
504
505
506
507
508
509
510
    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"),
511
          py::arg("calibration") = std::vector<migraphx::parameter_map>{},
Shucai Xiao's avatar
Shucai Xiao committed
512
          py::arg("ins_names")   = std::vector<std::string>{"dot", "convolution"});
Shucai Xiao's avatar
Shucai Xiao committed
513

Paul's avatar
Paul committed
514
515
516
517
#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);
518
    m.def("gpu_sync", [] { migraphx::gpu::gpu_sync(); });
Paul's avatar
Paul committed
519
520
#endif

Paul's avatar
Paul committed
521
522
523
524
525
526
#ifdef VERSION_INFO
    m.attr("__version__") = VERSION_INFO;
#else
    m.attr("__version__") = "dev";
#endif
}