concat_mkldnn_op.cc 8.7 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 33

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;
34 35
    const bool is_format_defined =
        input->format() != memory::format::format_undef;
M
Michal Gallus 已提交
36 37 38 39 40
    PADDLE_ENFORCE(is_layout_correct && is_format_defined,
                   "Wrong layout/format set for Input tensor");
  }
}

41
static memory::primitive_desc CreateMemPrimDesc(const Tensor& input,
42 43
                                                const mkldnn::engine& engine,
                                                const memory::data_type& dt) {
M
Michal Gallus 已提交
44 45
  const auto dims = paddle::framework::vectorize2int(input.dims());
  const auto format = input.format();
46
  auto description = memory::desc(dims, dt, format);
M
Michal Gallus 已提交
47 48 49 50
  auto mem_prim_desc = memory::primitive_desc(description, engine);
  return mem_prim_desc;
}

51 52 53 54 55
static mkldnn::memory::format GetDstMemFormat(
    const concat::primitive_desc& concat_pd) {
  return (memory::format)concat_pd.dst_primitive_desc().desc().data.format;
}

M
Michal Gallus 已提交
56 57 58 59 60 61 62 63
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 已提交
64
static const mkldnn::engine& GetMKLDNNEngine(
65 66 67
    const paddle::framework::ExecutionContext& ctx) {
  auto& dev_ctx = ctx.template device_context<platform::MKLDNNDeviceContext>();
  return dev_ctx.GetEngine();
M
Michal Gallus 已提交
68
}
M
Michal Gallus 已提交
69

70 71 72 73 74 75 76 77 78 79 80 81 82 83 84
std::string CreateKey(const paddle::framework::ExecutionContext& ctx,
                      const std::vector<const Tensor*> multi_input,
                      const int64_t& concat_axis, const memory::data_type& dt) {
  std::string key;
  key.reserve(platform::MKLDNNHandler::MaxKeyLength);
  for (size_t i = 0; i < multi_input.size(); i++) {
    platform::MKLDNNHandler::AppendKeyDims(
        &key, paddle::framework::vectorize2int(multi_input[i]->dims()));
  }
  platform::MKLDNNHandler::AppendKey(&key, std::to_string(concat_axis));
  platform::MKLDNNHandler::AppendKey(&key, ctx.op().Output("Out"));
  platform::MKLDNNHandler::AppendKey(&key, std::to_string(dt));
  return key;
}

M
Michal Gallus 已提交
85 86 87 88 89
template <typename T>
class ConcatPrimitiveFactory {
 public:
  concat::primitive_desc CreateConcatPrimDescriptor(
      const std::vector<const Tensor*> multi_input, Tensor* output,
90 91 92 93
      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 已提交
94 95
    return concat::primitive_desc(dst_desc, concat_axis, srcs_pd);
  }
M
Michal Gallus 已提交
96

M
Michal Gallus 已提交
97 98 99
  concat CreateConcatPrimitive(const concat::primitive_desc& concat_pd,
                               Tensor* output, platform::CPUPlace place) {
    CreateSourcePrimitiveAts();
100 101
    dst_mem = CreateDstMemory(concat_pd, output, place);
    return concat(concat_pd, inputs, dst_mem.get());
M
Michal Gallus 已提交
102 103
  }

104 105 106 107 108 109 110 111 112 113 114 115 116
  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 已提交
117
 private:
118 119
  memory::desc CreateDstMemDescriptor(Tensor* output,
                                      const memory::data_type& dt) {
M
Michal Gallus 已提交
120
    auto dst_dims = paddle::framework::vectorize2int(output->dims());
121
    return memory::desc(dst_dims, dt, memory::format::any);
M
Michal Gallus 已提交
122 123 124
  }

  mkldnn::memory CreateDstMemory(const concat::primitive_desc& concat_pd,
125 126
                                 Tensor* output,
                                 const platform::CPUPlace& place) {
M
Michal Gallus 已提交
127 128 129
    return memory(concat_pd.dst_primitive_desc(),
                  output->mutable_data<T>(place));
  }
M
Michal Gallus 已提交
130

M
Michal Gallus 已提交
131
  void CreateSourcesDescriptors(const std::vector<const Tensor*> multi_input,
132 133
                                const mkldnn::engine& mkldnn_engine,
                                const memory::data_type& dt) {
M
Michal Gallus 已提交
134
    for (size_t i = 0; i < multi_input.size(); i++) {
135 136
      auto mem_prim_desc =
          CreateMemPrimDesc(*multi_input[i], mkldnn_engine, dt);
137 138 139
      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 已提交
140
    }
M
Michal Gallus 已提交
141
  }
M
Michal Gallus 已提交
142

M
Michal Gallus 已提交
143
  void CreateSourcePrimitiveAts() {
M
Michal Gallus 已提交
144 145 146 147
    inputs.reserve(srcs.size());
    for (size_t i = 0; i < srcs.size(); i++) {
      inputs.push_back(srcs[i]);
    }
M
Michal Gallus 已提交
148 149 150 151 152 153
  }

 private:
  std::vector<memory::primitive_desc> srcs_pd;
  std::vector<memory> srcs;
  std::vector<primitive::at> inputs;
154 155
  boost::optional<memory> dst_mem;  // TODO(mgallus): change to std::optional
};                                  // upon introduction of C++17 to paddle
M
Michal Gallus 已提交
156 157 158 159 160 161 162 163 164

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"));
165 166 167 168 169 170
    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 已提交
171 172

    ConcatPrimitiveFactory<T> prim_creator;
173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211
    std::string key = CreateKey(ctx, multi_input, concat_axis, dt);
    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 已提交
212

213 214
    output->set_layout(DataLayout::kMKLDNN);
    output->set_format(GetDstMemFormat(*concat_pd));
M
Michal Gallus 已提交
215 216 217 218 219 220 221 222
  }
};
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OP_KERNEL(concat, MKLDNN, ::paddle::platform::CPUPlace,
223 224 225
                   ops::ConcatMKLDNNOpKernel<float>,
                   ops::ConcatMKLDNNOpKernel<int8_t>,
                   ops::ConcatMKLDNNOpKernel<uint8_t>);