From d3f09ad70262fdd91a111547472ba73d51233090 Mon Sep 17 00:00:00 2001
From: Jacek Czaja <jacek.czaja@intel.com>
Date: Tue, 23 Feb 2021 10:43:35 +0100
Subject: [PATCH] Update of onednn to 2.2 (#31067)

---
 cmake/external/mkldnn.cmake                     | 2 +-
 paddle/fluid/operators/mkldnn/conv_mkldnn_op.cc | 9 +++++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/cmake/external/mkldnn.cmake b/cmake/external/mkldnn.cmake
index 013c22c8a6c..e41d8fdb6da 100644
--- a/cmake/external/mkldnn.cmake
+++ b/cmake/external/mkldnn.cmake
@@ -20,7 +20,7 @@ SET(MKLDNN_SOURCE_DIR     ${THIRD_PARTY_PATH}/mkldnn/src/extern_mkldnn)
 SET(MKLDNN_INSTALL_DIR    ${THIRD_PARTY_PATH}/install/mkldnn)
 SET(MKLDNN_INC_DIR        "${MKLDNN_INSTALL_DIR}/include" CACHE PATH "mkldnn include directory." FORCE)
 SET(MKLDNN_REPOSITORY     ${GIT_URL}/oneapi-src/oneDNN.git)
-SET(MKLDNN_TAG            a18f78f1f058437e9efee403655d671633360f98)
+SET(MKLDNN_TAG            3d53cd3f17ce7ca365c980f0e1e50359751ca038)
 
 # Introduce variables:
 # * CMAKE_INSTALL_LIBDIR
diff --git a/paddle/fluid/operators/mkldnn/conv_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/conv_mkldnn_op.cc
index fc11951d743..2d5f560eac5 100644
--- a/paddle/fluid/operators/mkldnn/conv_mkldnn_op.cc
+++ b/paddle/fluid/operators/mkldnn/conv_mkldnn_op.cc
@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/framework/data_layout_transform.h"
 #include "paddle/fluid/operators/conv_op.h"
+#include "paddle/fluid/platform/cpu_info.h"
 #include "paddle/fluid/platform/mkldnn_reuse.h"
 
 namespace paddle {
@@ -976,8 +977,12 @@ class ConvMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
      * ('any') which lets a primitive (conv backward in this case) choose
      * the memory format preferred for best performance
      */
-
-    auto chosen_memory_format = MKLDNNMemoryFormat::any;
+    // TODO: NHWC is preferred starting from oneDNN 2.1 . Any may crash
+    auto chosen_memory_format =
+        platform::MayIUse(platform::cpu_isa_t::avx512_core) &&
+                is_conv3d == false
+            ? MKLDNNMemoryFormat::nhwc
+            : MKLDNNMemoryFormat::any;
     weights_format = MKLDNNMemoryFormat::any;
 
     auto src_md = platform::MKLDNNMemDesc(
-- 
GitLab