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 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(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