common.cpp 7.6 KB
Newer Older
M
Megvii Engine Team 已提交
1 2 3 4
/**
 * \file imperative/python/src/common.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
5
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
M
Megvii Engine Team 已提交
6 7 8 9 10 11
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

12 13 14 15 16 17 18 19 20 21 22 23 24 25
#include "./common.h"

#include <pybind11/operators.h>

#include "megbrain/comp_node.h"
#include "megbrain/graph.h"
#include "megbrain/imperative/physical_tensor.h"
#include "./numpy_dtypes.h"
#include "./helper.h"

namespace py = pybind11;
using namespace mgb;
using namespace imperative;

M
Megvii Engine Team 已提交
26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41
namespace {

template<typename XTensorND>
auto def_TensorND(py::object parent, const char* name) {
    return py::class_<XTensorND>(parent, name)
        .def_property_readonly("shape", py::overload_cast<>(&XTensorND::shape, py::const_))
        .def_property_readonly("dtype", py::overload_cast<>(&XTensorND::dtype, py::const_))
        .def_property_readonly("comp_node", py::overload_cast<>(&XTensorND::comp_node, py::const_))
        .def("copy_from", &XTensorND::template copy_from<DeviceTensorStorage>)
        .def("copy_from", &XTensorND::template copy_from<HostTensorStorage>)
        .def("copy_from_fixlayout", py::overload_cast<const DeviceTensorND&>(
            &XTensorND::template copy_from_fixlayout<DeviceTensorStorage>))
        .def("copy_from_fixlayout", py::overload_cast<const HostTensorND&>(
            &XTensorND::template copy_from_fixlayout<HostTensorStorage>));
}

42 43
std::string default_device = "xpux";

M
Megvii Engine Team 已提交
44 45
} // namespace

46 47 48 49 50 51 52 53
void set_default_device(const std::string &device) {
    default_device = device;
}

std::string get_default_device() {
    return default_device;
}

54
void init_common(py::module m) {
55
    auto PyCompNode = py::class_<CompNode>(m, "CompNode")
56 57
        .def(py::init())
        .def(py::init(py::overload_cast<const std::string&>(&CompNode::load)))
58 59 60
        .def_property_readonly("logical_name", [](const CompNode& cn) {
            return cn.to_string_logical();
        })
61 62 63
        .def_property_readonly("physical_name", [](const CompNode& cn) {
            return cn.to_string();
        })
64 65 66
        .def_property_readonly("get_mem_status_bytes", [](const CompNode& cn) {
            return cn.get_mem_status_bytes();
        })
M
Megvii Engine Team 已提交
67
        .def("create_event", &CompNode::create_event, py::arg("flags") = 0ul)
68 69
        .def_static("_set_default_device", &set_default_device)
        .def_static("_get_default_device", &get_default_device)
70
        .def("__str__", &CompNode::to_string_logical)
71
        .def("__repr__", [](const CompNode& cn) {
72 73 74
            return mgb::ssprintf("CompNode(\"%s\" from \"%s\")",
                    cn.to_string_physical().c_str(),
                    cn.to_string_logical().c_str());
75
        })
76
        .def("__hash__", [](CompNode cn){ return mgb::hash(cn); })
77 78 79 80 81 82 83 84 85 86 87 88
        .def_static("_sync_all", &CompNode::sync_all)
        .def(py::self == py::self)
        .def_static("_get_device_count", &CompNode::get_device_count,
                    "Get total number of specific devices on this system")
        .def(py::pickle(
                [](const CompNode& cn) {
                    return py::str(cn.to_string_logical());
                },
                [](py::str cn) {
                    return CompNode::load(cn);
                }));

M
Megvii Engine Team 已提交
89 90 91 92
    py::class_<CompNode::Event, std::shared_ptr<CompNode::Event>>(PyCompNode, "Event")
        .def("record", &CompNode::Event::record)
        .def("wait", &CompNode::Event::host_wait);

93 94
    py::implicitly_convertible<std::string, CompNode>();

M
Megvii Engine Team 已提交
95
    def_TensorND<DeviceTensorND>(m, "DeviceTensorND")
96 97 98 99 100 101
        .def("numpy", [](const DeviceTensorND& self) {
                HostTensorND hv;
                hv.copy_from(self).sync();
                return py::handle(npy::ndarray_from_tensor(hv, npy::ShareType::TRY_SHARE));
            });

M
Megvii Engine Team 已提交
102 103 104 105 106 107 108 109 110 111 112
    def_TensorND<HostTensorND>(m, "HostTensorND")
        .def(py::init([](py::array data, CompNode cn, DType dtype) {
                if (!cn.valid()) {
                    throw py::type_error("device must not be None");
                }
                return npy::np2tensor(data.ptr(), npy::Meth::borrow(cn), dtype);
            }))
        .def("numpy", [](const HostTensorND& self) {
                return py::reinterpret_steal<py::object>(npy::ndarray_from_tensor(self, npy::ShareType::TRY_SHARE));
            });

113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179
    py::class_<cg::OperatorNodeConfig>(m, "OperatorNodeConfig")
        .def(py::init())
        .def_property("name",
            [](const OperatorNodeConfig& config) -> py::object {
                auto name = config.name();
                if (name.valid()) {
                    return py::str(name.val());
                } else {
                    return py::none();
                }
            },
            [](OperatorNodeConfig& config, std::string name){
                config.name(std::move(name));
            })
        .def_property("dtype",
            [](const OperatorNodeConfig& config) {
                return config.output_dtype();
            },
            [](OperatorNodeConfig& config, DType dtype) {
                config.output_dtype(dtype);
            })
        .def_property("comp_node_arr",
            [](const OperatorNodeConfig& config) -> py::tuple {
                auto arr = config.comp_node();
                std::vector<CompNode> tmp(arr.begin(), arr.end());
                return py::cast(tmp);
            },
            [](OperatorNodeConfig& config, std::vector<CompNode> cns) {
                config.comp_node_arr({cns.begin(), cns.end()});
            })
        .def_property("comp_node",
            [](const OperatorNodeConfig& config) {
                auto arr = config.comp_node();
                if (arr.size() != 1) {
                    throw py::value_error("invalid number of comp_node");
                }
                return arr[0];
            },
            [](OperatorNodeConfig& config, CompNode cn) {
                OperatorNodeConfig::CompNodeArray arr{cn};
                config.comp_node_arr(arr);
            });

    py::class_<LogicalTensorDesc>(m, "TensorAttr")
        .def(py::init())
        .def(py::init([](const TensorShape& shape, const DType& dtype, const CompNode& comp_node){
                return LogicalTensorDesc{TensorLayout{shape, dtype}, comp_node};
            }))
        .def_property("shape",
            [](const LogicalTensorDesc& desc) {
                return static_cast<TensorShape>(desc.layout);
            },
            [](LogicalTensorDesc& desc, TensorShape shape) {
            })
        .def_property("dtype",
            [](const LogicalTensorDesc& desc) {
                return desc.layout.dtype;
            },
            [](LogicalTensorDesc& desc, DType dtype) {
                desc.layout.dtype = dtype;
            })
        .def_readwrite("comp_node", &LogicalTensorDesc::comp_node);

    py::enum_<CompNode::DeviceType>(m, "DeviceType")
            .value("UNSPEC", CompNode::DeviceType::UNSPEC)
            .value("CUDA", CompNode::DeviceType::CUDA)
            .value("CPU", CompNode::DeviceType::CPU)
180 181
            .value("CAMBRICON", CompNode::DeviceType::CAMBRICON)
            .value("ATLAS", CompNode::DeviceType::ATLAS)
182 183 184
            .value("MULTITHREAD", CompNode::DeviceType::MULTITHREAD)
            .value("MAX_DEVICE_ID", CompNode::DeviceType::MAX_DEVICE_ID);

185 186 187
    m.def("set_prealloc_config", &CompNode::set_prealloc_config, 
        "specifies how to pre-allocate from raw dev allocator");

188 189 190 191
    m.def("what_is_xpu", []{
        return CompNode::Locator::parse("xpux").to_physical().type;
    });

192 193
    init_npy_num_bfloat16(m);
    init_npy_num_intbx(m);
194
    init_dtypes(m);
195
}