migraphx_py.cpp 2.21 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
5
6
7
8

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <migraphx/program.hpp>
#include <migraphx/onnx.hpp>

namespace py = pybind11;

Paul's avatar
Paul committed
9
template <class F>
Paul's avatar
Paul committed
10
11
12
13
struct skip_half
{
    F f;

Paul's avatar
Paul committed
14
    template <class A>
Paul's avatar
Paul committed
15
16
17
18
19
20
21
22
    void operator()(A a) const
    {
        f(a);
    }

    void operator()(migraphx::shape::as<migraphx::half>) const
    {
        throw std::runtime_error("Half not supported in python yet.");
Paul's avatar
Paul committed
23
    }
Paul's avatar
Paul committed
24
25
};

Paul's avatar
Paul committed
26
template <class F>
Paul's avatar
Paul committed
27
28
29
30
31
void visit_type(const migraphx::shape& s, F f)
{
    s.visit_type(skip_half<F>{f});
}

Paul's avatar
Paul committed
32
template <class T>
Paul's avatar
Paul committed
33
34
35
36
37
py::buffer_info to_buffer_info(T& x)
{
    migraphx::shape s = x.get_shape();
    py::buffer_info b;
    visit_type(s, [&](auto as) {
Paul's avatar
Paul committed
38
39
40
41
42
43
        b = py::buffer_info(x.data(),
                            as.size(),
                            py::format_descriptor<decltype(as())>::format(),
                            s.lens().size(),
                            s.lens(),
                            s.strides());
Paul's avatar
Paul committed
44
45
46
47
    });
    return b;
}

Paul's avatar
Paul committed
48
49
PYBIND11_MODULE(migraphx, m)
{
Paul's avatar
Paul committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    py::class_<migraphx::shape>(m, "shape")
        .def(py::init<>())
        .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)
        .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)
        .def("scalar", &migraphx::shape::scalar);

    py::class_<migraphx::argument>(m, "argument", py::buffer_protocol())
Paul's avatar
Paul committed
65
        .def_buffer([](migraphx::argument& x) -> py::buffer_info { return to_buffer_info(x); });
Paul's avatar
Paul committed
66
67
68

    py::class_<migraphx::program>(m, "program")
        .def("get_parameter_shapes", &migraphx::program::get_parameter_shapes)
Paul's avatar
Paul committed
69
        .def("compile", [](migraphx::program& p, const migraphx::target& t) { p.compile(t); })
Paul's avatar
Paul committed
70
71
72
73
74
75
76
77
78
79
        .def("eval", &migraphx::program::eval);

    m.def("parse_onnx", &migraphx::parse_onnx);

#ifdef VERSION_INFO
    m.attr("__version__") = VERSION_INFO;
#else
    m.attr("__version__") = "dev";
#endif
}