Simplify CPU allocator arena usage helper function, fix unit tests th…

…at check old ifdefs. (#22876)
microsoft · Nov 19, 2024 · af0303f · af0303f
1 parent 13346fd
commit af0303f
Show file tree

Hide file tree

Showing 7 changed files with 24 additions and 25 deletions.
diff --git a/onnxruntime/core/framework/allocator_utils.cc b/onnxruntime/core/framework/allocator_utils.cc
@@ -77,7 +77,7 @@ AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) {
   }
 }
 
-bool ShouldCpuAllocatorUseArena([[maybe_unused]] bool is_arena_requested) {
+bool DoesCpuAllocatorSupportArenaUsage() {
 #if defined(USE_JEMALLOC) || defined(USE_MIMALLOC)
   // We use these allocators instead of the arena.
   return false;
@@ -89,7 +89,7 @@ bool ShouldCpuAllocatorUseArena([[maybe_unused]] bool is_arena_requested) {
   if constexpr (sizeof(void*) == 4) {
     return false;
   } else {
-    return is_arena_requested;
+    return true;
   }
 #endif
 }

diff --git a/onnxruntime/core/framework/allocator_utils.h b/onnxruntime/core/framework/allocator_utils.h
@@ -43,8 +43,8 @@ struct AllocatorCreationInfo {
 AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info);
 
 /**
- * Gets whether a CPU allocator should use an arena or not.
+ * Gets whether a CPU allocator supports arena usage.
  */
-bool ShouldCpuAllocatorUseArena(bool is_arena_requested);
+bool DoesCpuAllocatorSupportArenaUsage();
 
 }  // namespace onnxruntime
diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
@@ -31,8 +31,7 @@ CPUExecutionProvider::CPUExecutionProvider(const CPUExecutionProviderInfo& info)
     : IExecutionProvider{onnxruntime::kCpuExecutionProvider}, info_{info} {}
 
 std::vector<AllocatorPtr> CPUExecutionProvider::CreatePreferredAllocators() {
-  const bool is_arena_requested = info_.create_arena;
-  const bool create_arena = ShouldCpuAllocatorUseArena(is_arena_requested);
+  const bool create_arena = DoesCpuAllocatorSupportArenaUsage() ? info_.create_arena : false;
   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
@@ -117,8 +117,9 @@ Status Environment::CreateAndRegisterAllocator(const OrtMemoryInfo& mem_info, co
   }
 
   // determine if arena should be used
-  const bool is_arena_requested = mem_info.alloc_type == OrtArenaAllocator;
-  const bool create_arena = ShouldCpuAllocatorUseArena(is_arena_requested);
+  const bool create_arena = DoesCpuAllocatorSupportArenaUsage()
+                                ? (mem_info.alloc_type == OrtArenaAllocator)
+                                : false;
 
   AllocatorPtr allocator_ptr;
   // create appropriate DeviceAllocatorRegistrationInfo and allocator based on create_arena

diff --git a/onnxruntime/test/framework/allocator_test.cc b/onnxruntime/test/framework/allocator_test.cc
@@ -3,6 +3,7 @@
 #include <absl/base/config.h>
 
 #include "core/framework/allocator.h"
+#include "core/framework/allocator_utils.h"
 
 #include "test_utils.h"
 #include "gtest/gtest.h"
@@ -15,12 +16,10 @@ TEST(AllocatorTest, CPUAllocatorTest) {
   ASSERT_STREQ(cpu_arena->Info().name, CPU);
   EXPECT_EQ(cpu_arena->Info().id, 0);
 
-  // arena is disabled for CPUExecutionProvider on x86 and JEMalloc
-#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(__loongarch__) || defined(_M_ARM64)) && !defined(USE_JEMALLOC) && !defined(USE_MIMALLOC) && !defined(ABSL_HAVE_ADDRESS_SANITIZER)
-  EXPECT_EQ(cpu_arena->Info().alloc_type, OrtAllocatorType::OrtArenaAllocator);
-#else
-  EXPECT_EQ(cpu_arena->Info().alloc_type, OrtAllocatorType::OrtDeviceAllocator);
-#endif
+  const auto expected_allocator_type = DoesCpuAllocatorSupportArenaUsage()
+                                           ? OrtAllocatorType::OrtArenaAllocator
+                                           : OrtAllocatorType::OrtDeviceAllocator;
+  EXPECT_EQ(cpu_arena->Info().alloc_type, expected_allocator_type);
 
   size_t size = 1024;
   auto bytes = cpu_arena->Alloc(size);

diff --git a/onnxruntime/test/framework/session_state_test.cc b/onnxruntime/test/framework/session_state_test.cc
@@ -5,6 +5,7 @@
 #include <absl/base/config.h>
 
 #include "asserts.h"
+#include "core/framework/allocator_utils.h"
 #include "core/framework/execution_providers.h"
 #include "core/framework/graph_partitioner.h"
 #include "core/framework/kernel_registry.h"
@@ -216,10 +217,12 @@ TEST_P(SessionStateTestP, TestInitializerProcessing) {
 
 // Test that we allocate memory for an initializer from non-arena memory even if we provide an arena-based allocator
 // if the relevant session option config flag is set
-// For this test we need to enable the arena-based allocator which is not supported on x86 builds, so
-// enable this test only on x64 builds
-#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64)) && !defined(USE_MIMALLOC) && !defined(ABSL_HAVE_ADDRESS_SANITIZER)
 TEST(SessionStateTest, TestInitializerMemoryAllocatedUsingNonArenaMemory) {
+  // For this test we need to enable the arena-based allocator.
+  if (!DoesCpuAllocatorSupportArenaUsage()) {
+    GTEST_SKIP() << "CPU allocator does not support arena usage.";
+  }
+
   AllocatorPtr cpu_allocator = std::make_shared<CPUAllocator>();
   // Part 1: Feature turned ON (i.e.) allocate from non-arena memory
   {
@@ -348,8 +351,6 @@ TEST(SessionStateTest, TestInitializerMemoryAllocatedUsingNonArenaMemory) {
   }
 }
 
-#endif
-
 INSTANTIATE_TEST_SUITE_P(SessionStateTests, SessionStateTestP, testing::ValuesIn(param_list));
 
 #ifndef ENABLE_TRAINING_CORE

diff --git a/onnxruntime/test/framework/tensor_test.cc b/onnxruntime/test/framework/tensor_test.cc
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "core/framework/tensor.h"
+#include "core/framework/allocator_utils.h"
 #include "test_utils.h"
 
 #include "gmock/gmock.h"
@@ -137,12 +138,10 @@ TEST(TensorTest, EmptyTensorTest) {
   ASSERT_STREQ(location.name, CPU);
   EXPECT_EQ(location.id, 0);
 
-  // arena is disabled for CPUExecutionProvider on x86 and JEMalloc
-#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(__loongarch__) || defined(_M_ARM64)) && !defined(USE_JEMALLOC) && !defined(USE_MIMALLOC) && !defined(ABSL_HAVE_ADDRESS_SANITIZER)
-  EXPECT_EQ(location.alloc_type, OrtAllocatorType::OrtArenaAllocator);
-#else
-  EXPECT_EQ(location.alloc_type, OrtAllocatorType::OrtDeviceAllocator);
-#endif
+  const auto expected_allocator_type = DoesCpuAllocatorSupportArenaUsage()
+                                           ? OrtAllocatorType::OrtArenaAllocator
+                                           : OrtAllocatorType::OrtDeviceAllocator;
+  EXPECT_EQ(location.alloc_type, expected_allocator_type);
 }
 
 TEST(TensorTest, StringTensorTest) {