Move ShouldCpuAllocatorUseArena to allocator_utils.h/cc.

microsoft · Oct 16, 2024 · 13c77d7 · 13c77d7
1 parent 09d3aa6
commit 13c77d7
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 42 deletions.
diff --git a/onnxruntime/core/framework/allocator_utils.cc b/onnxruntime/core/framework/allocator_utils.cc
@@ -8,6 +8,8 @@
 #include <sstream>
 #include <unordered_map>
 
+#include <absl/base/config.h>
+
 #include "core/common/logging/logging.h"
 #include "core/common/narrow.h"
 #include "core/framework/bfc_arena.h"
@@ -75,4 +77,21 @@ AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) {
   }
 }
 
+bool ShouldCpuAllocatorUseArena([[maybe_unused]] bool is_arena_requested) {
+#if defined(USE_JEMALLOC) || defined(USE_MIMALLOC)
+  // We use these allocators instead of the arena.
+  return false;
+#elif defined(ABSL_HAVE_ADDRESS_SANITIZER)
+  // Using the arena may hide memory issues. Disable it in an ASan build.
+  return false;
+#else
+  // Disable the arena for 32-bit builds because it may run into an infinite loop on integer overflow.
+  if constexpr (sizeof(void*) == 4) {
+    return false;
+  } else {
+    return is_arena_requested;
+  }
+#endif
+}
+
 }  // namespace onnxruntime
diff --git a/onnxruntime/core/framework/allocator_utils.h b/onnxruntime/core/framework/allocator_utils.h
@@ -42,4 +42,9 @@ struct AllocatorCreationInfo {
 // Valid values can be found in onnxruntime_c_api.h.
 AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info);
 
+/**
+ * Gets whether a CPU allocator should use an arena or not.
+ */
+bool ShouldCpuAllocatorUseArena(bool is_arena_requested);
+
 }  // namespace onnxruntime
diff --git a/onnxruntime/core/framework/cpu_allocator_utils.cc b/onnxruntime/core/framework/cpu_allocator_utils.cc
diff --git a/onnxruntime/core/framework/cpu_allocator_utils.h b/onnxruntime/core/framework/cpu_allocator_utils.h
diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
@@ -3,7 +3,7 @@
 
 #include "core/providers/cpu/cpu_execution_provider.h"
 
-#include "core/framework/cpu_allocator_utils.h"
+#include "core/framework/allocator_utils.h"
 #include "core/framework/op_kernel.h"
 #include "core/framework/kernel_registry.h"
 #include "core/framework/int4.h"

diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc
@@ -4,7 +4,6 @@
 #include "core/session/environment.h"
 #include "core/session/allocator_adapters.h"
 #include "core/framework/allocator_utils.h"
-#include "core/framework/cpu_allocator_utils.h"
 #include "core/graph/constants.h"
 #include "core/graph/op.h"