From a5b4d2a8a783335c8a9fff42518d6b221960c28a Mon Sep 17 00:00:00 2001
From: Changming Sun <chasun@microsoft.com>
Date: Thu, 6 Apr 2023 00:43:13 -0700
Subject: [PATCH] XNNPack: allow users to choose whether enable CPU MEM arena
 or not (#15392)

### Description
XNNPack: allow users to choose whether enable CPU MEM arena or not.
Right now it is hardcoded to true and it is not impacted by the on/off
switch in SessionOption. We should make it work.

### Motivation and Context
As we have such a switch in SessionOption, it should work as expected.
---
 .../core/providers/xnnpack/xnnpack_execution_provider.cc    | 6 ++++--
 .../core/providers/xnnpack/xnnpack_execution_provider.h     | 1 +
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.cc b/onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.cc
index e1204944cf..ba20089667 100644
--- a/onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.cc
+++ b/onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.cc
@@ -123,7 +123,8 @@ std::unique_ptr<KernelRegistry> RegisterKernels() {
 using namespace xnnpack;
 
 XnnpackExecutionProvider::XnnpackExecutionProvider(const XnnpackExecutionProviderInfo& info)
-    : IExecutionProvider{kXnnpackExecutionProvider, true} {
+    : IExecutionProvider{kXnnpackExecutionProvider, true},
+      enable_cpu_mem_arena_(info.session_options ? info.session_options->enable_cpu_mem_arena : true) {
   int xnn_thread_pool_size = info.xnn_thread_pool_size;
   int ort_thread_pool_size = info.session_options ? info.session_options->intra_op_param.thread_pool_size : 1;
   bool allow_intra_op_spinning = (info.session_options == nullptr) ||
@@ -170,7 +171,8 @@ void XnnpackExecutionProvider::RegisterAllocator(AllocatorManager& allocator_man
             // lazy create the allocator
             return std::make_unique<CPUAllocator>(OrtMemoryInfo(kXnnpackExecutionProvider,
                                                                 OrtAllocatorType::OrtDeviceAllocator));
-          });
+          },
+          cpu_device.Id(), enable_cpu_mem_arena_);
       // only the first time we create the allocator do we pass in the xnn_allocator
       cpu_alloc = stored_allocator ? stored_allocator : CreateAllocator(allocator_info);
       // enable sharing of our allocator
diff --git a/onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.h b/onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.h
index d4eefdd759..50da147049 100644
--- a/onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.h
+++ b/onnxruntime/core/providers/xnnpack/xnnpack_execution_provider.h
@@ -52,6 +52,7 @@ class XnnpackExecutionProvider : public IExecutionProvider {
 
  private:
   pthreadpool* xnnpack_thread_pool_{nullptr};
+  const bool enable_cpu_mem_arena_;
 };
 
 }  // namespace onnxruntime