From 05b9440fceabc8b95b52173ba67330b4540ac5c1 Mon Sep 17 00:00:00 2001 From: Sreekanth Yalachigere <17345104+sreekanth-yalachigere@users.noreply.github.com> Date: Fri, 11 Jan 2019 08:48:20 -0800 Subject: [PATCH] mkldnn:Conv weight optimization (#256) * mkldnn:Conv weight optimization * weight optimization: review changes * lock_guard and mutex for thread safe * mutex added to provider * lock to ReOrder done only once * removed #ifndef mkldnn_hpp * keep re-ordered mem buffer in scope * applied clang format * review updates: map to unordered map * conv_mutex to mutex_ --- .../mkldnn/mkldnn_execution_provider.h | 36 ++++++++++++ onnxruntime/core/providers/mkldnn/nn/conv.cc | 57 ++++++++++++------- onnxruntime/core/providers/mkldnn/nn/conv.h | 5 ++ 3 files changed, 77 insertions(+), 21 deletions(-) diff --git a/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.h b/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.h index e5820f8951..5e70e7b075 100644 --- a/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.h +++ b/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.h @@ -4,11 +4,18 @@ #pragma once #include +#include +#include +#include #include "core/framework/allocatormgr.h" #include "core/framework/execution_provider.h" #include "core/graph/graph_transformer.h" +namespace mkldnn { +struct memory; +}; + namespace onnxruntime { // Information needed to construct MKL-DNN execution providers. @@ -37,6 +44,35 @@ class MKLDNNExecutionProvider : public IExecutionProvider { } virtual std::shared_ptr GetKernelRegistry() const override; + + std::shared_ptr GetWeightsMemoryBuffer(const std::string& weight_key) { + auto iter = weights_mem_map_.find(weight_key); + if (iter != weights_mem_map_.end()) + return iter->second; + return nullptr; + } + + void SetWeightsMemoryBuffer(const std::string& weight_key, + const std::shared_ptr& filter_dst_mem) { + weights_mem_map_.insert(std::make_pair(weight_key, filter_dst_mem)); + } + + std::mutex& GetMutex() { + return mutex_; + } + + void SaveAllocatedMemory(IAllocatorUniquePtr buffer) { + // keep reordered memory buffers in scope. + reordered_buffers_.push_back(std::move(buffer)); + } + + private: + // mkldnn weights(filer data) memory blocks from first iteration + // saved by weights name + std::unordered_map> weights_mem_map_; + // Save reordered memory buffers in list so that memory is not freed. + std::vector> reordered_buffers_; + std::mutex mutex_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/nn/conv.cc b/onnxruntime/core/providers/mkldnn/nn/conv.cc index f960c26fd9..b38324f38c 100644 --- a/onnxruntime/core/providers/mkldnn/nn/conv.cc +++ b/onnxruntime/core/providers/mkldnn/nn/conv.cc @@ -4,9 +4,11 @@ #ifdef _WIN32 #pragma warning(disable : 4244) #endif +#include +#include -#include "core/providers/mkldnn/mkldnn_common.h" #include "core/providers/mkldnn/nn/conv.h" +#include "core/providers/mkldnn/mkldnn_common.h" #include "core/providers/mkldnn/mkldnn_fwd.h" namespace onnxruntime { @@ -277,15 +279,15 @@ Status Conv::Compute(OpKernelContext* context) const { if (kernel_rank + 2 != W->Shape().NumDimensions()) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "kernel_shape num_dims is not compatible with W num_dims.", - " kernel_shape: ", TensorShape(kernel_shape).ToString().c_str(), - " W: ", W->Shape().ToString().c_str()); + " kernel_shape: ", TensorShape(kernel_shape).ToString().c_str(), + " W: ", W->Shape().ToString().c_str()); } for (size_t i = 0; i < kernel_rank; ++i) { if (kernel_shape[i] != W->Shape()[i + 2]) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "kernel_shape is not compatible with W shape.", - " kernel_shape: ", TensorShape(kernel_shape).ToString().c_str(), - " W: ", W->Shape().ToString().c_str()); + " kernel_shape: ", TensorShape(kernel_shape).ToString().c_str(), + " W: ", W->Shape().ToString().c_str()); } } @@ -336,7 +338,6 @@ Status Conv::Compute(OpKernelContext* context) const { AllocatorPtr alloc; ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); IAllocatorUniquePtr src_reorder_buffer; - IAllocatorUniquePtr filter_reorder_buffer; IAllocatorUniquePtr dst_reorder_buffer; const T* src_data = X->template Data(); @@ -402,22 +403,36 @@ Status Conv::Compute(OpKernelContext* context) const { src_data = static_cast(dst.get_data_handle()); } - // Reorder filter memory layout if necessary. - if (filter_format != conv_primitive->GetFilterMemoryFormat()) { - auto pd = mkldnn::memory::primitive_desc(mkldnn::memory::desc(filter_dims_mkl, - MklDnnType(), - filter_format), - cpu_engine); - mkldnn::memory src = mkldnn::memory(pd, (void*)filter_data); - // allocate the size queried from memory primitive desc. it may not match tensor logical size due to - // mkldnn using padding to allow use of blocked format. - filter_reorder_buffer = IAllocator::MakeUniquePtr(alloc, conv_primitive->GetFilterSize()); - mkldnn::memory dst = mkldnn::memory(conv_fwd_pd->weights_primitive_desc(), filter_reorder_buffer.get()); - MemoryReorderParams params(src, dst); - DoReorder(params); - filter_data = static_cast(dst.get_data_handle()); - } + // Reorder filter memory layout if necessary + // Avoid data reordering. Save filter memory in mkldnn format from first iteration + // in execution provider mapped by weight name. + { + // lock to make sure reordering is done only once + std::lock_guard lock(provider_->GetMutex()); + auto weight_name = OpKernel::Node().InputDefs()[1]->Name(); + std::shared_ptr filter_dst_mem = provider_->GetWeightsMemoryBuffer(weight_name); + if (filter_dst_mem == nullptr) { + if (filter_format != conv_primitive->GetFilterMemoryFormat()) { + auto pd = mkldnn::memory::primitive_desc(mkldnn::memory::desc( + filter_dims_mkl, MklDnnType(), filter_format), + cpu_engine); + mkldnn::memory src = mkldnn::memory(pd, (void*)filter_data); + IAllocatorUniquePtr filter_reorder_buffer = IAllocator::MakeUniquePtr(alloc, conv_primitive->GetFilterSize()); + filter_dst_mem.reset( + new mkldnn::memory(conv_fwd_pd->weights_primitive_desc(), filter_reorder_buffer.get())); + + MemoryReorderParams params(src, *filter_dst_mem); + DoReorder(params); + provider_->SaveAllocatedMemory(std::move(filter_reorder_buffer)); + + filter_data = static_cast(filter_dst_mem->get_data_handle()); + provider_->SetWeightsMemoryBuffer(weight_name, filter_dst_mem); + } + } else { + filter_data = static_cast(filter_dst_mem->get_data_handle()); + } + } // Allocate dst buffer if reorder is necessary if (dst_md.data.format != conv_primitive->GetDstMemoryFormat()) { // allocate the size queried from memory primitive desc. it may not match tensor logical size due to diff --git a/onnxruntime/core/providers/mkldnn/nn/conv.h b/onnxruntime/core/providers/mkldnn/nn/conv.h index 76bd48a4c8..e66566262e 100644 --- a/onnxruntime/core/providers/mkldnn/nn/conv.h +++ b/onnxruntime/core/providers/mkldnn/nn/conv.h @@ -4,18 +4,23 @@ #pragma once #include "core/framework/op_kernel.h" #include "core/providers/cpu/nn/conv.h" +#include "core/providers/mkldnn/mkldnn_execution_provider.h" namespace onnxruntime { namespace mkl_dnn { + template class Conv final : public onnxruntime::Conv { public: explicit Conv(const OpKernelInfo& info) : onnxruntime::Conv(info) { + provider_ = (const_cast( + dynamic_cast(info.GetExecutionProvider()))); } Status Compute(OpKernelContext* context) const override; private: + MKLDNNExecutionProvider* provider_; }; } // namespace mkl_dnn } // namespace onnxruntime