From 1d78681d0586a15bba19058c2762eaf3022ab654 Mon Sep 17 00:00:00 2001
From: WangZhen <23097963+0x45f@users.noreply.github.com>
Date: Thu, 15 Sep 2022 14:35:06 +0800
Subject: [PATCH] Support 0 shapes input Tensor for MKL slice (#45930)

Support 0 shapes input Tensor for MKL slice kernel
---
 .../fluid/framework/data_layout_transform.cc  |  4 ++-
 .../fluid/operators/mkldnn/slice_mkldnn_op.cc | 10 ++++++-
 .../kernels/funcs/data_layout_transform.cc    |  4 ++-
 .../fluid/tests/unittests/test_slice_op.py    | 27 +++++++++++++++++++
 4 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/framework/data_layout_transform.cc b/paddle/fluid/framework/data_layout_transform.cc
index 4bf81b46b34..3c6a89f2939 100644
--- a/paddle/fluid/framework/data_layout_transform.cc
+++ b/paddle/fluid/framework/data_layout_transform.cc
@@ -171,7 +171,9 @@ void innerTransDataLayoutFromMKLDNN(DataLayout in_layout,
   out->set_mem_desc(out_mem_desc);
   out->Resize(in.dims());
 
-  if ((in.mem_desc() != out->mem_desc()) || always_copy) {
+  // Note(0x45f): Using initialized() to support slice Tensors
+  // with shapes like [0, 0, 0].
+  if (in.initialized() && ((in.mem_desc() != out->mem_desc()) || always_copy)) {
     void* in_data = GetDataFromTensor(in, in_type);
 
     platform::ReorderMKLDNNHandler handler(
diff --git a/paddle/fluid/operators/mkldnn/slice_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/slice_mkldnn_op.cc
index e23971c86ad..a7c6bd28486 100644
--- a/paddle/fluid/operators/mkldnn/slice_mkldnn_op.cc
+++ b/paddle/fluid/operators/mkldnn/slice_mkldnn_op.cc
@@ -72,11 +72,19 @@ class SliceMKLDNNKernel : public framework::OpKernel<T> {
       ends[i] = ends[i] < 0 ? x_vec_dims[axes[i]] + ends[i]
                             : std::min(ends[i], x_vec_dims[axes[i]]);
       offsets[axes[i]] = starts[i];
-      slice_dims[axes[i]] = ends[i] - starts[i];
+      slice_dims[axes[i]] =
+          std::max(static_cast<int64_t>(0), ends[i] - starts[i]);
     }
 
     out->Resize(phi::make_ddim(slice_dims));
 
+    // Note(0x45f): To support slice Tensors with shapes like [0, 0, 0].
+    if (!x->initialized()) {
+      out->mutable_data(x->place(), x->dtype());
+      out->set_layout(experimental::DataLayout::kMKLDNN);
+      return;
+    }
+
     dnnl::memory::data_type x_type =
         framework::ToMKLDNNDataType(framework::TransToProtoVarType(x->dtype()));
 
diff --git a/paddle/phi/kernels/funcs/data_layout_transform.cc b/paddle/phi/kernels/funcs/data_layout_transform.cc
index 9b09d897eff..9d2d0bf3b5c 100644
--- a/paddle/phi/kernels/funcs/data_layout_transform.cc
+++ b/paddle/phi/kernels/funcs/data_layout_transform.cc
@@ -83,7 +83,9 @@ void innerTransDataLayoutFromOneDNN(DataLayout in_layout,
   out->set_mem_desc(out_mem_desc);
   out->Resize(in.dims());
 
-  if ((in.mem_desc() != out->mem_desc()) || always_copy) {
+  // Note(0x45f): Using initialized() to support slice Tensors
+  // with shapes like [0, 0, 0].
+  if (in.initialized() && ((in.mem_desc() != out->mem_desc()) || always_copy)) {
     void* in_data = GetDataFromTensor(in, in_type);
 
     ReorderOneDNNHandler handler(in_tz, in.dtype(), in_type, cpu_engine);
diff --git a/python/paddle/fluid/tests/unittests/test_slice_op.py b/python/paddle/fluid/tests/unittests/test_slice_op.py
index fee36c0d421..bc10e4f4fde 100644
--- a/python/paddle/fluid/tests/unittests/test_slice_op.py
+++ b/python/paddle/fluid/tests/unittests/test_slice_op.py
@@ -82,6 +82,33 @@ class TestCase2(TestSliceOp):
         self.out = self.input[-3:3, 0:100, :, 2:-1]
 
 
+class TestSliceZerosShapeTensor(OpTest):
+
+    def setUp(self):
+        self.op_type = "slice"
+        self.config()
+        self.inputs = {'Input': self.input}
+        self.outputs = {'Out': self.out}
+        self.attrs = {
+            'axes': self.axes,
+            'starts': self.starts,
+            'ends': self.ends,
+            'infer_flags': self.infer_flags,
+            'use_mkldnn': True
+        }
+
+    def config(self):
+        self.input = np.random.random([0, 0, 0]).astype("float32")
+        self.starts = [1]
+        self.ends = [2]
+        self.axes = [0]
+        self.infer_flags = []
+        self.out = self.input[1:2]
+
+    def test_check_output(self):
+        self.check_output_with_place(paddle.CPUPlace())
+
+
 # 1.2 with attr(decrease)
 class TestSliceOp_decs_dim(OpTest):
 
-- 
GitLab