diff --git a/onnxruntime/core/framework/bfc_arena.cc b/onnxruntime/core/framework/bfc_arena.cc index ea7fb9e34e..c6c1c2d537 100644 --- a/onnxruntime/core/framework/bfc_arena.cc +++ b/onnxruntime/core/framework/bfc_arena.cc @@ -114,7 +114,7 @@ bool BFCArena::Extend(size_t rounded_bytes) { ORT_THROW("Failed to allocate memory for requested buffer of size ", rounded_bytes); } - // if we didn't update already, default to growing by 2x next time + // we allocated the same number of bytes as the current region, so we have 2x that now 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 370cbfa010..526ccfb621 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/onnxruntime/core/framework/execution_frame.cc @@ -217,13 +217,9 @@ 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 = nullptr; - size_t peak_size = mem_patterns_->patterns[i].PeakSize(); - if (peak_size > 0) { - buffer = utils::AllocateBlock(*alloc, peak_size); - } - + void* buffer = mem_patterns_->patterns[i].PeakSize() > 0 + ? alloc->Alloc(mem_patterns_->patterns[i].PeakSize()) + : nullptr; buffers_[mem_patterns_->locations[i]] = BufferUniquePtr(buffer, alloc); } } @@ -381,7 +377,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 8f0a714f26..11b874f652 100644 --- a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h +++ b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h @@ -30,8 +30,7 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { "Failed to get allocator for location: " + location.ToString()); if (mem_patterns_.patterns[i].PeakSize() > 0) { - void* buffer = utils::AllocateBlock(*alloc, mem_patterns_.patterns[i].PeakSize()); - + void* buffer = alloc->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 821825ae69..64b6007719 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -6,7 +6,6 @@ #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" @@ -123,19 +122,6 @@ 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 35f2435fa5..4383d266e0 100644 --- a/onnxruntime/core/framework/utils.h +++ b/onnxruntime/core/framework/utils.h @@ -43,11 +43,6 @@ 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 5c6d15b7ac..7051f17246 100644 --- a/onnxruntime/test/framework/bfc_arena_test.cc +++ b/onnxruntime/test/framework/bfc_arena_test.cc @@ -2,10 +2,7 @@ // 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 { @@ -237,42 +234,5 @@ 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