From 4a4488baaead3479fec191fdd2d2686b6f243408 Mon Sep 17 00:00:00 2001 From: Chen Fu Date: Wed, 10 Mar 2021 14:07:20 -0800 Subject: [PATCH] Release buffers for prepacked tensors (#6820) Unsolved problems: 1. One test failure was caused by a bug in Cudnn rnn kernels, when they can allocate a buffer and partially initialize it, the garbage data near tail of the buffer caused problem in some of the hardware. To attack this problem in a broader sense, should we add code in our allocators, and during a memory fuzzing test, fill an allocated buffer with garbage before returning to the caller? 2. Prepacking is used more widely than we know. For instance, Cudnn rnn kernels also cache their weights. They mix several weight tensors together into a single buffer, and never touch the original weight tensor anymore. This is the same idea with pre-pack, but they didn't override the virtual function, and they never tried to release those weight tensors, leading to memory waste. It also seems to me that there are some other kernels have similar behavior. Wonder how much memory we can save if we try to cleanup those too. 3. Turning off memory pattern planning does increase memory fragmentation, leading to out of memory error in some training test cases. Perhaps we can revisit the idea of pushing kernels-creation stage earlier, and then during initializer deserialization, we only avoid tracing those that will be prepacked. --- .../core/framework/execution_providers.h | 6 +- onnxruntime/core/framework/session_state.cc | 21 +- .../core/framework/session_state_utils.cc | 115 ++++---- .../core/framework/session_state_utils.h | 2 +- .../core/framework/simple_tensor_allocator.cc | 28 +- .../core/framework/simple_tensor_allocator.h | 9 +- onnxruntime/core/framework/tensor_allocator.h | 21 +- .../tensor_allocator_with_mem_pattern.h | 11 +- .../core/framework/tensorprotoutils.cc | 268 +++++++++--------- onnxruntime/core/framework/tensorprotoutils.h | 18 +- .../optimizer/optimizer_execution_frame.cc | 6 +- .../core/providers/cuda/rnn/cudnn_rnn_base.cc | 7 + .../test/framework/test_tensor_loader.cc | 24 +- .../models/runner/training_runner.cc | 6 +- .../models/runner/training_util.cc | 6 +- .../test/framework/checkpointing_test.cc | 6 +- 16 files changed, 277 insertions(+), 277 deletions(-) diff --git a/onnxruntime/core/framework/execution_providers.h b/onnxruntime/core/framework/execution_providers.h index 8125bdd874..61a689b1b5 100644 --- a/onnxruntime/core/framework/execution_providers.h +++ b/onnxruntime/core/framework/execution_providers.h @@ -73,8 +73,12 @@ class ExecutionProviders { const_iterator begin() const noexcept { return exec_providers_.cbegin(); } const_iterator end() const noexcept { return exec_providers_.cend(); } + const AllocatorPtr GetDefaultCpuAllocator() const { + return Get(onnxruntime::kCpuExecutionProvider)->GetAllocator(0, OrtMemTypeDefault); + } + OrtMemoryInfo GetDefaultCpuMemoryInfo() const { - return Get(onnxruntime::kCpuExecutionProvider)->GetAllocator(0, OrtMemTypeDefault)->Info(); + return GetDefaultCpuAllocator()->Info(); } const std::vector& GetIds() const { return exec_provider_ids_; } diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc index 13f5ddf4f7..fc7062e657 100644 --- a/onnxruntime/core/framework/session_state.cc +++ b/onnxruntime/core/framework/session_state.cc @@ -1004,8 +1004,27 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string tensor_allocator( ITensorAllocator::Create(enable_mem_pattern_, *p_seq_exec_plan_, *this, weights_buffers_)); +#else + std::unique_ptr tensor_allocator( + ITensorAllocator::Create(false, *p_seq_exec_plan_, *this, weights_buffers_)); +#endif const auto& initializer_allocation_order = p_seq_exec_plan_->initializer_allocation_order; @@ -1013,7 +1032,7 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string Status { return AddInitializedTensor(idx, value, &d, constant); diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index f29389f600..93e12df7b9 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -31,58 +31,54 @@ namespace onnxruntime { namespace session_state_utils { static common::Status DeserializeTensorProto(const Env& env, const std::basic_string& proto_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, const MemBuffer& m, - const OrtMemoryInfo& default_cpu_memory_info, OrtValue& ort_value, - OrtCallback& deleter, - const DataTransferManager& data_transfer_mgr) { - const OrtMemoryInfo& alloc_info = m.GetAllocInfo(); - if (strcmp(alloc_info.name, CPU) == 0 || alloc_info.mem_type == OrtMemTypeCPUOutput) { - // deserialize directly to CPU tensor - return utils::TensorProtoToMLValue(env, proto_path.c_str(), tensor_proto, m, ort_value, deleter); + const ONNX_NAMESPACE::TensorProto& tensor_proto, const MemBuffer* m, + const AllocatorPtr& alloc, const AllocatorPtr& default_cpu_alloc, + OrtValue& ort_value, const DataTransferManager& data_transfer_mgr) { + if (bool(alloc) == (m != nullptr)) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, + "DeserializeTensorProto() takes either pre-allocated buffer or an allocator!"); } - // deserialize and copy. In the copy stage, it won't check if the buffer has enough room. - // The result tensor won't need a deleter because: - // 1. It mustn't be a string tensor - // 2. The memory is not memory-mapped. - deleter.f = nullptr; - deleter.param = nullptr; - if (tensor_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "string tensor is not supported for copying between allocators"); - } - - // deserialize to CPU first for non-CPU allocator, then alloc and copy - size_t cpu_tensor_length; - ORT_RETURN_IF_ERROR(utils::GetSizeInBytesFromTensorProto<0>(tensor_proto, &cpu_tensor_length)); - if (m.GetLen() < cpu_tensor_length) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Internal error. The preallocated buffer is too small. Requires ", - cpu_tensor_length, ", Got ", m.GetLen()); - } - - std::unique_ptr data(new char[cpu_tensor_length]); + // Get shape and type of the tensor, and allocate the empty tensor + TensorShape tensor_shape{utils::GetTensorShapeFromTensorProto(tensor_proto)}; + const DataTypeImpl* const type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type())->GetElementType(); std::unique_ptr p_tensor; - OrtValue tmp_ort_value; - OrtCallback d; - ORT_RETURN_IF_ERROR(utils::TensorProtoToMLValue(env, proto_path.c_str(), tensor_proto, - MemBuffer(data.get(), cpu_tensor_length, default_cpu_memory_info), - tmp_ort_value, d)); - - const Tensor& p_deserialize_tensor = tmp_ort_value.Get(); - - p_tensor = onnxruntime::make_unique(p_deserialize_tensor.DataType(), p_deserialize_tensor.Shape(), m.GetBuffer(), - m.GetAllocInfo()); - // TODO: does this function work for string tensor? - Status copy_status = data_transfer_mgr.CopyTensor(p_deserialize_tensor, *p_tensor); - if (d.f) d.f(d.param); - - if (!copy_status.IsOK()) { - if (copy_status.ErrorMessage().empty()) { - // The windows execution provider does not return any error message today for CopyTensor since it is - // not implemented yet. That's the reason we're adding our own error message so that we can debug better. - return Status(copy_status.Category(), copy_status.Code(), - "Failed to copy tensor to " + p_tensor->Location().ToString()); + if (m != nullptr) { + p_tensor = onnxruntime::make_unique(type, tensor_shape, m->GetBuffer(), m->GetAllocInfo()); + if (m->GetLen() < p_tensor->SizeInBytes()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Internal error. The preallocated buffer is too small. Requires ", + p_tensor->SizeInBytes(), ", Got ", m->GetLen()); + } + } else { + // tensor constructor should give us enough buffer size based on type and shape + p_tensor = onnxruntime::make_unique(type, tensor_shape, alloc); + } + + if (strcmp(p_tensor->Location().name, CPU) == 0) { + // deserialize directly to CPU tensor + ORT_RETURN_IF_ERROR(utils::TensorProtoToTensor(env, proto_path.c_str(), tensor_proto, *p_tensor)); + } else { + // non-cpu tensor + if (tensor_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "string tensor is not supported for copying between allocators"); + } + + // deserialize to CPU first for non-CPU allocator, then copy + std::unique_ptr p_deserialize_tensor = onnxruntime::make_unique(type, tensor_shape, default_cpu_alloc); + ORT_RETURN_IF_ERROR(utils::TensorProtoToTensor(env, proto_path.c_str(), tensor_proto, *p_deserialize_tensor)); + // TODO!! Need a temp buffer allocator for non-escape buffers that maybe too big for stack allocation. + + Status copy_status = data_transfer_mgr.CopyTensor(*p_deserialize_tensor, *p_tensor); + + if (!copy_status.IsOK()) { + if (copy_status.ErrorMessage().empty()) { + // The windows execution provider does not return any error message today for CopyTensor since it is + // not implemented yet. That's the reason we're adding our own error message so that we can debug better. + return Status(copy_status.Category(), copy_status.Code(), + "Failed to copy tensor to " + p_tensor->Location().ToString()); + } + return copy_status; } - return copy_status; } auto ml_tensor = DataTypeImpl::GetType(); @@ -92,7 +88,7 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st common::Status SaveInitializedTensors( const Env& env, const std::basic_string& graph_loc, - const GraphViewer& graph, const OrtMemoryInfo& default_cpu_memory_info, + const GraphViewer& graph, const AllocatorPtr& default_cpu_alloc, const OrtValueNameIdxMap& ort_value_name_idx_map, const std::vector& initializer_allocation_order, ITensorAllocator& planner, @@ -117,8 +113,8 @@ common::Status SaveInitializedTensors( if (!ort_value_name_idx_map.GetIdx(name, ort_value_index).IsOK()) { retval = false; } else { - auto planned_mem_info = exec_plan.GetLocation(ort_value_index); - auto user_mem_info = it->second->Get().Location(); + const auto& planned_mem_info = exec_plan.GetLocation(ort_value_index); + const auto& user_mem_info = it->second->Get().Location(); retval = user_mem_info.device == planned_mem_info.device; if (!retval) { LOGS(logger, WARNING) << "Cannot use user supplied initializer with name: (" @@ -149,7 +145,9 @@ common::Status SaveInitializedTensors( auto initialized_tensors_to_allocate = id_to_initialized_tensor; for (int ort_value_index : initializer_allocation_order) { const auto entry = initialized_tensors_to_allocate.find(ort_value_index); - ORT_ENFORCE(entry != initialized_tensors_to_allocate.end()); + // can not trace string tensor + ORT_ENFORCE(entry != initialized_tensors_to_allocate.end() + && entry->second->data_type() != ONNX_NAMESPACE::TensorProto_DataType_STRING); ORT_RETURN_IF_ERROR(planner.Trace(entry->first, entry->second)); initialized_tensors_to_allocate.erase(entry); } @@ -159,6 +157,10 @@ common::Status SaveInitializedTensors( if (user_supplied_initializer_ids.find(entry.first) != user_supplied_initializer_ids.end()) { continue; } + if (entry.second->data_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING) { + // do not trace string tensor + continue; + } ORT_RETURN_IF_ERROR(planner.Trace(entry.first, entry.second)); } //2. allocate weight buffer on different locations @@ -193,13 +195,10 @@ common::Status SaveInitializedTensors( const ONNX_NAMESPACE::TensorProto& tensor_proto = *(entry.second); std::unique_ptr m; + AllocatorPtr alloc; // TODO: if the tensor need be copied, does it have enough room? - ORT_RETURN_IF_ERROR(planner.GetPreallocatedBuffer(ort_value_index, name, m)); -#ifndef NDEBUG - ORT_ENFORCE(m != nullptr); - ORT_ENFORCE(m->GetBuffer() != nullptr || m->GetLen() == 0); -#endif - Status st = DeserializeTensorProto(env, graph_loc, tensor_proto, *m, default_cpu_memory_info, ort_value, deleter, + ORT_RETURN_IF_ERROR(planner.GetPreallocatedBuffer(ort_value_index, name, m, alloc)); + Status st = DeserializeTensorProto(env, graph_loc, tensor_proto, m.get(), alloc, default_cpu_alloc, ort_value, data_transfer_mgr); if (!st.IsOK()) { std::ostringstream oss; diff --git a/onnxruntime/core/framework/session_state_utils.h b/onnxruntime/core/framework/session_state_utils.h index eef60235bd..f74a9796e7 100644 --- a/onnxruntime/core/framework/session_state_utils.h +++ b/onnxruntime/core/framework/session_state_utils.h @@ -32,7 +32,7 @@ class Logger; namespace session_state_utils { common::Status SaveInitializedTensors( const Env& env, const std::basic_string& graph_loc, - const GraphViewer& graph, const OrtMemoryInfo& default_cpu_memory_info, + const GraphViewer& graph, const AllocatorPtr& default_cpu_memory_info, const OrtValueNameIdxMap& ort_value_name_idx_map, const std::vector& initializer_allocation_order, ITensorAllocator& planner, const std::function& save_tensor_func, diff --git a/onnxruntime/core/framework/simple_tensor_allocator.cc b/onnxruntime/core/framework/simple_tensor_allocator.cc index 360e052963..08ebb67284 100644 --- a/onnxruntime/core/framework/simple_tensor_allocator.cc +++ b/onnxruntime/core/framework/simple_tensor_allocator.cc @@ -5,32 +5,16 @@ #include "tensorprotoutils.h" namespace onnxruntime { -common::Status SimpleTensorAllocator::Trace(int id, const ONNX_NAMESPACE::TensorProto* value) { - values_[id] = value; +common::Status SimpleTensorAllocator::Trace(int /*id*/, const ONNX_NAMESPACE::TensorProto* /*value*/) { return Status::OK(); } -common::Status SimpleTensorAllocator::GetPreallocatedBuffer(int ort_value_index, const char* name, - std::unique_ptr& out) { - auto iter = values_.find(ort_value_index); - if (iter == values_.end()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "invalid ort_value_index:", ort_value_index); - } - - size_t len = 0; - ORT_RETURN_IF_ERROR(utils::GetSizeInBytesFromTensorProto(*iter->second, &len)); +common::Status SimpleTensorAllocator::GetPreallocatedBuffer(int ort_value_index, const char* /*name*/, + std::unique_ptr& /*buf_out*/, + AllocatorPtr& alloc_out) { const struct OrtMemoryInfo& location = seq_plan_.GetLocation(ort_value_index); - if (len == 0) { - out = onnxruntime::make_unique(nullptr, 0, location); + // just return allocator and let others handle it. + alloc_out = GetAllocator(location); return Status::OK(); - } - auto alloc = GetAllocator(location); - if (!alloc) - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to get allocator for initializer '", name, - "', location: ", location.ToString()); - void* buffer = alloc->Alloc(len); - weights_buffers_.push_back(BufferUniquePtr(buffer, alloc)); - out = onnxruntime::make_unique(buffer, len, location); - return Status::OK(); } } // namespace onnxruntime diff --git a/onnxruntime/core/framework/simple_tensor_allocator.h b/onnxruntime/core/framework/simple_tensor_allocator.h index 90ba5db35d..ac90d90c42 100644 --- a/onnxruntime/core/framework/simple_tensor_allocator.h +++ b/onnxruntime/core/framework/simple_tensor_allocator.h @@ -15,17 +15,12 @@ class ExecutionProviders; class SimpleTensorAllocator : public ITensorAllocator { private: MemoryPatternGroup mem_patterns_; - std::vector& weights_buffers_; const ExecutionPlanBase& seq_plan_; - private: - std::unordered_map values_; - public: SimpleTensorAllocator(const ExecutionPlanBase& execution_plan, const SessionState& session_state, - std::vector& weights_buffers) + std::vector& /*weights_buffers*/) : ITensorAllocator(session_state), - weights_buffers_(weights_buffers), seq_plan_(execution_plan) {} common::Status FinalizePlan(std::unordered_map& planned_memory_sizes_in_byte) override { @@ -34,7 +29,7 @@ class SimpleTensorAllocator : public ITensorAllocator { planned_memory_sizes_in_byte = std::unordered_map(); return Status::OK(); } - common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::unique_ptr& out) override; + common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::unique_ptr& buf_out, AllocatorPtr& alloc_out) override; common::Status Trace(int id, const ONNX_NAMESPACE::TensorProto* value) override; const MemoryPatternGroup& GetMemPatterns() override { return mem_patterns_; diff --git a/onnxruntime/core/framework/tensor_allocator.h b/onnxruntime/core/framework/tensor_allocator.h index cafff10021..081592f4e5 100644 --- a/onnxruntime/core/framework/tensor_allocator.h +++ b/onnxruntime/core/framework/tensor_allocator.h @@ -35,15 +35,20 @@ class ITensorAllocator { virtual common::Status FinalizePlan(std::unordered_map& planned_memory_sizes_in_byte) = 0; /** - * - * \param ort_value_index The index in planner - * \param name Tensor name. Only for logging purpose - * \param out The allocated buffer - * - * When it succeeded, p could be NULL if the tensor with 'ort_value_index' will not have any element - */ + * Handing out buffers reserved in @see #Trace() via parameter buf_out, + * or, in the case of not reserved tensor, returns an allocator so that + * the caller can take care of the dynamic buffer allocation. + * buf_out and alloc_out, one and only one can be non-null + * + * @param ort_value_index [In] int id of the tensor + * @param name [In] name of the tensor + * @param buf_out [Out] pre reserved buffer, if not null + * @param alloc_out [Out] allocator based on tensor's location, if not null + * @return + */ virtual common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, - std::unique_ptr& out) = 0; + std::unique_ptr& buf_out, + AllocatorPtr& alloc_out) = 0; virtual const MemoryPatternGroup& GetMemPatterns() = 0; /** diff --git a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h index 73c41cb497..bf05d184c2 100644 --- a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h +++ b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h @@ -73,7 +73,7 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { } common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, - std::unique_ptr& out) override { + std::unique_ptr& buf_out, AllocatorPtr& alloc_out) override { if (!is_sealed_) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Internal error."); } @@ -86,11 +86,16 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { // fall back to allocate separate buffer. // if it->second.get() is null, then fall back to the block not found case auto block = pattern->GetBlock(ort_value_index); + if (nullptr == block) { + // not traced, only return allocator + alloc_out = GetAllocator(location); + return Status::OK(); + } auto it = buffers_.find(location); if (it == buffers_.end()) { if (block != nullptr && block->size_ == 0) { // Because the size is 0, this miss find is expected. we won't allocate a buffer with size of zero. - out = onnxruntime::make_unique(nullptr, 0, location); + buf_out = onnxruntime::make_unique(nullptr, 0, location); return Status::OK(); } return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Weight buffer for initializer '", name, "' is not found"); @@ -100,7 +105,7 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Get preallocated buffer for initializer '", name, "' failed"); } - out = onnxruntime::make_unique(reinterpret_cast(it->second) + block->offset_, block->size_, location); + buf_out = onnxruntime::make_unique(reinterpret_cast(it->second) + block->offset_, block->size_, location); return Status::OK(); } common::Status Trace(int id, const ONNX_NAMESPACE::TensorProto* value) override { diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 015a2da9e1..c149fda6e9 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -86,16 +86,6 @@ bool operator!=(const ONNX_NAMESPACE::TensorShapeProto_Dimension& l, namespace { -std::vector GetTensorShapeFromTensorProto(const ONNX_NAMESPACE::TensorProto& tensor_proto) { - const auto& dims = tensor_proto.dims(); - std::vector tensor_shape_vec(static_cast(dims.size())); - for (int i = 0; i < dims.size(); ++i) { - tensor_shape_vec[i] = dims[i]; - } - - return tensor_shape_vec; -} - // This function doesn't support string tensors static Status UnpackTensorWithRawDataImpl(const void* raw_data, size_t raw_data_len, size_t expected_num_elements, size_t element_size, @@ -506,6 +496,16 @@ TensorShape GetTensorShapeFromTensorShapeProto(const ONNX_NAMESPACE::TensorShape return TensorShape(std::move(tensor_shape_vec)); } +std::vector GetTensorShapeFromTensorProto(const ONNX_NAMESPACE::TensorProto& tensor_proto) { + const auto& dims = tensor_proto.dims(); + std::vector tensor_shape_vec(static_cast(dims.size())); + for (int i = 0; i < dims.size(); ++i) { + tensor_shape_vec[i] = dims[i]; + } + + return tensor_shape_vec; +} + struct UnInitializeParam { void* preallocated; size_t preallocated_size; @@ -542,12 +542,6 @@ ORT_API(void, OrtUninitializeBuffer, _In_opt_ void* input, size_t input_len, enu } } -static void UnInitTensor(void* param) noexcept { - UnInitializeParam* p = reinterpret_cast(param); - OrtUninitializeBuffer(p->preallocated, p->preallocated_size, p->ele_type); - delete p; -} - class AutoDelete { public: OrtCallback d{nullptr, nullptr}; @@ -594,13 +588,6 @@ static Status GetFileContent( return Status::OK(); } -static void MoveOrtCallback(OrtCallback& from, OrtCallback& to) { - to.f = from.f; - to.param = from.param; - from.f = nullptr; - from.param = nullptr; -} - #define CASE_PROTO(X, Y) \ case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: \ ORT_RETURN_IF_ERROR( \ @@ -608,6 +595,110 @@ static void MoveOrtCallback(OrtCallback& from, OrtCallback& to) { (Y*)preallocated, static_cast(tensor_size))); \ break; +/** + * @brief Convert tensor_proto to tensor format and store it to pre-allocated tensor + * @param env + * @param model_path + * @param tensor_proto tensor data in protobuf format + * @param tensorp pre-allocated tensor object, where we store the data + * @return +*/ +Status TensorProtoToTensor(const Env& env, const ORTCHAR_T* model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, + Tensor& tensor) { + // Validate tensor compatibility + std::vector tensor_shape_vec = GetTensorShapeFromTensorProto(tensor_proto); + if (tensor_shape_vec != tensor.Shape().GetDims()) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "TensorProtoToTensor() tensor shape mismatch!"); + } + const DataTypeImpl* const source_type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type())->GetElementType(); + if (source_type->Size() > tensor.DataType()->Size()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "TensorProto type ", DataTypeImpl::ToString(source_type), + " can not be writen into Tensor type ", DataTypeImpl::ToString(tensor.DataType())); + } + + // find raw data in proto buf + void* raw_data = nullptr; + SafeInt raw_data_len = 0; + AutoDelete deleter_for_file_data; + + if (utils::HasExternalData(tensor_proto)) { + // Get the external data info + std::basic_string external_data_file_path; + FileOffsetType file_offset; + std::basic_string tensor_proto_dir; + if (model_path != nullptr) { + ORT_RETURN_IF_ERROR(GetDirNameFromFilePath(model_path, tensor_proto_dir)); + } + ORT_RETURN_IF_ERROR(GetExternalDataInfo( + tensor_proto, + tensor_proto_dir.size() == 0 ? nullptr : tensor_proto_dir.c_str(), + external_data_file_path, file_offset, raw_data_len)); + + // load the file + ORT_RETURN_IF_ERROR(GetFileContent( + env, external_data_file_path.c_str(), file_offset, raw_data_len, + raw_data, deleter_for_file_data.d)); + } else if (utils::HasRawData(tensor_proto)) { + raw_data = const_cast(tensor_proto.raw_data().data()); + // TODO The line above has const-correctness issues. Below is a possible fix which copies the tensor_proto data + // into a writeable buffer. However, it requires extra memory which may exceed the limit for certain tests. + //auto buffer = onnxruntime::make_unique(tensor_proto.raw_data().size()); + //std::memcpy(buffer.get(), tensor_proto.raw_data().data(), tensor_proto.raw_data().size()); + //deleter_for_file_data.d = OrtCallback{DeleteCharArray, buffer.get()}; + //raw_data = buffer.release(); + raw_data_len = tensor_proto.raw_data().size(); + } + + if (nullptr != raw_data && utils::IsPrimitiveDataType(source_type)) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "string tensor can not have raw data"); + } + + // unpacking tensor_proto data to preallocated tensor + void* preallocated = tensor.MutableDataRaw(); + int64_t tensor_size = 1; + { + for (auto i : tensor_proto.dims()) { + if (i < 0) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "tensor can't contain negative dims"); + } + tensor_size *= i; + } + } + // tensor_size could be zero. see test_slice_start_out_of_bounds\test_data_set_0\output_0.pb + if (static_cast(tensor_size) > SIZE_MAX) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "size overflow"); + } + switch (tensor_proto.data_type()) { + CASE_PROTO(FLOAT, float); + CASE_PROTO(DOUBLE, double); + CASE_PROTO(BOOL, bool); + CASE_PROTO(INT8, int8_t); + CASE_PROTO(INT16, int16_t); + CASE_PROTO(INT32, int32_t); + CASE_PROTO(INT64, int64_t); + CASE_PROTO(UINT8, uint8_t); + CASE_PROTO(UINT16, uint16_t); + CASE_PROTO(UINT32, uint32_t); + CASE_PROTO(UINT64, uint64_t); + CASE_PROTO(FLOAT16, MLFloat16); + CASE_PROTO(BFLOAT16, BFloat16); + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_STRING: + ORT_RETURN_IF_ERROR(UnpackTensor(tensor_proto, raw_data, raw_data_len, + static_cast(preallocated), + static_cast(tensor_size))); + break; + default: { + std::ostringstream ostr; + ostr << "Initialized tensor with unexpected type: " << tensor_proto.data_type(); + return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, ostr.str()); + } + } + + return Status::OK(); +} + + #ifdef _MSC_VER #pragma warning(push) #pragma warning(disable : 6239) @@ -615,120 +706,31 @@ static void MoveOrtCallback(OrtCallback& from, OrtCallback& to) { // TODO: Change the current interface to take Path object for model path // so that validating and manipulating path for reading external data becomes easy Status TensorProtoToMLValue(const Env& env, const ORTCHAR_T* model_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, const MemBuffer& m, OrtValue& value, - OrtCallback& deleter) { - const OrtMemoryInfo& allocator = m.GetAllocInfo(); - ONNXTensorElementDataType ele_type = utils::GetTensorElementType(tensor_proto); - deleter.f = nullptr; - deleter.param = nullptr; - void* raw_data = nullptr; - SafeInt raw_data_len = 0; - const DataTypeImpl* const type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type())->GetElementType(); - AutoDelete deleter_for_file_data; - void* tensor_data; - { - if (utils::HasExternalData(tensor_proto)) { - // Get the external data info - std::basic_string external_data_file_path; - FileOffsetType file_offset; - std::basic_string tensor_proto_dir; - if (model_path != nullptr) { - ORT_RETURN_IF_ERROR(GetDirNameFromFilePath(model_path, tensor_proto_dir)); - } - ORT_RETURN_IF_ERROR(GetExternalDataInfo( - tensor_proto, - tensor_proto_dir.size() == 0 ? nullptr : tensor_proto_dir.c_str(), - external_data_file_path, file_offset, raw_data_len)); - - // load the file - ORT_RETURN_IF_ERROR(GetFileContent( - env, external_data_file_path.c_str(), file_offset, raw_data_len, - raw_data, deleter_for_file_data.d)); - } else if (utils::HasRawData(tensor_proto)) { - if (ele_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING) - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "string tensor can not have raw data"); - raw_data = const_cast(tensor_proto.raw_data().data()); - // TODO The line above has const-correctness issues. Below is a possible fix which copies the tensor_proto data - // into a writeable buffer. However, it requires extra memory which may exceed the limit for certain tests. - //auto buffer = onnxruntime::make_unique(tensor_proto.raw_data().size()); - //std::memcpy(buffer.get(), tensor_proto.raw_data().data(), tensor_proto.raw_data().size()); - //deleter_for_file_data.d = OrtCallback{DeleteCharArray, buffer.get()}; - //raw_data = buffer.release(); - raw_data_len = tensor_proto.raw_data().size(); - } - if (endian::native == endian::little && raw_data != nullptr && deleter_for_file_data.d.f != nullptr) { - tensor_data = raw_data; - MoveOrtCallback(deleter_for_file_data.d, deleter); - } else { - void* preallocated = m.GetBuffer(); - size_t preallocated_size = m.GetLen(); - int64_t tensor_size = 1; - { - for (auto i : tensor_proto.dims()) { - if (i < 0) - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "tensor can't contain negative dims"); - tensor_size *= i; - } - } - // tensor_size could be zero. see test_slice_start_out_of_bounds\test_data_set_0\output_0.pb - if (static_cast(tensor_size) > SIZE_MAX) { - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "size overflow"); - } - size_t size_to_allocate; - if (!IAllocator::CalcMemSizeForArrayWithAlignment<0>(static_cast(tensor_size), type->Size(), - &size_to_allocate)) { - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "size overflow"); - } - - if (preallocated && preallocated_size < size_to_allocate) - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "The buffer planner is not consistent with tensor buffer size, expected ", - size_to_allocate, ", got ", preallocated_size); - switch (tensor_proto.data_type()) { - CASE_PROTO(FLOAT, float); - CASE_PROTO(DOUBLE, double); - CASE_PROTO(BOOL, bool); - CASE_PROTO(INT8, int8_t); - CASE_PROTO(INT16, int16_t); - CASE_PROTO(INT32, int32_t); - CASE_PROTO(INT64, int64_t); - CASE_PROTO(UINT8, uint8_t); - CASE_PROTO(UINT16, uint16_t); - CASE_PROTO(UINT32, uint32_t); - CASE_PROTO(UINT64, uint64_t); - CASE_PROTO(FLOAT16, MLFloat16); - CASE_PROTO(BFLOAT16, BFloat16); - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_STRING: - if (preallocated != nullptr) { - OrtStatus* status = OrtInitializeBufferForTensor(preallocated, preallocated_size, ele_type); - if (status != nullptr) { - OrtApis::ReleaseStatus(status); - return Status(common::ONNXRUNTIME, common::FAIL, "initialize preallocated buffer failed"); - } - - deleter.f = UnInitTensor; - deleter.param = new UnInitializeParam{preallocated, preallocated_size, ele_type}; - } - ORT_RETURN_IF_ERROR(UnpackTensor(tensor_proto, raw_data, raw_data_len, - static_cast(preallocated), - static_cast(tensor_size))); - break; - default: { - std::ostringstream ostr; - ostr << "Initialized tensor with unexpected type: " << tensor_proto.data_type(); - return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, ostr.str()); - } - } - tensor_data = preallocated; - } + const ONNX_NAMESPACE::TensorProto& tensor_proto, + const MemBuffer& m, OrtValue& value) { + if (m.GetBuffer() == nullptr) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, + "TensorProtoToMLValue() must take a pre-allocated MemBuffer!"); } - std::vector tensor_shape_vec = GetTensorShapeFromTensorProto(tensor_proto); + + ONNXTensorElementDataType ele_type = utils::GetTensorElementType(tensor_proto); + if (ele_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING) { + return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "string tensor can not use pre-allocated buffer"); + } + // Note: We permit an empty tensor_shape_vec, and treat it as a scalar (a tensor of size 1). - TensorShape tensor_shape{tensor_shape_vec}; + TensorShape tensor_shape{GetTensorShapeFromTensorProto(tensor_proto)}; + const DataTypeImpl* const type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type())->GetElementType(); + std::unique_ptr tensorp = onnxruntime::make_unique(type, tensor_shape, m.GetBuffer(), m.GetAllocInfo()); + if (tensorp->SizeInBytes() > m.GetLen()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "The preallocated buffer is too small. Requires ", + tensorp->SizeInBytes(), ", Got ", m.GetLen()); + } + + TensorProtoToTensor(env, model_path, tensor_proto, *tensorp); auto ml_tensor = DataTypeImpl::GetType(); - value.Init(new Tensor(type, tensor_shape, tensor_data, allocator), ml_tensor, - ml_tensor->GetDeleteFunc()); + value.Init(tensorp.release(), ml_tensor, ml_tensor->GetDeleteFunc()); return Status::OK(); } #ifdef _MSC_VER diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index 1f3f0c844a..f429349c5f 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -32,15 +32,27 @@ class Tensor; namespace utils { TensorShape GetTensorShapeFromTensorShapeProto(const ONNX_NAMESPACE::TensorShapeProto& tensor_shape_proto); -/** +std::vector GetTensorShapeFromTensorProto(const ONNX_NAMESPACE::TensorProto& tensor_proto); + + /** * deserialize a TensorProto into a preallocated memory buffer. * \param tensor_proto_path A local file path of where the 'input' was loaded from. Can be NULL if the tensor proto doesn't * have any external data or it was loaded from current working dir. This path could be either a * relative path or an absolute path. */ common::Status TensorProtoToMLValue(const Env& env, const ORTCHAR_T* tensor_proto_path, - const ONNX_NAMESPACE::TensorProto& input, const MemBuffer& m, OrtValue& value, - OrtCallback& deleter); + const ONNX_NAMESPACE::TensorProto& input, const MemBuffer& m, OrtValue& value); +/** + * @brief Deserialize a TensorProto into a preallocated empty Tensor + * @param env + * @param model_path + * @param tensor_proto source data + * @param tensorp destination empty tensor + * @return +*/ +common::Status TensorProtoToTensor(const Env& env, const ORTCHAR_T* model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, + Tensor& tensor); /** Creates a TensorProto from a Tensor. @param[in] tensor the Tensor whose data and shape will be used to create the TensorProto. diff --git a/onnxruntime/core/optimizer/optimizer_execution_frame.cc b/onnxruntime/core/optimizer/optimizer_execution_frame.cc index fcbe59d67d..4167cf70a0 100644 --- a/onnxruntime/core/optimizer/optimizer_execution_frame.cc +++ b/onnxruntime/core/optimizer/optimizer_execution_frame.cc @@ -41,18 +41,14 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, OrtValue ort_value; std::unique_ptr data(new char[cpu_tensor_length]); std::unique_ptr p_tensor; - OrtCallback d; ORT_RETURN_IF_ERROR(utils::TensorProtoToMLValue(Env::Default(), model_path.IsEmpty() ? nullptr : model_path.ToPathString().c_str(), tensor_proto, MemBuffer(data.get(), cpu_tensor_length, allocator_ptr_->Info()), - ort_value, - d)); + ort_value)); initializers_[idx] = ort_value; buffer_for_initialized_tensors_[idx] = std::move(data); - if (d.f != nullptr) - deleter_for_initialized_tensors_[idx] = d; } return Status::OK(); diff --git a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc index 01c237e0f5..7f73a142f3 100644 --- a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc +++ b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc @@ -96,6 +96,13 @@ Status CudnnRnnBase::ReorganizeWeights(const Tensor* W, const Tensor* R, cons // Prepare the weight data reorganized_w_data = GetScratchBuffer(w_size * sizeof(T)); + // In many cases, this allocation is bigger than needed, leaving part of + // the buffer unintialized. non-zero garbage data leads to wrong result + // in call to cudnnRNNForwardInference() + // TODO! refine allocation size for each case. + cudaMemset(reorganized_w_data.get(), 0, w_size * sizeof(T)); + + const T* W_data = W->template Data(); const T* R_data = R->template Data(); const T* B_data = B == nullptr ? nullptr : B->template Data(); diff --git a/onnxruntime/test/framework/test_tensor_loader.cc b/onnxruntime/test/framework/test_tensor_loader.cc index a2e979ba3a..940e3735d3 100644 --- a/onnxruntime/test/framework/test_tensor_loader.cc +++ b/onnxruntime/test/framework/test_tensor_loader.cc @@ -29,15 +29,11 @@ TEST(CApiTensorTest, load_simple_float_tensor_not_enough_space) { // deserialize it std::vector output(1); OrtValue value; - auto deleter = onnxruntime::make_unique(); OrtMemoryInfo cpu_memory_info(onnxruntime::CPU, OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeDefault); auto st = utils::TensorProtoToMLValue(Env::Default(), nullptr, p, - MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value, *deleter); + MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value); // check the result ASSERT_FALSE(st.IsOK()); - if (deleter->f) { - OrtRunCallback(deleter.release()); - } } TEST(CApiTensorTest, load_simple_float_tensor) { @@ -54,10 +50,9 @@ TEST(CApiTensorTest, load_simple_float_tensor) { // deserialize it std::vector output(3); OrtValue value; - auto deleter = onnxruntime::make_unique(); OrtMemoryInfo cpu_memory_info(onnxruntime::CPU, OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeDefault); auto st = utils::TensorProtoToMLValue(Env::Default(), nullptr, p, - MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value, *deleter); + MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value); ASSERT_TRUE(st.IsOK()) << st.ErrorMessage(); float* real_output; auto ort_st = g_ort->GetTensorMutableData(&value, (void**)&real_output); @@ -67,9 +62,6 @@ TEST(CApiTensorTest, load_simple_float_tensor) { ASSERT_EQ(real_output[1], 2.2f); ASSERT_EQ(real_output[2], 3.5f); g_ort->ReleaseStatus(ort_st); - if (deleter->f) { - OrtRunCallback(deleter.release()); - } } template @@ -113,10 +105,9 @@ static void run_external_data_test() { #endif } OrtValue value; - auto deleter = onnxruntime::make_unique(); OrtMemoryInfo cpu_memory_info(onnxruntime::CPU, OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeDefault); auto st = utils::TensorProtoToMLValue(Env::Default(), nullptr, p, - MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value, *deleter); + MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value); ASSERT_TRUE(st.IsOK()) << st.ErrorMessage(); float* real_output; auto ort_st = g_ort->GetTensorMutableData(&value, (void**)&real_output); @@ -126,9 +117,6 @@ static void run_external_data_test() { ASSERT_EQ(real_output[1], 2.2f); ASSERT_EQ(real_output[2], 3.5f); g_ort->ReleaseStatus(ort_st); - if (deleter->f) { - OrtRunCallback(deleter.release()); - } } TEST(CApiTensorTest, load_float_tensor_with_external_data) { @@ -167,10 +155,9 @@ TEST(CApiTensorTest, load_huge_tensor_with_external_data) { // deserialize it std::vector output(total_ele_count); OrtValue value; - auto deleter = onnxruntime::make_unique(); OrtMemoryInfo cpu_memory_info(onnxruntime::CPU, OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeDefault); auto st = utils::TensorProtoToMLValue(Env::Default(), nullptr, p, - MemBuffer(output.data(), output.size() * sizeof(int), cpu_memory_info), value, *deleter); + MemBuffer(output.data(), output.size() * sizeof(int), cpu_memory_info), value); // check the result ASSERT_TRUE(st.IsOK()) << "Error from TensorProtoToMLValue: " << st.ErrorMessage(); @@ -181,9 +168,6 @@ TEST(CApiTensorTest, load_huge_tensor_with_external_data) { ASSERT_EQ(1, buffer[i]); } g_ort->ReleaseStatus(ort_st); - if (deleter->f) { - OrtRunCallback(deleter.release()); - } } #endif #endif diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index c3df039a9c..688201dcc7 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -1225,7 +1225,6 @@ Status WithOrtValuesFromTensorProtos( NameMLValMap name_to_ort_value{}; std::vector> tensor_buffers{}; - std::vector tensor_deleters{}; for (const auto& tensor_proto : tensor_protos) { const auto* tensor_type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type()); @@ -1239,16 +1238,13 @@ Status WithOrtValuesFromTensorProtos( const MemBuffer mem_buffer{tensor_buffer.data(), tensor_buffer.size(), cpu_alloc_info}; OrtValue ort_value; - OrtCallback callback; ORT_RETURN_IF_ERROR(utils::TensorProtoToMLValue( Env::Default(), model_location.c_str(), tensor_proto, mem_buffer, - ort_value, callback)); - ScopedOrtCallbackInvoker callback_invoker{callback}; + ort_value)); name_to_ort_value.emplace(tensor_proto.name(), ort_value); tensor_buffers.emplace_back(std::move(tensor_buffer)); - tensor_deleters.emplace_back(std::move(callback_invoker)); } ORT_RETURN_IF_ERROR(use_name_to_ort_value_fn(name_to_ort_value)); diff --git a/orttraining/orttraining/models/runner/training_util.cc b/orttraining/orttraining/models/runner/training_util.cc index 66086dae43..5000aebe44 100644 --- a/orttraining/orttraining/models/runner/training_util.cc +++ b/orttraining/orttraining/models/runner/training_util.cc @@ -52,15 +52,11 @@ common::Status DataSet::AddData(const vector& featu OrtValue ort_value; OrtMemoryInfo info("Cpu", OrtDeviceAllocator, OrtDevice{}, 0, OrtMemTypeDefault); std::unique_ptr buffer(new char[cpu_tensor_length]); - OrtCallback deleter; ORT_RETURN_IF_ERROR(utils::TensorProtoToMLValue( - Env::Default(), nullptr, tensor_proto, MemBuffer(buffer.get(), cpu_tensor_length, info), ort_value, deleter)); + Env::Default(), nullptr, tensor_proto, MemBuffer(buffer.get(), cpu_tensor_length, info), ort_value)); sample->push_back(ort_value); ortvalue_buffers_.emplace_back(std::move(buffer)); - if (deleter.f != nullptr) { - ortvalue_deleters_.emplace_back(deleter); - } } data_.emplace_back(move(sample)); diff --git a/orttraining/orttraining/test/framework/checkpointing_test.cc b/orttraining/orttraining/test/framework/checkpointing_test.cc index 369780dcf8..1a2b64552b 100644 --- a/orttraining/orttraining/test/framework/checkpointing_test.cc +++ b/orttraining/orttraining/test/framework/checkpointing_test.cc @@ -53,7 +53,6 @@ void CompareOrtValuesToTensorProtoValues( NameMLValMap name_to_ort_value_from_tensor_proto{}; std::vector> tensor_buffers{}; - std::vector tensor_deleters{}; for (const auto& name_and_tensor_proto : name_to_tensor_proto) { const auto& name = name_and_tensor_proto.first; @@ -63,14 +62,11 @@ void CompareOrtValuesToTensorProtoValues( std::vector tensor_buffer(shape.Size() * sizeof(float)); MemBuffer m(tensor_buffer.data(), tensor_buffer.size(), cpu_alloc_info); OrtValue ort_value; - OrtCallback callback; ASSERT_STATUS_OK(utils::TensorProtoToMLValue( - Env::Default(), model_path.c_str(), tensor_proto, m, ort_value, callback)); - ScopedOrtCallbackInvoker callback_invoker{callback}; + Env::Default(), model_path.c_str(), tensor_proto, m, ort_value)); name_to_ort_value_from_tensor_proto.emplace(name, ort_value); tensor_buffers.emplace_back(std::move(tensor_buffer)); - tensor_deleters.emplace_back(std::move(callback_invoker)); } for (const auto& name_and_ort_value : name_to_ort_value) {