migraphx_py.cpp 22.6 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
#include <migraphx/float8.hpp>
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
#define MIGRAPHX_PYTHON_GENERATE_SHAPE_ENUM(x, t) .value(#x, migraphx::shape::type_t::x)
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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>());
    }
86
    else if(py::isinstance<py::int_>(x) or py::hasattr(x, "__index__"))
87
88
89
90
91
92
93
94
95
96
97
    {
        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>());
    }
98
99
100
101
    else if(py::isinstance<migraphx::shape::dynamic_dimension>(x))
    {
        f(migraphx::to_value(x.template cast<migraphx::shape::dynamic_dimension>()));
    }
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
128
129
130
131
132
    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
133
134
namespace pybind11 {
namespace detail {
Paul's avatar
Paul committed
135

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

147
148
149
150
151
152
153
154
155
156
157
158
template <>
struct npy_format_descriptor<migraphx::fp8::fp8e4m3fnuz>
{
    static std::string format()
    {
        // following: https://docs.python.org/3/library/struct.html#format-characters
        // TODO: need to figure out correct encoding
        return "z";
    }
    static constexpr auto name() { return _("fp8e4m3fnuz"); }
};

Shucai Xiao's avatar
Shucai Xiao committed
159
160
161
} // namespace detail
} // namespace pybind11

Paul's avatar
Paul committed
162
template <class F>
Paul's avatar
Paul committed
163
164
void visit_type(const migraphx::shape& s, F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
165
    s.visit_type(f);
Paul's avatar
Paul committed
166
167
}

Paul's avatar
Paul committed
168
169
170
template <class T, class F>
void visit(const migraphx::raw_data<T>& x, F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
171
    x.visit(f);
Paul's avatar
Paul committed
172
173
}

Paul's avatar
Paul committed
174
175
176
template <class F>
void visit_types(F f)
{
Shucai Xiao's avatar
Shucai Xiao committed
177
    migraphx::shape::visit_types(f);
Paul's avatar
Paul committed
178
179
}

Paul's avatar
Paul committed
180
template <class T>
Paul's avatar
Paul committed
181
182
183
py::buffer_info to_buffer_info(T& x)
{
    migraphx::shape s = x.get_shape();
184
185
186
187
    assert(s.type() != migraphx::shape::tuple_type);
    if(s.dynamic())
        MIGRAPHX_THROW("MIGRAPHX PYTHON: dynamic shape argument passed to to_buffer_info");
    auto strides = s.strides();
Paul's avatar
Paul committed
188
189
    std::transform(
        strides.begin(), strides.end(), strides.begin(), [&](auto i) { return i * s.type_size(); });
Paul's avatar
Paul committed
190
191
    py::buffer_info b;
    visit_type(s, [&](auto as) {
192
193
194
195
196
197
198
        // 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(),
199
                                s.ndim(),
200
201
202
203
204
205
206
207
                                s.lens(),
                                strides);
        }
        else
        {
            b = py::buffer_info(x.data(),
                                as.size(),
                                py::format_descriptor<decltype(as())>::format(),
208
                                s.ndim(),
209
210
211
                                s.lens(),
                                strides);
        }
Paul's avatar
Paul committed
212
213
214
215
    });
    return b;
}

Paul's avatar
Paul committed
216
217
218
migraphx::shape to_shape(const py::buffer_info& info)
{
    migraphx::shape::type_t t;
219
    std::size_t n = 0;
Paul's avatar
Paul committed
220
    visit_types([&](auto as) {
Shucai Xiao's avatar
Shucai Xiao committed
221
222
223
        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
224
        {
Paul's avatar
Paul committed
225
            t = as.type_enum();
226
227
            n = sizeof(as());
        }
Shucai Xiao's avatar
Shucai Xiao committed
228
229
230
231
232
        else if(info.format == "?" and py::format_descriptor<decltype(as())>::format() == "b")
        {
            t = migraphx::shape::bool_type;
            n = sizeof(bool);
        }
233
    });
234

Shucai Xiao's avatar
Shucai Xiao committed
235
    if(n == 0)
236
    {
Shucai Xiao's avatar
Shucai Xiao committed
237
        MIGRAPHX_THROW("MIGRAPHX PYTHON: Unsupported data type " + info.format);
238
239
    }

240
241
    auto strides = info.strides;
    std::transform(strides.begin(), strides.end(), strides.begin(), [&](auto i) -> std::size_t {
Paul's avatar
Paul committed
242
        return n > 0 ? i / n : 0;
Paul's avatar
Paul committed
243
    });
244
245
246
247
248
249
250
251
252
253

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

256
MIGRAPHX_PYBIND11_MODULE(migraphx, m)
Paul's avatar
Paul committed
257
{
258
259
    py::class_<migraphx::shape> shape_cls(m, "shape");
    shape_cls
260
        .def(py::init([](py::kwargs kwargs) {
261
262
263
264
265
266
267
268
269
            auto v = migraphx::to_value(kwargs);
            auto t = migraphx::shape::parse_type(v.get("type", "float"));
            if(v.contains("dyn_dims"))
            {
                auto dyn_dims =
                    migraphx::from_value<std::vector<migraphx::shape::dynamic_dimension>>(
                        v.at("dyn_dims"));
                return migraphx::shape(t, dyn_dims);
            }
270
271
272
            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>());
273
            else
274
                return migraphx::shape(t, lens);
275
        }))
Paul's avatar
Paul committed
276
277
278
        .def("type", &migraphx::shape::type)
        .def("lens", &migraphx::shape::lens)
        .def("strides", &migraphx::shape::strides)
279
        .def("ndim", &migraphx::shape::ndim)
Paul's avatar
Paul committed
280
281
        .def("elements", &migraphx::shape::elements)
        .def("bytes", &migraphx::shape::bytes)
282
        .def("type_string", &migraphx::shape::type_string)
Paul's avatar
Paul committed
283
        .def("type_size", &migraphx::shape::type_size)
284
        .def("dyn_dims", &migraphx::shape::dyn_dims)
Paul's avatar
Paul committed
285
286
287
288
        .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
289
        .def("scalar", &migraphx::shape::scalar)
290
        .def("dynamic", &migraphx::shape::dynamic)
Paul's avatar
Paul committed
291
292
        .def("__eq__", std::equal_to<migraphx::shape>{})
        .def("__ne__", std::not_equal_to<migraphx::shape>{})
Paul's avatar
Paul committed
293
        .def("__repr__", [](const migraphx::shape& s) { return migraphx::to_string(s); });
Paul's avatar
Paul committed
294

295
296
297
    py::enum_<migraphx::shape::type_t>(shape_cls, "type_t")
        MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_PYTHON_GENERATE_SHAPE_ENUM);

298
299
300
301
302
303
304
305
306
    py::class_<migraphx::shape::dynamic_dimension>(shape_cls, "dynamic_dimension")
        .def(py::init<>())
        .def(py::init<std::size_t, std::size_t>())
        .def(py::init<std::size_t, std::size_t, std::set<std::size_t>>())
        .def_readwrite("min", &migraphx::shape::dynamic_dimension::min)
        .def_readwrite("max", &migraphx::shape::dynamic_dimension::max)
        .def_readwrite("optimals", &migraphx::shape::dynamic_dimension::optimals)
        .def("is_fixed", &migraphx::shape::dynamic_dimension::is_fixed);

Paul's avatar
Paul committed
307
    py::class_<migraphx::argument>(m, "argument", py::buffer_protocol())
Paul's avatar
Paul committed
308
        .def_buffer([](migraphx::argument& x) -> py::buffer_info { return to_buffer_info(x); })
309
310
311
312
        .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
313
        .def("get_shape", &migraphx::argument::get_shape)
314
315
        .def("data_ptr",
             [](migraphx::argument& x) { return reinterpret_cast<std::uintptr_t>(x.data()); })
Paul's avatar
Paul committed
316
317
318
319
320
321
        .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
322
323
324
        .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
325

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

328
329
330
    py::class_<migraphx::instruction_ref>(m, "instruction_ref")
        .def("shape", [](migraphx::instruction_ref i) { return i->get_shape(); })
        .def("op", [](migraphx::instruction_ref i) { return i->get_operator(); });
331
332

    py::class_<migraphx::module, std::unique_ptr<migraphx::module, py::nodelete>>(m, "module")
Shucai Xiao's avatar
Shucai Xiao committed
333
        .def("print", [](const migraphx::module& mm) { std::cout << mm << std::endl; })
334
335
336
337
338
339
340
341
342
343
344
        .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*>{})
345
346
347
348
349
350
351
352
        .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"))
353
354
355
356
357
358
359
360
361
362
363
364
365
        .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
366
        .def("__repr__", [](const migraphx::module& mm) { return migraphx::to_string(mm); });
Shucai Xiao's avatar
Shucai Xiao committed
367

Paul's avatar
Paul committed
368
    py::class_<migraphx::program>(m, "program")
369
        .def(py::init([]() { return migraphx::program(); }))
370
        .def("get_parameter_names", &migraphx::program::get_parameter_names)
Paul's avatar
Paul committed
371
        .def("get_parameter_shapes", &migraphx::program::get_parameter_shapes)
372
        .def("get_output_shapes", &migraphx::program::get_output_shapes)
373
        .def("is_compiled", &migraphx::program::is_compiled)
kahmed10's avatar
kahmed10 committed
374
375
        .def(
            "compile",
376
377
378
379
380
            [](migraphx::program& p,
               const migraphx::target& t,
               bool offload_copy,
               bool fast_math,
               bool exhaustive_tune) {
kahmed10's avatar
kahmed10 committed
381
                migraphx::compile_options options;
382
383
384
                options.offload_copy    = offload_copy;
                options.fast_math       = fast_math;
                options.exhaustive_tune = exhaustive_tune;
kahmed10's avatar
kahmed10 committed
385
386
387
                p.compile(t, options);
            },
            py::arg("t"),
388
389
390
            py::arg("offload_copy")    = true,
            py::arg("fast_math")       = true,
            py::arg("exhaustive_tune") = false)
391
392
393
394
395
        .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"))
396
397
        .def("run",
             [](migraphx::program& p, py::dict params) {
398
                 migraphx::parameter_map pm;
399
400
401
402
403
404
405
406
407
                 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);
             })
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
        .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);
             })
425
        .def("sort", &migraphx::program::sort)
426
        .def("print", [](const migraphx::program& p) { std::cout << p << std::endl; })
Paul's avatar
Paul committed
427
428
        .def("__eq__", std::equal_to<migraphx::program>{})
        .def("__ne__", std::not_equal_to<migraphx::program>{})
Paul's avatar
Paul committed
429
        .def("__repr__", [](const migraphx::program& p) { return migraphx::to_string(p); });
Paul's avatar
Paul committed
430

431
432
433
434
435
436
437
438
439
    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);
      }))
440
441
        .def("name", &migraphx::operation::name);

442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
    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
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
    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,
481
           migraphx::shape::dynamic_dimension default_dyn_dim_value,
bpickrel's avatar
bpickrel committed
482
           std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims,
483
484
           std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>
               map_dyn_input_dims,
bpickrel's avatar
bpickrel committed
485
486
           bool skip_unknown_operators,
           bool print_program_on_error,
487
488
           int64_t max_loop_iterations,
           int64_t limit_max_iterations) {
bpickrel's avatar
bpickrel committed
489
490
            migraphx::onnx_options options;
            options.default_dim_value      = default_dim_value;
491
            options.default_dyn_dim_value  = default_dyn_dim_value;
bpickrel's avatar
bpickrel committed
492
            options.map_input_dims         = map_input_dims;
493
            options.map_dyn_input_dims     = map_dyn_input_dims;
bpickrel's avatar
bpickrel committed
494
495
496
            options.skip_unknown_operators = skip_unknown_operators;
            options.print_program_on_error = print_program_on_error;
            options.max_loop_iterations    = max_loop_iterations;
497
            options.limit_max_iterations   = limit_max_iterations;
bpickrel's avatar
bpickrel committed
498
499
500
501
            return migraphx::parse_onnx(filename, options);
        },
        "Parse onnx file",
        py::arg("filename"),
502
503
504
505
506
        py::arg("default_dim_value")     = 0,
        py::arg("default_dyn_dim_value") = migraphx::shape::dynamic_dimension{1, 1},
        py::arg("map_input_dims") = std::unordered_map<std::string, std::vector<std::size_t>>(),
        py::arg("map_dyn_input_dims") =
            std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>(),
bpickrel's avatar
bpickrel committed
507
508
        py::arg("skip_unknown_operators") = false,
        py::arg("print_program_on_error") = false,
509
510
        py::arg("max_loop_iterations")    = 10,
        py::arg("limit_max_iterations")   = std::numeric_limits<uint16_t>::max());
bpickrel's avatar
bpickrel committed
511
512
513
514
515

    m.def(
        "parse_onnx_buffer",
        [](const std::string& onnx_buffer,
           unsigned int default_dim_value,
516
           migraphx::shape::dynamic_dimension default_dyn_dim_value,
bpickrel's avatar
bpickrel committed
517
           std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims,
518
519
           std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>
               map_dyn_input_dims,
bpickrel's avatar
bpickrel committed
520
521
522
523
           bool skip_unknown_operators,
           bool print_program_on_error) {
            migraphx::onnx_options options;
            options.default_dim_value      = default_dim_value;
524
            options.default_dyn_dim_value  = default_dyn_dim_value;
bpickrel's avatar
bpickrel committed
525
            options.map_input_dims         = map_input_dims;
526
            options.map_dyn_input_dims     = map_dyn_input_dims;
bpickrel's avatar
bpickrel committed
527
528
529
530
531
532
            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"),
533
534
535
536
537
        py::arg("default_dim_value")     = 0,
        py::arg("default_dyn_dim_value") = migraphx::shape::dynamic_dimension{1, 1},
        py::arg("map_input_dims") = std::unordered_map<std::string, std::vector<std::size_t>>(),
        py::arg("map_dyn_input_dims") =
            std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>(),
bpickrel's avatar
bpickrel committed
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
        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
563

564
    m.def("get_target", &migraphx::make_target);
565
566
567
568
569
570
571
    m.def("create_argument", [](const migraphx::shape& s, const std::vector<double>& values) {
        if(values.size() != s.elements())
            MIGRAPHX_THROW("Values and shape elements do not match");
        migraphx::argument a{s};
        a.fill(values.begin(), values.end());
        return a;
    });
Paul's avatar
Paul committed
572
    m.def("generate_argument", &migraphx::generate_argument, py::arg("s"), py::arg("seed") = 0);
573
    m.def("fill_argument", &migraphx::fill_argument, py::arg("s"), py::arg("value"));
Shucai Xiao's avatar
Shucai Xiao committed
574
575
576
577
578
579
580
581
    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"),
582
          py::arg("calibration") = std::vector<migraphx::parameter_map>{},
Shucai Xiao's avatar
Shucai Xiao committed
583
          py::arg("ins_names")   = std::vector<std::string>{"dot", "convolution"});
Shucai Xiao's avatar
Shucai Xiao committed
584

Paul's avatar
Paul committed
585
586
587
588
#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);
589
    m.def("gpu_sync", [] { migraphx::gpu::gpu_sync(); });
Paul's avatar
Paul committed
590
591
#endif

Paul's avatar
Paul committed
592
593
594
595
596
597
#ifdef VERSION_INFO
    m.attr("__version__") = VERSION_INFO;
#else
    m.attr("__version__") = "dev";
#endif
}