From a0666b9d20c77d7ac1394de53a4214b35c351b37 Mon Sep 17 00:00:00 2001
From: jakpiase <62569058+jakpiase@users.noreply.github.com>
Date: Wed, 7 Jul 2021 06:48:48 +0200
Subject: [PATCH] Split op oneDNN AVX2 fix (#33944)

* added checking if md uses blocking format

* minor change

* removed unnecessary line
---
 paddle/fluid/operators/split_op.cc | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/operators/split_op.cc b/paddle/fluid/operators/split_op.cc
index 661e4ca727b..f81ac8882d1 100644
--- a/paddle/fluid/operators/split_op.cc
+++ b/paddle/fluid/operators/split_op.cc
@@ -78,9 +78,18 @@ class SplitOp : public framework::OperatorWithKernel {
 
 #ifdef PADDLE_WITH_MKLDNN
     if (this->CanMKLDNNBeUsed(ctx, input_data_type)) {
-      return framework::OpKernelType(input_data_type, ctx.GetPlace(),
-                                     framework::DataLayout::kMKLDNN,
-                                     framework::LibraryType::kMKLDNN);
+      // OneDNN uses blocking format, which cannot be always
+      // supported with reorders, because if blocked dimension is not divisible
+      // by
+      // 8 or 16(depending on which blocking format is used) submemory cannot be
+      // created, so in that scenario a fallback is needed
+      auto tmp_md = dnnl::memory::desc(
+          framework::vectorize(ctx.Input<Tensor>("X")->dims()),
+          dnnl::memory::data_type::f32, ctx.Input<Tensor>("X")->format());
+      if (tmp_md.data.format_desc.blocking.inner_nblks == 0)
+        return framework::OpKernelType(input_data_type, ctx.GetPlace(),
+                                       framework::DataLayout::kMKLDNN,
+                                       framework::LibraryType::kMKLDNN);
     }
 #endif
     return framework::OpKernelType(input_data_type, ctx.GetPlace());
-- 
GitLab