Use IArenaAllocator::Reserve for initializers and mem pattern planner blocks (#2835)

* Use IArenaAllocator::Reserve for initializers and mem pattern planner blocks.

onnxruntime/core/framework/bfc_arena.cc

@@ -106,7 +106,7 @@ bool BFCArena::Extend(size_t rounded_bytes) {
     return false;
   }
 
-  // we allocated the same number of bytes as the current region, so we have 2x that now
+  // if we didn't update already, default to growing by 2x next time
   if (!increased_allocation) {
     curr_region_allocation_bytes_ *= 2;
   }

onnxruntime/core/framework/execution_frame.cc

@@ -217,9 +217,13 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const
         for (size_t i = 0; i < mem_patterns_->locations.size(); i++) {
           ORT_ENFORCE(buffers_.find(mem_patterns_->locations[i]) == buffers_.end());
           AllocatorPtr alloc = GetAllocator(mem_patterns_->locations[i]);
-          void* buffer = mem_patterns_->patterns[i].PeakSize() > 0
-                             ? alloc->Alloc(mem_patterns_->patterns[i].PeakSize())
-                             : nullptr;
+
+          void* buffer = nullptr;
+          size_t peak_size = mem_patterns_->patterns[i].PeakSize();
+          if (peak_size > 0) {
+            buffer = utils::AllocateBlock(*alloc, peak_size);
+          }
+
           buffers_[mem_patterns_->locations[i]] = BufferUniquePtr(buffer, alloc);
         }
       }
@@ -374,7 +378,7 @@ static Status AllocateTraditionalMLValue(OrtValue& ort_value, const NonTensorTyp
   return Status::OK();
 }
 
-static Status AllocateTensorSequence (OrtValue& ort_value) {
+static Status AllocateTensorSequence(OrtValue& ort_value) {
   auto ml_tensor_sequence = DataTypeImpl::GetType<TensorSeq>();
   auto p_tensor_sequence = onnxruntime::make_unique<TensorSeq>();
   ort_value.Init(p_tensor_sequence.release(), ml_tensor_sequence, ml_tensor_sequence->GetDeleteFunc());

onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h

@@ -30,7 +30,8 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator {
                       "Failed to get allocator for location: " + location.ToString());
 
       if (mem_patterns_.patterns[i].PeakSize() > 0) {
-        void* buffer = alloc->Alloc(mem_patterns_.patterns[i].PeakSize());
+        void* buffer = utils::AllocateBlock(*alloc, mem_patterns_.patterns[i].PeakSize());
+
         weights_buffers_.push_back(BufferUniquePtr(buffer, alloc));
         auto kvp = buffers_.insert(std::make_pair(location, buffer));
         if (!kvp.second) {

onnxruntime/core/framework/utils.cc

@@ -6,6 +6,7 @@
 #include <iomanip>
 
 #include "core/graph/graph_viewer.h"
+#include "core/framework/arena.h"
 #include "core/framework/data_transfer_manager.h"
 #include "core/framework/execution_frame.h"
 #include "core/framework/execution_providers.h"
@@ -122,6 +123,19 @@ common::Status AllocateHelper(const IExecutionProvider& execution_provider, cons
   return Status::OK();
 }
 
+void* AllocateBlock(IAllocator& allocator, size_t size) {
+  // use the normal alloc for small allocations, or if there is no arena in use
+  if (size > 1024 * 1024) {
+    IArenaAllocator* arena_alloc = dynamic_cast<IArenaAllocator*>(&allocator);
+    if (arena_alloc) {
+      // note: Reserve allocates a block that is used directly by the caller
+      return arena_alloc->Reserve(size);
+    }
+  }
+
+  return allocator.Alloc(size);
+}
+
 const std::string& GetNodeInputProviderType(const SessionState::NodeInfo& info) {
   // the input index will be std::numeric_limits<size_t>::max() if it's an implicit input to a control flow node.
   // the input will be processed fully when executing the subgraph that consumes the implicit input.

onnxruntime/core/framework/utils.h

@@ -43,6 +43,11 @@ AllocatorPtr GetAllocator(const SessionState& session_state, const OrtMemoryInfo
 common::Status AllocateHelper(const IExecutionProvider& execution_provider, int device_id, const Tensor& fetched_tensor,
                               OrtValue& output_mlvalue);
 
+// Allocate a block using IArenaAllocator::Reserve if possible.
+// Use this for large blocks that are atypical and should not affect the growth rate of the arena
+// such as the initializers and memory pattern planner blocks.
+void* AllocateBlock(IAllocator& allocator, size_t size);
+
 const std::string& GetNodeInputProviderType(const SessionState::NodeInfo& info);
 
 common::Status CopyOneInputAcrossDevices(const SessionState& session_state, const std::string& input_name,

onnxruntime/test/framework/bfc_arena_test.cc

@@ -2,7 +2,10 @@
 // Licensed under the MIT License.
 
 #include "core/framework/bfc_arena.h"
+#include "core/framework/utils.h"
 #include "gtest/gtest.h"
+#include "test_utils.h"
+
 #include <cstdlib>
 
 namespace onnxruntime {
@@ -234,5 +237,42 @@ TEST(BFCArenaTest, TestReserve) {
   a.GetStats(&stats);
   EXPECT_EQ(stats.total_allocated_bytes, 1048576);
 }
+
+// arena is disabled for CPUExecutionProvider on x86 and JEMalloc
+#if (defined(__amd64__) || defined(_M_AMD64)) && !defined(USE_JEMALLOC)
+TEST(BFCArenaTest, UtilsAllocateBlockTest) {
+  // use a local CPUExecutionProvider so it has a clean arena.
+  CPUExecutionProviderInfo info;
+  info.create_arena = true;
+  CPUExecutionProvider cpu_provider(info);
+
+  auto cpu_arena = cpu_provider.GetAllocator(0, OrtMemTypeDefault);
+
+  EXPECT_EQ(cpu_arena->Info().alloc_type, OrtAllocatorType::OrtArenaAllocator);
+
+  // AllocateBlock should use Reserve so the growth rate of the arena doesn't change
+  size_t block_size = 4 * 1024 * 1024;
+  void* block = utils::AllocateBlock(*cpu_arena, block_size);
+  EXPECT_TRUE(block);
+
+  // if Reserve was called the initial allocation via Alloc should allocate 1MB of memory
+  // if Reserve was not called it would grow by another 4MB
+  size_t initial_extend_size = 1024 * 1024;
+  size_t size = 1024;
+  auto bytes = cpu_arena->Alloc(size);
+  EXPECT_TRUE(bytes);
+
+  AllocatorStats stats;
+  BFCArena* arena = dynamic_cast<BFCArena*>(&*cpu_arena);
+  EXPECT_TRUE(arena);
+  arena->GetStats(&stats);
+
+  EXPECT_EQ(size_t(stats.total_allocated_bytes), block_size + initial_extend_size);
+
+  cpu_arena->Free(block);
+  cpu_arena->Free(bytes);
+}
+#endif
+
 }  // namespace test
 }  // namespace onnxruntime