! * \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 */
| 435 | * \return the megbrain tensor, and whether memory is borrowed |
| 436 | */ |
| 437 | std::pair<HostTensorND, bool> np2tensor_try_borrow( |
| 438 | PyObject* obj, const npy::Meth& meth, DType dtype) { |
| 439 | auto dest_cn = meth.dest_cn_; |
| 440 | mgb_assert(dest_cn.valid()); |
| 441 | |
| 442 | PYTHON_GIL; |
| 443 | |
| 444 | PyArray_Descr* expected_descr = nullptr; |
| 445 | if (dtype.valid()) { |
| 446 | // The reference to expected_descr will be stealed later. |
| 447 | expected_descr = dtype_mgb2np_descr(dtype).release(); |
| 448 | } |
| 449 | |
| 450 | // make result from PyArrayObject; its reference may be stolen |
| 451 | auto make_from_arr = [&](PyArrayObject* input, bool allow_borrow) { |
| 452 | TensorLayout layout{{}, dtype_np2mgb_descr(PyArray_DESCR(input))}; |
| 453 | if (dtype.valid()) |
| 454 | mgb_assert(dtype == layout.dtype); |
| 455 | layout.ndim = PyArray_NDIM(input); |
| 456 | |
| 457 | if (layout.dtype.is_low_bit()) { |
| 458 | auto ret = lowbit_ndarray_to_host_tensor(dest_cn, layout, input); |
| 459 | if (meth.dest_tensor_) { |
| 460 | meth.dest_tensor_->copy_from(ret); |
| 461 | ret = *meth.dest_tensor_; |
| 462 | } |
| 463 | return std::make_pair(ret, false); |
| 464 | } |
| 465 | |
| 466 | auto data = reinterpret_cast<dt_byte*>(PyArray_DATA(input)); |
| 467 | if (!layout.ndim) { |
| 468 | // numpy scalar |
| 469 | mgb_assert(data, "can not convert from null numpy array"); |
| 470 | layout.init_contiguous_stride({1}); |
| 471 | } else { |
| 472 | mgb_assert( |
| 473 | layout.ndim && layout.ndim <= TensorShape::MAX_NDIM, |
| 474 | "unsupported ndim %zu", layout.ndim); |
| 475 | auto dsize = layout.dtype.size(); |
| 476 | bool is_empty = false; |
| 477 | for (size_t i = 0; i < layout.ndim; ++i) { |
| 478 | layout.shape[i] = PyArray_SHAPE(input)[i]; |
| 479 | layout.stride[i] = PyArray_STRIDE(input, i); |
| 480 | if (!layout.shape[i]) { |
| 481 | is_empty = true; |
| 482 | } |
| 483 | mgb_assert( |
| 484 | layout.stride[i] % dsize == 0, "bad stride %zd", |
| 485 | layout.stride[i]); |
| 486 | layout.stride[i] /= dsize; |
| 487 | } |
| 488 | mgb_assert(is_empty || layout.is_contiguous()); |
| 489 | } |
| 490 | |
| 491 | if (!meth.dest_tensor_ && allow_borrow) { |
| 492 | Py_INCREF(input); |
| 493 | PyObjRefKeeper ref_obj_cvt{reinterpret_cast<PyObject*>(input)}; |
| 494 | HostTensorStorage storage; |
no test coverage detected