This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch refactor-s2 in repository https://gitbox.apache.org/repos/asf/tvm.git
commit e82734e82cc88f25d7da91887e5c18551228178d Author: tqchen <[email protected]> AuthorDate: Wed Apr 23 09:11:15 2025 -0400 NDArray migrate to new mechanism --- include/tvm/runtime/container/shape_tuple.h | 1 + include/tvm/runtime/device_api.h | 18 +- include/tvm/runtime/memory/memory_manager.h | 6 - include/tvm/runtime/ndarray.h | 280 +++---------------------- python/tvm/_ffi/_cython/ndarray.pxi | 2 +- python/tvm/runtime/ndarray.py | 2 +- src/node/serialization.cc | 4 +- src/node/structural_equal.cc | 24 +-- src/node/structural_hash.cc | 18 +- src/runtime/contrib/random/mt_random_engine.cc | 4 +- src/runtime/memory/memory_manager.cc | 161 ++++++-------- src/runtime/ndarray.cc | 219 ++++--------------- src/runtime/opencl/opencl_device_api.cc | 3 +- src/runtime/relax_vm/vm.cc | 1 - src/runtime/rpc/rpc_module.cc | 39 ++-- 15 files changed, 194 insertions(+), 588 deletions(-) diff --git a/include/tvm/runtime/container/shape_tuple.h b/include/tvm/runtime/container/shape_tuple.h index 61f44d30be..c7a96b6623 100644 --- a/include/tvm/runtime/container/shape_tuple.h +++ b/include/tvm/runtime/container/shape_tuple.h @@ -35,6 +35,7 @@ namespace tvm { namespace runtime { +using Shape = tvm::ffi::Shape; using ShapeTuple = tvm::ffi::Shape; using ShapeTupleObj = tvm::ffi::ShapeObj; using IntTuple = ShapeTuple; diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h index 2ab9d32fcc..a4b53eb797 100644 --- a/include/tvm/runtime/device_api.h +++ b/include/tvm/runtime/device_api.h @@ -24,13 +24,18 @@ #ifndef TVM_RUNTIME_DEVICE_API_H_ #define TVM_RUNTIME_DEVICE_API_H_ +#include <tvm/ffi/any.h> +#include <tvm/ffi/optional.h> #include <tvm/runtime/c_runtime_api.h> -#include <tvm/runtime/ndarray.h> -#include <tvm/runtime/packed_func.h> +#include <tvm/runtime/logging.h> #include <string> namespace tvm { + +// alias DLDevice +using Device = DLDevice; + namespace runtime { /*! * \brief the query type into GetAttr @@ -96,7 +101,7 @@ class TVM_DLL DeviceAPI { * \param rv The return value. * \sa DeviceAttrKind */ - virtual void GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv) = 0; + virtual void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) = 0; /*! * \brief Get the physical memory size required. @@ -104,7 +109,8 @@ class TVM_DLL DeviceAPI { * \param mem_scope the memory scope if any * \return the memory size. */ - virtual size_t GetDataSize(const DLTensor& arr, Optional<String> mem_scope = NullOpt); + virtual size_t GetDataSize(const DLTensor& arr, + ffi::Optional<ffi::String> mem_scope = std::nullopt); /*! * \brief Query the device for specified properties. @@ -112,7 +118,7 @@ class TVM_DLL DeviceAPI { * This is used to expand "-from_device=N" in the target string to * all properties that can be determined from that device. */ - virtual void GetTargetProperty(Device dev, const std::string& property, TVMRetValue* rv) {} + virtual void GetTargetProperty(Device dev, const std::string& property, ffi::Any* rv) {} /*! * \brief Allocate a data space on device. @@ -135,7 +141,7 @@ class TVM_DLL DeviceAPI { * \return The allocated device pointer. */ virtual void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType dtype, - Optional<String> mem_scope = NullOpt); + ffi::Optional<ffi::String> mem_scope = std::nullopt); /*! * \brief Free a data space on device. * \param dev The device device to perform operation. diff --git a/include/tvm/runtime/memory/memory_manager.h b/include/tvm/runtime/memory/memory_manager.h index 3def5224ed..537beeb8fa 100644 --- a/include/tvm/runtime/memory/memory_manager.h +++ b/include/tvm/runtime/memory/memory_manager.h @@ -170,12 +170,6 @@ class StorageObj : public Object { TVM_DLL NDArray AllocNDArrayScoped(int64_t offset, ShapeTuple shape, DLDataType dtype, String scope = "global"); - /*! \brief The deleter for an NDArray when allocated from underlying storage. */ - static void ScopedDeleter(TVMFFIObject* ptr); - - /*! \brief The deleter for an NDArray when allocated from underlying storage. */ - static void Deleter(TVMFFIObject* ptr); - ~StorageObj() { if (allocator) { allocator->Free(buffer); diff --git a/include/tvm/runtime/ndarray.h b/include/tvm/runtime/ndarray.h index 227e319a7e..82c9b229ab 100644 --- a/include/tvm/runtime/ndarray.h +++ b/include/tvm/runtime/ndarray.h @@ -24,11 +24,13 @@ #ifndef TVM_RUNTIME_NDARRAY_H_ #define TVM_RUNTIME_NDARRAY_H_ +#include <tvm/ffi/container/ndarray.h> #include <tvm/runtime/c_runtime_api.h> #include <tvm/runtime/container/optional.h> #include <tvm/runtime/container/shape_tuple.h> #include <tvm/runtime/container/string.h> #include <tvm/runtime/data_type.h> +#include <tvm/runtime/device_api.h> #include <tvm/runtime/object.h> #include <tvm/runtime/serializer.h> @@ -38,43 +40,39 @@ #include <vector> namespace tvm { - -// alias DLDevice -using Device = DLDevice; - namespace runtime { +using ffi::GetDataSize; +using ffi::IsAligned; +using ffi::IsContiguous; + /*! * \brief Managed NDArray. * The array is backed by reference counted blocks. */ -class NDArray : public ObjectRef { +class NDArray : public tvm::ffi::NDArray { public: - /*! \brief ContainerBase used to back the TVMArrayHandle */ - class ContainerBase; - /*! \brief NDArray internal container type */ - class Container; - /*! \brief Container type for Object system. */ - using ContainerType = Container; - /*! \brief default constructor */ - NDArray() {} + using Container = ffi::NDArrayObj; + NDArray() = default; /*! * \brief constructor. * \param data ObjectPtr to the data container. */ - explicit NDArray(ObjectPtr<Object> data) : ObjectRef(data) {} + explicit NDArray(ObjectPtr<Object> data) : tvm::ffi::NDArray(data) {} + NDArray(ffi::NDArray&& other) : tvm::ffi::NDArray(std::move(other)) {} // NOLINT(*) + NDArray(const ffi::NDArray& other) : tvm::ffi::NDArray(other) {} // NOLINT(*) - /*! \brief reset the content of NDArray to be nullptr */ - inline void reset(); - /*! - * \return the reference counter - * \note this number is approximate in multi-threaded setting. - */ - inline int use_count() const; - /*! \return Pointer to content of DLTensor */ - inline const DLTensor* operator->() const; - /*! \return Whether the tensor is contiguous */ - inline bool IsContiguous() const; + ShapeTuple Shape() const { return this->shape(); } + runtime::DataType DataType() const { return runtime::DataType(this->dtype()); } + + // DLPack handling + static NDArray FromDLPack(DLManagedTensor* tensor) { + return tvm::ffi::NDArray::FromDLPack(tensor, kAllocAlignment, true); + } + + static NDArray FromDLPackVersioned(DLManagedTensorVersioned* tensor) { + return tvm::ffi::NDArray::FromDLPackVersioned(tensor, kAllocAlignment, true); + } /*! * \brief Copy data content from another array. * \param other The source array to be copied from. @@ -149,13 +147,6 @@ class NDArray : public ObjectRef { */ TVM_DLL NDArray CreateView(ShapeTuple shape, DLDataType dtype, uint64_t relative_byte_offset = 0) const; - - /*! - * \brief Create a reference view of NDArray that - * represents as DLManagedTensor. - * \return A DLManagedTensor - */ - TVM_DLL DLManagedTensor* ToDLPack() const; /*! * \brief Create an empty NDArray. * \param shape The shape of the new array. @@ -166,37 +157,6 @@ class NDArray : public ObjectRef { */ TVM_DLL static NDArray Empty(ShapeTuple shape, DLDataType dtype, Device dev, Optional<String> mem_scope = NullOpt); - /*! - * \brief Create a NDArray backed by an external DLTensor without memory copying. - * - * If DLTensor is not contiguous or has bad aligned data, It fails. - * This allows us to create a NDArray using the memory - * allocated by an external source. Responsibility for memory - * retaining lies with the external source. - * \param dl_tensor The DLTensor for NDArray base. - * \return The created NDArray view. - */ - TVM_DLL static NDArray FromExternalDLTensor(const DLTensor& dl_tensor); - /*! - * \brief Create new NDArray, data is copied from DLTensor. - * - * \param dl_tensor The DLTensor to copy from. - * \param dev device location of the created NDArray. - * \return The created NDArray view. - */ - TVM_DLL static NDArray NewFromDLTensor(DLTensor* dl_tensor, const Device& dev); - /*! - * \brief Create a NDArray backed by a dlpack tensor. - * - * This allows us to create a NDArray using the memory - * allocated by an external deep learning framework - * that is DLPack compatible. - * - * The memory is retained until the NDArray went out of scope. - * \param tensor The DLPack tensor to copy from. - * \return The created NDArray view. - */ - TVM_DLL static NDArray FromDLPack(DLManagedTensor* tensor); /*! * \brief Function to copy data from one array to another. * \param from The source array. @@ -206,42 +166,9 @@ class NDArray : public ObjectRef { TVM_DLL static void CopyFromTo(const DLTensor* from, DLTensor* to, TVMStreamHandle stream = nullptr); - TVM_DLL ShapeTuple Shape() const; - TVM_DLL runtime::DataType DataType() const; - /*! - * \brief Check conditions for construction NDArray over DLTensor without copying. - * There are three conditions to check: - * 1. Destination device is the same as DLTensor device - * 2. Destination device id is the same as DLTensor device id - * 3. Memory in DLTensor is aligned as expected for NDArray - * \param tensor the DLTensor. - * \param dev destination device. - * \return true if all conditions are satisfied. - */ - TVM_DLL static bool AbilityOfZeroCopyForDLTensor(DLTensor* tensor, const Device& dev); - // internal namespace struct Internal; - private: - TVM_DLL static bool IsAligned(const DLTensor& tensor); - protected: - /*! - * \brief Get mutable internal container pointer. - * \return a mutable container pointer. - */ - inline Container* get_mutable() const; - // Helper functions for FFI handling. - /*! - * \brief Construct NDArray's Data field from array handle in FFI. - * \param handle The array handle. - * \return The corresponding ObjectPtr to the constructed container object. - * - * \note We keep a special calling convention for NDArray by passing - * ContainerBase pointer in FFI. - * As a result, the argument is compatible to DLTensor*. - */ - inline static ObjectPtr<Object> FFIDataFromHandle(TVMArrayHandle handle); /*! * \brief DecRef resource managed by an FFI array handle. * \param handle The array handle. @@ -262,187 +189,44 @@ class NDArray : public ObjectRef { */ inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor); -/*! - * \brief The container base structure - * contains all the fields except for the Object header. - * - * \note We explicitly declare this structure in order to pass - * PackedFunc argument using ContainerBase*. - */ -class NDArray::ContainerBase { - public: - /*! - * \brief The corresponding dl_tensor field. - * \note it is important that the first field is DLTensor - * So that this data structure is DLTensor compatible. - * The head ptr of this struct can be viewed as DLTensor*. - */ - DLTensor dl_tensor; - - /*! - * \brief additional context, reserved for recycling - * \note We can attach additional content here - * which the current container depend on - * (e.g. reference to original memory when creating views). - */ - void* manager_ctx{nullptr}; - - protected: - /*! - * \brief The shape container, - * can be used for shape data. - */ - ShapeTuple shape_; -}; - -/*! - * \brief Object container class that backs NDArray. - * \note do not use this function directly, use NDArray. - */ -class NDArray::Container : public Object, public NDArray::ContainerBase { - public: - /*! \brief default constructor */ - Container() { - // Initialize the type index. - header_.type_index = Container::RuntimeTypeIndex(); - dl_tensor.data = nullptr; - dl_tensor.ndim = 0; - dl_tensor.shape = nullptr; - dl_tensor.strides = nullptr; - dl_tensor.byte_offset = 0; - } - - Container(void* data, ShapeTuple shape, DLDataType dtype, Device dev) { - // Initialize the type index. - header_.type_index = Container::RuntimeTypeIndex(); - dl_tensor.data = data; - shape_ = std::move(shape); - dl_tensor.ndim = static_cast<int>(shape_.size()); - dl_tensor.shape = const_cast<ShapeTuple::index_type*>(shape_.data()); - dl_tensor.dtype = dtype; - dl_tensor.strides = nullptr; - dl_tensor.byte_offset = 0; - dl_tensor.device = dev; - } - /*! - * \brief Set the deleter field. - * \param deleter The deleter. - */ - void SetDeleter(ffi::FObjectDeleter deleter) { header_.deleter = deleter; } - - // Expose DecRef and IncRef as public function - // NOTE: they are only for developer purposes only. - // using Object::DecRef; - // using Object::IncRef; - - // Information for object protocol. - static constexpr const uint32_t _type_index = ffi::TypeIndex::kTVMFFINDArray; - static constexpr const uint32_t _type_child_slots = 0; - static constexpr const uint32_t _type_child_slots_can_overflow = true; - static constexpr const char* _type_key = "runtime.NDArray"; - static const constexpr bool _type_final = true; - TVM_FFI_DECLARE_STATIC_OBJECT_INFO(NDArray::Container, Object); - - protected: - friend class RPCWrappedFunc; - friend class NDArray; -}; - -// implementations of inline functions -/*! - * \brief return the size of data the DLTensor hold, in term of number of bytes - * - * \param arr the input DLTensor - * \return number of bytes of data in the DLTensor. - */ -inline size_t GetDataSize(const DLTensor& arr) { - size_t size = 1; - for (tvm_index_t i = 0; i < arr.ndim; ++i) { - size *= static_cast<size_t>(arr.shape[i]); - } - size *= (arr.dtype.bits * arr.dtype.lanes + 7) / 8; - return size; -} - -/*! - * \brief check if a DLTensor is contiguous. - * \param arr The input DLTensor. - * \return The check result. - */ -static inline bool IsContiguous(const DLTensor& arr) { - if (arr.strides == nullptr) return true; - int64_t expected_stride = 1; - for (int32_t i = arr.ndim; i != 0; --i) { - int32_t k = i - 1; - if (arr.shape[k] == 1) { - // Skip stride check if shape[k] is 1, where the dimension is contiguous - // regardless of the value of stride. - // - // For example, PyTorch will normalize stride to 1 if shape is 1 when exporting - // to DLPack. - // More context: https://github.com/pytorch/pytorch/pull/83158 - continue; - } - if (arr.strides[k] != expected_stride) return false; - expected_stride *= arr.shape[k]; - } - return true; -} - -inline bool NDArray::IsContiguous() const { - return ::tvm::runtime::IsContiguous(get_mutable()->dl_tensor); -} - inline void NDArray::CopyFrom(const DLTensor* other) { ICHECK(data_ != nullptr); - CopyFromTo(other, &(get_mutable()->dl_tensor)); + CopyFromTo(other, get_mutable()); } inline void NDArray::CopyFrom(const NDArray& other) { ICHECK(data_ != nullptr); ICHECK(other.data_ != nullptr); - CopyFromTo(&(other.get_mutable()->dl_tensor), &(get_mutable()->dl_tensor)); + CopyFromTo(other.get_mutable(), get_mutable()); } inline void NDArray::CopyTo(DLTensor* other) const { ICHECK(data_ != nullptr); - CopyFromTo(&(get_mutable()->dl_tensor), other); + CopyFromTo(get_mutable(), other); } inline void NDArray::CopyTo(const NDArray& other) const { ICHECK(data_ != nullptr); ICHECK(other.data_ != nullptr); - CopyFromTo(&(get_mutable()->dl_tensor), &(other.get_mutable()->dl_tensor)); -} - -inline int NDArray::use_count() const { return data_.use_count(); } - -inline const DLTensor* NDArray::operator->() const { return &(get_mutable()->dl_tensor); } - -inline NDArray::Container* NDArray::get_mutable() const { - return static_cast<NDArray::Container*>(data_.get()); -} - -inline ObjectPtr<Object> NDArray::FFIDataFromHandle(TVMArrayHandle handle) { - return GetObjectPtr<Object>( - static_cast<NDArray::Container*>(reinterpret_cast<NDArray::ContainerBase*>(handle))); + CopyFromTo(get_mutable(), other.get_mutable()); } inline TVMArrayHandle NDArray::FFIGetHandle(const ObjectRef& nd) { // NOTE: it is necessary to cast to container then to base // so that the FFI handle uses the ContainerBase address. - auto ptr = reinterpret_cast<TVMArrayHandle>(static_cast<NDArray::ContainerBase*>( - static_cast<NDArray::Container*>(const_cast<Object*>(nd.get())))); + auto ptr = reinterpret_cast<TVMArrayHandle>( + TVMFFINDArrayGetDLTensorPtr(static_cast<ffi::NDArrayObj*>(const_cast<Object*>(nd.get())))); return ptr; } inline TVMArrayHandle ObjectHandleToTVMArrayHandle(Object* handle) { return reinterpret_cast<TVMArrayHandle>( - static_cast<NDArray::ContainerBase*>(static_cast<NDArray::Container*>(handle))); + TVMFFINDArrayGetDLTensorPtr(static_cast<ffi::NDArrayObj*>(handle))); } inline Object* TVMArrayHandleToObjectHandle(void* handle) { - return static_cast<NDArray::Container*>(reinterpret_cast<NDArray::ContainerBase*>(handle)); + // NOTE: legacy patch here for TFM FFI + return reinterpret_cast<ffi::NDArrayObj*>(reinterpret_cast<char*>(handle) - sizeof(TVMFFIObject)); } inline void NDArray::FFIDecRef(TVMArrayHandle handle) { diff --git a/python/tvm/_ffi/_cython/ndarray.pxi b/python/tvm/_ffi/_cython/ndarray.pxi index b88698319f..f220e866d3 100644 --- a/python/tvm/_ffi/_cython/ndarray.pxi +++ b/python/tvm/_ffi/_cython/ndarray.pxi @@ -26,7 +26,7 @@ cdef void _c_dlpack_deleter(object pycaps): cdef DLManagedTensor* dltensor if pycapsule.PyCapsule_IsValid(pycaps, _c_str_dltensor): dltensor = <DLManagedTensor*>pycapsule.PyCapsule_GetPointer(pycaps, _c_str_dltensor) - TVMDLManagedTensorCallDeleter(dltensor) + dltensor.deleter(dltensor) def _from_dlpack(object dltensor): diff --git a/python/tvm/runtime/ndarray.py b/python/tvm/runtime/ndarray.py index d001b671fc..eb72a01113 100644 --- a/python/tvm/runtime/ndarray.py +++ b/python/tvm/runtime/ndarray.py @@ -46,7 +46,7 @@ from tvm._ffi._cy3.core import ( from . import _ffi_api -@tvm._ffi.register_object("runtime.NDArray") +@tvm._ffi.register_object("object.NDArray") class NDArray(NDArrayBase): """Lightweight NDArray class of TVM runtime. diff --git a/src/node/serialization.cc b/src/node/serialization.cc index e98e27ea3a..51625aca36 100644 --- a/src/node/serialization.cc +++ b/src/node/serialization.cc @@ -82,7 +82,7 @@ class NodeIndexer : public AttrVisitor { void Visit(const char* key, DataType* value) final {} void Visit(const char* key, runtime::NDArray* value) final { - DLTensor* ptr = const_cast<DLTensor*>((*value).operator->()); + DLTensor* ptr = const_cast<ffi::NDArrayObj*>((*value).operator->()); if (tensor_index_.count(ptr)) return; ICHECK_EQ(tensor_index_.size(), tensor_list_.size()); tensor_index_[ptr] = tensor_list_.size(); @@ -238,7 +238,7 @@ class JSONAttrGetter : public AttrVisitor { void Visit(const char* key, DataType* value) final { node_->attrs[key] = Type2String(*value); } void Visit(const char* key, runtime::NDArray* value) final { node_->attrs[key] = - std::to_string(tensor_index_->at(const_cast<DLTensor*>((*value).operator->()))); + std::to_string(tensor_index_->at(const_cast<ffi::NDArrayObj*>((*value).operator->()))); } void Visit(const char* key, Optional<int64_t>* value) final { diff --git a/src/node/structural_equal.cc b/src/node/structural_equal.cc index a00040d4cd..260a02da8f 100644 --- a/src/node/structural_equal.cc +++ b/src/node/structural_equal.cc @@ -624,21 +624,21 @@ bool NDArrayEqual(const runtime::NDArray::Container* lhs, const runtime::NDArray SEqualReducer equal, bool compare_data) { if (lhs == rhs) return true; - auto ldt = lhs->dl_tensor.dtype; - auto rdt = rhs->dl_tensor.dtype; - ICHECK_EQ(lhs->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor"; - ICHECK_EQ(rhs->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor"; - ICHECK(runtime::IsContiguous(lhs->dl_tensor)) << "Can only compare contiguous tensor"; - ICHECK(runtime::IsContiguous(rhs->dl_tensor)) << "Can only compare contiguous tensor"; - - if (lhs->dl_tensor.ndim != rhs->dl_tensor.ndim) return false; - for (int i = 0; i < lhs->dl_tensor.ndim; ++i) { - if (!equal(lhs->dl_tensor.shape[i], rhs->dl_tensor.shape[i])) return false; + auto ldt = lhs->dtype; + auto rdt = rhs->dtype; + ICHECK_EQ(lhs->device.device_type, kDLCPU) << "can only compare CPU tensor"; + ICHECK_EQ(rhs->device.device_type, kDLCPU) << "can only compare CPU tensor"; + ICHECK(runtime::IsContiguous(*lhs)) << "Can only compare contiguous tensor"; + ICHECK(runtime::IsContiguous(*rhs)) << "Can only compare contiguous tensor"; + + if (lhs->ndim != rhs->ndim) return false; + for (int i = 0; i < lhs->ndim; ++i) { + if (!equal(lhs->shape[i], rhs->shape[i])) return false; } if (ldt.code == rdt.code && ldt.lanes == rdt.lanes && ldt.bits == rdt.bits) { - size_t data_size = runtime::GetDataSize(lhs->dl_tensor); + size_t data_size = runtime::GetDataSize(*lhs); if (compare_data) { - return std::memcmp(lhs->dl_tensor.data, rhs->dl_tensor.data, data_size) == 0; + return std::memcmp(lhs->data, rhs->data, data_size) == 0; } else { return true; } diff --git a/src/node/structural_hash.cc b/src/node/structural_hash.cc index 50b8cb4e9b..bd9f39d617 100644 --- a/src/node/structural_hash.cc +++ b/src/node/structural_hash.cc @@ -368,17 +368,17 @@ TVM_REGISTER_REFLECTION_VTABLE(runtime::ModuleNode, ModuleNodeTrait) void NDArrayHash(const runtime::NDArray::Container* arr, SHashReducer* hash_reduce, bool hash_data) { - ICHECK_EQ(arr->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor"; - ICHECK(runtime::IsContiguous(arr->dl_tensor)) << "Can only hash contiguous tensor"; - (*hash_reduce)(runtime::DataType(arr->dl_tensor.dtype)); - (*hash_reduce)(arr->dl_tensor.ndim); - for (int i = 0; i < arr->dl_tensor.ndim; ++i) { - (*hash_reduce)(arr->dl_tensor.shape[i]); + ICHECK_EQ(arr->device.device_type, kDLCPU) << "can only compare CPU tensor"; + ICHECK(runtime::IsContiguous(*arr)) << "Can only hash contiguous tensor"; + (*hash_reduce)(runtime::DataType(arr->dtype)); + (*hash_reduce)(arr->ndim); + for (int i = 0; i < arr->ndim; ++i) { + (*hash_reduce)(arr->shape[i]); } if (hash_data) { (*hash_reduce) - ->SHashReduceHashedValue(ffi::details::StableHashBytes( - static_cast<const char*>(arr->dl_tensor.data), runtime::GetDataSize(arr->dl_tensor))); + ->SHashReduceHashedValue(ffi::details::StableHashBytes(static_cast<const char*>(arr->data), + runtime::GetDataSize(*arr))); } } @@ -401,7 +401,7 @@ TVM_REGISTER_REFLECTION_VTABLE(runtime::NDArray::Container, NDArrayContainerTrai dmlc::MemoryStringStream mstrm(&blob); support::Base64OutStream b64strm(&mstrm); const auto* ndarray = static_cast<const runtime::NDArray::Container*>(n); - runtime::SaveDLTensor(&b64strm, &ndarray->dl_tensor); + runtime::SaveDLTensor(&b64strm, ndarray); b64strm.Finish(); return blob; }); diff --git a/src/runtime/contrib/random/mt_random_engine.cc b/src/runtime/contrib/random/mt_random_engine.cc index dc01114af0..04b53d74b4 100644 --- a/src/runtime/contrib/random/mt_random_engine.cc +++ b/src/runtime/contrib/random/mt_random_engine.cc @@ -124,7 +124,7 @@ class RandomEngine { } else { runtime::NDArray local = runtime::NDArray::Empty( std::vector<int64_t>{data->shape, data->shape + data->ndim}, data->dtype, {kDLCPU, 0}); - DLTensor* tensor = const_cast<DLTensor*>(local.operator->()); + DLTensor* tensor = const_cast<ffi::NDArrayObj*>(local.operator->()); FillData(tensor); runtime::NDArray::CopyFromTo(tensor, data); } @@ -136,7 +136,7 @@ class RandomEngine { } else { runtime::NDArray local = runtime::NDArray::Empty( std::vector<int64_t>{data->shape, data->shape + data->ndim}, data->dtype, {kDLCPU, 0}); - DLTensor* tensor = const_cast<DLTensor*>(local.operator->()); + DLTensor* tensor = const_cast<ffi::NDArrayObj*>(local.operator->()); FillDataForMeasure(tensor); runtime::NDArray::CopyFromTo(tensor, data); } diff --git a/src/runtime/memory/memory_manager.cc b/src/runtime/memory/memory_manager.cc index 97fcf0534f..60fd5db5b1 100644 --- a/src/runtime/memory/memory_manager.cc +++ b/src/runtime/memory/memory_manager.cc @@ -34,15 +34,6 @@ namespace tvm { namespace runtime { namespace memory { -static void BufferDeleter(TVMFFIObject* ptr_obj) { - auto* ptr = ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned<NDArray::Container>(ptr_obj); - ICHECK(ptr->manager_ctx != nullptr); - Buffer* buffer = reinterpret_cast<Buffer*>(ptr->manager_ctx); - MemoryManager::GetAllocator(buffer->device, buffer->alloc_type)->Free(*(buffer)); - delete buffer; - delete ptr; -} - Storage::Storage(Buffer buffer, Allocator* allocator) { auto n = make_object<StorageObj>(); n->buffer = std::move(buffer); @@ -50,23 +41,6 @@ Storage::Storage(Buffer buffer, Allocator* allocator) { data_ = std::move(n); } -void StorageObj::Deleter(TVMFFIObject* ptr_obj) { - auto* ptr = ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned<NDArray::Container>(ptr_obj); - // When invoking AllocNDArray we don't own the underlying allocation - // and should not delete the buffer, but instead let it be reclaimed - // by the storage object's destructor. - // - // We did bump the reference count by 1 to keep alive the StorageObj - // allocation in case this NDArray is the sole owner. - // - // We decrement the object allowing for the buffer to release our - // reference count from allocation. - StorageObj* storage = reinterpret_cast<StorageObj*>(ptr->manager_ctx); - // storage->DecRef(); - tvm::ffi::details::ObjectUnsafe::DecRefObjectHandle(storage); - delete ptr; -} - inline void VerifyDataType(DLDataType dtype) { ICHECK_GE(dtype.lanes, 1); if (dtype.code == kDLFloat) { @@ -79,81 +53,74 @@ inline void VerifyDataType(DLDataType dtype) { ICHECK_EQ(dtype.bits & (dtype.bits - 1), 0); } -inline size_t GetDataAlignment(const DLTensor& arr) { - size_t align = (arr.dtype.bits / 8) * arr.dtype.lanes; +inline size_t GetDataAlignment(const DLDataType& dtype) { + size_t align = dtype.lanes * dtype.bits / 8; if (align < kAllocAlignment) return kAllocAlignment; return align; } -void StorageObj::ScopedDeleter(TVMFFIObject* ptr_obj) { - auto* ptr = ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned<NDArray::Container>(ptr_obj); - StorageObj* storage = reinterpret_cast<StorageObj*>(ptr->manager_ctx); - - // Let the device handle proper cleanup of view - storage->allocator->FreeView(ptr->dl_tensor.device, ptr->dl_tensor.data); - // storage->DecRef(); - tvm::ffi::details::ObjectUnsafe::DecRefObjectHandle(storage); - delete ptr; -} - NDArray StorageObj::AllocNDArrayScoped(int64_t offset, ShapeTuple shape, DLDataType dtype, String scope) { if (scope == "global" || scope.empty()) { return AllocNDArray(offset, shape, dtype); } VerifyDataType(dtype); - void* data = this->allocator->CreateView(this->buffer, shape, dtype, scope); - NDArray::Container* container = new NDArray::Container(data, shape, dtype, this->buffer.device); - container->dl_tensor.byte_offset = offset; - container->SetDeleter(StorageObj::ScopedDeleter); - size_t needed_size = DeviceAPI::Get(this->buffer.device)->GetDataSize(container->dl_tensor); - // this->IncRef(); - tvm::ffi::details::ObjectUnsafe::IncRefObjectHandle(this); - container->manager_ctx = reinterpret_cast<void*>(this); - NDArray ret(GetObjectPtr<Object>(container)); - // RAII in effect, now run the check. + + struct StorageScopedAlloc { + public: + StorageScopedAlloc(Storage storage) : storage_(storage) {} + + void AllocData(DLTensor* tensor, const ffi::Shape& shape, const String& scope, + int64_t byte_offset) { + tensor->data = storage_->allocator->CreateView(storage_->buffer, shape, tensor->dtype, scope); + tensor->byte_offset = byte_offset; + } + void FreeData(DLTensor* tensor) { storage_->allocator->FreeView(tensor->device, tensor->data); } + + private: + Storage storage_; + }; + + size_t needed_size = ffi::GetPackedDataSize(shape.Product(), dtype); ICHECK(offset + needed_size <= this->buffer.size) << "storage allocation failure, attempted to allocate " << needed_size << " at offset " << offset << " in region that is " << this->buffer.size << "bytes"; - return ret; + + return NDArray::FromNDAlloc(StorageScopedAlloc(GetRef<Storage>(this)), shape, dtype, + this->buffer.device, shape, scope, offset); } NDArray StorageObj::AllocNDArray(int64_t offset, ShapeTuple shape, DLDataType dtype) { VerifyDataType(dtype); - // crtical zone: allocate header, cannot throw - NDArray::Container* container = - new NDArray::Container(this->buffer.data, shape, dtype, this->buffer.device); - container->dl_tensor.byte_offset = offset; - - container->SetDeleter(StorageObj::Deleter); - size_t needed_size = DeviceAPI::Get(this->buffer.device)->GetDataSize(container->dl_tensor); - // this->IncRef(); - tvm::ffi::details::ObjectUnsafe::IncRefObjectHandle(this); - // The manager context pointer must continue to point to the storage object - // which owns the backing memory, and keeps track of the reference count. - // - // When we free a container we extract the storage object, decrement its - // reference count, then destroy the container, but leave the underlying - // buffer intact. - container->manager_ctx = reinterpret_cast<void*>(this); - - if (this->buffer.device.device_type == kDLHexagon) { - // For Hexagon, non-zero offset support simply requires adjusting the - // beginning of data pointer - auto offset_ptr = reinterpret_cast<uint8_t*>(this->buffer.data) + offset; - container->dl_tensor.data = reinterpret_cast<void*>(offset_ptr); - container->dl_tensor.byte_offset = 0; - } - - NDArray ret(GetObjectPtr<Object>(container)); - // RAII in effect, now run the check. - + size_t needed_size = ffi::GetPackedDataSize(shape.Product(), dtype); ICHECK(offset + needed_size <= this->buffer.size) << "storage allocation failure, attempted to allocate " << needed_size << " at offset " << offset << " in region that is " << this->buffer.size << "bytes"; + struct StorageAlloc { + public: + StorageAlloc(Storage storage) : storage_(storage) {} + + void AllocData(DLTensor* tensor, int64_t offset) { + if (storage_->buffer.device.device_type == kDLHexagon) { + // For Hexagon, non-zero offset support simply requires adjusting the + // beginning of data pointer + auto offset_ptr = reinterpret_cast<uint8_t*>(storage_->buffer.data) + offset; + tensor->data = reinterpret_cast<void*>(offset_ptr); + tensor->byte_offset = 0; + } else { + tensor->data = storage_->buffer.data; + tensor->byte_offset = offset; + } + } + void FreeData(DLTensor* tensor) {} + + private: + Storage storage_; + }; - return ret; + return NDArray::FromNDAlloc(StorageAlloc(GetRef<Storage>(this)), shape, dtype, + this->buffer.device, offset); } MemoryManager* MemoryManager::Global() { @@ -248,19 +215,30 @@ void MemoryManager::Clear() { NDArray Allocator::Empty(ShapeTuple shape, DLDataType dtype, DLDevice dev, Optional<String> mem_scope) { VerifyDataType(dtype); - NDArray::Container* container = new NDArray::Container(nullptr, shape, dtype, dev); - container->SetDeleter(BufferDeleter); - size_t size = DeviceAPI::Get(dev)->GetDataSize(container->dl_tensor, mem_scope); - size_t alignment = GetDataAlignment(container->dl_tensor); - Buffer* buffer = new Buffer; + + struct BufferAlloc { + public: + BufferAlloc(Buffer buffer) : buffer_(buffer) {} + + void AllocData(DLTensor* tensor) { tensor->data = buffer_.data; } + void FreeData(DLTensor* tensor) { + MemoryManager::GetAllocator(buffer_.device, buffer_.alloc_type)->Free(buffer_); + } + + private: + Buffer buffer_; + }; + + size_t alignment = GetDataAlignment(dtype); + size_t size = ffi::GetPackedDataSize(shape.Product(), dtype); + + Buffer buffer; if (!mem_scope.defined() || mem_scope.value().empty() || mem_scope.value() == "global") { - *buffer = this->Alloc(dev, size, alignment, dtype); + buffer = this->Alloc(dev, size, alignment, dtype); } else { - *buffer = this->Alloc(dev, shape, dtype, mem_scope.value()); + buffer = this->Alloc(dev, shape, dtype, mem_scope.value()); } - container->manager_ctx = reinterpret_cast<void*>(buffer); - container->dl_tensor.data = buffer->data; - return NDArray(GetObjectPtr<Object>(container)); + return NDArray::FromNDAlloc(BufferAlloc(buffer), shape, dtype, dev); } bool Allocator::AllowMemoryScope(const std::string& mem_scope) const { @@ -271,9 +249,8 @@ Buffer Allocator::Alloc(Device dev, ShapeTuple shape, DLDataType type_hint, const std::string& mem_scope) { if (AllowMemoryScope(mem_scope)) { // by default, we can always redirect to the flat memory allocations - NDArray::Container container(nullptr, shape, type_hint, dev); - size_t size = DeviceAPI::Get(dev)->GetDataSize(container.dl_tensor); - size_t alignment = GetDataAlignment(container.dl_tensor); + size_t alignment = GetDataAlignment(type_hint); + size_t size = ffi::GetPackedDataSize(shape.Product(), type_hint); return Alloc(dev, size, alignment, type_hint); } LOG(FATAL) << "Allocator cannot allocate data space with " diff --git a/src/runtime/ndarray.cc b/src/runtime/ndarray.cc index ae7266e29e..379437b103 100644 --- a/src/runtime/ndarray.cc +++ b/src/runtime/ndarray.cc @@ -30,13 +30,6 @@ #include "runtime_base.h" #include "tvm/runtime/data_type.h" -extern "C" { -// C-mangled dlpack deleter. -static void TVMNDArrayDLPackDeleter(DLManagedTensor* tensor); -// helper function to get NDArray's type index, only used by ctypes. -TVM_DLL int TVMArrayGetTypeIndex(TVMArrayHandle handle, unsigned* out_tindex); -} - namespace tvm { namespace runtime { @@ -99,61 +92,21 @@ void ArrayCopyToBytes(const DLTensor* handle, void* data, size_t nbytes) { DeviceAPI::Get(handle->device)->StreamSync(handle->device, nullptr); } -struct NDArray::Internal { - // Default deleter for the container - static void DefaultDeleter(TVMFFIObject* ptr_obj) { - auto* ptr = ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned<NDArray::Container>(ptr_obj); - if (ptr->manager_ctx != nullptr) { - ffi::details::ObjectUnsafe::DecRefObjectHandle( - static_cast<NDArray::Container*>(ptr->manager_ctx)); - } else if (ptr->dl_tensor.data != nullptr) { - tvm::runtime::DeviceAPI::Get(ptr->dl_tensor.device) - ->FreeDataSpace(ptr->dl_tensor.device, ptr->dl_tensor.data); +NDArray NDArray::Empty(ShapeTuple shape, DLDataType dtype, Device dev, Optional<String> mem_scope) { + struct DeviceAPIAlloc { + void AllocData(DLTensor* tensor, ffi::Optional<ffi::String> mem_scope) { + tensor->data = DeviceAPI::Get(tensor->device) + ->AllocDataSpace(tensor->device, tensor->ndim, tensor->shape, + tensor->dtype, mem_scope); } - delete ptr; - } - // Deleter for NDArray converted from DLPack - // This is used from data which is passed from external DLPack(DLManagedTensor) - // that are not allocated inside of TVM. - // This enables us to create NDArray from memory allocated by other - // frameworks that are DLPack compatible - static void DLPackDeleter(TVMFFIObject* ptr_obj) { - auto* ptr = ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned<NDArray::Container>(ptr_obj); - DLManagedTensor* tensor = static_cast<DLManagedTensor*>(ptr->manager_ctx); - if (tensor->deleter != nullptr) { - (*tensor->deleter)(tensor); + void FreeData(DLTensor* tensor) { + DeviceAPI::Get(tensor->device)->FreeDataSpace(tensor->device, tensor->data); } - delete ptr; - } - // Deleter for NDArray based on external DLTensor - // The memory is allocated from outside and it is assumed that - // responsibility for its freeing is also outside - static void SelfDeleter(TVMFFIObject* ptr_obj) { - NDArray::Container* ptr = - ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned<NDArray::Container>(ptr_obj); - delete ptr; - } - // Local create function which allocates tensor metadata - // but does not allocate space for the data. - static NDArray Create(ShapeTuple shape, DLDataType dtype, Device dev) { - VerifyDataType(dtype); - - // critical zone: construct header - NDArray::Container* data = new NDArray::Container(); - data->SetDeleter(DefaultDeleter); - - // RAII now in effect - NDArray ret(GetObjectPtr<Object>(data)); - // setup shape - data->shape_ = std::move(shape); - data->dl_tensor.shape = const_cast<ShapeTuple::index_type*>(data->shape_.data()); - data->dl_tensor.ndim = static_cast<int>(data->shape_.size()); - // setup dtype - data->dl_tensor.dtype = dtype; - // setup device - data->dl_tensor.device = dev; - return ret; - } + }; + return ffi::NDArray::FromNDAlloc(DeviceAPIAlloc(), shape, dtype, dev, mem_scope); +} + +struct NDArray::Internal { // Implementation of API function static DLTensor* MoveToFFIHandle(NDArray arr) { DLTensor* handle = NDArray::FFIGetHandle(arr); @@ -162,35 +115,13 @@ struct NDArray::Internal { return handle; } static void FFIDecRef(TVMArrayHandle tensor) { NDArray::FFIDecRef(tensor); } - // Container to DLManagedTensor - static DLManagedTensor* ToDLPack(TVMArrayHandle handle) { - auto* from = - static_cast<NDArray::Container*>(reinterpret_cast<NDArray::ContainerBase*>(handle)); - return ToDLPack(from); - } - - static DLManagedTensor* ToDLPack(NDArray::Container* from) { - ICHECK(from != nullptr); - DLManagedTensor* ret = new DLManagedTensor(); - ret->dl_tensor = from->dl_tensor; - ret->manager_ctx = from; - tvm::ffi::details::ObjectUnsafe::IncRefObjectHandle(from); - ret->deleter = TVMNDArrayDLPackDeleter; - return ret; - } - // Delete dlpack object. - static void NDArrayDLPackDeleter(DLManagedTensor* tensor) { - ffi::details::ObjectUnsafe::DecRefObjectHandle( - static_cast<NDArray::Container*>(tensor->manager_ctx)); - delete tensor; - } }; NDArray NDArray::CreateView(ShapeTuple shape, DLDataType dtype, uint64_t relative_byte_offset) const { ICHECK(data_ != nullptr); - const DLTensor& orig = get_mutable()->dl_tensor; + const DLTensor& orig = *get_mutable(); CHECK(IsContiguous()) << [&orig]() { std::stringstream ss; ss << "Can only create view for compact tensor, but found strides "; @@ -211,13 +142,9 @@ NDArray NDArray::CreateView(ShapeTuple shape, DLDataType dtype, ss << "]"; return ss.str(); }(); - - const auto& curr_dl_tensor = get_mutable()->dl_tensor; - - NDArray ret = Internal::Create(shape, dtype, curr_dl_tensor.device); - - size_t curr_size = GetDataSize(this->get_mutable()->dl_tensor); - size_t view_size = GetDataSize(ret.get_mutable()->dl_tensor); + const auto& curr_dl_tensor = *get_mutable(); + size_t curr_size = GetDataSize(curr_dl_tensor); + size_t view_size = ffi::GetPackedDataSize(shape.Product(), dtype); CHECK_LE(relative_byte_offset + view_size, curr_size) << "ValueError: " << "View with shape " << shape << " and datatype " << dtype << " would have a size of " @@ -228,83 +155,36 @@ NDArray NDArray::CreateView(ShapeTuple shape, DLDataType dtype, << ShapeTuple(curr_dl_tensor.shape, curr_dl_tensor.shape + curr_dl_tensor.ndim) << ", dtype= " << curr_dl_tensor.dtype << ")."; - // increase ref count - // get_mutable()->IncRef(); - tvm::ffi::details::ObjectUnsafe::IncRefObjectHandle(get_mutable()); - ret.get_mutable()->manager_ctx = get_mutable(); - ret.get_mutable()->dl_tensor.data = get_mutable()->dl_tensor.data; - ret.get_mutable()->dl_tensor.byte_offset = - get_mutable()->dl_tensor.byte_offset + relative_byte_offset; - return ret; -} - -DLManagedTensor* NDArray::ToDLPack() const { return Internal::ToDLPack(get_mutable()); } - -NDArray NDArray::Empty(ShapeTuple shape, DLDataType dtype, Device dev, Optional<String> mem_scope) { - NDArray ret = Internal::Create(shape, dtype, dev); - ret.get_mutable()->dl_tensor.data = - DeviceAPI::Get(ret->device) - ->AllocDataSpace(ret->device, shape.size(), shape.data(), ret->dtype, mem_scope); - return ret; -} - -NDArray NDArray::FromExternalDLTensor(const DLTensor& dl_tensor) { - ICHECK(::tvm::runtime::IsContiguous(dl_tensor)) << "External DLTensor must be contiguous."; - ICHECK(IsAligned(dl_tensor)) << "Data in DLTensor is not aligned as required by NDArray"; - NDArray::Container* data = new NDArray::Container(); + // helper allocator class that retains ref count of original NDArray + class ViewBasedAlloc { + public: + ViewBasedAlloc(NDArray source) : source_(source) {} + void AllocData(DLTensor* tensor, int64_t byte_offset) { + tensor->data = source_.get_mutable()->data; + tensor->byte_offset = byte_offset; + } - data->SetDeleter(Internal::SelfDeleter); - data->dl_tensor = dl_tensor; - std::vector<ShapeTuple::index_type> shape; - shape.resize(data->dl_tensor.ndim); - shape.assign(data->dl_tensor.shape, data->dl_tensor.shape + data->dl_tensor.ndim); - data->shape_ = ShapeTuple(shape); - data->dl_tensor.shape = const_cast<ShapeTuple::index_type*>(data->shape_.data()); + void FreeData(DLTensor* tensor) {} - return NDArray(GetObjectPtr<Object>(data)); -} + private: + NDArray source_; + }; -NDArray NDArray::NewFromDLTensor(DLTensor* tensor, const Device& dev) { - ICHECK(::tvm::runtime::IsContiguous(*tensor)) - << "DLTensor is not contiguous. Copying from non-contiguous data is currently not supported"; - std::vector<int64_t> shape; - for (int64_t i = 0; i < tensor->ndim; i++) { - shape.push_back(tensor->shape[i]); - } - NDArray ary = NDArray::Empty(shape, tensor->dtype, dev); - ary.CopyFrom(tensor); - return ary; -} - -NDArray NDArray::FromDLPack(DLManagedTensor* tensor) { - NDArray::Container* data = new NDArray::Container(); - // construct header - data->SetDeleter(Internal::DLPackDeleter); - // fill up content. - data->manager_ctx = tensor; - ICHECK(::tvm::runtime::IsContiguous(tensor->dl_tensor)) << "DLManagedTensor must be contiguous."; - ICHECK(IsAligned(tensor->dl_tensor)) - << "Data in DLManagedTensor is not aligned as required by NDArray"; - data->dl_tensor = tensor->dl_tensor; - // update shape_ - std::vector<ShapeTuple::index_type> shape; - shape.resize(data->dl_tensor.ndim); - shape.assign(data->dl_tensor.shape, data->dl_tensor.shape + data->dl_tensor.ndim); - data->shape_ = ShapeTuple(shape); - data->dl_tensor.shape = const_cast<ShapeTuple::index_type*>(data->shape_.data()); - return NDArray(GetObjectPtr<Object>(data)); + NDArray ret = NDArray::FromNDAlloc(ViewBasedAlloc(NDArray(*this)), shape, dtype, (*this)->device, + curr_dl_tensor.byte_offset + relative_byte_offset); + return ret; } void NDArray::CopyToBytes(void* data, size_t nbytes) const { ICHECK(data != nullptr); ICHECK(data_ != nullptr); - ArrayCopyToBytes(&get_mutable()->dl_tensor, data, nbytes); + ArrayCopyToBytes(get_mutable(), data, nbytes); } void NDArray::CopyFromBytes(const void* data, size_t nbytes) { ICHECK(data != nullptr); ICHECK(data_ != nullptr); - ArrayCopyFromBytes(&get_mutable()->dl_tensor, data, nbytes); + ArrayCopyFromBytes(get_mutable(), data, nbytes); } NDArray NDArray::CopyTo(const Device& dev, Optional<String> mem_scope) const { @@ -337,38 +217,11 @@ void NDArray::CopyFromTo(const DLTensor* from, DLTensor* to, TVMStreamHandle str DeviceAPI::Get(dev)->CopyDataFromTo(const_cast<DLTensor*>(from), to, stream); } -ShapeTuple NDArray::Shape() const { - return static_cast<const NDArray::Container*>(data_.get())->shape_; -} - -runtime::DataType NDArray::DataType() const { - return runtime::DataType(get_mutable()->dl_tensor.dtype); -} - -bool NDArray::AbilityOfZeroCopyForDLTensor(DLTensor* tensor, const Device& dev) { - bool device_check = (dev.device_type == tensor->device.device_type); - bool device_id_check = (dev.device_id == tensor->device.device_id); - bool alignment_check = IsAligned(*tensor); - return device_check && device_id_check && alignment_check; -} - -bool NDArray::IsAligned(const DLTensor& tensor) { - return (reinterpret_cast<size_t>(static_cast<char*>(tensor.data) + tensor.byte_offset) % - tvm::runtime::kAllocAlignment == - 0); -} - -// TVM_REGISTER_OBJECT_TYPE(NDArray::Container); - } // namespace runtime } // namespace tvm using namespace tvm::runtime; -void TVMNDArrayDLPackDeleter(DLManagedTensor* tensor) { - NDArray::Internal::NDArrayDLPackDeleter(tensor); -} - int TVMArrayGetTypeIndex(TVMArrayHandle handle, unsigned* out_tindex) { API_BEGIN(); *out_tindex = @@ -416,12 +269,10 @@ int TVMArrayFromDLPack(DLManagedTensor* from, TVMArrayHandle* out) { int TVMArrayToDLPack(TVMArrayHandle from, DLManagedTensor** out) { API_BEGIN(); - *out = NDArray::Internal::ToDLPack(from); + *out = static_cast<tvm::ffi::NDArrayObj*>(TVMArrayHandleToObjectHandle(from))->ToDLPack(); API_END(); } -void TVMDLManagedTensorCallDeleter(DLManagedTensor* dltensor) { (*(dltensor->deleter))(dltensor); } - int TVMArrayCopyFromBytes(TVMArrayHandle handle, void* data, size_t nbytes) { API_BEGIN(); ArrayCopyFromBytes(handle, data, nbytes); diff --git a/src/runtime/opencl/opencl_device_api.cc b/src/runtime/opencl/opencl_device_api.cc index 416a1e2653..28cc5afd11 100644 --- a/src/runtime/opencl/opencl_device_api.cc +++ b/src/runtime/opencl/opencl_device_api.cc @@ -852,8 +852,7 @@ class OpenCLPooledAllocator final : public memory::PooledAllocator { Buffer Alloc(Device dev, ShapeTuple shape, DLDataType type_hint, const std::string& mem_scope) override { if (AllowMemoryScope(mem_scope)) { - NDArray::Container container(nullptr, shape, type_hint, dev); - size_t size = DeviceAPI::Get(dev)->GetDataSize(container.dl_tensor); + size_t size = ffi::GetPackedDataSize(shape.Product(), type_hint); Buffer buf; buf.device = dev; buf.size = size; diff --git a/src/runtime/relax_vm/vm.cc b/src/runtime/relax_vm/vm.cc index ec9e4e506d..4f518e4adb 100644 --- a/src/runtime/relax_vm/vm.cc +++ b/src/runtime/relax_vm/vm.cc @@ -111,7 +111,6 @@ Any ConvertObjectToDevice(Any src, const Device& dev, Allocator* alloc) { } TVMRetValue ConvertArgToDevice(AnyView input, Device dev, Allocator* alloc) { - // NOTE: NDArray::FromExternalDLTensor is not safe // in terms of memory-behavior. // To be extra careful, we copy DLTensor. // The developer can still explicitly allocate NDArray diff --git a/src/runtime/rpc/rpc_module.cc b/src/runtime/rpc/rpc_module.cc index e2a7694c5e..37c7fd9d3c 100644 --- a/src/runtime/rpc/rpc_module.cc +++ b/src/runtime/rpc/rpc_module.cc @@ -39,18 +39,6 @@ namespace tvm { namespace runtime { - -// deleter of RPC remote array -static void RemoteNDArrayDeleter(TVMFFIObject* ptr_obj) { - auto* ptr = ffi::details::ObjectUnsafe::RawObjectPtrFromUnowned<NDArray::Container>(ptr_obj); - RemoteSpace* space = static_cast<RemoteSpace*>(ptr->dl_tensor.data); - if (ptr->manager_ctx != nullptr) { - space->sess->FreeHandle(ptr->manager_ctx); - } - delete space; - delete ptr; -} - /*! * \brief Build a local NDArray with remote backing storage. * \param sess the RPCSession which owns the given handle. @@ -66,16 +54,23 @@ NDArray NDArrayFromRemoteOpaqueHandle(std::shared_ptr<RPCSession> sess, void* ha void* remote_ndarray_handle) { ICHECK_EQ(sess->table_index(), GetRPCSessionIndex(dev)) << "The Device given does not belong to the given session"; - RemoteSpace* space = new RemoteSpace(); - space->sess = sess; - space->data = handle; - std::vector<int64_t> shape_vec{template_tensor->shape, - template_tensor->shape + template_tensor->ndim}; - NDArray::Container* data = new NDArray::Container(static_cast<void*>(space), std::move(shape_vec), - template_tensor->dtype, dev); - data->manager_ctx = remote_ndarray_handle; - data->SetDeleter(RemoteNDArrayDeleter); - return NDArray(GetObjectPtr<Object>(data)); + class RemoteSpaceAlloc { + public: + RemoteSpaceAlloc(RemoteSpace space) : space_(space) {} + void AllocData(DLTensor* tensor) { + // the pointer to the remote space is passed in as the data pointer + tensor->data = &(space_); + } + void FreeData(DLTensor* tensor) { space_.sess->FreeHandle(space_.data); } + + private: + RemoteSpace space_; + }; + RemoteSpace space; + space.sess = sess; + space.data = handle; + ffi::Shape shape(template_tensor->shape, template_tensor->shape + template_tensor->ndim); + return NDArray::FromNDAlloc(RemoteSpaceAlloc(space), shape, template_tensor->dtype, dev); } /*!
