MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / lowbit_ndarray_to_host_tensor

Function lowbit_ndarray_to_host_tensor

imperative/python/src/helper.cpp:399–429  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

397}
398
399HostTensorND lowbit_ndarray_to_host_tensor(
400 CompNode comp_node, TensorLayout& layout, PyArrayObject* input) {
401 auto src_ptr = reinterpret_cast<dt_byte*>(PyArray_DATA(input));
402 if (!layout.ndim) {
403 // numpy scalar
404 mgb_assert(src_ptr, "can not convert from null numpy array");
405 layout.init_contiguous_stride({1});
406 } else {
407 mgb_assert(
408 layout.ndim && layout.ndim <= TensorShape::MAX_NDIM,
409 "unsupported ndim %zu", layout.ndim);
410 TensorLayout ly;
411 ly.ndim = layout.ndim;
412 for (size_t i = 0; i < layout.ndim; ++i) {
413 ly.shape[i] = layout.shape[i] = PyArray_SHAPE(input)[i];
414 ly.stride[i] = PyArray_STRIDE(input, i);
415 mgb_assert(layout.shape[i], "zero shape not supported");
416 }
417 mgb_assert(ly.is_physical_contiguous());
418 layout.init_contiguous_stride();
419 }
420 HostTensorND ret{comp_node, layout};
421 if (layout.format.is_lowbit_aligned()) {
422 mgb_assert(layout.is_contiguous());
423 lowbit_memcpy_byte2aligned(ret.raw_ptr(), src_ptr, layout);
424 } else {
425 lowbit_memcpy_byte2compact(
426 layout.dtype, ret.raw_ptr(), src_ptr, layout.total_nr_elems());
427 }
428 return ret;
429}
430
431/*!
432 * \brief convert a python object to tensor and try to borrow memory if the

Callers 1

np2tensor_try_borrowFunction · 0.85

Calls 6

is_lowbit_alignedMethod · 0.80
is_contiguousMethod · 0.45
raw_ptrMethod · 0.45
total_nr_elemsMethod · 0.45

Tested by

no test coverage detected