sum_mkldnn_op.cc 8.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27
//   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.

/*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/operators/sum_op.h"
J
Jacek Czaja 已提交
28
#include "paddle/fluid/platform/mkldnn_reuse.h"
29

W
wanghuancoder 已提交
30 31 32 33 34 35 36 37 38 39
namespace paddle {
namespace framework {
class Tensor;
}  // namespace framework
namespace platform {
class CPUDeviceContext;
class MKLDNNDeviceContext;
}  // namespace platform
}  // namespace paddle

40 41 42
namespace paddle {
namespace operators {

T
tangwei12 已提交
43 44
using paddle::platform::CPUDeviceContext;
using paddle::platform::MKLDNNDeviceContext;
45 46
using platform::to_void_cast;

J
Jacek Czaja 已提交
47 48 49 50 51 52 53 54 55 56
template <typename T>
class SumMKLDNNHandler : public platform::MKLDNNHandlerT<T, dnnl::sum> {
 public:
  SumMKLDNNHandler(const MKLDNNDeviceContext& dev_ctx,
                   platform::Place cpu_place,
                   const std::vector<framework::Variable*>& in_vars,
                   framework::LoDTensor* z, const std::string& uniq_name)

      : platform::MKLDNNHandlerT<T, dnnl::sum>(
            dev_ctx, dev_ctx.GetEngine(), cpu_place,
57 58
            platform::CreateKey(dev_ctx, framework::vectorize(z->dims()),
                                uniq_name)),
J
Jacek Czaja 已提交
59 60 61 62 63 64 65 66 67
        num_inputs_(0) {
    for (size_t i = 0; i < in_vars.size(); i++) {
      srcs_suffix_.push_back(std::string("-") + std::to_string(i));
    }

    if (!this->isCached()) {
      auto dst_tz = framework::vectorize<int64_t>(z->dims());
      auto src_tz = dst_tz;

J
Jacek Czaja 已提交
68
      std::vector<mkldnn::memory::desc> srcs_md;
J
Jacek Czaja 已提交
69 70 71 72 73 74
      for (size_t i = 0; i < in_vars.size(); i++) {
        auto& input_it = in_vars[i]->Get<framework::LoDTensor>();
        if (input_it.numel() == 0) {
          continue;
        }
        MKLDNNMemoryFormat input_format = input_it.format();
J
Jacek Czaja 已提交
75 76
        srcs_md.push_back(mkldnn::memory::desc(
            src_tz, platform::MKLDNNGetDataType<T>(), input_format));
J
Jacek Czaja 已提交
77 78 79 80
        ++num_inputs_;
      }
      std::vector<float> scales(num_inputs_, 1.0);

J
Jacek Czaja 已提交
81 82
      auto dst_md = mkldnn::memory::desc(
          dst_tz, platform::MKLDNNGetDataType<T>(), MKLDNNMemoryFormat::any);
J
Jacek Czaja 已提交
83 84 85 86 87 88 89 90

      this->AcquireForwardPrimitiveDescriptor(dst_md, scales, srcs_md);
    }
  }

  // (jczaja) sum oneDNN prim is not having .desc attribute so
  // we cannot use base AcquireForwardPrimitiveDescriptor
  void AcquireForwardPrimitiveDescriptor(
J
Jacek Czaja 已提交
91 92
      const mkldnn::memory::desc& dst_md, const std::vector<float>& scales,
      const std::vector<mkldnn::memory::desc>& srcs_md) {
J
Jacek Czaja 已提交
93 94 95 96 97
    // Sum op does not have backward so no passing from FWD to BWD is needed
    const std::string key_pd = this->key_ + "@fwd_pd";
    this->fwd_pd_ = std::static_pointer_cast<dnnl::sum::primitive_desc>(
        this->dev_ctx_.GetBlob(key_pd));
    if (this->fwd_pd_ == nullptr) {
J
Jacek Czaja 已提交
98 99
      this->fwd_pd_.reset(new dnnl::sum::primitive_desc(dst_md, scales, srcs_md,
                                                        this->engine_));
J
Jacek Czaja 已提交
100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136
      this->dev_ctx_.SetBlob(key_pd, this->fwd_pd_);
    }
  }

  std::shared_ptr<mkldnn::memory> AcquireSrcMemory(
      const framework::Tensor& input, int i) {
    const T* input_data = input.data<T>();
    return this->AcquireMemoryFromPrimitive(this->fwd_pd_->src_desc(i),
                                            to_void_cast<T>(input_data),
                                            "@src_mem_p" + srcs_suffix_[i]);
  }

  using platform::MKLDNNHandlerT<T, dnnl::sum>::AcquireDstMemory;

  std::shared_ptr<mkldnn::memory> AcquireDstMemory(void) {
    return this->AcquireMemoryFromPrimitive(this->fwd_pd_->dst_desc(),
                                            "@dst_mem_p");
  }

  inline int GetNumInputs(void) { return num_inputs_; }

 protected:
  // isCached need to be overloaded as base one works on key_common
  bool isCached() {
    const std::string key_pd = this->key_ + "@fwd_pd";
    this->fwd_pd_ = std::static_pointer_cast<dnnl::sum::primitive_desc>(
        this->dev_ctx_.GetBlob(key_pd));

    const std::string key_p = this->key_ + "@fwd_p";
    return (this->dev_ctx_.GetBlob(key_p) != nullptr);
  }

 private:
  int num_inputs_;
  std::vector<std::string> srcs_suffix_;
};

137 138 139 140
template <typename T>
class SumMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
 public:
  void Compute(const paddle::framework::ExecutionContext& ctx) const override {
141 142 143
    PADDLE_ENFORCE_EQ(platform::is_cpu_place(ctx.GetPlace()), true,
                      paddle::platform::errors::PreconditionNotMet(
                          "Operator DNNL Sum must use CPUPlace"));
144 145
    auto& dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
    auto in_vars = ctx.MultiInputVar("X");
146 147 148

    PADDLE_ENFORCE_NE(in_vars.empty(), true, platform::errors::InvalidArgument(
                                                 "Input variable is empty."));
J
Jacek Czaja 已提交
149
    auto& input0 = in_vars[0]->Get<LoDTensor>();
150
    LoDTensor* output = ctx.Output<LoDTensor>("Out");
151

J
Jacek Czaja 已提交
152
    bool in_place = (input0.numel() > 0) && input0.IsSharedBufferWith(*output);
153

J
Jacek Czaja 已提交
154 155
    SumMKLDNNHandler<T> handler(dev_ctx, ctx.GetPlace(), in_vars, output,
                                ctx.OutputName("Out"));
156

J
Jacek Czaja 已提交
157 158 159 160
    // Create list of SRC MEMs
    std::vector<std::shared_ptr<mkldnn::memory>> srcs_mem;
    srcs_mem.reserve(handler.GetNumInputs());
    int input_index = 0;
161
    for (size_t i = 0; i < in_vars.size(); i++) {
J
Jacek Czaja 已提交
162
      auto& input_it = in_vars[i]->Get<framework::LoDTensor>();
163 164
      if (input_it.numel() == 0) {
        continue;
A
Adam 已提交
165
      }
J
Jacek Czaja 已提交
166 167
      srcs_mem.push_back(handler.AcquireSrcMemory(input_it, input_index));
      ++input_index;
168
    }
169

J
Jacek Czaja 已提交
170 171
    auto dst_mem = in_place ? handler.AcquireDstMemory()
                            : handler.AcquireDstMemory(output);
172

J
Jacek Czaja 已提交
173
    auto sum_p = handler.AcquireForwardPrimitive();
174

J
Jacek Czaja 已提交
175
    std::unordered_map<int, mkldnn::memory> args;
176
    for (size_t i = 0; i < srcs_mem.size(); ++i) {
J
Jacek Czaja 已提交
177
      args.insert({MKLDNN_ARG_MULTIPLE_SRC + i, *(srcs_mem[i])});
178 179 180
    }
    args.insert({MKLDNN_ARG_DST, *dst_mem});

J
Jacek Czaja 已提交
181 182
    mkldnn::stream astream(dev_ctx.GetEngine());
    sum_p->execute(astream, args);
183 184
    astream.wait();

J
Jacek Czaja 已提交
185 186
    // For in-place execution which sum does not have we need to fake it
    // so from oneDNN dst memory we reorder data into input
187
    if (in_place) {
188 189 190
      const std::string reorder_key =
          platform::CreateKey(dev_ctx, framework::vectorize(output->dims()),
                              ctx.OutputName("Out") + "-I");
J
Jacek Czaja 已提交
191 192 193 194 195 196 197 198 199 200 201

      auto& in_out = in_vars[0]->Get<framework::LoDTensor>();
      auto output_tz = framework::vectorize<int64_t>(output->dims());
      platform::ReorderMKLDNNHandler reorder_handler(
          output_tz, output->type(), framework::ToMKLDNNDataType(in_out.type()),
          dev_ctx, dev_ctx.GetEngine(), reorder_key);

      auto target_mem = reorder_handler.AcquireDstMemory(
          output, in_out.format(), ctx.GetPlace());

      auto reorder_p = reorder_handler.AcquireReorder(target_mem, dst_mem);
202 203 204 205 206 207
      {
        platform::RecordEvent record_reorder("int_reorder",
                                             platform::EventRole::kUniqueOp);
        reorder_p->execute(astream, *dst_mem, *target_mem);
        astream.wait();
      }
208
    }
J
Jacek Czaja 已提交
209 210
    output->set_layout(framework::DataLayout::kMKLDNN);
    output->set_format(platform::GetMKLDNNFormat(*dst_mem));
211 212 213 214 215 216
  }
};

}  // namespace operators
}  // namespace paddle

J
Jacek Czaja 已提交
217 218 219 220
REGISTER_OP_KERNEL(
    sum, MKLDNN, ::paddle::platform::CPUPlace,
    paddle::operators::SumMKLDNNOpKernel<paddle::platform::bfloat16>,
    paddle::operators::SumMKLDNNOpKernel<float>);