mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-02 03:55:34 +00:00
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:
parent
ca68ab6126
commit
a5b4d2a8a7
2 changed files with 5 additions and 2 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@ class XnnpackExecutionProvider : public IExecutionProvider {
|
|||
|
||||
private:
|
||||
pthreadpool* xnnpack_thread_pool_{nullptr};
|
||||
const bool enable_cpu_mem_arena_;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
Loading…
Reference in a new issue