saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support plugging in custom user-defined allocators for sharing between sessions (#8059)

Browse source
This commit is contained in:
Hariharan Seshadri 2021-07-22 10:17:35 -07:00 committed by GitHub
parent 989491c333
commit 3360024a0b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
36 changed files with 652 additions and 424 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
include/onnxruntime/core/framework/allocator.h
View file

@ -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
*/

2
include/onnxruntime/core/framework/tensor.h
View file

@ -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.
*/

5
include/onnxruntime/core/session/environment.h
View file

@ -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);

49
include/onnxruntime/core/session/onnxruntime_c_api.h
View file

@ -12,7 +12,7 @@
#include <string.h>
// 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);
};
/*

2
onnxruntime/core/framework/allocator.cc
View file

@ -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) {

31
onnxruntime/core/framework/arena.h → onnxruntime/core/framework/allocator_stats.h
View file

@ -4,38 +4,9 @@
#pragma once
#include <string>
#include "core/common/common.h"
#include "core/framework/allocator.h"
#include <sstream>
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<IArenaAllocator>;
// Runtime statistics collected by an allocator.
struct AllocatorStats {

10
onnxruntime/core/framework/allocatormgr.cc
View file

@ -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 <mutex>
#include <sstream>
@ -48,11 +48,11 @@ AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) {
return nullptr;
}
#ifdef USE_MIMALLOC
return std::shared_ptr<IArenaAllocator>(
std::make_unique<MiMallocArena>(std::move(device_allocator), max_mem));
#ifdef USE_MIMALLOC_ARENA_ALLOCATOR
return std::shared_ptr<IAllocator>(
std::make_unique<MiMallocAllocator>(max_mem));
#else
return std::shared_ptr<IArenaAllocator>(
return std::shared_ptr<IAllocator>(
std::make_unique<BFCArena>(std::move(device_allocator),
max_mem,
arena_extend_str,

18
onnxruntime/core/framework/allocatormgr.h
View file

@ -19,14 +19,14 @@ using MemoryInfoSet = std::set<OrtMemoryInfo>;
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);

11
onnxruntime/core/framework/bfc_arena.cc
View file

@ -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 <type_traits>
@ -11,11 +12,11 @@ BFCArena::BFCArena(std::unique_ptr<IAllocator> 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),

16
onnxruntime/core/framework/bfc_arena.h
View file

@ -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 <intrin.h>
@ -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<size_t>(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);

73
onnxruntime/core/framework/execution_provider.cc
View file

@ -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<AllocatorManager> ) {
void IExecutionProvider::RegisterAllocator(std::shared_ptr<AllocatorManager>) {
return;
}

44
onnxruntime/core/framework/mimalloc_allocator.cc Normal file
View file

@ -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

31
onnxruntime/core/framework/mimalloc_allocator.h Normal file
View file

@ -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

50
onnxruntime/core/framework/mimalloc_arena.cc
View file

@ -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<IAllocator> 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

36
onnxruntime/core/framework/mimalloc_arena.h
View file

@ -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<IAllocator> 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

14
onnxruntime/core/framework/session_state_utils.cc
View file

@ -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<size_t> mem_size = 0;
if (!alloc->CalcMemSizeForArray(SafeInt<size_t>(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<IArenaAllocator*>(alloc.get())->Reserve(mem_size);
else
p_data = alloc->Alloc(mem_size);
if (!alloc->CalcMemSizeForArray(SafeInt<size_t>(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();

4
onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h
View file

@ -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<IArenaAllocator*>(alloc.get())->Reserve(peak_size);
buffer = static_cast<BFCArena*>(alloc.get())->Reserve(peak_size);
} else {
buffer = alloc->Alloc(peak_size);
}

131
onnxruntime/core/session/allocator_adapters.cc Normal file
View file

@ -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<OrtAllocatorImplWrappingIAllocator*>(this_)->Alloc(size); };
OrtAllocator::Free =
[](OrtAllocator* this_, void* p) { static_cast<OrtAllocatorImplWrappingIAllocator*>(this_)->Free(p); };
OrtAllocator::Info =
[](const OrtAllocator* this_) { return static_cast<const OrtAllocatorImplWrappingIAllocator*>(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<const ::onnxruntime::InferenceSession*>(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<IAllocator> i_alloc_ptr =
std::make_shared<onnxruntime::IAllocatorImplWrappingOrtAllocator>(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<onnxruntime::OrtAllocatorImpl*>(allocator);
}

53
onnxruntime/core/session/allocator_adapters.h Normal file
View file

@ -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

25
onnxruntime/core/session/allocator_impl.h
View file

@ -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

50
onnxruntime/core/session/default_cpu_allocator_c_api.cc
View file

@ -1,19 +1,14 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <atomic>
#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 <assert.h>
#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<OrtDefaultAllocator*>(this_)->Alloc(size); };
OrtAllocator::Free = [](OrtAllocator* this_, void* p) { static_cast<OrtDefaultAllocator*>(this_)->Free(p); };
OrtAllocator::Info = [](const OrtAllocator* this_) { return static_cast<const OrtDefaultAllocator*>(this_)->Info(); };
OrtAllocator::Alloc =
[](OrtAllocator* this_, size_t size) { return static_cast<OrtDefaultCpuAllocator*>(this_)->Alloc(size); };
OrtAllocator::Free =
[](OrtAllocator* this_, void* p) { static_cast<OrtDefaultCpuAllocator*>(this_)->Free(p); };
OrtAllocator::Info =
[](const OrtAllocator* this_) { return static_cast<const OrtDefaultCpuAllocator*>(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
}

58
onnxruntime/core/session/device_allocator.cc
View file

@ -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<const ::onnxruntime::InferenceSession*>(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<onnxruntime::OrtAllocatorForDevice*>(allocator);
}

45
onnxruntime/core/session/device_allocator.h
View file

@ -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 <assert.h>
#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<OrtAllocatorForDevice*>(this_)->Alloc(size); };
OrtAllocator::Free = [](OrtAllocator* this_, void* p) { static_cast<OrtAllocatorForDevice*>(this_)->Free(p); };
OrtAllocator::Info = [](const OrtAllocator* this_) { return static_cast<const OrtAllocatorForDevice*>(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

63
onnxruntime/core/session/environment.cc
View file

@ -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::LoggingManager> 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::LoggingManager> logging_manager,
const OrtThreadingOptions* tp_options,
bool create_global_thread_pools) {

8
onnxruntime/core/session/inference_session.cc
View file

@ -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<AllocatorPtr>& arenas_to_shrink) {
for (auto& alloc : arenas_to_shrink) {
auto status = static_cast<IArenaAllocator*>(alloc.get())->Shrink();
auto status = static_cast<BFCArena*>(alloc.get())->Shrink();
if (!status.IsOK()) {
LOGS(*session_logger_, WARNING) << "Unable to shrink arena: " << alloc->Info().ToString()

12
onnxruntime/core/session/onnxruntime_c_api.cc
View file

@ -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<IAllocator> alloc_ptr = std::make_shared<onnxruntime::AllocatorWrapper>(allocator);
std::shared_ptr<IAllocator> alloc_ptr = std::make_shared<onnxruntime::IAllocatorImplWrappingOrtAllocator>(allocator);
*out = std::make_unique<Tensor>(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<IAllocator> alloc_ptr = std::make_shared<onnxruntime::AllocatorWrapper>(allocator);
std::shared_ptr<IAllocator> alloc_ptr = std::make_shared<onnxruntime::IAllocatorImplWrappingOrtAllocator>(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)

2
onnxruntime/core/session/ort_apis.h
View file

@ -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

6
onnxruntime/core/session/ort_env.cc
View file

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

5
onnxruntime/core/session/ort_env.h
View file

@ -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_;

28
onnxruntime/test/framework/TestAllocatorManager.cc
View file

@ -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<IAllocator> 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);

5
onnxruntime/test/framework/TestAllocatorManager.h
View file

@ -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.

9
onnxruntime/test/framework/inference_session_test.cc
View file

@ -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<IOBinding> 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

154
onnxruntime/test/shared_lib/test_inference.cc
View file

@ -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<MockedOrtAllocator>();
@ -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<MockedOrtAllocator>();
// 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<Input> 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<MockedOrtAllocator>();
// prepare expected inputs and outputs
// prepare expected outputs
std::vector<int64_t> expected_dims_y = {3, 2};
std::vector<float> 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<OrtMemoryInfo, decltype(api.ReleaseMemoryInfo)> rel_info(mem_info, api.ReleaseMemoryInfo);
OrtArenaCfg* arena_cfg = nullptr;
ASSERT_TRUE(api.CreateArenaCfg(0, -1, -1, -1, &arena_cfg) == nullptr);
std::unique_ptr<OrtArenaCfg, decltype(api.ReleaseArenaCfg)> rel_arena_cfg(arena_cfg, api.ReleaseArenaCfg);
ASSERT_TRUE(api.CreateAndRegisterAllocator(env_ptr, mem_info, arena_cfg) == nullptr);
// test for duplicates
std::unique_ptr<OrtStatus, decltype(api.ReleaseStatus)> 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<MockedOrtAllocator>();
// 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<float>(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<float>(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<OrtMemoryInfo, decltype(api.ReleaseMemoryInfo)> rel_info(mem_info, api.ReleaseMemoryInfo);
OrtArenaCfg* arena_cfg = nullptr;
ASSERT_TRUE(api.CreateArenaCfg(0, -1, -1, -1, &arena_cfg) == nullptr);
std::unique_ptr<OrtArenaCfg, decltype(api.ReleaseArenaCfg)> 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<OrtStatus, decltype(api.ReleaseStatus)> 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<float>(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<float>(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<OrtStatus, decltype(api.ReleaseStatus)>
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<float>(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<float>(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) {

4
onnxruntime/test/testdata/custom_execution_provider_library/my_allocator.cc vendored
View file

@ -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) {

2
onnxruntime/test/util/include/test_allocator.h
View file

@ -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<size_t> memory_inuse{0};
std::atomic<size_t> num_allocations{0};
OrtMemoryInfo* cpu_memory_info;
};

5
onnxruntime/test/util/test_allocator.cc
View file

@ -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!!!");