From 7905c57f43eaccab9f7bf738d34c6c2dc8409082 Mon Sep 17 00:00:00 2001 From: Wei-Sheng Chin Date: Thu, 6 Aug 2020 18:42:05 -0700 Subject: [PATCH] Revert "Remove code which is not thread-safe. (#4454)" (#4712) * Revert "Remove code which is not thread-safe. (#4454)" This reverts commit 5222b2c6c0366033aaa78b95b2a0a49984984203. * Resolve race condition * More thread-safe changes * Remove unused lock Polish comments --- onnxruntime/core/framework/execution_frame.cc | 19 +++++++++- onnxruntime/core/framework/execution_frame.h | 36 +++++++++++++++++++ .../core/framework/finalize_session_state.cc | 8 ++++- .../core/framework/sequential_executor.cc | 10 ++++++ .../core/framework/simple_tensor_allocator.h | 5 ++- onnxruntime/core/framework/tensor_allocator.h | 5 ++- .../tensor_allocator_with_mem_pattern.h | 9 +++-- 7 files changed, 85 insertions(+), 7 deletions(-) diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index 4783c1f118..097c039323 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/onnxruntime/core/framework/execution_frame.cc @@ -261,9 +261,17 @@ ExecutionFrame::ExecutionFrame(const std::vector& feed_mlvalue_idxs, const // it's less efficient (the arena will add some overhead to coalesce individual allocations // back into blocks on 'free'), but better than failing completely. try { - // static_activation_memory_in_bytes_ is max virtual memory size the planner computes auto peak_size = mem_patterns_->patterns[i].PeakSize(); // Planning of one memory type should only happen once. + ORT_ENFORCE( + static_activation_memory_sizes_in_byte_.find(location.name) == + static_activation_memory_sizes_in_byte_.end(), + "Memory type ", + location.name, + " should only appear once."); + // static_activation_memory_in_bytes_ is max virtual memory size the planner computes. + // Memory dynamically allocated when executing kernels is not recorded using this field. + static_activation_memory_sizes_in_byte_[location.name] = peak_size; buffer = alloc->Alloc(peak_size); // handle allocator that doesn't throw if (buffer == nullptr) { @@ -387,6 +395,15 @@ Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(OrtValue& ort_va TraceAllocate(ort_value_index, size); } + + { + // This code block is not thread-safe. + // Dynamic activation size would be accessed by multiple threads + // if parallel executor is used. + std::unique_lock lock(mtx_); + dynamic_activation_memory_sizes_in_byte_[location.name] += size; + } + return Status::OK(); } diff --git a/onnxruntime/core/framework/execution_frame.h b/onnxruntime/core/framework/execution_frame.h index dc500314d3..91f234c5d2 100644 --- a/onnxruntime/core/framework/execution_frame.h +++ b/onnxruntime/core/framework/execution_frame.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include @@ -138,6 +139,26 @@ class ExecutionFrame final : public IExecutionFrame { // If the retrival is sucessful, this function returns true and false otherwise. bool TryGetInferredShape(int index, TensorShape& shape) const override; + // Return the size of virtual memory allocated in runtime. + // The memory is usually used for activations in forward and backward passes. + const std::unordered_map& GetDynamicMemorySizeInfo() { + // This function is not thread-safe. Please make sure dynamic_activation_memory_sizes_in_byte_ + // is not being changed when calling this function. + // If one day, race condition happens, please uncomment the following line: + // std::unique_lock lock(mtx_); + return dynamic_activation_memory_sizes_in_byte_; + } + + // Return the size of virtual memory allocated before computation. + // The memory is usually used for activations in forward and backward passes. + const std::unordered_map& GetStaticMemorySizeInfo() { + // This function is not thread-safe. Please make sure static_activation_memory_sizes_in_byte_ + // is not being changed when calling this function. + // If one day, race condition happens, please uncomment the following line: + // std::unique_lock lock(mtx_); + return static_activation_memory_sizes_in_byte_; + } + private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ExecutionFrame); @@ -183,5 +204,20 @@ class ExecutionFrame final : public IExecutionFrame { // by i, if the key i exists. // inferred_shapes_ is generated togehter with mem_patterns_. std::unordered_map inferred_shapes_; + + // Size of virtual memory allocated before any kernel execution. + // This field is not physical memory size. + // static_activation_memory_sizes_in_byte_[location] is the static memory consumption on "location". + std::unordered_map static_activation_memory_sizes_in_byte_; + + // Size of virtual memory allocated during kernel execution (i.e., inside a kernel, + // we may allocate some memory for its outputs, if not planned.). + // This field is not physical memory size. + // dynamic_activation_memory_sizes_in_byte_[location] is the dynamic memory consumption on "location". + std::unordered_map dynamic_activation_memory_sizes_in_byte_; + + // Mutex which should be acquired when executing non-thread-safe member functions. + // A current example is the tracker of dynamic memory allocation. + mutable std::mutex mtx_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/framework/finalize_session_state.cc b/onnxruntime/core/framework/finalize_session_state.cc index d889b7f360..e252d11b08 100644 --- a/onnxruntime/core/framework/finalize_session_state.cc +++ b/onnxruntime/core/framework/finalize_session_state.cc @@ -189,8 +189,14 @@ common::Status SaveInitializedTensors(const Env& env, const std::basic_string planned_initializers_memory_sizes_in_byte; ORT_RETURN_IF_ERROR( - planner.FinalizePlan()); + planner.FinalizePlan(planned_initializers_memory_sizes_in_byte)); + + for (auto i : planned_initializers_memory_sizes_in_byte) { + LOGS(logger, INFO) << "[Memory] SessionStateInitializer statically allocates " + << i.second << " bytes for " << i.first << std::endl; + } OrtCallback deleter; diff --git a/onnxruntime/core/framework/sequential_executor.cc b/onnxruntime/core/framework/sequential_executor.cc index 805f669ea3..7e4610d426 100644 --- a/onnxruntime/core/framework/sequential_executor.cc +++ b/onnxruntime/core/framework/sequential_executor.cc @@ -445,6 +445,16 @@ Status SequentialExecutor::Execute(const SessionState& session_state, const std: session_state.Profiler().EndTimeAndRecordEvent(profiling::SESSION_EVENT, "SequentialExecutor::Execute", tp); } + for (auto i: frame.GetStaticMemorySizeInfo()) { + LOGS(logger, INFO) << "[Memory] ExecutionFrame statically allocates " + << i.second << " bytes for " << i.first << std::endl; + } + + for (auto i: frame.GetDynamicMemorySizeInfo()) { + LOGS(logger, INFO) << "[Memory] ExecutionFrame dynamically allocates " + << i.second << " bytes for " << i.first << std::endl; + } + return Status::OK(); } diff --git a/onnxruntime/core/framework/simple_tensor_allocator.h b/onnxruntime/core/framework/simple_tensor_allocator.h index 0a6c888f5d..ce3dc1767a 100644 --- a/onnxruntime/core/framework/simple_tensor_allocator.h +++ b/onnxruntime/core/framework/simple_tensor_allocator.h @@ -28,7 +28,10 @@ class SimpleTensorAllocator : public ITensorAllocator { weights_buffers_(weights_buffers), seq_plan_(execution_plan) {} - common::Status FinalizePlan() override { + common::Status FinalizePlan(std::unordered_map& planned_memory_sizes_in_byte) override { + // There is no memory plan to allocate a big block of memory, so + // planned memory sizes in different locations are all empty. + 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; diff --git a/onnxruntime/core/framework/tensor_allocator.h b/onnxruntime/core/framework/tensor_allocator.h index 12e92bd90c..c6d7fbfb0c 100644 --- a/onnxruntime/core/framework/tensor_allocator.h +++ b/onnxruntime/core/framework/tensor_allocator.h @@ -25,10 +25,13 @@ class ITensorAllocator { AllocatorPtr GetAllocator(const OrtMemoryInfo& memory_info); /** + * + * \param planned_memory_sizes_in_byte The sizes of memory allocated inside FinalizePlan on different devices. + * * When there is no more tensor to trace, call this function to finalize the * allocation. */ - virtual common::Status FinalizePlan() = 0; + virtual common::Status FinalizePlan(std::unordered_map& planned_memory_sizes_in_byte) = 0; /** * diff --git a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h index 02b7f88ed4..4ab44d29ba 100644 --- a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h +++ b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h @@ -21,7 +21,8 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { bool is_sealed_ = false; const ExecutionPlanBase& seq_plan_; - common::Status AllocatePlannedBuffersAndReportTotalSize() { + common::Status AllocatePlannedBuffersAndReportTotalSize( + std::unordered_map& planned_memory_sizes_in_byte) { const size_t location_len = mem_patterns_.locations.size(); for (size_t i = 0; i < location_len; ++i) { auto& location = mem_patterns_.locations[i]; @@ -51,6 +52,8 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { alloc->Free(buffer); return Status(common::ONNXRUNTIME, common::FAIL, "duplicated location"); } + + planned_memory_sizes_in_byte[location.name] += peak_size; } return Status::OK(); } @@ -63,9 +66,9 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { weights_buffers_(weights_buffers), seq_plan_(execution_plan) {} - common::Status FinalizePlan() override { + common::Status FinalizePlan(std::unordered_map& planned_memory_sizes_in_byte) override { ORT_RETURN_IF_ERROR(planner_.GeneratePatterns(&mem_patterns_)); - ORT_RETURN_IF_ERROR(AllocatePlannedBuffersAndReportTotalSize()); + ORT_RETURN_IF_ERROR(AllocatePlannedBuffersAndReportTotalSize(planned_memory_sizes_in_byte)); is_sealed_ = true; return Status::OK(); }