concat_mkldnn_op.cc 7.8 KB
Newer Older
M
Michal Gallus 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* 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. */

M
Michal Gallus 已提交
15
#include <memory>
M
Michal Gallus 已提交
16 17
#include "paddle/fluid/operators/concat_op.h"
#include "paddle/fluid/platform/mkldnn_helper.h"
18
#include "paddle/fluid/platform/mkldnn_reuse.h"
M
Michal Gallus 已提交
19 20 21 22 23 24 25 26 27 28 29 30 31 32

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) {
33 34 35 36
    PADDLE_ENFORCE_EQ(input->layout(), DataLayout::kMKLDNN,
                      "Wrong layout set for Input tensor");
    PADDLE_ENFORCE_NE(input->format(), MKLDNNMemoryFormat::format_undef,
                      "Wrong format set for Input tensor");
M
Michal Gallus 已提交
37 38 39
  }
}

40
static memory::primitive_desc CreateMemPrimDesc(const Tensor& input,
41 42
                                                const mkldnn::engine& engine,
                                                const memory::data_type& dt) {
43
  const auto dims = paddle::framework::vectorize<int>(input.dims());
M
Michal Gallus 已提交
44
  const auto format = input.format();
45
  auto description = memory::desc(dims, dt, format);
M
Michal Gallus 已提交
46 47 48 49 50 51 52 53 54 55 56 57
  auto mem_prim_desc = memory::primitive_desc(description, engine);
  return mem_prim_desc;
}

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);
}

M
Michal Gallus 已提交
58
static const mkldnn::engine& GetMKLDNNEngine(
59 60 61
    const paddle::framework::ExecutionContext& ctx) {
  auto& dev_ctx = ctx.template device_context<platform::MKLDNNDeviceContext>();
  return dev_ctx.GetEngine();
M
Michal Gallus 已提交
62
}
M
Michal Gallus 已提交
63

M
Michal Gallus 已提交
64 65 66 67 68
template <typename T>
class ConcatPrimitiveFactory {
 public:
  concat::primitive_desc CreateConcatPrimDescriptor(
      const std::vector<const Tensor*> multi_input, Tensor* output,
69 70 71 72
      int concat_axis, const mkldnn::engine& mkldnn_engine,
      const memory::data_type& dt = memory::data_type::f32) {
    CreateSourcesDescriptors(multi_input, mkldnn_engine, dt);
    auto dst_desc = CreateDstMemDescriptor(output, dt);
M
Michal Gallus 已提交
73 74
    return concat::primitive_desc(dst_desc, concat_axis, srcs_pd);
  }
M
Michal Gallus 已提交
75

M
Michal Gallus 已提交
76 77 78
  concat CreateConcatPrimitive(const concat::primitive_desc& concat_pd,
                               Tensor* output, platform::CPUPlace place) {
    CreateSourcePrimitiveAts();
79 80
    dst_mem = CreateDstMemory(concat_pd, output, place);
    return concat(concat_pd, inputs, dst_mem.get());
M
Michal Gallus 已提交
81 82
  }

83 84 85 86 87 88 89 90 91 92 93 94 95
  void SetSrcDataHandleByIndex(const std::vector<memory>& srcs, const size_t& i,
                               void* handler) {
    srcs[i].set_data_handle(handler);
  }

  void SetDstDataHandle(const memory& dst_mem, void* handler) {
    dst_mem.set_data_handle(handler);
  }

  std::vector<memory> GetSrcs() { return srcs; }

  memory GetDst() { return dst_mem.get(); }

M
Michal Gallus 已提交
96
 private:
97 98
  memory::desc CreateDstMemDescriptor(Tensor* output,
                                      const memory::data_type& dt) {
99
    auto dst_dims = paddle::framework::vectorize<int>(output->dims());
100
    return memory::desc(dst_dims, dt, MKLDNNMemoryFormat::any);
M
Michal Gallus 已提交
101 102 103
  }

  mkldnn::memory CreateDstMemory(const concat::primitive_desc& concat_pd,
104 105
                                 Tensor* output,
                                 const platform::CPUPlace& place) {
M
Michal Gallus 已提交
106 107 108
    return memory(concat_pd.dst_primitive_desc(),
                  output->mutable_data<T>(place));
  }
M
Michal Gallus 已提交
109

M
Michal Gallus 已提交
110
  void CreateSourcesDescriptors(const std::vector<const Tensor*> multi_input,
111 112
                                const mkldnn::engine& mkldnn_engine,
                                const memory::data_type& dt) {
M
Michal Gallus 已提交
113
    for (size_t i = 0; i < multi_input.size(); i++) {
114 115
      auto mem_prim_desc =
          CreateMemPrimDesc(*multi_input[i], mkldnn_engine, dt);
116 117 118
      srcs_pd.push_back(mem_prim_desc);
      srcs.push_back(
          memory(mem_prim_desc, to_void_cast(multi_input[i]->data<T>())));
M
Michal Gallus 已提交
119
    }
M
Michal Gallus 已提交
120
  }
M
Michal Gallus 已提交
121

M
Michal Gallus 已提交
122
  void CreateSourcePrimitiveAts() {
M
Michal Gallus 已提交
123 124 125 126
    inputs.reserve(srcs.size());
    for (size_t i = 0; i < srcs.size(); i++) {
      inputs.push_back(srcs[i]);
    }
M
Michal Gallus 已提交
127 128 129 130 131 132
  }

 private:
  std::vector<memory::primitive_desc> srcs_pd;
  std::vector<memory> srcs;
  std::vector<primitive::at> inputs;
133 134
  boost::optional<memory> dst_mem;
};
M
Michal Gallus 已提交
135 136 137 138 139 140 141 142 143

template <typename T>
class ConcatMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
 public:
  void Compute(const paddle::framework::ExecutionContext& ctx) const override {
    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"));
144 145 146 147 148 149
    auto& dev_ctx =
        ctx.template device_context<paddle::platform::MKLDNNDeviceContext>();
    auto place = GetCpuPlace(ctx);

    memory::data_type dt =
        paddle::framework::ToMKLDNNDataType(multi_input[0]->type());
M
Michal Gallus 已提交
150 151

    ConcatPrimitiveFactory<T> prim_creator;
152 153 154 155
    std::string key = platform::CreateKey(
        paddle::framework::vectorize<int>(multi_input[0]->dims()), concat_axis,
        ctx.op().Output("Out"), dt, multi_input[0]->format(),
        platform::ThreadIDasStr());
156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193
    const std::string key_prim = key + "@concat_p";
    const std::string key_concat_pd = key + "@concat_pd";
    const std::string key_srcs = key + "@concat_srcs";
    const std::string key_dst = key + "@concat_dst";

    std::shared_ptr<concat::primitive_desc> concat_pd;
    std::shared_ptr<std::vector<memory>> srcs;
    std::shared_ptr<memory> dst_mem;
    auto concat_p = std::static_pointer_cast<concat>(dev_ctx.GetBlob(key_prim));

    if (concat_p == nullptr) {
      const auto& mkldnn_engine = dev_ctx.GetEngine();
      concat_pd = std::make_shared<concat::primitive_desc>(
          prim_creator.CreateConcatPrimDescriptor(multi_input, output,
                                                  static_cast<int>(concat_axis),
                                                  mkldnn_engine, dt));
      concat_p = std::make_shared<concat>(
          prim_creator.CreateConcatPrimitive(*concat_pd, output, place));
      srcs = std::make_shared<std::vector<memory>>(prim_creator.GetSrcs());
      dst_mem = std::make_shared<memory>(prim_creator.GetDst());
      dev_ctx.SetBlob(key_prim, concat_p);
      dev_ctx.SetBlob(key_concat_pd, concat_pd);
      dev_ctx.SetBlob(key_srcs, srcs);
      dev_ctx.SetBlob(key_dst, dst_mem);
    } else {
      srcs = std::static_pointer_cast<std::vector<memory>>(
          dev_ctx.GetBlob(key_srcs));
      dst_mem = std::static_pointer_cast<memory>(dev_ctx.GetBlob(key_dst));
      concat_pd = std::static_pointer_cast<concat::primitive_desc>(
          dev_ctx.GetBlob(key_concat_pd));
      for (size_t i = 0; i < multi_input.size(); i++) {
        prim_creator.SetSrcDataHandleByIndex(
            *srcs, i, to_void_cast<T>(multi_input[i]->data<T>()));
      }
      prim_creator.SetDstDataHandle(*dst_mem, output->mutable_data<T>(place));
    }

    stream(stream::kind::eager).submit({*concat_p}).wait();
M
Michal Gallus 已提交
194

195
    output->set_layout(DataLayout::kMKLDNN);
A
Adam 已提交
196
    output->set_format(platform::GetMKLDNNFormat(*dst_mem));
M
Michal Gallus 已提交
197 198 199 200 201 202 203 204
  }
};
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OP_KERNEL(concat, MKLDNN, ::paddle::platform::CPUPlace,
205 206 207
                   ops::ConcatMKLDNNOpKernel<float>,
                   ops::ConcatMKLDNNOpKernel<int8_t>,
                   ops::ConcatMKLDNNOpKernel<uint8_t>);