Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

upgrade onednn to v2.1.2 and fix coredump under poolmax op

This commit is contained in:
zuochuanyong 2021-04-12 16:58:37 +08:00
parent 0df1ef53bf
commit 0b8a3b71a8
3 changed files with 19 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
cmake/external_libs/mkl_dnn.cmake
View File

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

9
mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.cc
View File

@ -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<size_t> 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<dnnl::pooling_forward>(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<kernel::AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> & /*workspace*/,
const std::vector<kernel::AddressPtr> &workspace,
const std::vector<kernel::AddressPtr> &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;
}

6
mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/pooling_cpu_kernel.h
View File

@ -31,6 +31,12 @@ class PoolingCPUKernel : public MKLCPUKernel {
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
const std::vector<AddressPtr> &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),