saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Release buffers for prepacked tensors (#6820)

Browse source 
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.
This commit is contained in:
Chen Fu 2021-03-10 14:07:20 -08:00 committed by GitHub
parent 2f307dd223
commit 4a4488baae
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 277 additions and 277 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
onnxruntime/core/framework/execution_providers.h
View file

@ -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<std::string>& GetIds() const { return exec_provider_ids_; }

21
onnxruntime/core/framework/session_state.cc
View file

@ -1004,8 +1004,27 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
MemoryInfo::GenerateTensorMap(GetExecutionPlan(), GetOrtValueNameIdxMap());
#endif
// Memory pattern tracer allocates all initializers on a single continous
// buffer. This has the effect of reducing memory fragementation.
// Further more, NCCL kernels require initializers to be allocated
// continously.
//
// In inferencing scenarios, however, we often want to pre-process and then
// release some initializers. See OpKernel::PrePack(). Letting all initializers
// sharing a single buffer makes it hard to release individual ones, leading
// to memory waste.
//
// TODO!! disabling memory pattern tracer increases fragementation, leading to
// out of memory error in some training tests. Need to create kernel first,
// and let the kernel tells us whether the initalizer needs to be traced.
//
#if defined(ENABLE_TRAINING)
std::unique_ptr<ITensorAllocator> tensor_allocator(
ITensorAllocator::Create(enable_mem_pattern_, *p_seq_exec_plan_, *this, weights_buffers_));
#else
std::unique_ptr<ITensorAllocator> 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<PATH_CHAR_
ORT_RETURN_IF_ERROR(
session_state_utils::SaveInitializedTensors(
Env::Default(), graph_location, *graph_viewer_,
execution_providers_.GetDefaultCpuMemoryInfo(),
execution_providers_.GetDefaultCpuAllocator(),
ort_value_name_idx_map_, initializer_allocation_order, *tensor_allocator,
[this](int idx, const OrtValue& value, const OrtCallback& d, bool constant) -> Status {
return AddInitializedTensor(idx, value, &d, constant);

115
onnxruntime/core/framework/session_state_utils.cc
View file

@ -31,58 +31,54 @@ namespace onnxruntime {
namespace session_state_utils {
static common::Status DeserializeTensorProto(const Env& env, const std::basic_string<PATH_CHAR_TYPE>& 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<char[]> 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<Tensor> 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<Tensor>();
p_tensor = onnxruntime::make_unique<Tensor>(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<Tensor>(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<Tensor>(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<Tensor> p_deserialize_tensor = onnxruntime::make_unique<Tensor>(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<Tensor>();
@ -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<PATH_CHAR_TYPE>& 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<OrtValueIndex>& 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<Tensor>().Location();
const auto& planned_mem_info = exec_plan.GetLocation(ort_value_index);
const auto& user_mem_info = it->second->Get<Tensor>().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<MemBuffer> 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;

2
onnxruntime/core/framework/session_state_utils.h
View file

@ -32,7 +32,7 @@ class Logger;
namespace session_state_utils {
common::Status SaveInitializedTensors(
const Env& env, const std::basic_string<PATH_CHAR_TYPE>& 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<OrtValueIndex>& initializer_allocation_order,
ITensorAllocator& planner,
const std::function<Status(int idx, const OrtValue& value, const OrtCallback& d, bool constant)>& save_tensor_func,

28
onnxruntime/core/framework/simple_tensor_allocator.cc
View file

@ -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<MemBuffer>& 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<kAllocAlignment>(*iter->second, &len));
common::Status SimpleTensorAllocator::GetPreallocatedBuffer(int ort_value_index, const char* /*name*/,
std::unique_ptr<MemBuffer>& /*buf_out*/,
AllocatorPtr& alloc_out) {
const struct OrtMemoryInfo& location = seq_plan_.GetLocation(ort_value_index);
if (len == 0) {
out = onnxruntime::make_unique<MemBuffer>(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<MemBuffer>(buffer, len, location);
return Status::OK();
}
} // namespace onnxruntime

9
onnxruntime/core/framework/simple_tensor_allocator.h
View file

@ -15,17 +15,12 @@ class ExecutionProviders;
class SimpleTensorAllocator : public ITensorAllocator {
private:
MemoryPatternGroup mem_patterns_;
std::vector<BufferUniquePtr>& weights_buffers_;
const ExecutionPlanBase& seq_plan_;
private:
std::unordered_map<int, const ONNX_NAMESPACE::TensorProto*> values_;
public:
SimpleTensorAllocator(const ExecutionPlanBase& execution_plan, const SessionState& session_state,
std::vector<BufferUniquePtr>& weights_buffers)
std::vector<BufferUniquePtr>& /*weights_buffers*/)
: ITensorAllocator(session_state),
weights_buffers_(weights_buffers),
seq_plan_(execution_plan) {}
common::Status FinalizePlan(std::unordered_map<std::string, size_t>& planned_memory_sizes_in_byte) override {
@ -34,7 +29,7 @@ class SimpleTensorAllocator : public ITensorAllocator {
planned_memory_sizes_in_byte = std::unordered_map<std::string, size_t>();
return Status::OK();
}
common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::unique_ptr<MemBuffer>& out) override;
common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::unique_ptr<MemBuffer>& buf_out, AllocatorPtr& alloc_out) override;
common::Status Trace(int id, const ONNX_NAMESPACE::TensorProto* value) override;
const MemoryPatternGroup& GetMemPatterns() override {
return mem_patterns_;

21
onnxruntime/core/framework/tensor_allocator.h
View file

@ -35,15 +35,20 @@ class ITensorAllocator {
virtual common::Status FinalizePlan(std::unordered_map<std::string, size_t>& 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<MemBuffer>& out) = 0;
std::unique_ptr<MemBuffer>& buf_out,
AllocatorPtr& alloc_out) = 0;
virtual const MemoryPatternGroup& GetMemPatterns() = 0;
/**

11
onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h
View file

@ -73,7 +73,7 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator {
}
common::Status GetPreallocatedBuffer(int ort_value_index, const char* name,
std::unique_ptr<MemBuffer>& out) override {
std::unique_ptr<MemBuffer>& 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<MemBuffer>(nullptr, 0, location);
buf_out = onnxruntime::make_unique<MemBuffer>(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<MemBuffer>(reinterpret_cast<char*>(it->second) + block->offset_, block->size_, location);
buf_out = onnxruntime::make_unique<MemBuffer>(reinterpret_cast<char*>(it->second) + block->offset_, block->size_, location);
return Status::OK();
}
common::Status Trace(int id, const ONNX_NAMESPACE::TensorProto* value) override {

268
onnxruntime/core/framework/tensorprotoutils.cc
View file

@ -86,16 +86,6 @@ bool operator!=(const ONNX_NAMESPACE::TensorShapeProto_Dimension& l,
namespace {
std::vector<int64_t> GetTensorShapeFromTensorProto(const ONNX_NAMESPACE::TensorProto& tensor_proto) {
const auto& dims = tensor_proto.dims();
std::vector<int64_t> tensor_shape_vec(static_cast<size_t>(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<int64_t> GetTensorShapeFromTensorProto(const ONNX_NAMESPACE::TensorProto& tensor_proto) {
const auto& dims = tensor_proto.dims();
std::vector<int64_t> tensor_shape_vec(static_cast<size_t>(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<UnInitializeParam*>(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<size_t>(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<int64_t> 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<size_t> raw_data_len = 0;
AutoDelete deleter_for_file_data;
if (utils::HasExternalData(tensor_proto)) {
// Get the external data info
std::basic_string<ORTCHAR_T> external_data_file_path;
FileOffsetType file_offset;
std::basic_string<ORTCHAR_T> 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<char*>(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<char[]>(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<std::string>(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<uint64_t>(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<std::string>(tensor_proto, raw_data, raw_data_len,
static_cast<std::string*>(preallocated),
static_cast<size_t>(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<size_t> 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<ORTCHAR_T> external_data_file_path;
FileOffsetType file_offset;
std::basic_string<ORTCHAR_T> 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<char*>(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<char[]>(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<uint64_t>(tensor_size) > SIZE_MAX) {
return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "size overflow");
}
size_t size_to_allocate;
if (!IAllocator::CalcMemSizeForArrayWithAlignment<0>(static_cast<size_t>(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<std::string>(tensor_proto, raw_data, raw_data_len,
static_cast<std::string*>(preallocated),
static_cast<size_t>(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<int64_t> 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<Tensor> tensorp = onnxruntime::make_unique<Tensor>(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<Tensor>();
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

18
onnxruntime/core/framework/tensorprotoutils.h
View file

@ -32,15 +32,27 @@ class Tensor;
namespace utils {
TensorShape GetTensorShapeFromTensorShapeProto(const ONNX_NAMESPACE::TensorShapeProto& tensor_shape_proto);
/**
std::vector<int64_t> 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.

6
onnxruntime/core/optimizer/optimizer_execution_frame.cc
View file

@ -41,18 +41,14 @@ OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
OrtValue ort_value;
std::unique_ptr<char[]> data(new char[cpu_tensor_length]);
std::unique_ptr<Tensor> 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();

7
onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc
View file

@ -96,6 +96,13 @@ Status CudnnRnnBase<T>::ReorganizeWeights(const Tensor* W, const Tensor* R, cons
// Prepare the weight data
reorganized_w_data = GetScratchBuffer<void>(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<T>();
const T* R_data = R->template Data<T>();
const T* B_data = B == nullptr ? nullptr : B->template Data<T>();

24
onnxruntime/test/framework/test_tensor_loader.cc
View file

@ -29,15 +29,11 @@ TEST(CApiTensorTest, load_simple_float_tensor_not_enough_space) {
// deserialize it
std::vector<float> output(1);
OrtValue value;
auto deleter = onnxruntime::make_unique<onnxruntime::OrtCallback>();
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<float> output(3);
OrtValue value;
auto deleter = onnxruntime::make_unique<onnxruntime::OrtCallback>();
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 <bool use_current_dir>
@ -113,10 +105,9 @@ static void run_external_data_test() {
#endif
}
OrtValue value;
auto deleter = onnxruntime::make_unique<onnxruntime::OrtCallback>();
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<int> output(total_ele_count);
OrtValue value;
auto deleter = onnxruntime::make_unique<onnxruntime::OrtCallback>();
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

6
orttraining/orttraining/models/runner/training_runner.cc
View file

@ -1225,7 +1225,6 @@ Status WithOrtValuesFromTensorProtos(
NameMLValMap name_to_ort_value{};
std::vector<std::vector<char>> tensor_buffers{};
std::vector<ScopedOrtCallbackInvoker> 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));

6
orttraining/orttraining/models/runner/training_util.cc
View file

@ -52,15 +52,11 @@ common::Status DataSet::AddData(const vector<ONNX_NAMESPACE::TensorProto>& featu
OrtValue ort_value;
OrtMemoryInfo info("Cpu", OrtDeviceAllocator, OrtDevice{}, 0, OrtMemTypeDefault);
std::unique_ptr<char[]> 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));

6
orttraining/orttraining/test/framework/checkpointing_test.cc
View file

@ -53,7 +53,6 @@ void CompareOrtValuesToTensorProtoValues(
NameMLValMap name_to_ort_value_from_tensor_proto{};
std::vector<std::vector<char>> tensor_buffers{};
std::vector<ScopedOrtCallbackInvoker> 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<char> 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) {