MegEngine/imperative/python/src/helper.cpp

/**
 * \file imperative/python/src/helper.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 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 "./helper.h"

#include <pybind11/eval.h>

#include "megbrain/graph/cg.h"
#include "megbrain/graph/event.h"
#include "megbrain/graph/exc_extra_info.h"
#include "megbrain/tensor.h"
#include "megbrain/utils/mempool.h"

#include "./numpy_dtypes.h"

namespace py = pybind11;

PyTaskDipatcher py_task_q = {};

py::module submodule(py::module parent, const char* name, const char* doc) {
    auto m = parent.def_submodule(name, doc);
    m.attr("__package__") = parent.attr("__name__");
    m.attr("__builtins__") = py::module::import("builtins");
    return m;
}

py::module rel_import(py::str name, py::module m, int level) {
    py::object import = py::module::import("builtins").attr("__import__");
    return import(name, m.attr("__dict__"), py::arg("level") = level);
}

/*
 * demangle typeid, see
 * http://stackoverflow.com/questions/281818/unmangling-the-result-of-stdtype-infoname
 */
#ifdef __GNUG__
#include <cxxabi.h>
#include <cstdlib>
#include <memory>

namespace {

std::string demangle_typeid(const char* name) {
    int status = -4;  // some arbitrary value to eliminate the compiler warning

    // enable c++11 by passing the flag -std=c++11 to g++
    std::unique_ptr<char, void (*)(void*)> res{
            abi::__cxa_demangle(name, nullptr, nullptr, &status), std::free};

    return (status == 0) ? res.get() : name;
}
}  // namespace
#else

namespace {
// does nothing if not g++
std::string demangle_typeid(const char* name) {
    return name;
}
}  // namespace

#endif

using namespace mgb;
using namespace cg;

namespace {

std::string repr_pyobj(PyObject* obj) {
    if (!obj)
        return "<null PyObject>";
    PYTHON_GIL;
    auto str = PyObject_Repr(obj);
    if (!str)
        return ssprintf("<PyObject at %p (repr failed)>", obj);
    std::string ret{PyUnicode_AsUTF8(str)};
    Py_DECREF(str);
    return ret;
}

template <typename T>
std::string typeid_name(const T& t) {
    return demangle_typeid(typeid(t).name());
}

}  // anonymous namespace

/* ============== PyExceptionForward ============== */

PyExceptionForward::~PyExceptionForward() {
    PYTHON_GIL;
    PyObjRefKeeper::deleter(m_type);
    PyObjRefKeeper::deleter(m_value);
    PyObjRefKeeper::deleter(m_traceback);
}

void PyExceptionForward::restore() {
    PyErr_Restore(m_type, m_value, m_traceback);
    m_type = m_value = m_traceback = nullptr;
}

void PyExceptionForward::throw_() {
    PyObject *etype, *obj, *trace;
    PyErr_Fetch(&etype, &obj, &trace);
    PyErr_NormalizeException(&etype, &obj, &trace);

    std::string msg{"python exception"};
    bool succ = false;
    if (etype && obj && trace) {
        auto run = [&]() {
#define DEF(name, expr)        \
    PyObjRefKeeper name{expr}; \
    if (!name.get())           \
    return
            DEF(mod, PyImport_ImportModule("traceback"));
            DEF(result,
                PyObject_CallMethod(
                        mod.get(), "format_exception", "(OOO)", etype, obj, trace));
            if (!PyList_Check(result.get()))
                return;
            auto size = PyList_Size(result.get());
            msg.append(":\n");
            for (Py_ssize_t i = 0; i < size; ++i) {
                msg.append("  ");
                msg.append(PyUnicode_AsUTF8(PyList_GetItem(result.get(), i)));
            }
            msg.pop_back();  // remove last \n
            succ = true;
#undef DEF
        };
        run();
    }
    if (!succ) {
        PyObject* obj_str_py;
        if (obj && (obj_str_py = PyObject_Repr(obj))) {
            msg.append(" with message ");
            msg.append(PyUnicode_AsUTF8(obj_str_py));
            Py_DECREF(obj_str_py);
        } else {
            msg.append(" with unknown message");
        }
    }
    // throwing exception may cause abort due to unknown reasons; so we first
    // log the message
    mgb_log_error("caught exception from python callback: %s", msg.c_str());
    fflush(stdout);
    fflush(stderr);
    throw PyExceptionForward{etype, obj, trace, msg};
}

/* ============== namespace npy ============== */

namespace npy {

int to_mgb_supported_dtype_raw(int dtype) {
    if (dtype == NPY_INT64)
        return NPY_INT32;
    if (dtype == NPY_FLOAT64)
        return NPY_FLOAT32;
    return dtype;
}

#define FOREACH_NPY_DTYPE_PAIR(cb)                                             \
    cb(Uint8, NPY_UINT8) cb(Int8, NPY_INT8) cb(Uint16, NPY_UINT16)             \
            cb(Int16, NPY_INT16) cb(Int32, NPY_INT32) cb(Float16, NPY_FLOAT16) \
                    cb(Float32, NPY_FLOAT32) cb(Bool, NPY_BOOL)

#define FOREACH_NPY_MGB_DTYPE_PAIR(cb) \
    FOREACH_NPY_DTYPE_PAIR(cb)         \
    FOREACH_MGB_DTYPE_PAIR(cb)

//! convert megbrain dtype to numpy dtype
int dtype_mgb2np_raw(DType dtype) {
    mgb_assert(dtype.valid(), "attempt to convert from invalid dtype");
    switch (dtype.enumv()) {
#define cb(_m, _n)      \
    case DTypeEnum::_m: \
        return _n;
        FOREACH_NPY_MGB_DTYPE_PAIR(cb)
#undef cb
        default:
            break;
    }
    throw ConversionError(
            ssprintf("can not convert dtype %s to numpy dtype", dtype.name()));
}

//! Convert MegBrain DType to NumPy DType descriptor, the caller receives a new
//! reference to the descriptor.
std::unique_ptr<PyArray_Descr, PyArrayDescrDeleter> dtype_mgb2np_descr(DType dtype) {
    PYTHON_GIL;
    mgb_assert(dtype.valid(), "attempt to convert from invalid dtype");
    auto build_mgb_dtype_dict =
            [](const char* name,
               const std::vector<std::pair<const char*, PyObject*>>& data) {
                PyObject* metadata = PyDict_New();
                PyObject* mgb_dtype_metadata = PyDict_New();
                PyDict_SetItemString(
                        mgb_dtype_metadata, "name", PyUnicode_FromString(name));
                for (const auto& d : data) {
                    PyDict_SetItemString(mgb_dtype_metadata, d.first, d.second);
                }
                PyDict_SetItemString(metadata, "mgb_dtype", mgb_dtype_metadata);
                return metadata;
            };
    if (dtype.has_param()) {
        PyArray_Descr* type_descr;
        switch (dtype.enumv()) {
            case DTypeEnum::QuantizedS1: {
                auto& param = dtype.param<dtype::QuantizedS1>();
                type_descr = PyArray_DescrNewFromType(NPY_INT8);
                type_descr->metadata = build_mgb_dtype_dict(
                        DTypeTrait<dtype::QuantizedS1>::name,
                        {{"scale", PyFloat_FromDouble(param.scale)}});
                break;
            }
            case DTypeEnum::Quantized4Asymm: {
                auto& param = dtype.param<dtype::Quantized4Asymm>();
                type_descr = PyArray_DescrNewFromType(NPY_UINT8);
                type_descr->metadata = build_mgb_dtype_dict(
                        DTypeTrait<dtype::Quantized4Asymm>::name,
                        {{"scale", PyFloat_FromDouble(param.scale)},
                         {"zero_point", PyLong_FromLong(param.zero_point)}});
                break;
            }
            case DTypeEnum::QuantizedS4: {
                auto& param = dtype.param<dtype::QuantizedS4>();
                type_descr = PyArray_DescrNewFromType(NPY_INT8);
                type_descr->metadata = build_mgb_dtype_dict(
                        DTypeTrait<dtype::QuantizedS4>::name,
                        {{"scale", PyFloat_FromDouble(param.scale)}});
                break;
            }
            case DTypeEnum::Quantized8Asymm: {
                auto& param = dtype.param<dtype::Quantized8Asymm>();
                type_descr = PyArray_DescrNewFromType(NPY_UINT8);
                type_descr->metadata = build_mgb_dtype_dict(
                        DTypeTrait<dtype::Quantized8Asymm>::name,
                        {{"scale", PyFloat_FromDouble(param.scale)},
                         {"zero_point", PyLong_FromLong(param.zero_point)}});
                break;
            }
            case DTypeEnum::QuantizedS8: {
                auto& param = dtype.param<dtype::QuantizedS8>();
                type_descr = PyArray_DescrNewFromType(NPY_INT8);
                type_descr->metadata = build_mgb_dtype_dict(
                        DTypeTrait<dtype::QuantizedS8>::name,
                        {{"scale", PyFloat_FromDouble(param.scale)}});
                break;
            }
            case DTypeEnum::QuantizedS32: {
                auto& param = dtype.param<dtype::QuantizedS32>();
                type_descr = PyArray_DescrNewFromType(NPY_INT32);
                type_descr->metadata = build_mgb_dtype_dict(
                        DTypeTrait<dtype::QuantizedS32>::name,
                        {{"scale", PyFloat_FromDouble(param.scale)}});
                break;
            }
            default:
                mgb_throw(
                        ConversionError, "unhandled parameterized DType %s",
                        dtype.name());
        }
        return std::unique_ptr<PyArray_Descr, PyArrayDescrDeleter>(type_descr);
    }
    PyArray_Descr* basic_descr = PyArray_DescrFromType(dtype_mgb2np_raw(dtype));
    mgb_assert(
            basic_descr != nullptr,
            "failed to convert expected dtype to numpy type descriptor");
    return std::unique_ptr<PyArray_Descr, PyArrayDescrDeleter>(basic_descr);
}

DType dtype_np2mgb_raw(int npt) {
    switch (npt) {
#define cb(_m, _n) \
    case _n:       \
        return dtype::_m();
        FOREACH_NPY_DTYPE_PAIR(cb)
#undef cb
    }
#define cb(_m, _n) \
    if (_n == npt) \
        return dtype::_m();
    FOREACH_MGB_DTYPE_PAIR(cb)
#undef cb

    PYTHON_GIL;
    std::string msg;
    auto py_obj = PyArray_TypeObjectFromType(npt);
    if (!py_obj) {
        msg = ssprintf("unknown numpy dtype enum %d", npt);
    } else {
        msg = ssprintf("unsupported numpy dtype %s", repr_pyobj(py_obj).c_str());
    }
    Py_DECREF(py_obj);
    throw ConversionError(msg);
}

DType dtype_np2mgb_descr(PyArray_Descr* descr) {
    PYTHON_GIL;
    auto handle_parameterized_dtype = [](PyObject* metadata) -> DType {
        mgb_assert(
                PyDict_Check(metadata),
                "Invalid parameterized DType metadata: should be a dict");
        PyObject* dtype_name_py = PyDict_GetItemString(metadata, "name");
        mgb_assert(
                PyUnicode_Check(dtype_name_py),
                "Invalid parameterized DType metadata: name should be a str");
        std::string dtype_name(PyUnicode_AsUTF8(dtype_name_py));
        if (dtype_name == "Quantized8Asymm") {
            PyObject* scale_py = PyDict_GetItemString(metadata, "scale");
            PyObject* zero_point_py = PyDict_GetItemString(metadata, "zero_point");
            mgb_assert(
                    scale_py && zero_point_py,
                    "Invalid Quantized8Asymm metadata: missing scale or "
                    "zero_point.");
            mgb_assert(
                    PyFloat_Check(scale_py),
                    "Invalid Quantized8Asymm metadata: scale should be float");
            mgb_assert(
                    PyLong_Check(zero_point_py),
                    "Invalid Quantized8Asymm metadata: zero_point should be "
                    "integer");
            auto zero_point = PyLong_AS_LONG(zero_point_py);
            mgb_assert(
                    zero_point >= 0 && zero_point < 256,
                    "Invalid Quantized8Asymm metadata: zero_point should be "
                    "in [0, 256)");
            return dtype::Quantized8Asymm(
                    static_cast<float>(PyFloat_AS_DOUBLE(scale_py)),
                    static_cast<uint8_t>(zero_point));
        }
        if (dtype_name == "Quantized4Asymm") {
            PyObject* scale_py = PyDict_GetItemString(metadata, "scale");
            PyObject* zero_point_py = PyDict_GetItemString(metadata, "zero_point");
            mgb_assert(
                    scale_py && zero_point_py,
                    "Invalid Quantized4Asymm metadata: missing scale or "
                    "zero_point.");
            mgb_assert(
                    PyFloat_Check(scale_py),
                    "Invalid Quantized4Asymm metadata: scale should be float");
            mgb_assert(
                    PyLong_Check(zero_point_py),
                    "Invalid Quantized4Asymm metadata: zero_point should be "
                    "integer");
            auto zero_point = PyLong_AS_LONG(zero_point_py);
            mgb_assert(
                    zero_point >= 0 && zero_point < 15,
                    "Invalid Quantized4Asymm metadata: zero_point should be "
                    "in [0, 15)");
            return dtype::Quantized4Asymm(
                    static_cast<float>(PyFloat_AS_DOUBLE(scale_py)),
                    static_cast<uint8_t>(zero_point));
        }
        if (dtype_name == "QuantizedS32" || dtype_name == "QuantizedS8" ||
            dtype_name == "QuantizedS4" || dtype_name == "QuantizedS1") {
            PyObject* scale_py = PyDict_GetItemString(metadata, "scale");
            mgb_assert(scale_py, "Invalid metadata: missing scale");
            mgb_assert(
                    PyFloat_Check(scale_py), "Invalid metadata: scale should be float");
            float scale = static_cast<float>(PyFloat_AS_DOUBLE(scale_py));
            if (dtype_name == "QuantizedS32") {
                return dtype::QuantizedS32(scale);
            } else if (dtype_name == "QuantizedS8") {
                return dtype::QuantizedS8(scale);
            } else if (dtype_name == "QuantizedS4") {
                return dtype::QuantizedS4(scale);
            } else if (dtype_name == "QuantizedS1") {
                return dtype::QuantizedS1(scale);
            }
        }
        throw ConversionError(
                ssprintf("Unknown parameterized DType: %s", dtype_name.c_str())
                        .c_str());
    };
    PyObject* dtype_metadata;
    if (descr->metadata && PyDict_Check(descr->metadata) &&
        (dtype_metadata = PyDict_GetItemString(descr->metadata, "mgb_dtype"))) {
        return handle_parameterized_dtype(dtype_metadata);
    }
    return dtype_np2mgb_raw(descr->type_num);
}

HostTensorND lowbit_ndarray_to_host_tensor(
        CompNode comp_node, TensorLayout& layout, PyArrayObject* input) {
    auto src_ptr = reinterpret_cast<dt_byte*>(PyArray_DATA(input));
    if (!layout.ndim) {
        // numpy scalar
        mgb_assert(src_ptr, "can not convert from null numpy array");
        layout.init_contiguous_stride({1});
    } else {
        mgb_assert(
                layout.ndim && layout.ndim <= TensorShape::MAX_NDIM,
                "unsupported ndim %zu", layout.ndim);
        TensorLayout ly;
        ly.ndim = layout.ndim;
        for (size_t i = 0; i < layout.ndim; ++i) {
            ly.shape[i] = layout.shape[i] = PyArray_SHAPE(input)[i];
            ly.stride[i] = PyArray_STRIDE(input, i);
            mgb_assert(layout.shape[i], "zero shape not supported");
        }
        mgb_assert(ly.is_physical_contiguous());
        layout.init_contiguous_stride();
    }
    HostTensorND ret{comp_node, layout};
    if (layout.format.is_lowbit_aligned()) {
        mgb_assert(layout.is_contiguous());
        lowbit_memcpy_byte2aligned(ret.raw_ptr(), src_ptr, layout);
    } else {
        lowbit_memcpy_byte2compact(
                layout.dtype, ret.raw_ptr(), src_ptr, layout.total_nr_elems());
    }
    return ret;
}

/*!
 * \brief convert a python object to tensor and try to borrow memory if the
 *      original object is a contiguous numpy array
 * \param dtype see np2tensor
 * \return the megbrain tensor, and whether memory is borrowed
 */
std::pair<HostTensorND, bool> np2tensor_try_borrow(
        PyObject* obj, const npy::Meth& meth, DType dtype) {
    auto dest_cn = meth.dest_cn_;
    mgb_assert(dest_cn.valid());

    PYTHON_GIL;

    PyArray_Descr* expected_descr = nullptr;
    if (dtype.valid()) {
        // The reference to expected_descr will be stealed later.
        expected_descr = dtype_mgb2np_descr(dtype).release();
    }

    // make result from PyArrayObject; its reference may be stolen
    auto make_from_arr = [&](PyArrayObject* input, bool allow_borrow) {
        TensorLayout layout{{}, dtype_np2mgb_descr(PyArray_DESCR(input))};
        if (dtype.valid())
            mgb_assert(dtype == layout.dtype);
        layout.ndim = PyArray_NDIM(input);

        if (layout.dtype.is_low_bit()) {
            auto ret = lowbit_ndarray_to_host_tensor(dest_cn, layout, input);
            if (meth.dest_tensor_) {
                meth.dest_tensor_->copy_from(ret);
                ret = *meth.dest_tensor_;
            }
            return std::make_pair(ret, false);
        }

        auto data = reinterpret_cast<dt_byte*>(PyArray_DATA(input));
        if (!layout.ndim) {
            // numpy scalar
            mgb_assert(data, "can not convert from null numpy array");
            layout.init_contiguous_stride({1});
        } else {
            mgb_assert(
                    layout.ndim && layout.ndim <= TensorShape::MAX_NDIM,
                    "unsupported ndim %zu", layout.ndim);
            auto dsize = layout.dtype.size();
            bool is_empty = false;
            for (size_t i = 0; i < layout.ndim; ++i) {
                layout.shape[i] = PyArray_SHAPE(input)[i];
                layout.stride[i] = PyArray_STRIDE(input, i);
                if (!layout.shape[i]) {
                    is_empty = true;
                }
                mgb_assert(
                        layout.stride[i] % dsize == 0, "bad stride %zd",
                        layout.stride[i]);
                layout.stride[i] /= dsize;
            }
            mgb_assert(is_empty || layout.is_contiguous());
        }

        if (!meth.dest_tensor_ && allow_borrow) {
            Py_INCREF(input);
            PyObjRefKeeper ref_obj_cvt{reinterpret_cast<PyObject*>(input)};
            HostTensorStorage storage;
            auto input_ptr = ref_obj_cvt.make_shared(data);
            storage.reset(dest_cn, layout.span().high_byte, input_ptr);
            HostTensorND ret;
            ret.reset(storage, layout);
            return std::make_pair(ret, true);
        } else {
            auto storage = HostTensorStorage(dest_cn);
            storage.ensure_size(layout.span().dist_byte());
            memcpy(storage.ptr(), data, layout.span().dist_byte());
            HostTensorND ret{dest_cn, layout.dtype};
            if (meth.dest_tensor_) {
                meth.dest_tensor_->reset(storage, layout);
                return std::make_pair(*meth.dest_tensor_, false);
            } else {
                HostTensorND ret;
                ret.reset(storage, layout);
                return std::make_pair(ret, false);
            }
        }
    };

    PyArrayObject* obj_as_arr = nullptr;
    do {
        // check contiguous and dtype, and borrow mem if ok
        if (!PyArray_Check(obj))
            break;
        obj_as_arr = reinterpret_cast<PyArrayObject*>(obj);
        int typenum = PyArray_DTYPE(obj_as_arr)->type_num;
        // We have to check dtype.valid() and typenum first to avoid
        // accidentally trigger ConversionError on incompatible dtypes which can
        // be automatically converted into comptaible ones (e.g. float64).
        if (dtype.valid() && (expected_descr->type_num != typenum ||
                              dtype_np2mgb_descr(PyArray_DTYPE(obj_as_arr)) != dtype))
            break;
        if (typenum != to_mgb_supported_dtype_raw(typenum)) {
            mgb_assert(!dtype.valid() && expected_descr == nullptr);
            expected_descr = PyArray_DescrFromType(to_mgb_supported_dtype_raw(typenum));
            break;
        }
        if (PyArray_ISCARRAY_RO(obj_as_arr)) {
            return make_from_arr(obj_as_arr, true);
        }
    } while (0);

    constexpr auto NP_FLAGS = NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_FORCECAST;
    PyObject* obj_cvt;
    if (obj_as_arr) {
        obj_cvt = PyArray_FromArray(obj_as_arr, expected_descr, NP_FLAGS);
    } else {
        obj_cvt = PyArray_FromAny(obj, expected_descr, 0, 0, NP_FLAGS, nullptr);
    }

    if (obj_cvt) {
        // convert to mgb supported dtype
        auto arr = reinterpret_cast<PyArrayObject*>(obj_cvt);
        int dt0 = PyArray_TYPE(arr), dt1 = to_mgb_supported_dtype_raw(dt0);
        if (dt0 != dt1) {
            mgb_assert(expected_descr == nullptr);
            expected_descr = PyArray_DescrFromType(dt1);
            mgb_assert(expected_descr);
            auto obj_cvt_new =
                    PyArray_FromAny(obj_cvt, expected_descr, 0, 0, NP_FLAGS, nullptr);
            Py_DECREF(obj_cvt);
            obj_cvt = obj_cvt_new;
        }
    }

    if (!obj_cvt) {
        if (PyErr_Occurred()) {
            PyExceptionForward::throw_();
        }
        throw ConversionError(ssprintf(
                "can not convert to numpy array from %s", repr_pyobj(obj).c_str()));
    }

    auto ret = make_from_arr(reinterpret_cast<PyArrayObject*>(obj_cvt), false);
    Py_DECREF(obj_cvt);
    return ret;
}

//! hold a reference to HostTensorND
class HostTensorNDRefHolder final : public NonCopyableObj {
    HostTensorND m_val;
    static MemPool<HostTensorNDRefHolder> sm_mem_pool;

    friend class MemPool<HostTensorNDRefHolder>;

    HostTensorNDRefHolder(const HostTensorND& v) : m_val{v} {}

public:
    static HostTensorNDRefHolder* alloc(const HostTensorND& v) {
        return sm_mem_pool.alloc(v);
    }

    static void free(HostTensorNDRefHolder* p) { return sm_mem_pool.free(p); }
};
MemPool<HostTensorNDRefHolder> HostTensorNDRefHolder::sm_mem_pool;

void ndarray_shared_from_tensor_py_capsule_dtor(PyObject* cap) {
    auto ptr = PyCapsule_GetPointer(cap, "HostTensorND");
    mgb_assert(ptr, "not a PyCapsule: %s", repr_pyobj(cap).c_str());
    HostTensorNDRefHolder::free(static_cast<HostTensorNDRefHolder*>(ptr));
}

PyObject* ndarray_from_tensor(const HostTensorND& val, ShareType share_type) {
    if (!val.layout().is_contiguous() && !val.shape().is_empty()) {
        mgb_assert(share_type != ShareType::MUST_SHARE);
        HostTensorND contig;
        contig.copy_from(val);
        return ndarray_from_tensor(contig, ShareType::TRY_SHARE);
    }
    PYTHON_GIL;
    npy_intp dims[TensorLayout::MAX_NDIM];
    for (size_t i = 0; i < val.layout().ndim; ++i)
        dims[i] = val.shape()[i];
    PyObject* ret = nullptr;

    auto alloc_new_ret = [&]() {
        mgb_assert(!ret);
        ret = PyArray_NewFromDescr(
                &PyArray_Type, dtype_mgb2np_descr(val.dtype()).release(),
                val.layout().ndim, dims, nullptr, nullptr, 0, nullptr);
        mgb_assert(ret, "failed to allocate array");
        mgb_assert(PyArray_Check(ret));
        return PyArray_DATA(reinterpret_cast<PyArrayObject*>(ret));
    };
    if (val.dtype().is_low_bit()) {
        mgb_assert(
                share_type != ShareType::MUST_SHARE,
                "can not share memory for lowbit dtype");
        const auto& layout = val.layout();
        if (layout.format.is_lowbit_aligned()) {
            lowbit_memcpy_aligned2byte(alloc_new_ret(), val.raw_ptr(), val.layout());
        } else {
            lowbit_memcpy_compact2byte(
                    val.dtype(), alloc_new_ret(), val.raw_ptr(),
                    val.layout().total_nr_elems());
        }
    } else if (share_type == ShareType::MUST_UNSHARE) {
        memcpy(alloc_new_ret(), val.raw_ptr(), val.layout().span().dist_byte());
    } else {
        // share data
        ret = PyArray_NewFromDescr(
                &PyArray_Type, dtype_mgb2np_descr(val.dtype()).release(),
                val.layout().ndim, dims, nullptr, const_cast<dt_byte*>(val.raw_ptr()),
                0, nullptr);
        mgb_assert(ret, "failed to alloc ndarray");
        auto capsule = PyCapsule_New(
                HostTensorNDRefHolder::alloc(val), "HostTensorND",
                ndarray_shared_from_tensor_py_capsule_dtor);
        mgb_assert(capsule, "failed to create PyCapsule");
        auto err =
                PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(ret), capsule);
        mgb_assert(!err);
    }
    return ret;
}

HostTensorND np2tensor(PyObject* obj, const Meth& meth, DType dtype) {
    auto ret_full = np2tensor_try_borrow(obj, meth, dtype);
    if (meth.must_borrow_) {
        mgb_assert(
                ret_full.second,
                "can not borrow from numpy array as contig array with dtype "
                "%s; src=%s",
                dtype.name(), repr_pyobj(obj).c_str());
    }
    return ret_full.first;
}

PyObject* dtype_mgb2np(mgb::DType dtype) {
    PYTHON_GIL;
    // According to
    // https://docs.scipy.org/doc/numpy/reference/c-api.array.html#c.PyArray_TypeObjectFromType
    // the following is equivalent to PyArray_TypeObjectFromType for built-in
    // types.
    if (!dtype.valid()) {
        Py_XINCREF(Py_None);
        return Py_None;
    }
    auto descr = dtype_mgb2np_descr(dtype);
    if (descr == nullptr) {
        Py_XINCREF(Py_None);
        return Py_None;
    }
    if (dtype.has_param()) {
        return reinterpret_cast<PyObject*>(descr.release());
    }
    PyObject* typeobj = reinterpret_cast<PyObject*>(descr->typeobj);
    Py_XINCREF(typeobj);
    return typeobj;
}

mgb::DType dtype_np2mgb(PyObject* obj) {
    mgb_assert(obj && obj != Py_None, "can not convert null PyObject to numpy dtype");
    // see
    // http://stackoverflow.com/questions/8477122/numpy-c-api-convert-type-object-to-type-number
    PYTHON_GIL;

    PyArray_Descr* dtype;
    if (!PyArray_DescrConverter(obj, &dtype)) {
        throw ConversionError(ssprintf(
                "can not convert to np.dtype from %s", repr_pyobj(obj).c_str()));
    }

    mgb::DType result = dtype_np2mgb_descr(dtype);
    Py_DECREF(dtype);
    return result;
}

PyObject* to_mgb_supported_dtype(PyObject* dtype) {
    PYTHON_GIL;

    PyArray_Descr* descr;
    if (!PyArray_DescrConverter(dtype, &descr)) {
        throw ConversionError(ssprintf(
                "can not convert to np.dtype from %s", repr_pyobj(dtype).c_str()));
    }
    mgb_assert(
            !descr->metadata,
            "unexpected metadata in dtype: "
            "dtype_obj=%s metadata=%s",
            repr_pyobj(dtype).c_str(), repr_pyobj(descr->metadata).c_str());
    int type_num = to_mgb_supported_dtype_raw(descr->type_num);
    return PyArray_TypeObjectFromType(type_num);
}

TensorShape vec2shape(const std::vector<size_t>& vec) {
    TensorShape shape;
    mgb_assert(
            vec.size() <= TensorShape::MAX_NDIM, "dim too large: %zd (max %zd)",
            vec.size(), TensorShape::MAX_NDIM);
    shape.ndim = vec.size();
    for (size_t i = 0; i < vec.size(); i++) {
        if (!vec[i]) {
            shape.ndim = 0;
            break;
        }
        shape[i] = vec[i];
    }
    mgb_assert(shape.ndim, "shape should not be empty");
    return shape;
}

}  // namespace npy