Added shape (U)INT8/BF16/FP32 oneDNN kernel

paddle/fluid/operators/mkldnn/shape_mkldnn_op.cc

@@ -0,0 +1,58 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/shape_op.h"
+#include "paddle/fluid/platform/mkldnn_helper.h"
+#include "paddle/fluid/platform/mkldnn_reuse.h"
+
+namespace paddle {
+namespace operators {
+
+using paddle::framework::Tensor;
+using paddle::framework::LoDTensor;
+
+template <typename T>
+class ShapeMKLDNNKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    auto* in_var = ctx.InputVar("Input");
+    auto* out = ctx.Output<Tensor>("Out");
+
+    framework::DDim in_dims;
+    if (in_var->IsType<SelectedRows>()) {
+      in_dims = in_var->Get<SelectedRows>().value().dims();
+    } else {
+      in_dims = in_var->Get<LoDTensor>().dims();
+    }
+
+    out->Resize({in_dims.size()});
+    auto out_data = out->mutable_data<int32_t>(platform::CPUPlace());
+    for (int i = 0; i < in_dims.size(); ++i) {
+      out_data[i] = in_dims[i];
+    }
+
+    out->set_layout(framework::DataLayout::kMKLDNN);
+    out->set_format(platform::GetPlainMKLDNNFormat(in_dims.size()));
+  }
+};
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OP_KERNEL(shape, MKLDNN, paddle::platform::CPUPlace,
+                   ops::ShapeMKLDNNKernel<float>,
+                   ops::ShapeMKLDNNKernel<paddle::platform::bfloat16>,
+                   ops::ShapeMKLDNNKernel<int8_t>,
+                   ops::ShapeMKLDNNKernel<uint8_t>);

paddle/fluid/operators/shape_op.cc

@@ -35,6 +35,21 @@ class ShapeOp : public framework::OperatorWithKernel {
     ctx->SetOutputDim("Out", {in_dim.size()});
   }
 
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext &ctx) const override {
+    auto input_data_type =
+        framework::OperatorWithKernel::IndicateVarDataType(ctx, "Input");
+
+#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);
+    }
+#endif
+    return framework::OpKernelType(input_data_type, ctx.GetPlace());
+  }
+
  protected:
   framework::OpKernelType GetKernelTypeForVar(
       const std::string &var_name, const framework::Tensor &tensor,
@@ -58,6 +73,14 @@ Shape Operator.
 
 Return the shape of the input.
 )DOC");
+    AddAttr<bool>("use_mkldnn",
+                  "(bool, default false) Only used in mkldnn kernel")
+        .SetDefault(false);
+    AddAttr<std::string>(
+        "mkldnn_data_type",
+        "(string, default \"float32\"). Data type of mkldnn kernel")
+        .SetDefault("float32")
+        .InEnum({"float32", "bfloat16", "int8"});
   }
 };
 

paddle/fluid/platform/mkldnn_helper.h

@@ -328,6 +328,23 @@ inline mkldnn::memory::format_tag GetMKLDNNFormat(
   return mkldnn::memory::format_tag::undef;
 }
 
+inline mkldnn::memory::format_tag GetPlainMKLDNNFormat(int tensor_rank) {
+  switch (tensor_rank) {
+    case 1:
+      return mkldnn::memory::format_tag::a;
+    case 2:
+      return mkldnn::memory::format_tag::ab;
+    case 3:
+      return mkldnn::memory::format_tag::abc;
+    case 4:
+      return mkldnn::memory::format_tag::abcd;
+    case 5:
+      return mkldnn::memory::format_tag::abcde;
+  }
+
+  return mkldnn::memory::format_tag::abcdef;
+}
+
 inline mkldnn::memory::format_tag GetMKLDNNFormat(const mkldnn::memory memory) {
   auto mem_desc = memory.get_desc();
   return GetMKLDNNFormat(mem_desc);

python/paddle/fluid/tests/unittests/mkldnn/test_shape_mkldnn_op.py

@@ -0,0 +1,85 @@
+#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import unittest
+import numpy as np
+from paddle.fluid.tests.unittests.op_test import OpTest, OpTestTool
+import paddle
+from paddle.fluid import core
+from paddle.fluid.op import Operator
+
+
+@OpTestTool.skip_if_not_cpu_bf16()
+class TestShapeOneDNNOp(OpTest):
+    def setUp(self):
+        self.op_type = "shape"
+        self.config()
+        self.attrs = {'use_mkldnn': True}
+        self.inputs = {'Input': np.zeros(self.shape).astype(self.dtype)}
+        self.outputs = {'Out': np.array(self.shape)}
+
+    def config(self):
+        self.shape = [5, 7, 4]
+        self.dtype = np.float32
+
+    def test_check_output(self):
+        self.check_output_with_place(core.CPUPlace())
+
+
+class TestShape1DOneDNNOp(TestShapeOneDNNOp):
+    def config(self):
+        self.shape = [2]
+
+
+class TestShape4DBF16OneDNNOp(TestShapeOneDNNOp):
+    def config(self):
+        self.shape = [10, 2, 3, 5]
+        self.dtype = np.uint16
+
+
+class TestShape6DBF16OneDNNOp(TestShapeOneDNNOp):
+    def config(self):
+        self.shape = [10, 2, 3, 4, 5, 2]
+        self.dtype = np.uint16
+
+
+class TestShape3DINT8OneDNNOp(TestShapeOneDNNOp):
+    def config(self):
+        self.shape = [10, 2, 3]
+        self.dtype = np.int8
+
+
+class TestShape5DINT8OneDNNOp(TestShapeOneDNNOp):
+    def config(self):
+        self.shape = [10, 2, 3, 4, 3]
+        self.dtype = np.int8
+
+
+class TestShape2DUINT8OneDNNOp(TestShapeOneDNNOp):
+    def config(self):
+        self.shape = [7, 11]
+        self.dtype = np.uint8
+
+
+class TestShape3DUINT8OneDNNOp(TestShapeOneDNNOp):
+    def config(self):
+        self.shape = [2, 7, 11]
+        self.dtype = np.uint8
+
+
+if __name__ == '__main__':
+    paddle.enable_static()
+    unittest.main()