Revert "Remove code which is not thread-safe. (#4454)" (#4712)

* Revert "Remove code which is not thread-safe. (#4454)"

This reverts commit 5222b2c6c0.

* Resolve race condition

* More thread-safe changes

* Remove unused lock

Polish comments

onnxruntime/core/framework/execution_frame.cc

@@ -261,9 +261,17 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& 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<std::mutex> lock(mtx_);
+    dynamic_activation_memory_sizes_in_byte_[location.name] += size;
+  }
+
   return Status::OK();
 }
 

onnxruntime/core/framework/execution_frame.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <mutex>
 #include <vector>
 #include <unordered_map>
 
@@ -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<std::string, size_t>& 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<std::mutex> 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<std::string, size_t>& 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<std::mutex> 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<int, TensorShape> 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<std::string, size_t> 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<std::string, size_t> 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

onnxruntime/core/framework/finalize_session_state.cc

@@ -189,8 +189,14 @@ common::Status SaveInitializedTensors(const Env& env, const std::basic_string<PA
   //2. allocate weight buffer on different locations
   // planned_initializers_memory_size_in_byte is not actual physical size.
   // It's the virtual size computed by planner.
+  std::unordered_map<std::string, size_t> 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;
 

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();
 }
 

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<std::string, size_t>& 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<std::string, size_t>();
     return Status::OK();
   }
   common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::unique_ptr<MemBuffer>& out) override;

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<std::string, size_t>& planned_memory_sizes_in_byte) = 0;
 
   /**
    *

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<std::string, size_t>& 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<std::string, size_t>& 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();
   }