Consolidate CPU allocator arena creation checks into a helper function.

onnxruntime/core/framework/cpu_allocator_utils.cc

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/framework/cpu_allocator_utils.h"
+
+#include <absl/base/config.h>
+
+namespace onnxruntime {
+
+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

onnxruntime/core/framework/cpu_allocator_utils.h

@@ -0,0 +1,13 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+namespace onnxruntime {
+
+/**
+ * Gets whether a CPU allocator should use an arena or not.
+ */
+bool ShouldCpuAllocatorUseArena(bool is_arena_requested);
+
+}  // namespace onnxruntime

onnxruntime/core/providers/cpu/cpu_execution_provider.cc

@@ -2,7 +2,8 @@
 // Licensed under the MIT License.
 
 #include "core/providers/cpu/cpu_execution_provider.h"
-#include <absl/base/config.h>
+
+#include "core/framework/cpu_allocator_utils.h"
 #include "core/framework/op_kernel.h"
 #include "core/framework/kernel_registry.h"
 #include "core/framework/int4.h"
@@ -30,14 +31,8 @@ CPUExecutionProvider::CPUExecutionProvider(const CPUExecutionProviderInfo& info)
     : IExecutionProvider{onnxruntime::kCpuExecutionProvider}, info_{info} {}
 
 std::vector<AllocatorPtr> CPUExecutionProvider::CreatePreferredAllocators() {
-  bool create_arena = info_.create_arena;
-#if defined(USE_JEMALLOC) || defined(USE_MIMALLOC) || defined(ABSL_HAVE_ADDRESS_SANITIZER)
-  // JEMalloc/mimalloc already have memory pool, so just use device allocator.
-  create_arena = false;
-#elif !(defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64))
-  // Disable Arena allocator for x86_32 build because it may run into infinite loop when integer overflow happens
-  create_arena = false;
-#endif
+  const bool is_arena_requested = info_.create_arena;
+  const bool create_arena = ShouldCpuAllocatorUseArena(is_arena_requested);
   AllocatorCreationInfo device_info{[](int) { return std::make_unique<CPUAllocator>(); },
                                     DEFAULT_CPU_ALLOCATOR_DEVICE_ID, create_arena};
 

onnxruntime/core/session/environment.cc

@@ -4,6 +4,7 @@
 #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"
 
@@ -117,19 +118,8 @@ Status Environment::CreateAndRegisterAllocator(const OrtMemoryInfo& mem_info, co
   }
 
   // determine if arena should be used
-  const bool create_arena = [&]() -> bool {
-#if defined(USE_JEMALLOC) || defined(USE_MIMALLOC)
-    // We use these allocators instead of the arena
-    return false;
-#else
-    // Disable Arena allocator for 32-bit builds because it may run into infinite loop when integer overflow happens
-    if constexpr (sizeof(void*) == 4) {
-      return false;
-    } else {
-      return mem_info.alloc_type == OrtArenaAllocator;
-    }
-#endif
-  }();
+  const bool is_arena_requested = mem_info.alloc_type == OrtArenaAllocator;
+  const bool create_arena = ShouldCpuAllocatorUseArena(is_arena_requested);
 
   AllocatorPtr allocator_ptr;
   // create appropriate DeviceAllocatorRegistrationInfo and allocator based on create_arena