Merge pull request #14732 from Sand3r-/mgallus/mkldnn-concat

[MKL-DNN] Concat Layer

paddle/fluid/operators/concat_mkldnn_op.cc

+/* Copyright (c) 2018 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 <memory>
+#include "paddle/fluid/operators/concat_op.h"
+#include "paddle/fluid/platform/mkldnn_helper.h"
+
+namespace paddle {
+namespace operators {
+
+using framework::DataLayout;
+using framework::Tensor;
+using mkldnn::memory;
+using mkldnn::primitive;
+using mkldnn::concat;
+using mkldnn::stream;
+using platform::to_void_cast;
+
+static void EnforceLayouts(const std::vector<const Tensor*> inputs) {
+  for (auto* input : inputs) {
+    const bool is_layout_correct = input->layout() == DataLayout::kMKLDNN;
+    const bool is_format_defined =
+        input->format() != memory::format::format_undef;
+    PADDLE_ENFORCE(is_layout_correct && is_format_defined,
+                   "Wrong layout/format set for Input tensor");
+  }
+}
+
+static memory::primitive_desc CreateMemPrimDesc(const Tensor& input,
+                                                const mkldnn::engine& engine) {
+  constexpr auto data_type = mkldnn::memory::f32;
+  const auto dims = paddle::framework::vectorize2int(input.dims());
+  const auto format = input.format();
+  auto description = memory::desc(dims, data_type, format);
+  auto mem_prim_desc = memory::primitive_desc(description, engine);
+  return mem_prim_desc;
+}
+
+static mkldnn::memory::format GetDstMemFormat(
+    const concat::primitive_desc& concat_pd) {
+  return (memory::format)concat_pd.dst_primitive_desc().desc().data.format;
+}
+
+static platform::CPUPlace GetCpuPlace(
+    const paddle::framework::ExecutionContext& ctx) {
+  auto place = ctx.GetPlace();
+  PADDLE_ENFORCE(paddle::platform::is_cpu_place(place),
+                 "It must use CPUPlace.");
+  return boost::get<platform::CPUPlace>(place);
+}
+
+static const mkldnn::engine& GetMKLDNNEngine(
+    const paddle::framework::ExecutionContext& ctx) {
+  auto& dev_ctx = ctx.template device_context<platform::MKLDNNDeviceContext>();
+  return dev_ctx.GetEngine();
+}
+
+template <typename T>
+class ConcatPrimitiveFactory {
+ public:
+  concat::primitive_desc CreateConcatPrimDescriptor(
+      const std::vector<const Tensor*> multi_input, Tensor* output,
+      int concat_axis, const mkldnn::engine& mkldnn_engine) {
+    CreateSourcesDescriptors(multi_input, mkldnn_engine);
+    auto dst_desc = CreateDstMemDescriptor(output);
+    return concat::primitive_desc(dst_desc, concat_axis, srcs_pd);
+  }
+
+  concat CreateConcatPrimitive(const concat::primitive_desc& concat_pd,
+                               Tensor* output, platform::CPUPlace place) {
+    CreateSourcePrimitiveAts();
+    dst_mem = CreateDstMemory(concat_pd, output, place);
+    return concat(concat_pd, inputs, dst_mem.get());
+  }
+
+ private:
+  memory::desc CreateDstMemDescriptor(Tensor* output) {
+    auto dst_dims = paddle::framework::vectorize2int(output->dims());
+    return memory::desc(dst_dims, platform::MKLDNNGetDataType<T>(),
+                        memory::format::any);
+  }
+
+  mkldnn::memory CreateDstMemory(const concat::primitive_desc& concat_pd,
+                                 Tensor* output, platform::CPUPlace place) {
+    return memory(concat_pd.dst_primitive_desc(),
+                  output->mutable_data<T>(place));
+  }
+
+  void CreateSourcesDescriptors(const std::vector<const Tensor*> multi_input,
+                                const mkldnn::engine& mkldnn_engine) {
+    for (size_t i = 0; i < multi_input.size(); i++) {
+      auto mem_prim_desc = CreateMemPrimDesc(*multi_input[i], mkldnn_engine);
+      srcs_pd.push_back(mem_prim_desc);
+      srcs.push_back(
+          memory(mem_prim_desc, to_void_cast(multi_input[i]->data<T>())));
+    }
+  }
+
+  void CreateSourcePrimitiveAts() {
+    inputs.reserve(srcs.size());
+    for (size_t i = 0; i < srcs.size(); i++) {
+      inputs.push_back(srcs[i]);
+    }
+  }
+
+ private:
+  std::vector<memory::primitive_desc> srcs_pd;
+  std::vector<memory> srcs;
+  std::vector<primitive::at> inputs;
+  boost::optional<memory> dst_mem;  // TODO(mgallus): change to std::optional
+};                                  // upon introduction of C++17 to paddle
+
+template <typename T>
+class ConcatMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
+ public:
+  void Compute(const paddle::framework::ExecutionContext& ctx) const override {
+    auto place = GetCpuPlace(ctx);
+    const auto& mkldnn_engine = GetMKLDNNEngine(ctx);
+
+    auto multi_input = ctx.MultiInput<Tensor>("X");
+    EnforceLayouts(multi_input);
+    Tensor* output = ctx.Output<Tensor>("Out");
+    int64_t concat_axis = static_cast<int64_t>(ctx.Attr<int>("axis"));
+
+    ConcatPrimitiveFactory<T> prim_creator;
+    auto concat_pd = prim_creator.CreateConcatPrimDescriptor(
+        multi_input, output, static_cast<int>(concat_axis), mkldnn_engine);
+    auto concat = prim_creator.CreateConcatPrimitive(concat_pd, output, place);
+    stream(stream::kind::eager).submit({concat}).wait();
+
+    output->set_layout(DataLayout::kMKLDNN);
+    output->set_format(GetDstMemFormat(concat_pd));
+  }
+};
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+
+REGISTER_OP_KERNEL(concat, MKLDNN, ::paddle::platform::CPUPlace,
+                   ops::ConcatMKLDNNOpKernel<float>)

paddle/fluid/operators/concat_op.cc

@@ -13,10 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/concat_op.h"
-
 #include <string>
 #include <vector>
 
+#ifdef PADDLE_WITH_MKLDNN
+#include <paddle/fluid/platform/mkldnn_helper.h>
+#endif
+
 namespace paddle {
 namespace operators {
 using framework::Tensor;
@@ -59,6 +62,22 @@ class ConcatOp : public framework::OperatorWithKernel {
     ctx->SetOutputDim("Out", out_dims);
     ctx->ShareLoD("X", /*->*/ "Out");
   }
+
+ protected:
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext &ctx) const override {
+    auto input_data_type =
+        framework::GetDataTypeOfVar(ctx.MultiInputVar("X")[0]);
+
+#ifdef PADDLE_WITH_MKLDNN
+    if (platform::CanMKLDNNBeUsed(ctx)) {
+      return framework::OpKernelType(input_data_type, ctx.GetPlace(),
+                                     framework::DataLayout::kMKLDNN,
+                                     framework::LibraryType::kMKLDNN);
+    }
+#endif
+    return framework::OpKernelType(input_data_type, ctx.GetPlace());
+  }
 };
 
 class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
@@ -66,6 +85,10 @@ class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
   void Make() override {
     AddInput("X", "Input tensors of concat operator.").AsDuplicable();
     AddOutput("Out", "Output tensor of concat operator.");
+    AddAttr<bool>(
+        "use_mkldnn",
+        "(bool, default false) Indicates if MKL-DNN kernel will be used")
+        .SetDefault(false);
     AddAttr<int>("axis",
                  "The axis along which the input tensors will be concatenated.")
         .SetDefault(0);

python/paddle/fluid/tests/unittests/test_concat_mkldnn_op.py

+#   Copyright (c) 2018 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
+from test_concat_op import TestConcatOp, TestConcatOp2, TestConcatOp3
+
+
+class TestMKLDNNConcatOp(TestConcatOp):
+    def setUp(self):
+        super(TestMKLDNNConcatOp, self).setUp()
+        self.attrs["use_mkldnn"] = True
+        self._cpu_only = True
+
+    def test_check_grad(self):
+        pass
+
+    def init_kernel_type(self):
+        self.use_mkldnn = True
+
+
+class TestMKLDNNConcatOp2(TestConcatOp2):
+    def setUp(self):
+        super(TestMKLDNNConcatOp2, self).setUp()
+        self.attrs["use_mkldnn"] = True
+        self._cpu_only = True
+
+    def test_check_grad(self):
+        pass
+
+    def init_kernel_type(self):
+        self.use_mkldnn = True
+
+
+class TestMKLDNNConcatOp3(TestConcatOp3):
+    def setUp(self):
+        super(TestMKLDNNConcatOp3, self).setUp()
+        self.attrs["use_mkldnn"] = True
+        self._cpu_only = True
+
+    def test_check_grad(self):
+        pass
+
+    def init_kernel_type(self):
+        self.use_mkldnn = True
+
+
+if __name__ == '__main__':
+    unittest.main()