From 09d3aa6e38c4918a9f30ad53b40cf548ef391597 Mon Sep 17 00:00:00 2001
From: edgchen1 <18449977+edgchen1@users.noreply.github.com>
Date: Tue, 15 Oct 2024 20:03:00 -0700
Subject: [PATCH 1/2] Consolidate CPU allocator arena creation check into a
 helper function.

---
 .../core/framework/cpu_allocator_utils.cc     | 27 +++++++++++++++++++
 .../core/framework/cpu_allocator_utils.h      | 13 +++++++++
 .../providers/cpu/cpu_execution_provider.cc   | 13 +++------
 onnxruntime/core/session/environment.cc       | 16 +++--------
 4 files changed, 47 insertions(+), 22 deletions(-)
 create mode 100644 onnxruntime/core/framework/cpu_allocator_utils.cc
 create mode 100644 onnxruntime/core/framework/cpu_allocator_utils.h

diff --git a/onnxruntime/core/framework/cpu_allocator_utils.cc b/onnxruntime/core/framework/cpu_allocator_utils.cc
new file mode 100644
index 0000000000000..b35407cb674f3
--- /dev/null
+++ b/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
diff --git a/onnxruntime/core/framework/cpu_allocator_utils.h b/onnxruntime/core/framework/cpu_allocator_utils.h
new file mode 100644
index 0000000000000..86e39055b8f0c
--- /dev/null
+++ b/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
diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
index 424bee63511ad..b9c60970f742a 100644
--- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
+++ b/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};
 
diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc
index e5e718fb8d1de..7837975f6bd67 100644
--- a/onnxruntime/core/session/environment.cc
+++ b/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

From 13c77d7d0210772d7fed91e7aed9f276dea54d88 Mon Sep 17 00:00:00 2001
From: edgchen1 <18449977+edgchen1@users.noreply.github.com>
Date: Wed, 16 Oct 2024 09:57:40 -0700
Subject: [PATCH 2/2] Move ShouldCpuAllocatorUseArena to allocator_utils.h/cc.

---
 onnxruntime/core/framework/allocator_utils.cc | 19 +++++++++++++
 onnxruntime/core/framework/allocator_utils.h  |  5 ++++
 .../core/framework/cpu_allocator_utils.cc     | 27 -------------------
 .../core/framework/cpu_allocator_utils.h      | 13 ---------
 .../providers/cpu/cpu_execution_provider.cc   |  2 +-
 onnxruntime/core/session/environment.cc       |  1 -
 6 files changed, 25 insertions(+), 42 deletions(-)
 delete mode 100644 onnxruntime/core/framework/cpu_allocator_utils.cc
 delete mode 100644 onnxruntime/core/framework/cpu_allocator_utils.h

diff --git a/onnxruntime/core/framework/allocator_utils.cc b/onnxruntime/core/framework/allocator_utils.cc
index 7493ac7d0a4e8..797b6e1606f97 100644
--- 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
index 7dda1d1a6fd8f..4035a0cc349e4 100644
--- 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
deleted file mode 100644
index b35407cb674f3..0000000000000
--- a/onnxruntime/core/framework/cpu_allocator_utils.cc
+++ /dev/null
@@ -1,27 +0,0 @@
-// 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
diff --git a/onnxruntime/core/framework/cpu_allocator_utils.h b/onnxruntime/core/framework/cpu_allocator_utils.h
deleted file mode 100644
index 86e39055b8f0c..0000000000000
--- a/onnxruntime/core/framework/cpu_allocator_utils.h
+++ /dev/null
@@ -1,13 +0,0 @@
-// 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
diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
index b9c60970f742a..f880a39188a06 100644
--- 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
index 7837975f6bd67..5f929d3760a95 100644
--- 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"