Skip to content

Commit cb0c5e9

Browse files
authored
Fix cleanup of Environment class. (microsoft#25743)
### Description <!-- Describe your changes. --> The clearing of shared_allocators_ invalidates all entries in shared_ort_allocators_. Remove unused shared_arena_allocators_. That became unnecessary by providing EPs an example implementation for an OrtAllocator based stream-aware arena that they can use directly. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> Fix access violation (swallowed as it happens during shutdown) in dtor.
1 parent a61fb39 commit cb0c5e9

File tree

2 files changed

+4
-27
lines changed

2 files changed

+4
-27
lines changed

include/onnxruntime/core/session/environment.h

Lines changed: 0 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -199,23 +199,6 @@ class Environment {
199199

200200
using OrtAllocatorUniquePtr = std::unique_ptr<OrtAllocator, std::function<void(OrtAllocator*)>>;
201201

202-
// if the user calls CreateSharedAllocator and wraps the plugin EP's allocator with an arena we end up with
203-
// OrtAllocator from EP -> wrapped in IAllocatorImplWrappingOrtAllocator -> inside a BFCArena IAllocator.
204-
// we can put that in shared_allocators_ for sessions to use, but to have an OrtAllocator available in
205-
// shared_ort_allocators_ that can be used outside of a session we need to additionally wrap that in an
206-
// OrtAllocatorImplWrappingIAllocator. way too many levels of indirection but that is what it is currently.
207-
// we need something to own that final OrtAllocator, so we add it to arena_ort_allocators_.
208-
//
209-
// TODO: we could split out the BFCArena implementation so it can be plugged into either an IAllocator
210-
// or an OrtAllocator instance to reduce the indirection a little.
211-
// with that we get an OrtAllocator from the EP, wrap it with an OrtAllocator based BFCArena, and wrap that with the
212-
// IAllocatorImplWrappingOrtAllocator which takes ownership of the OrtAllocator and is in shared_allocators_.
213-
//
214-
// Alternatively we can disable wrapping an EP's allocator with a BFCArena and say the EP should provide the arena
215-
// implementation directly. They're free to copy BFCArena as it came from TF originally. Or we could provide a
216-
// cut-and-paste BFCArena implementation that works using the EP API that can be included in the EP source.
217-
std::unordered_map<const OrtMemoryInfo*, std::unique_ptr<OrtAllocatorImplWrappingIAllocator>> arena_ort_allocators_;
218-
219202
#if !defined(ORT_MINIMAL_BUILD)
220203
// register EPs that are built into the ORT binary so they can take part in AutoEP selection
221204
// added to ep_libraries

onnxruntime/core/session/environment.cc

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -182,11 +182,6 @@ Status Environment::UnregisterAllocatorImpl(const OrtMemoryInfo& mem_info, bool
182182
shared_ort_allocators_.erase(it2);
183183
}
184184

185-
// also remove an arena wrapped allocator from an EP if the user called CreateSharedAllocator to create one
186-
if (auto it3 = arena_ort_allocators_.find(&mem_info); it3 != arena_ort_allocators_.end()) {
187-
arena_ort_allocators_.erase(it3);
188-
}
189-
190185
if (found_shared_allocator) {
191186
shared_allocators_.erase(it);
192187
}
@@ -436,6 +431,10 @@ Environment::~Environment() {
436431
// instance and will call Release on it. If the plugin EP has been freed the Release will fail.
437432
shared_allocators_.clear();
438433

434+
// and as any OrtAllocator instances in shared_ort_allocators_ were owned by values in shared_allocators_ and have
435+
// now been released we need to clear that too before calling UnregisterExecutionProviderLibrary().
436+
shared_ort_allocators_.clear();
437+
439438
#if !defined(ORT_MINIMAL_BUILD)
440439
// unregister any remaining EP libraries so they're cleaned up in a determistic way.
441440
while (!ep_libraries_.empty()) {
@@ -673,11 +672,6 @@ Status Environment::CreateSharedAllocatorImpl(const OrtEpDevice& ep_device,
673672
shared_ort_allocators_.erase(it);
674673
}
675674

676-
// if a previous call created an arena wrapped allocator for the EP's memory_info we also need to remove that
677-
if (auto it = arena_ort_allocators_.find(&memory_info); it != arena_ort_allocators_.end()) {
678-
arena_ort_allocators_.erase(it);
679-
}
680-
681675
// we only want one shared allocator for an OrtDevice in the shared_allocators_ so that it's deterministic which
682676
// one will be used for an inference session. ignore the name so that is the case.
683677
if (auto it = FindExistingAllocator(shared_allocators_, memory_info, /*match_name*/ false);

0 commit comments

Comments
 (0)