[Intel MKL] Fix the issue which fails to create memory descriptor in mkl concat

tensorflow/core/kernels/mkl_concat_op.cc

@@ -17,7 +17,6 @@ limitations under the License.
 #include <vector>
 
 #include "mkldnn.hpp"
-#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
@@ -30,6 +29,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/mkl_util.h"
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 
 using mkldnn::concat;
 using mkldnn::stream;
@@ -461,8 +461,14 @@ class MklConcatOp : public OpKernel {
               dst_dims, MklDnnDataFormatToTFDataFormat(orig_tf_format));
           // Set the output format same as the most common format of inputs
           // to avoid layout conversions.
-          dst_md = memory::desc(dst_dims_in_nchw, MklDnnType<T>(),
-                                mkl_common_format);
+          if (mkl_common_format == memory::format::blocked) {
+            VLOG(1) << "mkl_common_format == memory::format::blocked";
+            dst_md = MklDnnData<T>::CreateBlockedMemDesc(
+                dst_dims_in_nchw, CalculateTFStrides(dst_dims_in_nchw));
+          } else {
+            dst_md = memory::desc(dst_dims_in_nchw, MklDnnType<T>(),
+                                  mkl_common_format);
+          }
         } else if (dst_dims.size() == 2 &&
                    mkl_common_format == memory::format::nc) {
           // When memory::format::nc, dst_dims are already in MKL-DNN order

tensorflow/python/kernel_tests/concat_op_test.py

@@ -29,6 +29,7 @@
 from tensorflow.python.ops import gradient_checker
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
 
@@ -714,6 +715,26 @@ def testNegativeDim(self):
       ans = self.evaluate(off)
       self.assertAllEqual(ans, [[0, 0, 0], [2, 0, 0], [3, 0, 0]])
 
+  def testCreateMemDecBlockedFormat(self):
+    """Try to create the mkl concat operation
+    when one of the input's memory descriptor is in blocked format"""
+    if test_util.IsMklEnabled():
+      s0 = np.ones((1, 8188, 4092, 1), dtype=np.uint8).astype(np.float32)
+      s1 = array_ops.strided_slice(s0, [0, 1, 1, 0], [0, -1, -1, 0],
+                                   [1, 1, 1, 1], begin_mask=9, end_mask=9)
+      s2 = array_ops.slice(s1, [0, 0, 0, 0], [-1, -1, -1, 1])
+      s3_1 = array_ops.slice(s2, [0, 4, 4, 0], [-1, 8178, 4082, 1])
+      s3_2 = array_ops.slice(s2, [0, 4, 4, 0], [-1, 8178, 4082, 1])
+      filter4_1 = constant_op.constant([[[[1.18, -0.51]]]])
+      s4_1 = nn_ops.conv2d(s3_1, filter4_1,
+                           strides=[1, 1, 1, 1], padding="VALID")
+      filter4_2 = constant_op.constant([[[[1.38, -0.11]]]])
+      s4_2 = nn_ops.conv2d(s3_2, filter4_2,
+                           strides=[1, 1, 1, 1], padding="VALID")
+      s5_1 = array_ops.slice(s4_1, [0, 6, 6, 0], [-1, 1, 1, -1])
+      s5_2 = array_ops.slice(s4_2, [0, 6, 6, 0], [-1, 1, 1, -1])
+      x_concat = array_ops.concat([s5_1, s5_2], 3)
+      self.evaluate(x_concat)  # This test is only meant to check the creation is not crashed
 
 if __name__ == "__main__":
   test.main()