From 0b8a3b71a82b214bbc27919b71ae93a40aa2a6af Mon Sep 17 00:00:00 2001 From: zuochuanyong Date: Mon, 12 Apr 2021 16:58:37 +0800 Subject: [PATCH] upgrade onednn to v2.1.2 and fix coredump under poolmax op --- cmake/external_libs/mkl_dnn.cmake | 10 +++++----- .../kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.cc | 9 ++++++++- .../kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.h | 6 ++++++ 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/cmake/external_libs/mkl_dnn.cmake b/cmake/external_libs/mkl_dnn.cmake index 9942902ffe2..8b93346adad 100644 --- a/cmake/external_libs/mkl_dnn.cmake +++ b/cmake/external_libs/mkl_dnn.cmake @@ -10,14 +10,14 @@ if(CMAKE_SYSTEM_NAME MATCHES "Windows") MD5 fe660e34e9f73ab13a65987819a0712e) else() if(ENABLE_GITEE) - set(REQ_URL "https://gitee.com/mirrors/MKL-DNN/repository/archive/v1.6.tar.gz") - set(MD5 "44da423a3b6848990a907f99a65b26e7") + set(REQ_URL "https://gitee.com/mirrors/MKL-DNN/repository/archive/v2.1.2.tar.gz") + set(MD5 "d98f171d7e66e252c79e2e167ba4a8e8") else() - set(REQ_URL "https://github.com/oneapi-src/oneDNN/archive/v1.6.tar.gz") - set(MD5 "7cf251209f774ae6d61489ad7c2c3bea") + set(REQ_URL "https://github.com/oneapi-src/oneDNN/archive/v2.1.2.tar.gz") + set(MD5 "1df4f16f650b7ea08610a10af013faa3") endif() mindspore_add_pkg(onednn - VER 1.6 + VER 2.1.2 LIBS dnnl mkldnn URL ${REQ_URL} MD5 ${MD5} diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.cc index d927189af7f..ce74b83fbd6 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.cc @@ -22,6 +22,11 @@ namespace mindspore { namespace kernel { +void PoolingCPUKernel::InitInputOutputSize(const CNodePtr &kernel_node) { + CPUKernel::InitInputOutputSize(kernel_node); + workspace_size_list_.emplace_back(workspace_size_); +} + void PoolingCPUKernel::InitKernel(const CNodePtr &kernel_node) { MS_EXCEPTION_IF_NULL(kernel_node); std::vector src_shape = AnfAlgo::GetInputDeviceShape(kernel_node, 0); @@ -62,6 +67,7 @@ void PoolingCPUKernel::InitKernel(const CNodePtr &kernel_node) { dst_desc, strides_dims, kernels_dims, padding_l, padding_r); } auto prim_desc = dnnl::pooling_forward::primitive_desc(desc, MKLKernelEngine::Get().engine()); + workspace_size_ = prim_desc.workspace_desc().get_size(); primitive_ = std::make_shared(prim_desc); AddArgument(DNNL_ARG_SRC, src_desc); AddArgument(DNNL_ARG_DST, dst_desc); @@ -69,13 +75,14 @@ void PoolingCPUKernel::InitKernel(const CNodePtr &kernel_node) { } bool PoolingCPUKernel::Launch(const std::vector &inputs, - const std::vector & /*workspace*/, + const std::vector &workspace, const std::vector &outputs) { if (inputs.empty() || outputs.empty()) { MS_LOG(EXCEPTION) << "error input output size!"; } SetArgumentHandle(DNNL_ARG_SRC, inputs[0]->addr); SetArgumentHandle(DNNL_ARG_DST, outputs[0]->addr); + SetArgumentHandle(DNNL_ARG_WORKSPACE, workspace[0]->addr); ExecutePrimitive(); return true; } diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.h index fef6ea08a2e..4dfb11070a4 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.h @@ -31,6 +31,12 @@ class PoolingCPUKernel : public MKLCPUKernel { bool Launch(const std::vector &inputs, const std::vector &workspace, const std::vector &outputs) override; + + protected: + void InitInputOutputSize(const CNodePtr &kernel_node) override; + + private: + size_t workspace_size_{0}; }; MS_REG_CPU_KERNEL(MaxPool, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),