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.
This commit is contained in:
Changming Sun 2023-04-06 00:43:13 -07:00 committed by GitHub
parent ca68ab6126
commit a5b4d2a8a7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 5 additions and 2 deletions

View file

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

View file

@ -52,6 +52,7 @@ class XnnpackExecutionProvider : public IExecutionProvider {
private:
pthreadpool* xnnpack_thread_pool_{nullptr};
const bool enable_cpu_mem_arena_;
};
} // namespace onnxruntime