diff --git a/onnxruntime/core/framework/bfc_arena.cc b/onnxruntime/core/framework/bfc_arena.cc index 4b6fdb647c..c6a569e5d6 100644 --- a/onnxruntime/core/framework/bfc_arena.cc +++ b/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; } diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index 2cb6cd001c..7d6bf2385a 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/onnxruntime/core/framework/execution_frame.cc @@ -217,9 +217,13 @@ ExecutionFrame::ExecutionFrame(const std::vector& 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(); auto p_tensor_sequence = onnxruntime::make_unique(); ort_value.Init(p_tensor_sequence.release(), ml_tensor_sequence, ml_tensor_sequence->GetDeleteFunc()); diff --git a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h index 11b874f652..8f0a714f26 100644 --- a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h +++ b/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) { diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index 64b6007719..821825ae69 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -6,6 +6,7 @@ #include #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(&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::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. diff --git a/onnxruntime/core/framework/utils.h b/onnxruntime/core/framework/utils.h index 4383d266e0..35f2435fa5 100644 --- a/onnxruntime/core/framework/utils.h +++ b/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, diff --git a/onnxruntime/test/framework/bfc_arena_test.cc b/onnxruntime/test/framework/bfc_arena_test.cc index 62e87f5b56..2ace528da3 100644 --- a/onnxruntime/test/framework/bfc_arena_test.cc +++ b/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 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(&*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