From dd9f54cdfa2d63e2cd5ecfe0f489672364b01281 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Thu, 31 Dec 2020 10:37:19 +0800 Subject: [PATCH] refactor(mge): migrate to new core implementation * remove dispatcher/interpreter python wrapper * rename tensor_wrapper to array_method GitOrigin-RevId: b8a402c2be58a3e1f990802d772e8fc80ce23006 --- .../tensor/{tensor_wrapper.py => array_method.py} | 0 imperative/python/megengine/functional/tensor.py | 2 +- imperative/python/megengine/jit/tracing.py | 2 +- imperative/python/megengine/tensor.py | 3 +- imperative/python/megengine/utils/profiler.py | 2 +- .../python/megengine/utils/tensor_sanity_check.py | 2 +- imperative/python/src/dispatcher.cpp | 229 --------------------- imperative/python/src/dispatcher.h | 16 -- imperative/python/src/imperative_rt.cpp | 54 ----- imperative/python/src/module.cpp | 4 - imperative/python/test/unit/core/test_autodiff.py | 2 +- .../python/test/unit/core/test_imperative_rt.py | 8 +- 12 files changed, 10 insertions(+), 314 deletions(-) rename imperative/python/megengine/core/tensor/{tensor_wrapper.py => array_method.py} (100%) delete mode 100644 imperative/python/src/dispatcher.cpp delete mode 100644 imperative/python/src/dispatcher.h diff --git a/imperative/python/megengine/core/tensor/tensor_wrapper.py b/imperative/python/megengine/core/tensor/array_method.py similarity index 100% rename from imperative/python/megengine/core/tensor/tensor_wrapper.py rename to imperative/python/megengine/core/tensor/array_method.py diff --git a/imperative/python/megengine/functional/tensor.py b/imperative/python/megengine/functional/tensor.py index 307d5f49..dac5c7d0 100644 --- a/imperative/python/megengine/functional/tensor.py +++ b/imperative/python/megengine/functional/tensor.py @@ -18,7 +18,7 @@ from ..core._imperative_rt.core2 import apply from ..core._wrap import device as as_device from ..core.ops import builtin from ..core.ops.special import Const -from ..core.tensor.tensor_wrapper import _broadcast, _remove_axis +from ..core.tensor.array_method import _broadcast, _remove_axis from ..core.tensor.utils import ( astensor1d, convert_inputs, diff --git a/imperative/python/megengine/jit/tracing.py b/imperative/python/megengine/jit/tracing.py index 7beb82a2..ad02403b 100644 --- a/imperative/python/megengine/jit/tracing.py +++ b/imperative/python/megengine/jit/tracing.py @@ -18,7 +18,7 @@ import weakref import numpy as np -from ..core._imperative_rt import GraphProfiler, common, put +from ..core._imperative_rt import GraphProfiler, common from ..core._imperative_rt.core2 import Tensor as RawTensor from ..core._imperative_rt.core2 import TensorWeakRef from ..core._imperative_rt.core2 import __make_empty_tensor as make_empty_tensor diff --git a/imperative/python/megengine/tensor.py b/imperative/python/megengine/tensor.py index aadd7a9f..fb263b2f 100644 --- a/imperative/python/megengine/tensor.py +++ b/imperative/python/megengine/tensor.py @@ -18,7 +18,7 @@ from .core._imperative_rt.core2 import apply from .core._trace_option import use_symbolic_shape from .core._wrap import device as as_device from .core.ops.builtin import Copy, GetVarShape -from .core.tensor.tensor_wrapper import ArrayMethodMixin +from .core.tensor.array_method import ArrayMethodMixin from .device import _valid_device, get_default_device from .utils.deprecation import deprecated @@ -42,7 +42,6 @@ class Tensor(_Tensor, ArrayMethodMixin): else: cn = device._cn - # import pdb; pdb.set_trace() if isinstance(data, _Tensor): obj = _Tensor.__new__(cls, data) else: diff --git a/imperative/python/megengine/utils/profiler.py b/imperative/python/megengine/utils/profiler.py index 2a362e41..5c67ee70 100644 --- a/imperative/python/megengine/utils/profiler.py +++ b/imperative/python/megengine/utils/profiler.py @@ -14,7 +14,7 @@ from typing import Iterable, List, Optional from ..core._imperative_rt import OperatorNodeConfig, ProfileEntry from ..core._imperative_rt import ProfilerImpl as _Profiler -from ..core._imperative_rt.imperative import sync +from ..core._imperative_rt.core2 import sync from ..core._imperative_rt.ops import CollectiveComm diff --git a/imperative/python/megengine/utils/tensor_sanity_check.py b/imperative/python/megengine/utils/tensor_sanity_check.py index 818dfa6f..8fae8ac7 100644 --- a/imperative/python/megengine/utils/tensor_sanity_check.py +++ b/imperative/python/megengine/utils/tensor_sanity_check.py @@ -1,5 +1,5 @@ from ..core._imperative_rt import TensorSanityCheckImpl -from ..core._imperative_rt.imperative import sync +from ..core._imperative_rt.core2 import sync class TensorSanityCheck: diff --git a/imperative/python/src/dispatcher.cpp b/imperative/python/src/dispatcher.cpp deleted file mode 100644 index 616e1791..00000000 --- a/imperative/python/src/dispatcher.cpp +++ /dev/null @@ -1,229 +0,0 @@ -/** - * \file imperative/python/src/dispatcher.cpp - * MegEngine is Licensed under the Apache License, Version 2.0 (the "License") - * - * Copyright (c) 2014-2020 Megvii Inc. All rights reserved. - * - * 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. - */ - -#include "./dispatcher.h" -#include "./pyext17.h" -#include "megbrain/exception.h" -#include "megbrain/utils/hash.h" -#include "megbrain/utils/small_vector.h" - -#include -#include - -namespace py = pybind11; -namespace pyx = pyext17; - -namespace { - -struct Handler { - PyObject* func; // borrowed - bool enabled; - - Handler() = default; - Handler(PyObject* func_, bool enable = true) : func(func_), enabled(enable) {} -}; - -using FastSig = mgb::SmallVector; -using MRO = std::vector; - -struct Frame { - MRO* mro; - size_t mro_offset; - - Frame() = default; - Frame(MRO* mro_, size_t mro_offset_ = 0) : mro(mro_), mro_offset(mro_offset_) {} -}; - -struct FastSigHash { - size_t operator()(const FastSig& sig) const { - auto* ptr = &sig.front(); - return mgb::XXHash() - .update(ptr, sig.size() * sizeof(FastSig::value_type)) - .digest(); - } -}; - -struct ObjectIdHash : std::hash { - size_t operator()(const py::handle& h) const { - return std::hash::operator()(h.ptr()); - } -}; - -namespace { -using Container = std::vector; -struct DispatcherStack: Container { - constexpr static size_t MAX_RECURSIVE_DEPTH = 1024u; - DispatcherStack() { reserve(MAX_RECURSIVE_DEPTH); } - - template - auto&& emplace_back_safely(Args&& ...args) { - mgb_throw_if(size() >= MAX_RECURSIVE_DEPTH, mgb::MegBrainError, - "recursion depth %zu is greater than the MAX_RECURSIVE_DEPTH(%zu)", - size(), MAX_RECURSIVE_DEPTH); - return emplace_back(std::forward(args)...); - } -}; -} // anonymous namespace - -struct Dispatcher { - std::unordered_map, FastSigHash> cache; - DispatcherStack stack; - std::unordered_map, ObjectIdHash> registry; - - inline py::handle self() { - return pyx::wrap::pycast(this); - } - - bool prepare_call(PyObject*const* args, Py_ssize_t nargs) { - FastSig sig(nargs); - for (Py_ssize_t i = 0; i < nargs; ++i) { - sig[i] = Py_TYPE(args[i]); - } - auto it = cache.find(sig); - if (it == cache.end()) { - if (auto mro = resolve(sig)) { - it = cache.emplace(std::move(sig), std::move(mro)).first; - } else { - return false; - } - } - stack.emplace_back_safely(it->second.get()); - return true; - } - - template - PyObject* do_call(T&& caller) { - auto& frame = stack.back(); - auto& mro = *frame.mro; - auto& i = frame.mro_offset; - if (!mro.size()) { - PyErr_SetString(PyExc_NotImplementedError, "function not registered in dispatcher"); - return nullptr; - } - for (; i < mro.size(); ++i) { - if (mro[i]->enabled) { - auto ret = caller(mro[i]->func); - if (ret != Py_NotImplemented) { - stack.pop_back(); - return ret; - } - Py_DECREF(ret); - } - } - PyErr_SetString(PyExc_NotImplementedError, "mro exhausted"); - stack.pop_back(); - return nullptr; - } - - std::unique_ptr resolve(const FastSig& sig) { - try { - py::tuple args(sig.size()); - for (size_t i = 0; i < sig.size(); ++i) { - args[i] = (PyObject*)sig[i]; - } - auto mro_iter = self().attr("dispatch_iter")(*args); - auto ret = std::make_unique(); - for (auto i : mro_iter) { - auto it = registry.find(py::reinterpret_borrow(i)); - if (it == registry.end()) { - PyErr_SetString(PyExc_RuntimeError, "resolved to unregistered function"); - return nullptr; - } - ret->push_back(it->second.get()); - } - return ret; - } catch (py::error_already_set& e) { - e.restore(); - } catch (std::runtime_error& e) { - PyErr_SetString(PyExc_RuntimeError, e.what()); - } - return nullptr; - } - -public: - static constexpr auto tp_name = "Dispatcher"; - - PyObject* tp_call(PyObject* args, PyObject* kwargs) { - if (!prepare_call(&PyTuple_GET_ITEM(args, 0), PyTuple_GET_SIZE(args))) return nullptr; - return do_call([=](PyObject* func){return PyObject_Call(func, args, kwargs);}); - } - -#if PY_MINOR_VERSION >= 6 - PyObject* tp_vectorcall(PyObject*const* args, Py_ssize_t nargs) { - if (!prepare_call(args, nargs)) return nullptr; - return do_call([=](PyObject* func){return _PyObject_FastCall(func, const_cast(args), nargs);}); - } -#endif - -#if PY_MINOR_VERSION >= 6 - PyObject* super(PyObject*const* args, Py_ssize_t nargs) { - if (stack.empty()) { - PyErr_SetString(PyExc_RuntimeError, "super called at top level"); - return nullptr; - } - stack.emplace_back_safely(stack.back()).mro_offset++; - return do_call([=](PyObject* func){return _PyObject_FastCall(func, const_cast(args), nargs);}); - } -#else - PyObject* super(PyObject* args, PyObject* kwargs) { - if (stack.empty()) { - PyErr_SetString(PyExc_RuntimeError, "super called at top level"); - return nullptr; - } - stack.emplace_back_safely(stack.back()).mro_offset++; - return do_call([=](PyObject* func){return PyObject_Call(func, args, kwargs);}); - } -#endif - - void enable(PyObject* func) { - auto obj = py::reinterpret_borrow(func); - auto it = registry.find(obj); - if (it != registry.end()) { - it->second->enabled = true; - } else { - registry.emplace(std::move(obj), std::make_unique(func)); - } - } - - PyObject* disable(PyObject* func) { - auto obj = py::reinterpret_borrow(func); - auto it = registry.find(obj); - if (it == registry.end()) { - PyErr_SetString(PyExc_ValueError, "function not registered"); - return nullptr; - } else { - it->second->enabled = false; - } - Py_RETURN_NONE; - } - - void clear_cache() { - cache.clear(); - } -}; - -} // namespace - -void init_dispatcher(py::module m) { - auto* dispatcher_type = pyx::wrap::type() - .def<&Dispatcher::enable>("enable") - .def<&Dispatcher::disable>("disable") - .def<&Dispatcher::clear_cache>("clear_cache") -#if PY_MINOR_VERSION >= 6 - .def<&Dispatcher::tp_vectorcall>("call") -#else - .def<&Dispatcher::tp_call>("call") -#endif - .def<&Dispatcher::super>("super") - .finalize(); - if (!dispatcher_type) throw py::error_already_set(); - m.attr("Dispatcher") = dispatcher_type; -} diff --git a/imperative/python/src/dispatcher.h b/imperative/python/src/dispatcher.h deleted file mode 100644 index 5bb53311..00000000 --- a/imperative/python/src/dispatcher.h +++ /dev/null @@ -1,16 +0,0 @@ -/** - * \file imperative/python/src/dispatcher.h - * MegEngine is Licensed under the Apache License, Version 2.0 (the "License") - * - * Copyright (c) 2014-2020 Megvii Inc. All rights reserved. - * - * 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. - */ - -#pragma once - -#include - -void init_dispatcher(pybind11::module); diff --git a/imperative/python/src/imperative_rt.cpp b/imperative/python/src/imperative_rt.cpp index 8baa379f..acf20758 100644 --- a/imperative/python/src/imperative_rt.cpp +++ b/imperative/python/src/imperative_rt.cpp @@ -51,59 +51,5 @@ make_backward_graph( } // namespace void init_imperative_rt(py::module m) { - py::class_(m, "Interpreter") - .def("put", [](Interpreter::Channel& self, py::array data, DType dtype, CompNode cn) { - if (!cn.valid()) { - cn = CompNode::load(get_default_device()); - } - constexpr int size_threshhold = TensorShape::MAX_NDIM; - if (data.size() > size_threshhold) { - return self.put(npy::np2tensor(data.ptr(), npy::Meth::borrow(cn), dtype)); - } else { - HostTensorND ret(cn); - return self.put(npy::np2tensor(data.ptr(), npy::Meth::copy_into(&ret), dtype)); - } - }, py::arg(), py::arg("dtype") = py::none(), py::arg("device") = py::none()) - .def("put", py::overload_cast(&Interpreter::Channel::put)) - .def("delete", [](Interpreter::Channel& self, Interpreter::Handle handle) { - return self.del(handle); - }) - .def("_swap_in", [](Interpreter::Channel& self, Interpreter::Handle handle) { - self.swap_in(handle); - }) - .def("_swap_out", [](Interpreter::Channel& self, Interpreter::Handle handle) { - self.swap_out(handle); - }) - .def("_drop", [](Interpreter::Channel& self, Interpreter::Handle handle) { - self.drop(handle); - }) - .def("get_value", [](Interpreter::Channel& self, Interpreter::Handle handle) { - PyObject* optr = npy::ndarray_from_tensor(self.get_value(handle), npy::ShareType::TRY_SHARE); - return py::reinterpret_steal(optr); - }) - .def("get_dtype", &Interpreter::Channel::get_dtype) - .def("get_device", &Interpreter::Channel::get_device) - .def("get_shape", &Interpreter::Channel::get_shape) - .def("_get_dev_tensor", &Interpreter::Channel::get_dev_tensor) - .def("_set_swap_flag", &Interpreter::Channel::set_swap_flag) - .def("_set_drop_flag", &Interpreter::Channel::set_drop_flag) - .def("apply_op", &Interpreter::Channel::apply_op) - .def("config_async_level", &Interpreter::Channel::config_async_level) - .def("get_async_level", &Interpreter::Channel::get_async_level) - .def("sync", &Interpreter::Channel::sync, py::call_guard()); - - std::unique_ptr ch = Interpreter::inst().create_channel(); - m.attr("interpreter") = py::detail::make_caster::cast( - std::move(ch), py::return_value_policy::move, {}); - for (auto name : {"put", "delete", "get_value", "get_dtype", "get_device", "get_shape", "_get_dev_tensor", "apply_op", "config_async_level", "get_async_level", "_drop", "_swap_in", "_swap_out", "_set_drop_flag", "_set_swap_flag"}) { - m.attr(name) = m.attr("interpreter").attr(name); - } - - m.def("sync", [m]() { - m.attr("interpreter").attr("sync")(); - py::gil_scoped_release _; - py_task_q.wait_all_task_finish(); - }); - m.def("make_backward_graph", &make_backward_graph); } diff --git a/imperative/python/src/module.cpp b/imperative/python/src/module.cpp index b2eab035..4ffed371 100644 --- a/imperative/python/src/module.cpp +++ b/imperative/python/src/module.cpp @@ -21,8 +21,6 @@ #include "./graph_rt.h" #include "./ops.h" -#include "./dispatcher.h" - #include "./tensor.h" namespace py = pybind11; @@ -70,7 +68,5 @@ PYBIND11_MODULE(MODULE_NAME, m) { )", py::getattr(m, "__dict__")); - init_dispatcher(submodule(m, "dispatcher")); - init_tensor(submodule(m, "core2")); } diff --git a/imperative/python/test/unit/core/test_autodiff.py b/imperative/python/test/unit/core/test_autodiff.py index 9ea6cf97..05d1526b 100644 --- a/imperative/python/test/unit/core/test_autodiff.py +++ b/imperative/python/test/unit/core/test_autodiff.py @@ -16,7 +16,7 @@ import pytest import megengine as mge import megengine.distributed as dist import megengine.functional as F -from megengine.core._imperative_rt import CompNode, TensorAttr, core2, imperative +from megengine.core._imperative_rt import CompNode, TensorAttr, imperative from megengine.core._imperative_rt.core2 import TensorWeakRef, apply, sync from megengine.core.autodiff.grad import Grad from megengine.core.ops.builtin import Elemwise diff --git a/imperative/python/test/unit/core/test_imperative_rt.py b/imperative/python/test/unit/core/test_imperative_rt.py index b68287b7..efc781e5 100644 --- a/imperative/python/test/unit/core/test_imperative_rt.py +++ b/imperative/python/test/unit/core/test_imperative_rt.py @@ -54,10 +54,10 @@ def test_simple_arith(): def test_tensor_on_device(): device = megengine.core._imperative_rt.CompNode("cpu0:1") x = np.random.rand(10).astype("float32") - xx = megengine.core._imperative_rt.put(x, device=device) - assert str(megengine.core._imperative_rt.get_device(xx)) == "cpu0:1" - np.testing.assert_equal(x, megengine.core._imperative_rt.get_value(xx)) - megengine.core._imperative_rt.delete(xx) + xx = megengine.tensor(x, device=device) + assert str(xx.device) == "cpu0:1" + np.testing.assert_equal(x, xx.numpy()) + del xx def test_raw_tensor():