/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file ndarray.cc * \brief NDArray container infratructure. */ #include <dmlc/logging.h> #include <tvm/runtime/ndarray.h> #include <tvm/runtime/c_runtime_api.h> #include <tvm/runtime/device_api.h> #include "runtime_base.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 { inline void VerifyDataType(DLDataType dtype) { CHECK_GE(dtype.lanes, 1); if (dtype.code == kDLFloat) { CHECK_EQ(dtype.bits % 8, 0); } else { // allow uint1 as a special flag for bool. if (dtype.bits == 1 && dtype.code == kDLUInt) return; // allow int1/uint4/int4 else if (dtype.bits == 1 && dtype.code == kDLInt) return; else if (dtype.bits == 4 && dtype.code == kDLUInt) return; else if (dtype.bits == 4 && dtype.code == kDLInt) return; else CHECK_EQ(dtype.bits % 8, 0); } CHECK_EQ(dtype.bits & (dtype.bits - 1), 0); } inline size_t GetDataAlignment(const DLTensor& arr) { size_t align = (arr.dtype.bits / 8) * arr.dtype.lanes; if (align < kAllocAlignment) return kAllocAlignment; return align; } void ArrayCopyFromBytes(DLTensor* handle, const void* data, size_t nbytes) { TVMContext cpu_ctx; cpu_ctx.device_type = kDLCPU; cpu_ctx.device_id = 0; size_t arr_size = GetDataSize(*handle); CHECK_EQ(arr_size, nbytes) << "ArrayCopyFromBytes: size mismatch"; DeviceAPI::Get(handle->ctx)->CopyDataFromTo( data, 0, handle->data, static_cast<size_t>(handle->byte_offset), nbytes, cpu_ctx, handle->ctx, handle->dtype, nullptr); } void ArrayCopyToBytes(const DLTensor* handle, void* data, size_t nbytes) { TVMContext cpu_ctx; cpu_ctx.device_type = kDLCPU; cpu_ctx.device_id = 0; size_t arr_size = GetDataSize(*handle); CHECK_EQ(arr_size, nbytes) << "ArrayCopyToBytes: size mismatch"; DeviceAPI::Get(handle->ctx)->CopyDataFromTo( handle->data, static_cast<size_t>(handle->byte_offset), data, 0, nbytes, handle->ctx, cpu_ctx, handle->dtype, nullptr); } struct NDArray::Internal { // Default deleter for the container static void DefaultDeleter(Object* ptr_obj) { auto* ptr = static_cast<NDArray::Container*>(ptr_obj); if (ptr->manager_ctx != nullptr) { static_cast<NDArray::Container*>(ptr->manager_ctx)->DecRef(); } else if (ptr->dl_tensor.data != nullptr) { tvm::runtime::DeviceAPI::Get(ptr->dl_tensor.ctx)->FreeDataSpace( ptr->dl_tensor.ctx, ptr->dl_tensor.data); } 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(Object* ptr_obj) { auto* ptr = static_cast<NDArray::Container*>(ptr_obj); DLManagedTensor* tensor = static_cast<DLManagedTensor*>(ptr->manager_ctx); if (tensor->deleter != nullptr) { (*tensor->deleter)(tensor); } delete ptr; } // Local create function which allocates tensor metadata // but does not allocate space for the data. static NDArray Create(std::vector<int64_t> shape, DLDataType dtype, DLContext ctx) { 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 = dmlc::BeginPtr(data->shape_); data->dl_tensor.ndim = static_cast<int>(data->shape_.size()); // setup dtype data->dl_tensor.dtype = dtype; // setup ctx data->dl_tensor.ctx = ctx; return ret; } // Implementation of API function static DLTensor* MoveToFFIHandle(NDArray arr) { DLTensor* handle = NDArray::FFIGetHandle(arr); ObjectRef::FFIClearAfterMove(&arr); 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) { CHECK(from != nullptr); DLManagedTensor* ret = new DLManagedTensor(); ret->dl_tensor = from->dl_tensor; ret->manager_ctx = from; from->IncRef(); ret->deleter = TVMNDArrayDLPackDeleter; return ret; } // Delete dlpack object. static void NDArrayDLPackDeleter(DLManagedTensor* tensor) { static_cast<NDArray::Container*>(tensor->manager_ctx)->DecRef(); delete tensor; } }; NDArray NDArray::CreateView(std::vector<int64_t> shape, DLDataType dtype) { CHECK(data_ != nullptr); CHECK(get_mutable()->dl_tensor.strides == nullptr) << "Can only create view for compact tensor"; NDArray ret = Internal::Create(shape, dtype, get_mutable()->dl_tensor.ctx); ret.get_mutable()->dl_tensor.byte_offset = this->get_mutable()->dl_tensor.byte_offset; size_t curr_size = GetDataSize(this->get_mutable()->dl_tensor); size_t view_size = GetDataSize(ret.get_mutable()->dl_tensor); CHECK_LE(view_size, curr_size) << "Tries to create a view that has bigger memory than current one"; // increase ref count get_mutable()->IncRef(); ret.get_mutable()->manager_ctx = get_mutable(); ret.get_mutable()->dl_tensor.data = get_mutable()->dl_tensor.data; return ret; } DLManagedTensor* NDArray::ToDLPack() const { return Internal::ToDLPack(get_mutable()); } NDArray NDArray::Empty(std::vector<int64_t> shape, DLDataType dtype, DLContext ctx) { NDArray ret = Internal::Create(shape, dtype, ctx); // setup memory content size_t size = GetDataSize(ret.get_mutable()->dl_tensor); size_t alignment = GetDataAlignment(ret.get_mutable()->dl_tensor); ret.get_mutable()->dl_tensor.data = DeviceAPI::Get(ret->ctx)->AllocDataSpace( ret->ctx, size, alignment, ret->dtype); return ret; } NDArray NDArray::FromDLPack(DLManagedTensor* tensor) { NDArray::Container* data = new NDArray::Container(); // construct header data->SetDeleter(Internal::DLPackDeleter); // fill up content. data->manager_ctx = tensor; data->dl_tensor = tensor->dl_tensor; // update shape_ data->shape_.resize(data->dl_tensor.ndim); data->shape_.assign(data->dl_tensor.shape, data->dl_tensor.shape + data->dl_tensor.ndim); data->dl_tensor.shape = data->shape_.data(); return NDArray(GetObjectPtr<Object>(data)); } void NDArray::CopyToBytes(void* data, size_t nbytes) const { CHECK(data != nullptr); CHECK(data_ != nullptr); ArrayCopyToBytes(&get_mutable()->dl_tensor, data, nbytes); } void NDArray::CopyFromBytes(const void* data, size_t nbytes) { CHECK(data != nullptr); CHECK(data_ != nullptr); ArrayCopyFromBytes(&get_mutable()->dl_tensor, data, nbytes); } void NDArray::CopyFromTo(const DLTensor* from, DLTensor* to, TVMStreamHandle stream) { size_t from_size = GetDataSize(*from); size_t to_size = GetDataSize(*to); CHECK_EQ(from_size, to_size) << "TVMArrayCopyFromTo: The size must exactly match"; CHECK(from->ctx.device_type == to->ctx.device_type || from->ctx.device_type == kDLCPU || to->ctx.device_type == kDLCPU || from->ctx.device_type == kDLCPUPinned || to->ctx.device_type == kDLCPUPinned) << "Can not copy across different ctx types directly"; // Use the context that is *not* a cpu context to get the correct device // api manager. TVMContext ctx = from->ctx.device_type != kDLCPU ? from->ctx : to->ctx; DeviceAPI::Get(ctx)->CopyDataFromTo( from->data, static_cast<size_t>(from->byte_offset), to->data, static_cast<size_t>(to->byte_offset), from_size, from->ctx, to->ctx, from->dtype, stream); } std::vector<int64_t> NDArray::Shape() const { return get_mutable()->shape_; } 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 = TVMArrayHandleToObjectHandle(handle)->type_index(); API_END(); } int TVMArrayAlloc(const tvm_index_t* shape, int ndim, int dtype_code, int dtype_bits, int dtype_lanes, int device_type, int device_id, TVMArrayHandle* out) { API_BEGIN(); DLDataType dtype; dtype.code = static_cast<uint8_t>(dtype_code); dtype.bits = static_cast<uint8_t>(dtype_bits); dtype.lanes = static_cast<uint16_t>(dtype_lanes); DLContext ctx; ctx.device_type = static_cast<DLDeviceType>(device_type); ctx.device_id = device_id; *out = NDArray::Internal::MoveToFFIHandle( NDArray::Empty(std::vector<int64_t>(shape, shape + ndim), dtype, ctx)); API_END(); } int TVMArrayFree(TVMArrayHandle handle) { API_BEGIN(); NDArray::Internal::FFIDecRef(handle); API_END(); } int TVMArrayCopyFromTo(TVMArrayHandle from, TVMArrayHandle to, TVMStreamHandle stream) { API_BEGIN(); NDArray::CopyFromTo(from, to, stream); API_END(); } int TVMArrayFromDLPack(DLManagedTensor* from, TVMArrayHandle* out) { API_BEGIN(); *out = NDArray::Internal::MoveToFFIHandle(NDArray::FromDLPack(from)); API_END(); } int TVMArrayToDLPack(TVMArrayHandle from, DLManagedTensor** out) { API_BEGIN(); *out = NDArray::Internal::ToDLPack(from); API_END(); } void TVMDLManagedTensorCallDeleter(DLManagedTensor* dltensor) { (*(dltensor->deleter))(dltensor); } int TVMArrayCopyFromBytes(TVMArrayHandle handle, void* data, size_t nbytes) { API_BEGIN(); ArrayCopyFromBytes(handle, data, nbytes); API_END(); } int TVMArrayCopyToBytes(TVMArrayHandle handle, void* data, size_t nbytes) { API_BEGIN(); ArrayCopyToBytes(handle, data, nbytes); API_END(); }