Share allocator across sessions

microsoft · Sep 20, 2024 · 0291560 · 0291560
1 parent 99ea5ae
commit 0291560
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/generators.cpp b/src/generators.cpp
@@ -13,7 +13,12 @@ namespace Generators {
 
 static bool _ = (Ort::InitApi(), false);
 
-OrtGlobals::OrtGlobals() : env_{OrtEnv::Create(OrtLoggingLevel::ORT_LOGGING_LEVEL_ERROR)} {}
+OrtGlobals::OrtGlobals()
+    : env_{OrtEnv::Create(OrtLoggingLevel::ORT_LOGGING_LEVEL_ERROR)},
+      arena_config_{OrtArenaCfg::Create(0, -1, -1, -1)} {
+  Ort::Allocator& allocator_cpu_{Ort::Allocator::GetWithDefaultOptions()};
+  env_->CreateAndRegisterAllocator(allocator_cpu_.GetInfo(), *arena_config_);
+}
 
 // Ensure Shutdown() has been called before process exit
 struct ValidateShutdown {

diff --git a/src/generators.h b/src/generators.h
@@ -146,6 +146,7 @@ struct OrtGlobals {
   OrtGlobals();
 
   std::unique_ptr<OrtEnv> env_;
+  std::unique_ptr<OrtArenaCfg> arena_config_;
 #if USE_CUDA
   std::unique_ptr<OrtMemoryInfo> memory_info_cuda_;
   std::unique_ptr<Ort::Allocator> allocator_cuda_;