diff --git a/include/onnxruntime/core/session/environment.h b/include/onnxruntime/core/session/environment.h index 89467f5238fa9..59ca1a1df762e 100644 --- a/include/onnxruntime/core/session/environment.h +++ b/include/onnxruntime/core/session/environment.h @@ -199,23 +199,6 @@ class Environment { using OrtAllocatorUniquePtr = std::unique_ptr>; - // if the user calls CreateSharedAllocator and wraps the plugin EP's allocator with an arena we end up with - // OrtAllocator from EP -> wrapped in IAllocatorImplWrappingOrtAllocator -> inside a BFCArena IAllocator. - // we can put that in shared_allocators_ for sessions to use, but to have an OrtAllocator available in - // shared_ort_allocators_ that can be used outside of a session we need to additionally wrap that in an - // OrtAllocatorImplWrappingIAllocator. way too many levels of indirection but that is what it is currently. - // we need something to own that final OrtAllocator, so we add it to arena_ort_allocators_. - // - // TODO: we could split out the BFCArena implementation so it can be plugged into either an IAllocator - // or an OrtAllocator instance to reduce the indirection a little. - // with that we get an OrtAllocator from the EP, wrap it with an OrtAllocator based BFCArena, and wrap that with the - // IAllocatorImplWrappingOrtAllocator which takes ownership of the OrtAllocator and is in shared_allocators_. - // - // Alternatively we can disable wrapping an EP's allocator with a BFCArena and say the EP should provide the arena - // implementation directly. They're free to copy BFCArena as it came from TF originally. Or we could provide a - // cut-and-paste BFCArena implementation that works using the EP API that can be included in the EP source. - std::unordered_map> arena_ort_allocators_; - #if !defined(ORT_MINIMAL_BUILD) // register EPs that are built into the ORT binary so they can take part in AutoEP selection // added to ep_libraries diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index dfb2e33f8cb32..39b785c327d56 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -182,11 +182,6 @@ Status Environment::UnregisterAllocatorImpl(const OrtMemoryInfo& mem_info, bool shared_ort_allocators_.erase(it2); } - // also remove an arena wrapped allocator from an EP if the user called CreateSharedAllocator to create one - if (auto it3 = arena_ort_allocators_.find(&mem_info); it3 != arena_ort_allocators_.end()) { - arena_ort_allocators_.erase(it3); - } - if (found_shared_allocator) { shared_allocators_.erase(it); } @@ -436,6 +431,10 @@ Environment::~Environment() { // instance and will call Release on it. If the plugin EP has been freed the Release will fail. shared_allocators_.clear(); + // and as any OrtAllocator instances in shared_ort_allocators_ were owned by values in shared_allocators_ and have + // now been released we need to clear that too before calling UnregisterExecutionProviderLibrary(). + shared_ort_allocators_.clear(); + #if !defined(ORT_MINIMAL_BUILD) // unregister any remaining EP libraries so they're cleaned up in a determistic way. while (!ep_libraries_.empty()) { @@ -673,11 +672,6 @@ Status Environment::CreateSharedAllocatorImpl(const OrtEpDevice& ep_device, shared_ort_allocators_.erase(it); } - // if a previous call created an arena wrapped allocator for the EP's memory_info we also need to remove that - if (auto it = arena_ort_allocators_.find(&memory_info); it != arena_ort_allocators_.end()) { - arena_ort_allocators_.erase(it); - } - // we only want one shared allocator for an OrtDevice in the shared_allocators_ so that it's deterministic which // one will be used for an inference session. ignore the name so that is the case. if (auto it = FindExistingAllocator(shared_allocators_, memory_info, /*match_name*/ false);