From 3360024a0b2a55a9093ef477b097e4bbaa612cb5 Mon Sep 17 00:00:00 2001 From: Hariharan Seshadri Date: Thu, 22 Jul 2021 10:17:35 -0700 Subject: [PATCH] Support plugging in custom user-defined allocators for sharing between sessions (#8059) --- .../onnxruntime/core/framework/allocator.h | 15 ++ include/onnxruntime/core/framework/tensor.h | 2 +- .../onnxruntime/core/session/environment.h | 5 + .../core/session/onnxruntime_c_api.h | 49 ++++-- onnxruntime/core/framework/allocator.cc | 2 - .../framework/{arena.h => allocator_stats.h} | 31 +--- onnxruntime/core/framework/allocatormgr.cc | 10 +- onnxruntime/core/framework/allocatormgr.h | 18 +- onnxruntime/core/framework/bfc_arena.cc | 11 +- onnxruntime/core/framework/bfc_arena.h | 16 +- .../core/framework/execution_provider.cc | 73 ++++++++- .../core/framework/mimalloc_allocator.cc | 44 +++++ .../core/framework/mimalloc_allocator.h | 31 ++++ onnxruntime/core/framework/mimalloc_arena.cc | 50 ------ onnxruntime/core/framework/mimalloc_arena.h | 36 ---- .../core/framework/session_state_utils.cc | 14 +- .../tensor_allocator_with_mem_pattern.h | 4 +- .../core/session/allocator_adapters.cc | 131 +++++++++++++++ onnxruntime/core/session/allocator_adapters.h | 53 ++++++ onnxruntime/core/session/allocator_impl.h | 25 --- .../session/default_cpu_allocator_c_api.cc | 50 ++---- onnxruntime/core/session/device_allocator.cc | 58 ------- onnxruntime/core/session/device_allocator.h | 45 ----- onnxruntime/core/session/environment.cc | 63 ++++++- onnxruntime/core/session/inference_session.cc | 8 +- onnxruntime/core/session/onnxruntime_c_api.cc | 12 +- onnxruntime/core/session/ort_apis.h | 2 + onnxruntime/core/session/ort_env.cc | 6 +- onnxruntime/core/session/ort_env.h | 5 + .../test/framework/TestAllocatorManager.cc | 28 +--- .../test/framework/TestAllocatorManager.h | 5 +- .../test/framework/inference_session_test.cc | 9 +- onnxruntime/test/shared_lib/test_inference.cc | 154 +++++++++++++----- .../my_allocator.cc | 4 +- .../test/util/include/test_allocator.h | 2 + onnxruntime/test/util/test_allocator.cc | 5 + 36 files changed, 652 insertions(+), 424 deletions(-) rename onnxruntime/core/framework/{arena.h => allocator_stats.h} (63%) create mode 100644 onnxruntime/core/framework/mimalloc_allocator.cc create mode 100644 onnxruntime/core/framework/mimalloc_allocator.h delete mode 100644 onnxruntime/core/framework/mimalloc_arena.cc delete mode 100644 onnxruntime/core/framework/mimalloc_arena.h create mode 100644 onnxruntime/core/session/allocator_adapters.cc create mode 100644 onnxruntime/core/session/allocator_adapters.h delete mode 100644 onnxruntime/core/session/allocator_impl.h delete mode 100644 onnxruntime/core/session/device_allocator.cc delete mode 100644 onnxruntime/core/session/device_allocator.h diff --git a/include/onnxruntime/core/framework/allocator.h b/include/onnxruntime/core/framework/allocator.h index 3820d79606..689ab2fb05 100644 --- a/include/onnxruntime/core/framework/allocator.h +++ b/include/onnxruntime/core/framework/allocator.h @@ -5,6 +5,7 @@ #include "core/common/common.h" #include "core/framework/fence.h" +#include "core/framework/allocator_stats.h" #include "core/session/onnxruntime_c_api.h" #include "ortdevice.h" #include "ortmemoryinfo.h" @@ -55,9 +56,23 @@ class IAllocator { @remarks Use SafeInt when calculating the size of memory to allocate using Alloc. */ virtual void* Alloc(size_t size) = 0; + virtual void Free(void* p) = 0; + + // TODO: Find a better name than Reserve() and update in all places. + // Reserve() is an interface exposed for an implementation of IAllocator + // to optionally implement some allocation logic that by-passes any arena-based + // logic that may be housed in the Alloc() implementation. + // There are SessionOptions config(s) that allow users to allocate some memory + // by-passing arena-based logic. + // By default, the base implementation just calls Alloc(). + virtual void* Reserve(size_t size) { return Alloc(size); } + const OrtMemoryInfo& Info() const { return memory_info_; }; + // Each implementation of IAllocator can override and provide their own implementation + virtual void GetStats(AllocatorStats* /*stats*/) { return; } + /** optional CreateFence interface, as provider like DML has its own fence */ diff --git a/include/onnxruntime/core/framework/tensor.h b/include/onnxruntime/core/framework/tensor.h index 4b09260856..56e149dc44 100644 --- a/include/onnxruntime/core/framework/tensor.h +++ b/include/onnxruntime/core/framework/tensor.h @@ -69,7 +69,7 @@ class Tensor final { * \param p_type Data type of the tensor * \param shape Shape of the tensor * \param p_data A preallocated buffer. Can be NULL if the shape is empty. - * Tensor does not own the data and will not delete it + * Tensor will own the memory and will delete it when the tensor instance is destructed. * \param deleter Allocator used to free the pre-allocated memory * \param offset Offset in bytes to start of Tensor within p_data. */ diff --git a/include/onnxruntime/core/session/environment.h b/include/onnxruntime/core/session/environment.h index 5f7f6e2ba1..f7a4638371 100644 --- a/include/onnxruntime/core/session/environment.h +++ b/include/onnxruntime/core/session/environment.h @@ -74,6 +74,11 @@ class Environment { return shared_allocators_; } + /** + * Removes registered allocator that was previously registered for sharing between multiple sessions. + */ + Status UnregisterAllocator(const OrtMemoryInfo& mem_info); + private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Environment); diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 7b2be8b88d..10bd0f7db4 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -12,7 +12,7 @@ #include // This value is used in structures passed to ORT so that a newer version of ORT will still work with them -#define ORT_API_VERSION 8 +#define ORT_API_VERSION 9 #ifdef __cplusplus extern "C" { @@ -675,9 +675,9 @@ struct OrtApi { ORT_API2_STATUS(AllocatorFree, _Inout_ OrtAllocator* ptr, void* p); ORT_API2_STATUS(AllocatorGetInfo, _In_ const OrtAllocator* ptr, _Outptr_ const struct OrtMemoryInfo** out); + // This API returns a CPU non-arena based allocator // The returned pointer doesn't have to be freed. // Always returns the same instance on every invocation. - // Please note that this is a non-arena based allocator. ORT_API2_STATUS(GetAllocatorWithDefaultOptions, _Outptr_ OrtAllocator** out); // Override symbolic dimensions (by specific denotation strings) with actual values if known at session initialization time to enable @@ -1009,11 +1009,15 @@ struct OrtApi { ORT_API2_STATUS(AddSessionConfigEntry, _Inout_ OrtSessionOptions* options, _In_z_ const char* config_key, _In_z_ const char* config_value); - /** + /** + * This API returns an allocator bound to the provided OrtSession instance according + * to the spec within mem_info if successful * \param sess valid OrtSession instance * \param mem_info - valid OrtMemoryInfo instance - * \param - out a ptr to a new instance of OrtAllocator according to the spec within mem_info - * if successful + * \param - out a ptr to an instance of OrtAllocator which wraps the allocator + bound to the OrtSession instance + Freeing the returned pointer only frees the OrtAllocator instance and not + the wrapped session owned allocator itself. * \return OrtStatus or nullptr if successful */ ORT_API2_STATUS(CreateAllocator, _In_ const OrtSession* sess, _In_ const OrtMemoryInfo* mem_info, @@ -1124,7 +1128,8 @@ struct OrtApi { * sharing between multiple sessions that use the same env instance. * Lifetime of the created allocator will be valid for the duration of the environment. * Returns an error if an allocator with the same OrtMemoryInfo is already registered. - * \param mem_info must be non-null. + * \param env OrtEnv instance (must be non-null). + * \param mem_info (must be non-null). * \param arena_cfg if nullptr defaults will be used. * See docs/C_API.md for details. */ @@ -1390,7 +1395,7 @@ struct OrtApi { _In_ const OrtSessionOptions* options, _Inout_ OrtPrepackedWeightsContainer* prepacked_weights_container, _Outptr_ OrtSession** out); - /** + /* * Append TensorRT execution provider to the session options with TensorRT provider options. * If TensorRT is not available (due to a non TensorRT enabled build), this function will return failure. * Note: this API is slightly different than SessionOptionsAppendExecutionProvider_TensorRT. @@ -1425,9 +1430,9 @@ struct OrtApi { * \param num_keys - number of keys */ ORT_API2_STATUS(UpdateTensorRTProviderOptions, _Inout_ OrtTensorRTProviderOptionsV2* tensorrt_options, - _In_reads_(num_keys) const char* const* provider_options_keys, - _In_reads_(num_keys) const char* const* provider_options_values, - _In_ size_t num_keys); + _In_reads_(num_keys) const char* const* provider_options_keys, + _In_reads_(num_keys) const char* const* provider_options_values, + _In_ size_t num_keys); /** * Get serialized TensorRT provider options string. @@ -1446,10 +1451,32 @@ struct OrtApi { */ ORT_CLASS_RELEASE2(TensorRTProviderOptions); - /** + /* * Enable custom operators in onnxruntime-extensions: https://github.com/microsoft/onnxruntime-extensions.git */ ORT_API2_STATUS(EnableOrtCustomOps, _Inout_ OrtSessionOptions* options); + + /** + * Registers a custom allocator instance with the env to enable + * sharing between multiple sessions that use the same env instance. + * Returns an error if an allocator with the same OrtMemoryInfo is already registered. + * \param env OrtEnv instance (must be non-null). + * \param allocator user provided allocator (must be non-null). + * The behavior of this API is exactly the same as CreateAndRegisterAllocator() except + * instead of ORT creating an allocator based on provided info, in this case + * ORT uses the user-provided custom allocator. + * See docs/C_API.md for details. + */ + ORT_API2_STATUS(RegisterAllocator, _Inout_ OrtEnv* env, _In_ OrtAllocator* allocator); + + /** + * Unregisters a registered allocator for sharing across sessions + * based on provided OrtMemoryInfo. + * It is an error if you provide an OrtmemoryInfo not corresponding to any + * registered allocators for sharing. + */ + ORT_API2_STATUS(UnregisterAllocator, _Inout_ OrtEnv* env, + _In_ const OrtMemoryInfo* mem_info); }; /* diff --git a/onnxruntime/core/framework/allocator.cc b/onnxruntime/core/framework/allocator.cc index 6967d57f8d..adc2718d85 100644 --- a/onnxruntime/core/framework/allocator.cc +++ b/onnxruntime/core/framework/allocator.cc @@ -42,8 +42,6 @@ void* MiMallocAllocator::Alloc(size_t size) { void MiMallocAllocator::Free(void* p) { mi_free(p); } - -const OrtMemoryInfo& MiMallocAllocator::Info() const { return *memory_info_; } #endif void* CPUAllocator::Alloc(size_t size) { diff --git a/onnxruntime/core/framework/arena.h b/onnxruntime/core/framework/allocator_stats.h similarity index 63% rename from onnxruntime/core/framework/arena.h rename to onnxruntime/core/framework/allocator_stats.h index 7befd0e7e8..e3c942fec1 100644 --- a/onnxruntime/core/framework/arena.h +++ b/onnxruntime/core/framework/allocator_stats.h @@ -4,38 +4,9 @@ #pragma once #include - -#include "core/common/common.h" -#include "core/framework/allocator.h" +#include namespace onnxruntime { -// The interface for arena which manage memory allocations -// Arena will hold a pool of pre-allocate memories and manage their lifecycle. -// Need an underline IResourceAllocator to allocate memories. -// The setting like max_chunk_size is init by IDeviceDescriptor from resource allocator -class IArenaAllocator : public IAllocator { - public: - IArenaAllocator(const OrtMemoryInfo& info) : IAllocator(info) {} - ~IArenaAllocator() override = default; - // Alloc call needs to be thread safe. - void* Alloc(size_t size) override = 0; - // The chunk allocated by Reserve call won't be reused with other request - // (i.e.) it is not maintained by the arena and - // it will be return to the devices when it is freed. - // Reserve call needs to be thread safe. - virtual void* Reserve(size_t size) = 0; - // Free call needs to be thread safe. - void Free(void* p) override = 0; - // All unused device allocations maintained by the arena - // (i.e.) physical allocations with no chunks in use will be de-allocated. - // Shrink call needs to be thread safe. - virtual Status Shrink() = 0; - virtual size_t Used() const = 0; - virtual size_t Max() const = 0; - // allocate host pinned memory? -}; - -using ArenaPtr = std::shared_ptr; // Runtime statistics collected by an allocator. struct AllocatorStats { diff --git a/onnxruntime/core/framework/allocatormgr.cc b/onnxruntime/core/framework/allocatormgr.cc index 9157130bcc..125a450726 100644 --- a/onnxruntime/core/framework/allocatormgr.cc +++ b/onnxruntime/core/framework/allocatormgr.cc @@ -3,7 +3,7 @@ #include "core/framework/allocatormgr.h" #include "core/framework/bfc_arena.h" -#include "core/framework/mimalloc_arena.h" +#include "core/framework/mimalloc_allocator.h" #include "core/common/logging/logging.h" #include #include @@ -48,11 +48,11 @@ AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) { return nullptr; } -#ifdef USE_MIMALLOC - return std::shared_ptr( - std::make_unique(std::move(device_allocator), max_mem)); +#ifdef USE_MIMALLOC_ARENA_ALLOCATOR + return std::shared_ptr( + std::make_unique(max_mem)); #else - return std::shared_ptr( + return std::shared_ptr( std::make_unique(std::move(device_allocator), max_mem, arena_extend_str, diff --git a/onnxruntime/core/framework/allocatormgr.h b/onnxruntime/core/framework/allocatormgr.h index 2401fe3101..14ae1b62f6 100644 --- a/onnxruntime/core/framework/allocatormgr.h +++ b/onnxruntime/core/framework/allocatormgr.h @@ -19,14 +19,14 @@ using MemoryInfoSet = std::set; const int DEFAULT_CPU_ALLOCATOR_DEVICE_ID = 0; struct AllocatorCreationInfo { - AllocatorCreationInfo(AllocatorFactory device_alloc_factory0, - OrtDevice::DeviceId device_id0 = 0, - bool use_arena0 = true, - OrtArenaCfg arena_cfg0 = {0, -1, -1, -1, -1}) - : device_alloc_factory(device_alloc_factory0), - device_id(device_id0), - use_arena(use_arena0), - arena_cfg(arena_cfg0) { + AllocatorCreationInfo(AllocatorFactory device_alloc_factory, + OrtDevice::DeviceId device_id = 0, + bool use_arena = true, + OrtArenaCfg arena_cfg = {0, -1, -1, -1, -1}) + : device_alloc_factory(device_alloc_factory), + device_id(device_id), + use_arena(use_arena), + arena_cfg(arena_cfg) { } AllocatorFactory device_alloc_factory; @@ -35,7 +35,7 @@ struct AllocatorCreationInfo { OrtArenaCfg arena_cfg; }; -// Returns an allocator based on the creation info provided. +// Returns an allocator (an instance of IAllocator) based on the creation info provided. // Returns nullptr if an invalid value of info.arena_cfg.arena_extend_strategy is supplied. // Valid values can be found in onnxruntime_c_api.h. AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info); diff --git a/onnxruntime/core/framework/bfc_arena.cc b/onnxruntime/core/framework/bfc_arena.cc index eab62c08b0..57e4fea00b 100644 --- a/onnxruntime/core/framework/bfc_arena.cc +++ b/onnxruntime/core/framework/bfc_arena.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include "core/framework/allocator.h" #include "core/framework/bfc_arena.h" #include @@ -11,11 +12,11 @@ BFCArena::BFCArena(std::unique_ptr resource_allocator, int initial_chunk_size_bytes, int max_dead_bytes_per_chunk, int initial_growth_chunk_size_bytes) - : IArenaAllocator(OrtMemoryInfo(resource_allocator->Info().name, - OrtAllocatorType::OrtArenaAllocator, - resource_allocator->Info().device, - resource_allocator->Info().id, - resource_allocator->Info().mem_type)), + : IAllocator(OrtMemoryInfo(resource_allocator->Info().name, + OrtAllocatorType::OrtArenaAllocator, + resource_allocator->Info().device, + resource_allocator->Info().id, + resource_allocator->Info().mem_type)), device_allocator_(std::move(resource_allocator)), free_chunks_list_(kInvalidChunkHandle), next_allocation_id_(1), diff --git a/onnxruntime/core/framework/bfc_arena.h b/onnxruntime/core/framework/bfc_arena.h index 914af6899b..f84b9038fb 100644 --- a/onnxruntime/core/framework/bfc_arena.h +++ b/onnxruntime/core/framework/bfc_arena.h @@ -28,8 +28,8 @@ limitations under the License. #include "core/common/safeint.h" #include "core/platform/ort_mutex.h" -#include "core/framework/arena.h" #include "core/framework/arena_extend_strategy.h" +#include "core/framework/allocator.h" #if defined(PLATFORM_WINDOWS) #include @@ -50,7 +50,7 @@ namespace onnxruntime { // coalescing. One assumption we make is that the process using this // allocator owns pretty much all of the memory, and that nearly // all requests to allocate memory go through this interface. -class BFCArena : public IArenaAllocator { +class BFCArena : public IAllocator { public: static const ArenaExtendStrategy DEFAULT_ARENA_EXTEND_STRATEGY = ArenaExtendStrategy::kNextPowerOfTwo; static const int DEFAULT_INITIAL_CHUNK_SIZE_BYTES = 1 * 1024 * 1024; @@ -81,24 +81,16 @@ class BFCArena : public IArenaAllocator { // `initial_growth_chunk_size_bytes_` but ultimately all // future allocation sizes are determined by the arena growth strategy // and the allocation request. - Status Shrink() override; + Status Shrink(); void* Reserve(size_t size) override; - size_t Used() const override { - return static_cast(stats_.bytes_in_use); - } - - size_t Max() const override { - return memory_limit_; - } - FencePtr CreateFence(const SessionState* session_state) override { // arena always rely on its device allocator to create fence return device_allocator_->CreateFence(session_state); } - void GetStats(AllocatorStats* stats); + void GetStats(AllocatorStats* stats) override; size_t RequestedSize(const void* ptr); diff --git a/onnxruntime/core/framework/execution_provider.cc b/onnxruntime/core/framework/execution_provider.cc index dba534f194..e40313c7b4 100644 --- a/onnxruntime/core/framework/execution_provider.cc +++ b/onnxruntime/core/framework/execution_provider.cc @@ -52,13 +52,76 @@ IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, #endif } +// Returns true if an allocator was found and replaced +static bool FindAndReplaceAllocator(const OrtMemoryInfo& mem_info, + const MemoryInfoSet& mem_info_set, + AllocatorMap& allocators, + AllocatorPtr replacing_allocator) { + auto ite = mem_info_set.find(mem_info); + + if (ite != mem_info_set.end()) { + const int key = MakeKey(mem_info.id, mem_info.mem_type); + allocators[key] = replacing_allocator; + return true; + } + + return false; +} + // Update allocator in the provider if already present; ignore if not. void IExecutionProvider::ReplaceAllocator(AllocatorPtr allocator) { const auto& info = allocator->Info(); - auto ite = mem_info_set_.find(info); - if (ite != mem_info_set_.end()) { - const int key = MakeKey(info.id, info.mem_type); - allocators_[key] = allocator; + + if (FindAndReplaceAllocator(info, mem_info_set_, allocators_, allocator)) { + // We found an allocator corresponding to the provided + // allocator's OrtMemoryInfo and we replaced it with the + // provided allocator. + // We return back. + return; + } + + else { + // If we can't find an allocator registered with the exact OrtMemoryInfo + // as that of the replacing allocator, we do a "loosened" check + // (i.e.) check if there is an allocator registered with OrtAllocatorType + // as OrtArenaAllocator because for external user provided allocator + // we only accept OrtAllocatorType as OrtDeviceAllocator. + // If we do find such a registered allocator, we can safely go ahead + // and replace that with the provided allocator. This may seem like + // we are replacing an arena allocator with a non-arena allocator + // but in reality any user provided allocator may still be an arena + // allocator. We don't allow users to use OrtAllocatorType as + // OrtArenaAllocator for their allocators because we reserve its usage + // for our internal BFCArena. + // TODO: Should we remove the OrtAllocatorType field from OrtMemoryInfo to + // avoid such problems and also remove the unintuitive phenomenon of binding + // the allocator type info to OrtMemoryInfo (which loosely is just device info) ? + const auto& original_info = allocator->Info(); + + // If the alloc_type was OrtArenaAllocator already, then it is a no-op + if (original_info.alloc_type == OrtAllocatorType::OrtArenaAllocator) { + return; + } + + auto check_info = original_info; + + // Mutate the alloc_type + check_info.alloc_type = OrtAllocatorType::OrtArenaAllocator; + + if (FindAndReplaceAllocator(check_info, mem_info_set_, + allocators_, allocator)) { + // We found an allocator corresponding to the mutated OrtMemoryInfo + // and we replaced it with the provided allocator. + // Before we return back, we need to do some house-keeping + // (i.e.) update the EP's OrtMemoryInfo set + + // Delete the existing OrtMemoryInfo corresponding to the allocator + // that was replaced + mem_info_set_.erase(check_info); + + // Replace it with the provided allocator's OrtMemoryInfo + mem_info_set_.insert(allocator->Info()); + } } } @@ -84,7 +147,7 @@ void IExecutionProvider::TryInsertAllocator(AllocatorPtr allocator) { InsertAllocator(allocator); } -void IExecutionProvider::RegisterAllocator(std::shared_ptr ) { +void IExecutionProvider::RegisterAllocator(std::shared_ptr) { return; } diff --git a/onnxruntime/core/framework/mimalloc_allocator.cc b/onnxruntime/core/framework/mimalloc_allocator.cc new file mode 100644 index 0000000000..17a8855e87 --- /dev/null +++ b/onnxruntime/core/framework/mimalloc_allocator.cc @@ -0,0 +1,44 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if defined(USE_MIMALLOC_ARENA_ALLOCATOR) + +#include "mimalloc.h" +#include "core/framework/mimalloc_allocator.h" + +namespace onnxruntime { + +MiMallocAllocator::MiMallocAllocator(size_t total_memory) + : IAllocator(OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)) { + stats_.bytes_limit = total_memory; +} + +void* MiMallocAllocator::Alloc(size_t size) { +#if (MI_STAT > 1) + stats_.num_allocs++; +#endif + return mi_malloc(size); +} + +void MiMallocAllocator::Free(void* p) { + mi_free(p); +} + +// mimalloc only maintains stats when compiled under debug (which in turn sets MI_STAT) +void MiMallocAllocator::GetStats(AllocatorStats* stats) { +#if (MI_STAT > 1) + auto current_stats = mi_heap_get_default()->tld->stats; + stats_.bytes_in_use = current_stats.malloc.current; + stats_.total_allocated_bytes = current_stats.reserved.current; + stats_.max_bytes_in_use = current_stats.reserved.peak; +#endif + *stats = stats_; +} + +size_t MiMallocAllocator::AllocatedSize(const void* ptr) { + return mi_usable_size(ptr); +} + +} // namespace onnxruntime + +#endif diff --git a/onnxruntime/core/framework/mimalloc_allocator.h b/onnxruntime/core/framework/mimalloc_allocator.h new file mode 100644 index 0000000000..f2357e894f --- /dev/null +++ b/onnxruntime/core/framework/mimalloc_allocator.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if defined(USE_MIMALLOC_ARENA_ALLOCATOR) + +#include "core/common/common.h" +#include "core/framework/allocator.h" +#include "onnxruntime_config.h" + +namespace onnxruntime { + +class MiMallocAllocator : public IAllocator { + public: + explicit MiMallocAllocator(size_t total_memory); + + void* Alloc(size_t size) override; + + void Free(void* p) override; + + // mimalloc only maintains stats when compiled under debug, or when MI_STAT >= 2 + void GetStats(AllocatorStats* stats) override; + + size_t AllocatedSize(const void* ptr); + + private: + AllocatorStats stats_; +}; + +} // namespace onnxruntime + +#endif diff --git a/onnxruntime/core/framework/mimalloc_arena.cc b/onnxruntime/core/framework/mimalloc_arena.cc deleted file mode 100644 index 460f31e9f1..0000000000 --- a/onnxruntime/core/framework/mimalloc_arena.cc +++ /dev/null @@ -1,50 +0,0 @@ -#if defined(USE_MIMALLOC_ARENA_ALLOCATOR) -#include "mimalloc.h" -#include "core/framework/mimalloc_arena.h" - -namespace onnxruntime { -MiMallocArena::MiMallocArena(std::unique_ptr resource_allocator, - size_t total_memory) - : info_(resource_allocator->Info().name, OrtAllocatorType::OrtArenaAllocator, resource_allocator->Info().device, resource_allocator->Info().id, resource_allocator->Info().mem_type) { - stats_.bytes_limit = total_memory; -} - -void* MiMallocArena::Alloc(size_t size) { -#if (MI_STAT > 1) - stats_.num_allocs++; -#endif - return mi_malloc(size); -} - -void MiMallocArena::Free(void* p) { - mi_free(p); -} - -void* MiMallocArena::Reserve(size_t size) { - return mi_malloc(size); -} - -// mimalloc only maintains stats when compiled under debug (which in turn sets MI_STAT) -void MiMallocArena::GetStats(AllocatorStats* stats) { -#if (MI_STAT > 1) - auto current_stats = mi_heap_get_default()->tld->stats; - stats_.bytes_in_use = current_stats.malloc.current; - stats_.total_allocated_bytes = current_stats.reserved.current; - stats_.max_bytes_in_use = current_stats.reserved.peak; -#endif - *stats = stats_; -} - -size_t MiMallocArena::Used() const { -#if (MI_STAT > 1) - return mi_heap_get_default()->tld->stats.malloc.current; -#else - return 0; -#endif -} - -size_t MiMallocArena::AllocatedSize(const void* ptr) { - return mi_usable_size(ptr); -} -} // namespace onnxruntime -#endif diff --git a/onnxruntime/core/framework/mimalloc_arena.h b/onnxruntime/core/framework/mimalloc_arena.h deleted file mode 100644 index 9c478bddc2..0000000000 --- a/onnxruntime/core/framework/mimalloc_arena.h +++ /dev/null @@ -1,36 +0,0 @@ -#if defined(USE_MIMALLOC_ARENA_ALLOCATOR) -#include "core/common/common.h" -#include "core/framework/arena.h" -#include "onnxruntime_config.h" - -namespace onnxruntime { -class MiMallocArena : public IArenaAllocator { - public: - MiMallocArena(std::unique_ptr resource_allocator, size_t total_memory); - - void* Alloc(size_t size) override; - - void Free(void* p) override; - - // mimalloc only maintains stats when compiled under debug, or when MI_STAT >= 2 - void GetStats(AllocatorStats* stats); - - void* Reserve(size_t size) override; - - size_t Used() const override; - - size_t Max() const override { - return stats_.bytes_limit; - } - - const OrtMemoryInfo& Info() const override { - return info_; - } - - size_t AllocatedSize(const void* ptr); - - OrtMemoryInfo info_; - AllocatorStats stats_; -}; -} // namespace onnxruntime -#endif diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index 4448ae8b55..55e079a06d 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -21,7 +21,7 @@ #include "core/framework/session_state.h" #include "core/framework/tensorprotoutils.h" #include "core/framework/utils.h" -#include "core/framework/arena.h" +#include "core/framework/bfc_arena.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "core/framework/mem_buffer.h" #include "core/framework/tensor_allocator.h" @@ -43,14 +43,12 @@ static common::Status AllocateBufferUsingDeviceAllocatorFromShapeAndType(const T p_data = nullptr; if (shape_size > 0) { SafeInt mem_size = 0; - if (!alloc->CalcMemSizeForArray(SafeInt(shape_size), type->Size(), &mem_size)) - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed memory size calculation"); - if (alloc->Info().alloc_type == OrtArenaAllocator) - // Reserve() uses the device allocator to make the allocation - p_data = static_cast(alloc.get())->Reserve(mem_size); - else - p_data = alloc->Alloc(mem_size); + if (!alloc->CalcMemSizeForArray(SafeInt(shape_size), type->Size(), &mem_size)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed memory size calculation"); + } + + p_data = alloc->Reserve(mem_size); } return Status::OK(); diff --git a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h index f2a0da7d16..4a354bc609 100644 --- a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h +++ b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h @@ -8,7 +8,7 @@ #include "ort_value_pattern_planner.h" #include "utils.h" #include "tensorprotoutils.h" -#include "arena.h" +#include "bfc_arena.h" namespace onnxruntime { @@ -41,7 +41,7 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { if (alloc->Info().alloc_type == OrtArenaAllocator) { // Arena has a specific way to store static memory. // Arena does not reuse static memory allocated by Reserve. - buffer = static_cast(alloc.get())->Reserve(peak_size); + buffer = static_cast(alloc.get())->Reserve(peak_size); } else { buffer = alloc->Alloc(peak_size); } diff --git a/onnxruntime/core/session/allocator_adapters.cc b/onnxruntime/core/session/allocator_adapters.cc new file mode 100644 index 0000000000..1053d0d1e8 --- /dev/null +++ b/onnxruntime/core/session/allocator_adapters.cc @@ -0,0 +1,131 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "allocator_adapters.h" +#include "core/session/inference_session.h" +#include "core/session/ort_env.h" +#include "core/session/ort_apis.h" +#include "core/framework/error_code_helper.h" + +namespace onnxruntime { +OrtAllocatorImplWrappingIAllocator::OrtAllocatorImplWrappingIAllocator(onnxruntime::AllocatorPtr&& i_allocator) + : i_allocator_(std::move(i_allocator)) { + OrtAllocator::version = ORT_API_VERSION; + OrtAllocator::Alloc = + [](OrtAllocator* this_, size_t size) { return static_cast(this_)->Alloc(size); }; + OrtAllocator::Free = + [](OrtAllocator* this_, void* p) { static_cast(this_)->Free(p); }; + OrtAllocator::Info = + [](const OrtAllocator* this_) { return static_cast(this_)->Info(); }; +} + +void* OrtAllocatorImplWrappingIAllocator::Alloc(size_t size) { + return i_allocator_->Alloc(size); +} + +void OrtAllocatorImplWrappingIAllocator::Free(void* p) { + i_allocator_->Free(p); +} + +const OrtMemoryInfo* OrtAllocatorImplWrappingIAllocator::Info() const { + return &i_allocator_->Info(); +} + +IAllocatorImplWrappingOrtAllocator::IAllocatorImplWrappingOrtAllocator(OrtAllocator* ort_allocator) + : IAllocator(*ort_allocator->Info(ort_allocator)), ort_allocator_(ort_allocator) {} + +void* IAllocatorImplWrappingOrtAllocator::Alloc(size_t size) { + return ort_allocator_->Alloc(ort_allocator_, size); +} + +void IAllocatorImplWrappingOrtAllocator::Free(void* p) { + return ort_allocator_->Free(ort_allocator_, p); +} + +} // namespace onnxruntime + +ORT_API_STATUS_IMPL(OrtApis::CreateAllocator, const OrtSession* sess, + const OrtMemoryInfo* mem_info, _Outptr_ OrtAllocator** out) { + API_IMPL_BEGIN + auto* session = reinterpret_cast(sess); + auto allocator_ptr = session->GetAllocator(*mem_info); + if (!allocator_ptr) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "No requested allocator available"); + } + *out = new onnxruntime::OrtAllocatorImplWrappingIAllocator(std::move(allocator_ptr)); + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::CreateAndRegisterAllocator, _Inout_ OrtEnv* env, + _In_ const OrtMemoryInfo* mem_info, + _In_ const OrtArenaCfg* arena_cfg) { + using namespace onnxruntime; + if (!env) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Env is null"); + } + + if (!mem_info) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "OrtMemoryInfo is null"); + } + + auto st = env->CreateAndRegisterAllocator(*mem_info, arena_cfg); + + if (!st.IsOK()) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, st.ErrorMessage().c_str()); + } + return nullptr; +} + +ORT_API_STATUS_IMPL(OrtApis::RegisterAllocator, _Inout_ OrtEnv* env, + _In_ OrtAllocator* allocator) { + using namespace onnxruntime; + if (!env) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Env is null"); + } + + if (!allocator) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Provided allocator is null"); + } + + if (allocator->Info(allocator)->alloc_type == OrtAllocatorType::OrtArenaAllocator) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, + "Please register the allocator as OrtDeviceAllocator " + "even if the provided allocator has arena logic built-in. " + "OrtArenaAllocator is reserved for internal arena logic based " + "allocators only."); + } + + std::shared_ptr i_alloc_ptr = + std::make_shared(allocator); + + auto st = env->RegisterAllocator(i_alloc_ptr); + + if (!st.IsOK()) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, st.ErrorMessage().c_str()); + } + return nullptr; +} + +ORT_API_STATUS_IMPL(OrtApis::UnregisterAllocator, _Inout_ OrtEnv* env, + _In_ const OrtMemoryInfo* mem_info) { + using namespace onnxruntime; + if (!env) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Env is null"); + } + + if (!mem_info) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Provided OrtMemoryInfo is null"); + } + + auto st = env->UnregisterAllocator(*mem_info); + + if (!st.IsOK()) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, st.ErrorMessage().c_str()); + } + return nullptr; +} + +ORT_API(void, OrtApis::ReleaseAllocator, _Frees_ptr_opt_ OrtAllocator* allocator) { + delete static_cast(allocator); +} diff --git a/onnxruntime/core/session/allocator_adapters.h b/onnxruntime/core/session/allocator_adapters.h new file mode 100644 index 0000000000..6b39dd98d1 --- /dev/null +++ b/onnxruntime/core/session/allocator_adapters.h @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/framework/allocator.h" +#include "core/session/onnxruntime_cxx_api.h" + +namespace onnxruntime { + +// Since all allocators are of type 'OrtAllocator' and there is a single +// OrtApis:ReleaseAllocator function, we need to have a common base type that lets us delete them. +struct OrtAllocatorImpl : public OrtAllocator { + virtual ~OrtAllocatorImpl() = default; +}; + +// The following are "adapters" to allow using an IAllocator implementation wrapped as an OrtAllocator +// and vice versa to plug into any ORT internal code/ API implementation as necessary + +struct OrtAllocatorImplWrappingIAllocator final : public OrtAllocatorImpl { + explicit OrtAllocatorImplWrappingIAllocator(onnxruntime::AllocatorPtr&& i_allocator); + + ~OrtAllocatorImplWrappingIAllocator() override = default; + + void* Alloc(size_t size); + + void Free(void* p); + + const OrtMemoryInfo* Info() const; + + ORT_DISALLOW_COPY_AND_ASSIGNMENT(OrtAllocatorImplWrappingIAllocator); + + private: + onnxruntime::AllocatorPtr i_allocator_; +}; + +class IAllocatorImplWrappingOrtAllocator final : public IAllocator { + public: + explicit IAllocatorImplWrappingOrtAllocator(OrtAllocator* ort_allocator); + + ~IAllocatorImplWrappingOrtAllocator() override = default; + + void* Alloc(size_t size) override; + + void Free(void* p) override; + + ORT_DISALLOW_COPY_AND_ASSIGNMENT(IAllocatorImplWrappingOrtAllocator); + + private: + OrtAllocator* ort_allocator_ = nullptr; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/session/allocator_impl.h b/onnxruntime/core/session/allocator_impl.h deleted file mode 100644 index 4c0086af57..0000000000 --- a/onnxruntime/core/session/allocator_impl.h +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once -#include "core/session/onnxruntime_c_api.h" -#include "core/framework/allocator.h" -#include "core/framework/arena.h" -#include "core/session/device_allocator.h" - -namespace onnxruntime { -class AllocatorWrapper : public IAllocator { - public: - AllocatorWrapper(OrtAllocator* impl) : IAllocator(*impl->Info(impl)), impl_(impl) {} - void* Alloc(size_t size) override { - return impl_->Alloc(impl_, size); - } - void Free(void* p) override { - return impl_->Free(impl_, p); - } - - private: - OrtAllocator* impl_; -}; - -} // namespace onnxruntime diff --git a/onnxruntime/core/session/default_cpu_allocator_c_api.cc b/onnxruntime/core/session/default_cpu_allocator_c_api.cc index 5df27d7412..1cdac1ee26 100644 --- a/onnxruntime/core/session/default_cpu_allocator_c_api.cc +++ b/onnxruntime/core/session/default_cpu_allocator_c_api.cc @@ -1,19 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include #include "core/framework/utils.h" +#include "core/framework/allocator.h" +#include "core/session/allocator_adapters.h" #include "core/session/onnxruntime_cxx_api.h" #include "core/session/ort_apis.h" -#include +#include "core/framework/error_code_helper.h" -// In the future we'll have more than one allocator type. Since all allocators are of type 'OrtAllocator' and there is a single -// OrtReleaseAllocator function, we need to have a common base type that lets us delete them. -struct OrtAllocatorImpl : OrtAllocator { - virtual ~OrtAllocatorImpl() = default; -}; - -void ThrowOnError(OrtStatus* status) { +static void ThrowOnError(OrtStatus* status) { if (status) { std::string ort_error_message = OrtApis::GetErrorMessage(status); OrtErrorCode ort_error_code = OrtApis::GetErrorCode(status); @@ -22,16 +17,19 @@ void ThrowOnError(OrtStatus* status) { } } -struct OrtDefaultAllocator : OrtAllocatorImpl { - OrtDefaultAllocator() { +struct OrtDefaultCpuAllocator : onnxruntime::OrtAllocatorImpl { + OrtDefaultCpuAllocator() { OrtAllocator::version = ORT_API_VERSION; - OrtAllocator::Alloc = [](OrtAllocator* this_, size_t size) { return static_cast(this_)->Alloc(size); }; - OrtAllocator::Free = [](OrtAllocator* this_, void* p) { static_cast(this_)->Free(p); }; - OrtAllocator::Info = [](const OrtAllocator* this_) { return static_cast(this_)->Info(); }; + OrtAllocator::Alloc = + [](OrtAllocator* this_, size_t size) { return static_cast(this_)->Alloc(size); }; + OrtAllocator::Free = + [](OrtAllocator* this_, void* p) { static_cast(this_)->Free(p); }; + OrtAllocator::Info = + [](const OrtAllocator* this_) { return static_cast(this_)->Info(); }; ThrowOnError(OrtApis::CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &cpu_memory_info)); } - ~OrtDefaultAllocator() override { OrtApis::ReleaseMemoryInfo(cpu_memory_info); } + ~OrtDefaultCpuAllocator() { OrtApis::ReleaseMemoryInfo(cpu_memory_info); } void* Alloc(size_t size) { return onnxruntime::utils::DefaultAlloc(size); @@ -43,32 +41,16 @@ struct OrtDefaultAllocator : OrtAllocatorImpl { return cpu_memory_info; } - ORT_DISALLOW_COPY_AND_ASSIGNMENT(OrtDefaultAllocator); + ORT_DISALLOW_COPY_AND_ASSIGNMENT(OrtDefaultCpuAllocator); private: OrtMemoryInfo* cpu_memory_info; }; -#ifndef ORT_NO_EXCEPTIONS - -#define API_IMPL_BEGIN try { -#define API_IMPL_END \ - } \ - catch (std::exception & ex) { \ - return OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, ex.what()); \ - } - -#else - -#define API_IMPL_BEGIN { -#define API_IMPL_END } - -#endif - ORT_API_STATUS_IMPL(OrtApis::GetAllocatorWithDefaultOptions, _Outptr_ OrtAllocator** out) { API_IMPL_BEGIN - static OrtDefaultAllocator ort_default_allocator; - *out = &ort_default_allocator; + static OrtDefaultCpuAllocator ort_default_cpu_allocator; + *out = &ort_default_cpu_allocator; return nullptr; API_IMPL_END } diff --git a/onnxruntime/core/session/device_allocator.cc b/onnxruntime/core/session/device_allocator.cc deleted file mode 100644 index 72e1a21416..0000000000 --- a/onnxruntime/core/session/device_allocator.cc +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "device_allocator.h" -#include "core/session/inference_session.h" -#include "core/session/ort_env.h" -#include "core/session/allocator_impl.h" - -#ifndef ORT_NO_EXCEPTIONS - -#define API_IMPL_BEGIN try { -#define API_IMPL_END \ - } \ - catch (const std::exception& ex) { \ - return OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, ex.what()); \ - } - -#else - -#define API_IMPL_BEGIN { -#define API_IMPL_END } - -#endif - -ORT_API_STATUS_IMPL(OrtApis::CreateAllocator, const OrtSession* sess, const OrtMemoryInfo* mem_info, _Outptr_ OrtAllocator** out) { - API_IMPL_BEGIN - auto session = reinterpret_cast(sess); - auto allocator_ptr = session->GetAllocator(*mem_info); - if (!allocator_ptr) { - return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "No requested allocator available"); - } - *out = new onnxruntime::OrtAllocatorForDevice(std::move(allocator_ptr)); - return nullptr; - API_IMPL_END -} - -ORT_API_STATUS_IMPL(OrtApis::CreateAndRegisterAllocator, _Inout_ OrtEnv* env, _In_ const OrtMemoryInfo* mem_info, - _In_ const OrtArenaCfg* arena_cfg) { - using namespace onnxruntime; - if (!env) { - return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Env is null"); - } - - if (!mem_info) { - return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "OrtMemoryInfo is null"); - } - - auto st = env->CreateAndRegisterAllocator(*mem_info, arena_cfg); - - if (!st.IsOK()) { - return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, st.ErrorMessage().c_str()); - } - return nullptr; -} - -ORT_API(void, OrtApis::ReleaseAllocator, _Frees_ptr_opt_ OrtAllocator* allocator) { - delete reinterpret_cast(allocator); -} diff --git a/onnxruntime/core/session/device_allocator.h b/onnxruntime/core/session/device_allocator.h deleted file mode 100644 index 89e29b1ce1..0000000000 --- a/onnxruntime/core/session/device_allocator.h +++ /dev/null @@ -1,45 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. -#pragma once - -#include "core/framework/allocator.h" -#include "core/framework/utils.h" -#include "core/session/onnxruntime_cxx_api.h" -#include "core/session/ort_apis.h" -#include -#include "core/framework/allocatormgr.h" - -namespace onnxruntime { -struct OrtAllocatorForDevice : public OrtAllocator { - explicit OrtAllocatorForDevice(onnxruntime::AllocatorPtr&& dev_allocator) - : device_allocator_(std::move(dev_allocator)) { - OrtAllocator::version = ORT_API_VERSION; - OrtAllocator::Alloc = [](OrtAllocator* this_, size_t size) { return static_cast(this_)->Alloc(size); }; - OrtAllocator::Free = [](OrtAllocator* this_, void* p) { static_cast(this_)->Free(p); }; - OrtAllocator::Info = [](const OrtAllocator* this_) { return static_cast(this_)->Info(); }; - } - - ~OrtAllocatorForDevice() = default; - - void* Alloc(size_t size) const { - return device_allocator_->Alloc(size); - } - void Free(void* p) const { - device_allocator_->Free(p); - } - - const OrtMemoryInfo* Info() const { - return &device_allocator_->Info(); - } - - OrtAllocatorForDevice(const OrtAllocatorForDevice&) = delete; - OrtAllocatorForDevice& operator=(const OrtAllocatorForDevice&) = delete; - - onnxruntime::IAllocator* GetAllocator() { - return device_allocator_.get(); - } - - private: - onnxruntime::AllocatorPtr device_allocator_; -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index 12d3932581..14af0c9887 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/session/environment.h" +#include "core/session/allocator_adapters.h" #include "core/framework/allocatormgr.h" #include "core/graph/constants.h" #include "core/graph/op.h" @@ -25,7 +26,6 @@ #include "core/platform/env.h" #include "core/util/thread_utils.h" -#include "core/session/allocator_impl.h" #ifdef ONNXRUNTIME_ENABLE_INSTRUMENT #include "core/platform/tracing.h" @@ -56,16 +56,54 @@ Status Environment::Create(std::unique_ptr logging_mana return status; } +// Ugly but necessary for instances where we want to check equality of two OrtMemoryInfos +// without accounting for OrtAllocatorType in the equality checking process. +// TODO: Should we remove the OrtAllocatorType field from the OrtMemoryInfo struct to +// avoid such problems and also remove the unintuitive phenomenon of binding an allocator +// type to OrtMemoryInfo (which loosely is just device info) ? +static bool AreOrtMemoryInfosEquivalent( + const OrtMemoryInfo& left, const OrtMemoryInfo& right, + bool include_allocator_type_for_equivalence_checking = true) { + if (include_allocator_type_for_equivalence_checking) { + return left == right; + } else { + return left.mem_type == right.mem_type && + left.id == right.id && + strcmp(left.name, right.name) == 0; + } +} + Status Environment::RegisterAllocator(AllocatorPtr allocator) { const auto& mem_info = allocator->Info(); + + if (mem_info.device.Type() != OrtDevice::CPU) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Only CPU allocators can be shared between " + "multiple sessions for now."); + } + // We don't expect millions of allocators getting registered. Hence linear search should be fine. auto ite = std::find_if(std::begin(shared_allocators_), std::end(shared_allocators_), - [&mem_info](const AllocatorPtr& alloc_ptr) { return alloc_ptr->Info() == mem_info; }); + [&mem_info](const AllocatorPtr& alloc_ptr) { + // We want to do the equality checking of 2 OrtMemoryInfos sans the OrtAllocatorType field. + // This is because we want to avoid registering two allocators for the same device that just + // differ on OrtAllocatorType. + // To be more specific, we want to avoid the scenario where the user calls CreateAndRegiserAllocator() + // and registers the ORT-internal arena allocator and then tries to register their own custom + // allocator using RegisterAllocator() for the same device that has an OrtAllocatorType as + // OrtDeviceAllocator (which is the only accepted value while registering a custom allocator). + // If we allowed this, it could potentially cause a lot of confusion as to which shared allocator + // to use for that device and we want to avoid having any ugly logic around this. + return AreOrtMemoryInfosEquivalent(alloc_ptr->Info(), mem_info, false); + }); + if (ite != shared_allocators_.end()) { - return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Allocator with this OrtMemoryInfo is already registered."); + return Status(ONNXRUNTIME, INVALID_ARGUMENT, "An allocator for this device has already been registered for sharing."); } + shared_allocators_.insert(ite, allocator); + return Status::OK(); } @@ -133,6 +171,25 @@ Status Environment::CreateAndRegisterAllocator(const OrtMemoryInfo& mem_info, co return RegisterAllocator(allocator_ptr); } +Status Environment::UnregisterAllocator(const OrtMemoryInfo& mem_info) { + auto ite = std::find_if(std::begin(shared_allocators_), + std::end(shared_allocators_), + [&mem_info](const AllocatorPtr& alloc_ptr) { + // See comment in RegisterAllocator() as to why we + // use this method of OrtMemoryInfo equality checking + return AreOrtMemoryInfosEquivalent(alloc_ptr->Info(), mem_info, false); + }); + + if (ite == shared_allocators_.end()) { + return Status(ONNXRUNTIME, INVALID_ARGUMENT, + "No allocator for this device has been registered for sharing."); + } + + shared_allocators_.erase(ite); + + return Status::OK(); +} + Status Environment::Initialize(std::unique_ptr logging_manager, const OrtThreadingOptions* tp_options, bool create_global_thread_pools) { diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 8ed0212d7c..17d0742d31 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -14,7 +14,7 @@ #include "core/common/denormal.h" #include "core/common/logging/logging.h" #include "core/common/parse_string.h" -#include "core/framework/arena.h" +#include "core/framework/bfc_arena.h" #include "core/framework/allocatormgr.h" #include "core/framework/error_code_helper.h" #include "core/framework/execution_frame.h" @@ -1164,10 +1164,10 @@ common::Status InferenceSession::Initialize() { onnxruntime::Graph& graph = model_->MainGraph(); #ifdef DISABLE_EXTERNAL_INITIALIZERS const InitializedTensorSet& initializers = graph.GetAllInitializedTensors(); - for (const auto& it: initializers) { + for (const auto& it : initializers) { if (utils::HasExternalData(*it.second)) { return common::Status(common::ONNXRUNTIME, common::FAIL, - "Initializer tensors with external data is not allowed."); + "Initializer tensors with external data is not allowed."); } } #endif @@ -1947,7 +1947,7 @@ common::Status InferenceSession::ValidateAndParseShrinkArenaString(const std::st void InferenceSession::ShrinkMemoryArenas(const std::vector& arenas_to_shrink) { for (auto& alloc : arenas_to_shrink) { - auto status = static_cast(alloc.get())->Shrink(); + auto status = static_cast(alloc.get())->Shrink(); if (!status.IsOK()) { LOGS(*session_logger_, WARNING) << "Unable to shrink arena: " << alloc->Info().ToString() diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index 5c3b464f33..40a1dcc656 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -2,7 +2,7 @@ // Licensed under the MIT License. #include "core/session/onnxruntime_c_api.h" -#include "core/session/allocator_impl.h" +#include "core/session/allocator_adapters.h" #include "core/session/inference_session_utils.h" #include "core/session/IOBinding.h" #include "core/framework/allocator.h" @@ -157,7 +157,7 @@ ORT_STATUS_PTR CreateTensorImpl(MLDataType ml_type, const int64_t* shape, size_t for (size_t i = 0; i != shape_len; ++i) { shapes[i] = shape[i]; } - std::shared_ptr alloc_ptr = std::make_shared(allocator); + std::shared_ptr alloc_ptr = std::make_shared(allocator); *out = std::make_unique(ml_type, onnxruntime::TensorShape(shapes), alloc_ptr); return nullptr; } @@ -174,7 +174,7 @@ ORT_STATUS_PTR CreateTensorImplForSeq(MLDataType elem_type, const int64_t* shape if (st) { return st; } - std::shared_ptr alloc_ptr = std::make_shared(allocator); + std::shared_ptr alloc_ptr = std::make_shared(allocator); out = Tensor(elem_type, onnxruntime::TensorShape(shapes), alloc_ptr); return nullptr; } @@ -409,9 +409,9 @@ ORT_API_STATUS_IMPL(OrtApis::EnableOrtCustomOps, _Inout_ OrtSessionOptions* opti if (options) { #ifdef ENABLE_EXTENSION_CUSTOM_OPS - return RegisterCustomOps(options, OrtGetApiBase()); + return RegisterCustomOps(options, OrtGetApiBase()); #else - return OrtApis::CreateStatus(ORT_FAIL, "EnableOrtCustomOps: Custom operators in onnxruntime-extensions are not enabled"); + return OrtApis::CreateStatus(ORT_FAIL, "EnableOrtCustomOps: Custom operators in onnxruntime-extensions are not enabled"); #endif } return nullptr; @@ -2292,6 +2292,8 @@ static constexpr OrtApi ort_api_1_to_9 = { &OrtApis::GetTensorRTProviderOptionsAsString, &OrtApis::ReleaseTensorRTProviderOptions, &OrtApis::EnableOrtCustomOps, + &OrtApis::RegisterAllocator, + &OrtApis::UnregisterAllocator, }; // Assert to do a limited check to ensure Version 1 of OrtApi never changes (will detect an addition or deletion but not if they cancel out each other) diff --git a/onnxruntime/core/session/ort_apis.h b/onnxruntime/core/session/ort_apis.h index 15c49592ec..080431028f 100644 --- a/onnxruntime/core/session/ort_apis.h +++ b/onnxruntime/core/session/ort_apis.h @@ -286,4 +286,6 @@ ORT_API_STATUS_IMPL(UpdateTensorRTProviderOptions, _Inout_ OrtTensorRTProviderOp ORT_API_STATUS_IMPL(GetTensorRTProviderOptionsAsString, _In_ const OrtTensorRTProviderOptionsV2* tensorrt_options, _Inout_ OrtAllocator* allocator, _Outptr_ char** ptr); ORT_API(void, ReleaseTensorRTProviderOptions, _Frees_ptr_opt_ OrtTensorRTProviderOptionsV2*); ORT_API_STATUS_IMPL(EnableOrtCustomOps, _Inout_ OrtSessionOptions* options); +ORT_API_STATUS_IMPL(RegisterAllocator, _Inout_ OrtEnv* env, _In_ OrtAllocator* allocator); +ORT_API_STATUS_IMPL(UnregisterAllocator, _Inout_ OrtEnv* env, _In_ const OrtMemoryInfo* mem_info); } // namespace OrtApis diff --git a/onnxruntime/core/session/ort_env.cc b/onnxruntime/core/session/ort_env.cc index 4645b55fb6..834bbbeae9 100644 --- a/onnxruntime/core/session/ort_env.cc +++ b/onnxruntime/core/session/ort_env.cc @@ -8,7 +8,7 @@ #include "ort_env.h" #include "core/session/ort_apis.h" #include "core/session/environment.h" -#include "core/session/allocator_impl.h" +#include "core/session/allocator_adapters.h" #include "core/common/logging/logging.h" #include "core/framework/provider_shutdown.h" #include "core/platform/logging/make_platform_default_log_sink.h" @@ -113,3 +113,7 @@ onnxruntime::common::Status OrtEnv::CreateAndRegisterAllocator(const OrtMemoryIn auto status = value_->CreateAndRegisterAllocator(mem_info, arena_cfg); return status; } + +onnxruntime::common::Status OrtEnv::UnregisterAllocator(const OrtMemoryInfo& mem_info) { + return value_->UnregisterAllocator(mem_info); +} diff --git a/onnxruntime/core/session/ort_env.h b/onnxruntime/core/session/ort_env.h index f1967f3b04..8d85c1a4b3 100644 --- a/onnxruntime/core/session/ort_env.h +++ b/onnxruntime/core/session/ort_env.h @@ -69,6 +69,11 @@ struct OrtEnv { onnxruntime::common::Status CreateAndRegisterAllocator(const OrtMemoryInfo& mem_info, const OrtArenaCfg* arena_cfg = nullptr); + /** + * Removes registered allocator that was previously registered for sharing between multiple sessions. + */ + onnxruntime::common::Status UnregisterAllocator(const OrtMemoryInfo& mem_info); + private: static OrtEnv* p_instance_; static onnxruntime::OrtMutex m_; diff --git a/onnxruntime/test/framework/TestAllocatorManager.cc b/onnxruntime/test/framework/TestAllocatorManager.cc index be7e842037..97cc86e8a1 100644 --- a/onnxruntime/test/framework/TestAllocatorManager.cc +++ b/onnxruntime/test/framework/TestAllocatorManager.cc @@ -8,14 +8,14 @@ namespace onnxruntime { namespace test { // Dummy Arena which just call underline device allocator directly. -class DummyArena : public IArenaAllocator { +class DummyArena : public IAllocator { public: explicit DummyArena(std::unique_ptr resource_allocator) - : IArenaAllocator(OrtMemoryInfo(resource_allocator->Info().name, - OrtAllocatorType::OrtArenaAllocator, - resource_allocator->Info().device, - resource_allocator->Info().id, - resource_allocator->Info().mem_type)), + : IAllocator(OrtMemoryInfo(resource_allocator->Info().name, + OrtAllocatorType::OrtDeviceAllocator, + resource_allocator->Info().device, + resource_allocator->Info().id, + resource_allocator->Info().mem_type)), allocator_(std::move(resource_allocator)) { } @@ -31,22 +31,6 @@ class DummyArena : public IArenaAllocator { allocator_->Free(p); } - void* Reserve(size_t size) override { - return Alloc(size); - } - - Status Shrink() override { - ORT_NOT_IMPLEMENTED(__FUNCTION__, " is not implemented"); - } - - size_t Used() const override { - ORT_NOT_IMPLEMENTED(__FUNCTION__, " is not implemented"); - } - - size_t Max() const override { - ORT_NOT_IMPLEMENTED(__FUNCTION__, " is not implemented"); - } - private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(DummyArena); diff --git a/onnxruntime/test/framework/TestAllocatorManager.h b/onnxruntime/test/framework/TestAllocatorManager.h index 6278d558e5..d1b67fe72a 100644 --- a/onnxruntime/test/framework/TestAllocatorManager.h +++ b/onnxruntime/test/framework/TestAllocatorManager.h @@ -2,9 +2,12 @@ // Licensed under the MIT License. #pragma once -#include "core/framework/arena.h" + +#include "core/framework/allocator.h" + namespace onnxruntime { namespace test { + class AllocatorManager { public: // the allocator manager is a just for onnx runner to allocate space for input/output tensors. diff --git a/onnxruntime/test/framework/inference_session_test.cc b/onnxruntime/test/framework/inference_session_test.cc index 7c41528949..3b3cb822a1 100644 --- a/onnxruntime/test/framework/inference_session_test.cc +++ b/onnxruntime/test/framework/inference_session_test.cc @@ -39,8 +39,7 @@ #endif #include "core/session/environment.h" #include "core/session/IOBinding.h" -#include "core/session/device_allocator.h" -#include "core/session/allocator_impl.h" +#include "core/session/inference_session_utils.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "core/session/onnxruntime_run_options_config_keys.h" #include "dummy_provider.h" @@ -265,7 +264,7 @@ void RunModelWithBindingMatMul(InferenceSession& session_object, ProviderType bind_provider_type, bool is_preallocate_output_vec, ProviderType allocation_provider, - IExecutionProvider *gpu_provider, + IExecutionProvider* gpu_provider, OrtDevice* output_device) { unique_ptr io_binding; Status st = session_object.NewIOBinding(&io_binding); @@ -849,11 +848,11 @@ static void TestBindHelper(const std::string& log_str, so.session_log_verbosity_level = 1; // change to 1 for detailed logging InferenceSession session_object{so, GetEnvironment()}; - IExecutionProvider *gpu_provider{}; + IExecutionProvider* gpu_provider{}; if (bind_provider_type == kCudaExecutionProvider || bind_provider_type == kRocmExecutionProvider) { #ifdef USE_CUDA - auto provider = DefaultCudaExecutionProvider(); + auto provider = DefaultCudaExecutionProvider(); gpu_provider = provider.get(); ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(std::move(provider))); #elif USE_ROCM diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 2ecf013f6a..4c43a575b2 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -276,7 +276,6 @@ TEST(CApiTest, custom_op_handler) { #ifdef ENABLE_EXTENSION_CUSTOM_OPS // test enabled ort-customops negpos TEST(CApiTest, test_enable_ort_customops_negpos) { - Ort::MemoryInfo info("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault); auto allocator = std::make_unique(); @@ -316,7 +315,6 @@ TEST(CApiTest, test_enable_ort_customops_negpos) { // test enabled ort-customops stringlower TEST(CApiTest, test_enable_ort_customops_stringlower) { - auto allocator = std::make_unique(); // Create Inputs @@ -1364,63 +1362,131 @@ TEST(CApiTest, get_available_providers_cpp) { ASSERT_TRUE(std::find(providers.begin(), providers.end(), "CUDAExecutionProvider") != providers.end()); #endif } +TEST(CApiTest, TestSharedAllocators) { + OrtEnv* env_ptr = (OrtEnv*)(*ort_env); -// This test uses the CreateAndRegisterAllocator API to register an allocator with the env, -// creates 2 sessions and then runs those 2 sessions one after another -TEST(CApiTest, TestSharedAllocatorUsingCreateAndRegisterAllocator) { - // simple inference test // prepare inputs std::vector inputs(1); Input& input = inputs.back(); input.name = "X"; input.dims = {3, 2}; input.values = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + auto allocator_for_input_memory_allocation = std::make_unique(); - // prepare expected inputs and outputs + // prepare expected outputs std::vector expected_dims_y = {3, 2}; std::vector expected_values_y = {1.0f, 4.0f, 9.0f, 16.0f, 25.0f, 36.0f}; - OrtEnv* env_ptr = (OrtEnv*)(*ort_env); - - OrtMemoryInfo* mem_info = nullptr; - const auto& api = Ort::GetApi(); - ASSERT_TRUE(api.CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &mem_info) == nullptr); - std::unique_ptr rel_info(mem_info, api.ReleaseMemoryInfo); - - OrtArenaCfg* arena_cfg = nullptr; - ASSERT_TRUE(api.CreateArenaCfg(0, -1, -1, -1, &arena_cfg) == nullptr); - std::unique_ptr rel_arena_cfg(arena_cfg, api.ReleaseArenaCfg); - - ASSERT_TRUE(api.CreateAndRegisterAllocator(env_ptr, mem_info, arena_cfg) == nullptr); - - // test for duplicates - std::unique_ptr status_releaser( - api.CreateAndRegisterAllocator(env_ptr, mem_info, arena_cfg), - api.ReleaseStatus); - ASSERT_FALSE(status_releaser.get() == nullptr); + // Create session options and configure it appropriately Ort::SessionOptions session_options; - auto default_allocator = std::make_unique(); + // Turn on sharing of the allocator between sessions session_options.AddConfigEntry(kOrtSessionOptionsConfigUseEnvAllocators, "1"); - // create session 1 - Ort::Session session1(*ort_env, MODEL_URI, session_options); - RunSession(default_allocator.get(), - session1, - inputs, - "Y", - expected_dims_y, - expected_values_y, - nullptr); + const auto& api = Ort::GetApi(); - // create session 2 - Ort::Session session2(*ort_env, MODEL_URI, session_options); - RunSession(default_allocator.get(), - session2, - inputs, - "Y", - expected_dims_y, - expected_values_y, - nullptr); + // CASE 1: We test creating and registering an ORT-internal allocator implementation instance + // for sharing between sessions + { + OrtMemoryInfo* mem_info = nullptr; + ASSERT_TRUE(api.CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &mem_info) == nullptr); + std::unique_ptr rel_info(mem_info, api.ReleaseMemoryInfo); + + OrtArenaCfg* arena_cfg = nullptr; + ASSERT_TRUE(api.CreateArenaCfg(0, -1, -1, -1, &arena_cfg) == nullptr); + std::unique_ptr rel_arena_cfg(arena_cfg, api.ReleaseArenaCfg); + + // This creates an ORT-internal allocator instance and registers it in the environment for sharing + // NOTE: On x86 builds arenas are not supported and will default to using non-arena based allocator + ASSERT_TRUE(api.CreateAndRegisterAllocator(env_ptr, mem_info, arena_cfg) == nullptr); + + // Test that duplicates are handled + std::unique_ptr status_releaser( + api.CreateAndRegisterAllocator(env_ptr, mem_info, arena_cfg), + api.ReleaseStatus); + ASSERT_FALSE(status_releaser.get() == nullptr); + + { + // create session 1 + Ort::Session session1(*ort_env, MODEL_URI, session_options); + RunSession(allocator_for_input_memory_allocation.get(), + session1, + inputs, + "Y", + expected_dims_y, + expected_values_y, + nullptr); + + // create session 2 + Ort::Session session2(*ort_env, MODEL_URI, session_options); + RunSession(allocator_for_input_memory_allocation.get(), + session2, + inputs, + "Y", + expected_dims_y, + expected_values_y, + nullptr); + } + + // Remove the registered shared allocator for part 2 of this test + // where-in we will register a custom allocator for the same device. + ASSERT_TRUE(api.UnregisterAllocator(env_ptr, mem_info) == nullptr); + } + + // CASE 2: We test registering a custom allocator implementation + // for sharing between sessions + { + // This creates a custom allocator instance and registers it in the environment for sharing + // NOTE: This is a very basic allocator implementation. For optimal performance, allocations + // need to be aligned for certain devices/build configurations/math libraries. + // See docs/C_API.md for details. + MockedOrtAllocator custom_allocator; + ASSERT_TRUE(api.RegisterAllocator(env_ptr, &custom_allocator) == nullptr); + + // Test that duplicates are handled + std::unique_ptr + status_releaser( + api.RegisterAllocator(env_ptr, &custom_allocator), + api.ReleaseStatus); + ASSERT_FALSE(status_releaser.get() == nullptr); + + { + // Keep this scoped to destroy the underlying sessions after use + // This should trigger frees in our custom allocator + + // create session 1 + Ort::Session session1(*ort_env, MODEL_URI, session_options); + RunSession(allocator_for_input_memory_allocation.get(), + session1, + inputs, + "Y", + expected_dims_y, + expected_values_y, + nullptr); + + // create session 2 + Ort::Session session2(*ort_env, MODEL_URI, session_options); + RunSession(allocator_for_input_memory_allocation.get(), + session2, + inputs, + "Y", + expected_dims_y, + expected_values_y, + nullptr); + } + + // Remove the registered shared allocator from the global environment + // (common to all tests) to prevent its accidental usage elsewhere + ASSERT_TRUE(api.UnregisterAllocator(env_ptr, custom_allocator.Info()) == nullptr); + + // Ensure that the registered custom allocator was indeed used for both sessions + // We should have seen 2 allocations per session (one for the sole initializer + // and one for the output). So, for two sessions, we should have seen 4 allocations. + size_t num_allocations = custom_allocator.NumAllocations(); + ASSERT_TRUE(num_allocations == 4); + + // Ensure that there was no leak + custom_allocator.LeakCheck(); + } } TEST(CApiTest, TestSharingOfInitializerAndItsPrepackedVersion) { diff --git a/onnxruntime/test/testdata/custom_execution_provider_library/my_allocator.cc b/onnxruntime/test/testdata/custom_execution_provider_library/my_allocator.cc index 97bcdc8d52..c5f2a8e831 100644 --- a/onnxruntime/test/testdata/custom_execution_provider_library/my_allocator.cc +++ b/onnxruntime/test/testdata/custom_execution_provider_library/my_allocator.cc @@ -4,7 +4,9 @@ #include "my_allocator.h" namespace onnxruntime { -MyEPAllocator::MyEPAllocator(OrtDevice::DeviceId device_id) : IAllocator(OrtMemoryInfo(MyEP, OrtAllocatorType::OrtArenaAllocator, OrtDevice(MyEPDevice, OrtDevice::MemType::DEFAULT, device_id))) { +MyEPAllocator::MyEPAllocator(OrtDevice::DeviceId device_id) + : IAllocator(OrtMemoryInfo(MyEP, OrtAllocatorType::OrtArenaAllocator, + OrtDevice(MyEPDevice, OrtDevice::MemType::DEFAULT, device_id))) { } void* MyEPAllocator::Alloc(size_t size) { diff --git a/onnxruntime/test/util/include/test_allocator.h b/onnxruntime/test/util/include/test_allocator.h index 4bf5beb958..985d97329f 100644 --- a/onnxruntime/test/util/include/test_allocator.h +++ b/onnxruntime/test/util/include/test_allocator.h @@ -14,6 +14,7 @@ struct MockedOrtAllocator : OrtAllocator { void* Alloc(size_t size); void Free(void* p); const OrtMemoryInfo* Info() const; + size_t NumAllocations() const; void LeakCheck(); @@ -22,5 +23,6 @@ struct MockedOrtAllocator : OrtAllocator { MockedOrtAllocator& operator=(const MockedOrtAllocator&) = delete; std::atomic memory_inuse{0}; + std::atomic num_allocations{0}; OrtMemoryInfo* cpu_memory_info; }; diff --git a/onnxruntime/test/util/test_allocator.cc b/onnxruntime/test/util/test_allocator.cc index 73a4754dd8..a62eae4af4 100644 --- a/onnxruntime/test/util/test_allocator.cc +++ b/onnxruntime/test/util/test_allocator.cc @@ -22,6 +22,7 @@ void* MockedOrtAllocator::Alloc(size_t size) { void* p = ::malloc(size); if (p == nullptr) return p; + num_allocations.fetch_add(1); *(size_t*)p = size; return (char*)p + extra_len; } @@ -39,6 +40,10 @@ const OrtMemoryInfo* MockedOrtAllocator::Info() const { return cpu_memory_info; } +size_t MockedOrtAllocator::NumAllocations() const { + return num_allocations.load(); +} + void MockedOrtAllocator::LeakCheck() { if (memory_inuse.load()) ORT_THROW("memory leak!!!");