mkldnn_reuse.h 60.1 KB
Newer Older
J
Jacek Czaja 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/* Copyright (c) 2017 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. */
#pragma once

16
#include <memory>
17
#include <sstream>
J
Jacek Czaja 已提交
18
#include <string>
19
#include <utility>
J
Jacek Czaja 已提交
20
#include <vector>
21
#include "boost/optional.hpp"
X
xiaoli.liu@intel.com 已提交
22
#include "paddle/fluid/framework/data_layout_transform.h"
J
Jacek Czaja 已提交
23
#include "paddle/fluid/framework/operator.h"
24
#include "paddle/fluid/operators/pool_op.h"
J
Jacek Czaja 已提交
25 26 27 28 29 30
#include "paddle/fluid/platform/mkldnn_helper.h"
#include "paddle/fluid/platform/place.h"

namespace paddle {
namespace platform {

31 32
using framework::DataLayout;
using framework::Tensor;
J
Jacek Czaja 已提交
33
using user_function = std::function<std::shared_ptr<float>(const float*)>;
34
using memory = mkldnn::memory;
J
Jacek Czaja 已提交
35

36 37
template <typename T, typename TForward,
          typename TBackward = mkldnn_dummy_primitive>
38 39 40 41 42 43 44 45
class MKLDNNHandlerT {
 public:
  MKLDNNHandlerT(const MKLDNNDeviceContext& dev_ctx, mkldnn::engine engine,
                 platform::Place cpu_place, const std::string& base_key)
      : dev_ctx_(dev_ctx),
        engine_(engine),
        place_(cpu_place),
        key_common_(base_key),
46
        key_(platform::ExtendKeyWithThreadInfoIfNeeded(dev_ctx, base_key)),
47
        fwd_pd_(nullptr),
48 49 50
        bwd_pd_(nullptr) {
    platform::MKLDNNDeviceContext::tls().log_lib_version();
  }
51

A
Adam 已提交
52
  std::shared_ptr<TForward> AcquireForwardPrimitive() {
53
    const std::string key_p = key_ + "@fwd_p";
54 55 56
    auto forward_p =
        std::static_pointer_cast<TForward>(dev_ctx_.GetBlob(key_p));
    if (forward_p == nullptr) {
A
Adam 已提交
57
      forward_p = std::make_shared<TForward>(*fwd_pd_);
58 59 60 61 62
      dev_ctx_.SetBlob(key_p, forward_p);
    }
    return forward_p;
  }

A
Adam 已提交
63
  std::shared_ptr<TBackward> AcquireBackwardPrimitive() {
64
    const std::string key_p = key_ + "@bwd_p";
65 66 67
    auto backward_p =
        std::static_pointer_cast<TBackward>(dev_ctx_.GetBlob(key_p));
    if (backward_p == nullptr) {
A
Adam 已提交
68
      backward_p = std::make_shared<TBackward>(*bwd_pd_);
69 70 71 72 73
      dev_ctx_.SetBlob(key_p, backward_p);
    }
    return backward_p;
  }

74 75 76
  std::shared_ptr<mkldnn::memory> AcquireSrcMemory(
      const framework::Tensor* input) {
    const T* input_data = input->data<T>();
A
Adam 已提交
77 78
    return this->AcquireMemoryFromPrimitive(
        fwd_pd_->src_desc(), to_void_cast<T>(input_data), "@src_mem_p");
79 80
  }

81
  template <typename T_out = T>
82
  std::shared_ptr<mkldnn::memory> AcquireDstMemory(framework::Tensor* output) {
83 84
    T_out* ptr =
        output->mutable_data<T_out>(place_, fwd_pd_->dst_desc().get_size());
A
Adam 已提交
85
    return this->AcquireMemoryFromPrimitive(fwd_pd_->dst_desc(), ptr,
86 87 88
                                            "@dst_mem_p");
  }

89
  template <typename T_out = T>
90 91
  std::shared_ptr<mkldnn::memory> AcquireDstMemory(
      const framework::Tensor* output) {
92 93 94 95
    const T_out* output_data = output->data<T_out>();
    return this->AcquireMemoryFromPrimitive(bwd_pd_->dst_desc(),
                                            to_void_cast<T_out>(output_data),
                                            "@bwd-dst_mem_p");
96 97 98 99 100
  }

  std::shared_ptr<mkldnn::memory> AcquireDiffDstMemory(
      const framework::Tensor* diffdst) {
    const T* ptr = diffdst->data<T>();
A
Adam 已提交
101 102
    return this->AcquireMemoryFromPrimitive(
        bwd_pd_->diff_dst_desc(), to_void_cast<T>(ptr), "@diff_dst_mem_p");
103 104 105 106
  }

  std::shared_ptr<mkldnn::memory> AcquireDiffSrcMemory(
      framework::Tensor* diffsrc) {
A
Adam 已提交
107 108 109 110
    T* ptr =
        diffsrc->mutable_data<T>(place_, bwd_pd_->diff_src_desc().get_size());
    return this->AcquireMemoryFromPrimitive(bwd_pd_->diff_src_desc(), ptr,
                                            "@diff_src_mem_p");
111 112
  }

113
 protected:
114
  bool isCached() {
115
    const std::string key_pd = key_common_ + "@fwd_pd";
116 117
    fwd_pd_ = std::static_pointer_cast<typename TForward::primitive_desc>(
        dev_ctx_.GetBlob(key_pd));
118

119
    const std::string key_p = key_ + "@fwd_p";
120
    return (dev_ctx_.GetBlob(key_p) != nullptr);
121 122
  }

123 124 125 126 127 128
  // If your primitive descriptor requires attributes, pass them as a
  // first argument and paramters to descriptor constructor in the following
  // arguments. Otherwise, all arguments will be forwarded to descriptor
  // constructor, including the first one.
  template <typename Arg, typename... Args>
  void AcquireForwardPrimitiveDescriptor(Arg&& first_arg, Args&&... args) {
129 130 131
    // Forward PD has to be passed to Grad op that
    // may be executed by diffrent thread, hence
    // for that one we use key that does not contain TID
132
    const std::string key_pd = key_common_ + "@fwd_pd";
133 134 135 136 137 138 139 140 141
    fwd_pd_ = std::static_pointer_cast<typename TForward::primitive_desc>(
        dev_ctx_.GetBlob(key_pd));
    if (fwd_pd_ == nullptr) {
      static std::mutex acquire_barrier;
      std::lock_guard<std::mutex> block_threads_until_finish_this_job(
          acquire_barrier);
      fwd_pd_ = std::static_pointer_cast<typename TForward::primitive_desc>(
          dev_ctx_.GetBlob(key_pd));
      if (fwd_pd_ == nullptr) {
142 143
        CreateForwardPrimitiveDescriptor(first_arg,
                                         std::forward<Args>(args)...);
144 145 146 147 148
        dev_ctx_.SetBlob(key_pd, fwd_pd_);
      }
    }
  }

149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169
  // Using sfinae to specialise variadic function. Workaround for not having
  // if constexpr in C++ 11.
  template <class First, class... Args>
  typename std::enable_if<std::is_same<typename std::decay<First>::type,
                                       dnnl::primitive_attr>::value>::type
  CreateForwardPrimitiveDescriptor(First&& first, Args&&... args) {
    auto fwd_desc = typename TForward::desc(std::forward<Args>(args)...);
    fwd_pd_ = std::make_shared<typename TForward::primitive_desc>(
        fwd_desc, first, engine_);
  }

  template <class First, class... Args>
  typename std::enable_if<!std::is_same<typename std::decay<First>::type,
                                        dnnl::primitive_attr>::value>::type
  CreateForwardPrimitiveDescriptor(First&& first, Args&&... args) {
    auto fwd_desc = typename TForward::desc(std::forward<First>(first),
                                            std::forward<Args>(args)...);
    fwd_pd_ =
        std::make_shared<typename TForward::primitive_desc>(fwd_desc, engine_);
  }

170 171
  template <typename... Args>
  void AcquireBackwardPrimitiveDescriptor(Args&&... args) {
172
    const std::string key_fwd_pd = key_common_ + "@fwd_pd";
173 174
    fwd_pd_ = std::static_pointer_cast<typename TForward::primitive_desc>(
        dev_ctx_.GetBlob(key_fwd_pd));
G
GaoWei8 已提交
175 176 177
    PADDLE_ENFORCE_NOT_NULL(
        fwd_pd_, platform::errors::Unavailable(
                     "Get MKLDNN Forward primitive %s failed.", key_fwd_pd));
178
    const std::string key_pd = key_ + "@bwd_pd";
179 180 181 182 183 184 185 186 187 188
    bwd_pd_ = std::static_pointer_cast<typename TBackward::primitive_desc>(
        dev_ctx_.GetBlob(key_pd));
    if (bwd_pd_ == nullptr) {
      auto bwd_desc = typename TBackward::desc(std::forward<Args>(args)...);
      bwd_pd_ = std::make_shared<typename TBackward::primitive_desc>(
          bwd_desc, engine_, *fwd_pd_);
      dev_ctx_.SetBlob(key_pd, bwd_pd_);
    }
  }

189 190 191 192 193 194
  std::shared_ptr<mkldnn::memory> AcquireMemoryFromPrimitive(
      const std::string& suffix) {
    return std::static_pointer_cast<mkldnn::memory>(
        dev_ctx_.GetBlob(key_ + suffix));
  }

195
  std::shared_ptr<mkldnn::memory> AcquireMemoryFromPrimitive(
A
Adam 已提交
196
      mkldnn::memory::desc md, void* ptr, const std::string& suffix) {
197
    const auto local_key = key_ + suffix;
198 199 200
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
A
Adam 已提交
201
      mem_p = std::make_shared<mkldnn::memory>(md, engine_, ptr);
202 203 204 205 206 207 208
      dev_ctx_.SetBlob(local_key, mem_p);
    } else {
      mem_p->set_data_handle(ptr);
    }
    return mem_p;
  }

209 210 211 212 213 214 215 216 217 218 219 220
  std::shared_ptr<mkldnn::memory> AcquireMemoryFromPrimitive(
      mkldnn::memory::desc md, const std::string& suffix) {
    const auto local_key = key_ + suffix;
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
      mem_p = std::make_shared<mkldnn::memory>(md, engine_);
      dev_ctx_.SetBlob(local_key, mem_p);
    }
    return mem_p;
  }

221 222 223 224 225 226 227 228 229 230 231 232 233 234 235
  void AcquireReorder(const std::shared_ptr<mkldnn::memory>& user_memory_p,
                      const std::shared_ptr<mkldnn::memory>& target_memory_p,
                      const std::string& suffix) {
    const auto key_reorder_p = key_ + suffix + "reorder_p";

    auto reorder_p = std::static_pointer_cast<mkldnn::reorder>(
        dev_ctx_.GetBlob(key_reorder_p));

    if (reorder_p == nullptr) {
      reorder_p =
          std::make_shared<mkldnn::reorder>(*user_memory_p, *target_memory_p);
      dev_ctx_.SetBlob(key_reorder_p, reorder_p);
    }

    mkldnn::stream astream(engine_);
236 237 238

    platform::RecordEvent record_reorder("int_reorder",
                                         platform::EventRole::kUniqueOp);
239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264
    reorder_p->execute(astream, {{MKLDNN_ARG_FROM, *user_memory_p},
                                 {MKLDNN_ARG_TO, *target_memory_p}});
    astream.wait();
  }

  std::shared_ptr<mkldnn::memory> AcquireMemoryWithReorder(
      const mkldnn::memory::desc& user_md,
      const mkldnn::memory::desc& target_md, void* ptr,
      const std::string& suffix, bool is_persistent = false) {
    const auto target_key = key_ + suffix + "_target";
    const auto key_reorder_p = key_ + suffix + "reorder_p";
    const auto user_key = key_ + suffix + "_user";

    auto target_memory_p =
        std::static_pointer_cast<dnnl::memory>(dev_ctx_.GetBlob(target_key));

    if (target_memory_p == nullptr) {
      auto user_memory_p =
          std::make_shared<dnnl::memory>(user_md, engine_, ptr);
      if (user_md != target_md) {
        target_memory_p = std::make_shared<mkldnn::memory>(target_md, engine_);
        auto reorder_p =
            std::make_shared<dnnl::reorder>(*user_memory_p, *target_memory_p);
        dev_ctx_.SetBlob(key_reorder_p, reorder_p);

        mkldnn::stream astream(engine_);
265 266
        platform::RecordEvent record_reorder("int_reorder",
                                             platform::EventRole::kUniqueOp);
267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284
        reorder_p->execute(astream, {{MKLDNN_ARG_FROM, *user_memory_p},
                                     {MKLDNN_ARG_TO, *target_memory_p}});
        astream.wait();
      } else {
        target_memory_p = user_memory_p;
      }
      dev_ctx_.SetBlob(user_key, user_memory_p);
      dev_ctx_.SetBlob(target_key, target_memory_p);
    } else if (!is_persistent) {
      mkldnn::stream astream(engine_);

      auto user_memory_p =
          std::static_pointer_cast<dnnl::memory>(dev_ctx_.GetBlob(user_key));
      user_memory_p->set_data_handle(ptr);

      auto reorder_p = std::static_pointer_cast<mkldnn::reorder>(
          dev_ctx_.GetBlob(key_reorder_p));
      if (reorder_p != nullptr) {
285 286
        platform::RecordEvent record_reorder("int_reorder",
                                             platform::EventRole::kUniqueOp);
287 288 289 290 291 292 293 294
        reorder_p->execute(astream, {{MKLDNN_ARG_FROM, *user_memory_p},
                                     {MKLDNN_ARG_TO, *target_memory_p}});
        astream.wait();
      }
    }
    return target_memory_p;
  }

295 296 297 298 299 300
  std::shared_ptr<mkldnn::memory> AcquireMemory(const std::string& suffix) {
    const auto local_key = key_ + suffix;
    return std::static_pointer_cast<mkldnn::memory>(
        dev_ctx_.GetBlob(local_key));
  }

301 302 303 304
  const MKLDNNDeviceContext& dev_ctx_;
  mkldnn::engine engine_;
  platform::Place place_;
  std::string key_common_;
305
  std::string key_;
306 307 308 309 310
  std::shared_ptr<typename TForward::primitive_desc> fwd_pd_;
  std::shared_ptr<typename TBackward::primitive_desc> bwd_pd_;
};

// TODO(grygielski) this class will be deleted later.
J
Jacek Czaja 已提交
311 312 313 314
class MKLDNNHandler {
 public:
  MKLDNNHandler(const MKLDNNDeviceContext& dev_ctx, mkldnn::engine engine,
                const std::string& base_key)
315 316 317
      : dev_ctx_(dev_ctx),
        engine_(engine),
        key_common_(base_key),
318 319 320
        key_(platform::ExtendKeyWithThreadInfoIfNeeded(dev_ctx, base_key)) {
    platform::MKLDNNDeviceContext::tls().log_lib_version();
  }
J
Jacek Czaja 已提交
321 322 323 324 325 326 327 328 329 330 331

  std::shared_ptr<mkldnn::memory> AcquireSrcMemory(
      const mkldnn::memory::desc& md, void* ptr) {
    return this->AcquireMemory(md, ptr, "@user_src_mem_p");
  }

  std::shared_ptr<mkldnn::memory> AcquireDstMemory(
      const mkldnn::memory::desc& md, void* ptr) {
    return this->AcquireMemory(md, ptr, "@user_dst_mem_p");
  }

A
Adam 已提交
332
  std::shared_ptr<mkldnn::memory> AcquireDiffSrcMemory(
J
Jacek Czaja 已提交
333
      const mkldnn::memory::desc& md, void* ptr) {
A
Adam 已提交
334
    return this->AcquireMemory(md, ptr, "@user_diff_src_mem_p");
J
Jacek Czaja 已提交
335 336
  }

A
Adam 已提交
337
  std::shared_ptr<mkldnn::memory> AcquireDiffDstMemory(
J
Jacek Czaja 已提交
338
      const mkldnn::memory::desc& md, void* ptr) {
A
Adam 已提交
339
    return this->AcquireMemory(md, ptr, "@user_diff_dst_mem_p");
J
Jacek Czaja 已提交
340 341 342
  }

  std::shared_ptr<mkldnn::memory> AcquireMemoryFromPrimitive(
A
Adam 已提交
343
      mkldnn::memory::desc md, void* ptr, const std::string& suffix) {
J
Jacek Czaja 已提交
344 345 346 347
    auto local_key = key_ + suffix;
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
A
Adam 已提交
348
      mem_p = std::make_shared<mkldnn::memory>(md, engine_, ptr);
J
Jacek Czaja 已提交
349 350 351 352 353 354 355
      dev_ctx_.SetBlob(local_key, mem_p);
    } else {
      mem_p->set_data_handle(ptr);
    }
    return mem_p;
  }

356 357 358 359 360 361 362 363 364 365 366 367
  std::shared_ptr<mkldnn::memory> AcquireMemoryFromPrimitive(
      mkldnn::memory::desc md, const std::string& suffix) {
    const auto local_key = key_ + suffix;
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
      mem_p = std::make_shared<mkldnn::memory>(md, engine_);
      dev_ctx_.SetBlob(local_key, mem_p);
    }
    return mem_p;
  }

J
Jacek Czaja 已提交
368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384
  // This incarnation of AcquireMemory can call user function eg. custom reorder
  // or preprocessing routine if needed
  std::shared_ptr<mkldnn::memory> AcquireMemory(
      const mkldnn::memory::desc& md, void* ptr, const std::string& suffix,
      user_function custom_func = {}) {
    /*Generate key*/
    auto local_key = key_ + suffix;
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
      // Call custom reorder/preprocessing func if available
      if (custom_func) {
        auto reordered_data = custom_func(reinterpret_cast<const float*>(ptr));
        dev_ctx_.SetBlob(local_key + "-custom_reorder", reordered_data);
        ptr = reinterpret_cast<void*>(reordered_data.get());
      }

A
Adam 已提交
385
      mem_p = std::make_shared<mkldnn::memory>(md, engine_, ptr);
J
Jacek Czaja 已提交
386 387 388 389 390 391 392
      dev_ctx_.SetBlob(local_key, mem_p);
    } else {
      mem_p->set_data_handle(ptr);
    }
    return mem_p;
  }

393
  std::shared_ptr<mkldnn::memory> AcquireMemory(
A
Adam 已提交
394
      const std::vector<int64_t>& dims, const mkldnn::memory::data_type dtype,
395
      const MKLDNNMemoryFormat& fmt, void* ptr, const std::string& suffix) {
396 397 398 399 400 401 402
    /*Generate key*/
    auto local_key = key_ + suffix;
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
      auto md = mkldnn::memory::desc(dims, dtype, fmt);

A
Adam 已提交
403
      mem_p = std::make_shared<mkldnn::memory>(md, engine_, ptr);
404 405 406 407 408 409 410
      dev_ctx_.SetBlob(local_key, mem_p);
    } else {
      mem_p->set_data_handle(ptr);
    }
    return mem_p;
  }

J
Jacek Czaja 已提交
411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427
  std::shared_ptr<mkldnn::memory> AcquireMemory(
      const std::shared_ptr<mkldnn::memory>& user_memory_p,
      const std::shared_ptr<mkldnn::memory>& target_memory_p,
      const std::string& suffix,
      std::vector<mkldnn::primitive>& pipeline) {  // NOLINT
    auto local_key = key_ + suffix;
    auto key_reorder_p = key_ + suffix + "reorder_p";

    auto stored_reorder_p = std::static_pointer_cast<mkldnn::reorder>(
        dev_ctx_.GetBlob(key_reorder_p));

    if (stored_reorder_p) {
      pipeline.push_back(*stored_reorder_p);
    } else {
      auto reorder_p =
          std::make_shared<mkldnn::reorder>(*user_memory_p, *target_memory_p);
      dev_ctx_.SetBlob(key_reorder_p, reorder_p);
A
Adam 已提交
428
      mkldnn::stream astream(engine_);
429 430
      platform::RecordEvent record_reorder("int_reorder",
                                           platform::EventRole::kUniqueOp);
A
Adam 已提交
431 432 433
      reorder_p->execute(astream, {{MKLDNN_ARG_FROM, *user_memory_p},
                                   {MKLDNN_ARG_TO, *target_memory_p}});
      astream.wait();
J
Jacek Czaja 已提交
434 435 436 437 438 439
    }

    return target_memory_p;
  }

  std::shared_ptr<mkldnn::memory> AcquireMemory(
A
Adam 已提交
440 441
      mkldnn::memory::desc& md,       // NOLINT
      mkldnn::memory::desc& user_md,  // NOLINT
J
Jacek Czaja 已提交
442 443 444
      const std::shared_ptr<mkldnn::memory> user_memory_p,
      const std::string& suffix,
      std::vector<mkldnn::primitive>& pipeline,  // NOLINT
445 446
      bool is_persistent = false, bool is_INT8 = false,
      std::vector<float> scale_data = {1.0f}, int mask = 0) {
J
Jacek Czaja 已提交
447 448 449 450 451 452
    // create reorder primitive if the input format is not the preferred one
    auto local_key = key_ + suffix;
    auto key_reorder_p = key_ + suffix + "reorder_p";

    auto target_memory_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
A
Adam 已提交
453 454 455

    mkldnn::stream astream(engine_);

J
Jacek Czaja 已提交
456 457
    if (target_memory_p == nullptr) {
      target_memory_p = user_memory_p;
A
Adam 已提交
458 459 460
      if (md != user_md) {
        target_memory_p = std::make_shared<mkldnn::memory>(md, engine_);
        std::shared_ptr<mkldnn::reorder::primitive_desc> reorder_pd;
461 462 463 464 465
        if (is_INT8) {
          mkldnn::primitive_attr
              attri;  // attribute for int8 weights and bias data reorder.
          attri.set_output_scales(mask, scale_data);

A
Adam 已提交
466 467 468
          reorder_pd = std::shared_ptr<mkldnn::reorder::primitive_desc>(
              new mkldnn::reorder::primitive_desc(*user_memory_p,
                                                  *target_memory_p, attri));
469
        } else {
A
Adam 已提交
470 471 472
          reorder_pd = std::shared_ptr<mkldnn::reorder::primitive_desc>(
              new mkldnn::reorder::primitive_desc(*user_memory_p,
                                                  *target_memory_p));
473
        }
A
Adam 已提交
474 475
        auto reorder_p =
            std::shared_ptr<mkldnn::reorder>(new mkldnn::reorder(*reorder_pd));
J
Jacek Czaja 已提交
476
        dev_ctx_.SetBlob(key_reorder_p, reorder_p);
A
Adam 已提交
477

478 479
        platform::RecordEvent record_reorder("int_reorder",
                                             platform::EventRole::kUniqueOp);
A
Adam 已提交
480 481 482
        reorder_p->execute(astream, {{MKLDNN_ARG_FROM, *user_memory_p},
                                     {MKLDNN_ARG_TO, *target_memory_p}});
        astream.wait();
J
Jacek Czaja 已提交
483 484 485 486 487 488 489
      }
      dev_ctx_.SetBlob(local_key, target_memory_p);
    } else if (!is_persistent) {
      // Make reorder if needed
      auto reorder_p = std::static_pointer_cast<mkldnn::reorder>(
          dev_ctx_.GetBlob(key_reorder_p));
      if (reorder_p != nullptr) {
490 491
        platform::RecordEvent record_reorder("int_reorder",
                                             platform::EventRole::kUniqueOp);
A
Adam 已提交
492 493 494
        reorder_p->execute(astream, {{MKLDNN_ARG_FROM, *user_memory_p},
                                     {MKLDNN_ARG_TO, *target_memory_p}});
        astream.wait();
J
Jacek Czaja 已提交
495 496 497 498 499 500 501 502
      }
    }
    return target_memory_p;
  }

 protected:
  const MKLDNNDeviceContext& dev_ctx_;
  mkldnn::engine engine_;
503
  std::string key_common_;
504
  std::string key_;
J
Jacek Czaja 已提交
505 506
};

507 508 509
template <typename T>
class BinaryMKLDNNHandler : public platform::MKLDNNHandlerT<T, dnnl::binary> {
 public:
510 511
  BinaryMKLDNNHandler(const dnnl::algorithm algo, const int axis,
                      const MKLDNNDeviceContext& dev_ctx,
512 513
                      const mkldnn::engine engine, platform::Place cpu_place,
                      const Tensor* x, const Tensor* y, Tensor* z,
514
                      float scale_x, float scale_y, float scale_z,
515
                      const std::string& uniq_name)
516
      : platform::MKLDNNHandlerT<T, dnnl::binary>(
517
            dev_ctx, engine, cpu_place,
518
            platform::CreateKey(
519 520
                dev_ctx, framework::vectorize(x->dims()), uniq_name,
                (algo == dnnl::algorithm::binary_mul ? "M" : ""))) {
521
    // bradcasting combined with in-place may require
522 523
    auto rankdiff = x->dims().size() - y->dims().size();
    if (rankdiff > 0) {
524 525 526
      auto suffix = std::to_string(rankdiff);
      this->key_ += suffix;
      this->key_common_ += suffix;
527 528
    }

529 530 531
    if (!this->isCached()) {
      PADDLE_ENFORCE_EQ(
          x->layout(), DataLayout::kMKLDNN,
G
GaoWei8 已提交
532
          platform::errors::InvalidArgument("Wrong layout set for X tensor."));
533 534
      PADDLE_ENFORCE_NE(
          x->format(), MKLDNNMemoryFormat::undef,
G
GaoWei8 已提交
535
          platform::errors::InvalidArgument("Wrong format set for X tensor."));
536 537 538

      PADDLE_ENFORCE_EQ(
          y->layout(), DataLayout::kMKLDNN,
G
GaoWei8 已提交
539
          platform::errors::InvalidArgument("Wrong layout set for Y tensor."));
540 541
      PADDLE_ENFORCE_NE(
          y->format(), MKLDNNMemoryFormat::undef,
G
GaoWei8 已提交
542
          platform::errors::InvalidArgument("Wrong format set for Y tensor."));
543 544 545 546 547 548 549

      const auto src_x_tz = framework::vectorize(x->dims());
      const auto src_y_tz = framework::vectorize(y->dims());
      const auto dst_tz = framework::vectorize(z->dims());

      const auto src0_md = dnnl::memory::desc(
          src_x_tz, platform::MKLDNNGetDataType<T>(), x->format());
550
      auto src1_md = dnnl::memory::desc(
551
          src_y_tz, platform::MKLDNNGetDataType<T>(), y->format());
552
      if (rankdiff > 0) {
553 554 555
        std::vector<int64_t> dims1_ex(rankdiff, 1);
        dims1_ex.insert(next(dims1_ex.begin(), (axis == -1 ? rankdiff : axis)),
                        src_y_tz.begin(), src_y_tz.end());
556 557
        src1_md = src1_md.reshape(dims1_ex);
      }
558 559 560
      const auto dst_md = memory::desc(dst_tz, platform::MKLDNNGetDataType<T>(),
                                       MKLDNNMemoryFormat::any);

561 562 563
      auto attributes = CreateAttributes(algo, scale_x, scale_y, scale_z);
      this->AcquireForwardPrimitiveDescriptor(attributes, algo, src0_md,
                                              src1_md, dst_md);
564
    }
565 566 567 568 569 570
  }

  std::shared_ptr<mkldnn::memory> AcquireSecondSrcMemory(
      const framework::Tensor* input) {
    const T* input_data = input->data<T>();
    return this->AcquireMemoryFromPrimitive(
571
        this->fwd_pd_->src1_desc(), to_void_cast<T>(input_data), "@src1_mem_p");
572
  }
573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604

 private:
  static inline dnnl::primitive_attr CreateAttributes(dnnl::algorithm op,
                                                      float scale_x,
                                                      float scale_y,
                                                      float scale_z) {
    // Scales set in attributes for inputs contibute to the output equation
    // in the following way (assuming no broadcasting takes place):
    // output_i = scale_0 * x_i <+ or *> scale_1 * y_i;
    // Hence we have to create scales that will:
    // 1. Dequantize both values, by multiplying with (1.0 / scale_x_or_y)
    // 2. Quantize their result to output scale range, by multiplying with
    // (scale_z)
    // If we combine these two, we end up with following equation
    // output = scale_out * (1/scale_x * x <* or +> 1/scale_y * y)
    // Hence, to mimic such behaviour using provided interface,
    // For add operation the equation is equal to:
    // output = (scale_out / scale_x) * x + (scale_out / scale_y) * y
    //                <scale_0>                  <scale_1>
    // For mul operation on the other hand
    // output = (scale_out / scale_x) * x * (1.0 / scale_y) * y
    //                <scale_0>                 <scale_1>
    float scale_0 = scale_z / scale_x;
    float scale_1 =
        op == dnnl::algorithm::binary_add ? scale_z / scale_y : 1.0 / scale_y;
    dnnl::primitive_attr attributes;
    attributes.set_scales(/* input_x_id = */ DNNL_ARG_SRC_0, /* mask = */ 0,
                          {scale_0});
    attributes.set_scales(/* input_y_id = */ DNNL_ARG_SRC_1, /* mask = */ 0,
                          {scale_1});
    return attributes;
  }
605 606
};

607
template <typename T>
608 609 610
class ActivationMKLDNNHandler
    : public MKLDNNHandlerT<T, mkldnn::eltwise_forward,
                            mkldnn::eltwise_backward> {
611
 public:
A
Adam 已提交
612
  ActivationMKLDNNHandler(const std::vector<int64_t>& dims,
613
                          mkldnn::algorithm algorithm, float alpha, float beta,
614
                          const MKLDNNMemoryFormat fmt,
615 616 617 618
                          const platform::MKLDNNDeviceContext& dev_ctx,
                          platform::Place cpu_place,
                          const std::string& unique_name)

619 620 621
      : platform::MKLDNNHandlerT<T, mkldnn::eltwise_forward,
                                 mkldnn::eltwise_backward>(
            dev_ctx, dev_ctx.GetEngine(), cpu_place,
622
            platform::CreateKey(dev_ctx, dims, "a", algorithm, unique_name)) {
623 624
    auto md = mkldnn::memory::desc(dims, platform::MKLDNNGetDataType<T>(), fmt);

625 626
    this->AcquireForwardPrimitiveDescriptor(mkldnn::prop_kind::forward_training,
                                            algorithm, md, alpha, beta);
627
  }
628

A
Adam 已提交
629
  ActivationMKLDNNHandler(const std::vector<int64_t>& dims,
630 631 632 633 634 635 636
                          mkldnn::algorithm algorithm, float alpha, float beta,
                          const MKLDNNMemoryFormat fmt,
                          const MKLDNNMemoryFormat diff_fmt,
                          const platform::MKLDNNDeviceContext& dev_ctx,
                          platform::Place cpu_place,
                          const std::string& unique_name)

637 638 639
      : platform::MKLDNNHandlerT<T, mkldnn::eltwise_forward,
                                 mkldnn::eltwise_backward>(
            dev_ctx, dev_ctx.GetEngine(), cpu_place,
640
            platform::CreateKey(dev_ctx, dims, "a", algorithm, unique_name)) {
641 642 643 644 645 646 647
    auto diff_dst_md = platform::MKLDNNMemDesc(
        dims, platform::MKLDNNGetDataType<T>(), diff_fmt);
    auto src_md =
        platform::MKLDNNMemDesc(dims, platform::MKLDNNGetDataType<T>(), fmt);

    this->AcquireBackwardPrimitiveDescriptor(algorithm, diff_dst_md, src_md,
                                             alpha, beta);
648
  }
649

650 651 652
  std::shared_ptr<mkldnn::memory> AcquireBackwardSrcMemory(
      const framework::Tensor* input) {
    const T* input_data = input->data<T>();
A
Adam 已提交
653
    return this->AcquireMemoryFromPrimitive(this->bwd_pd_->src_desc(),
654 655
                                            to_void_cast<T>(input_data),
                                            "@bwd-src_mem_p");
656 657 658
  }
};

J
Jacek Czaja 已提交
659 660 661
template <typename T>
class LRNMKLDNNHandler
    : public MKLDNNHandlerT<T, mkldnn::lrn_forward, mkldnn::lrn_backward> {
662
 public:
663
  LRNMKLDNNHandler(const paddle::framework::ExecutionContext& ctx,
J
Jacek Czaja 已提交
664
                   const platform::MKLDNNDeviceContext& dev_ctx,
665 666 667
                   const mkldnn::engine mkldnn_engine,
                   platform::Place cpu_place, const Tensor* input,
                   const std::string& unique_name)
668

J
Jacek Czaja 已提交
669
      : platform::MKLDNNHandlerT<T, mkldnn::lrn_forward, mkldnn::lrn_backward>(
670
            dev_ctx, mkldnn_engine, cpu_place,
671
            platform::CreateKey(dev_ctx, framework::vectorize(input->dims()),
672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695
                                unique_name)) {
    if (!this->isCached()) {
      const int n = ctx.Attr<int>("n");
      // MKL-DNN implements LRN in a caffe way:
      // http://caffe.berkeleyvision.org/tutorial/layers/lrn.html
      // Where sum of squares is divided by size of normalization window
      // this is not the case for PaddlePaddle LRN.
      // Hence we need to compensate for this diffrence by
      // multipliing alpha by size of window(n)
      const float alpha = ctx.Attr<float>("alpha") * static_cast<float>(n);
      const float beta = ctx.Attr<float>("beta");
      const float k = ctx.Attr<float>("k");
      bool is_test = ctx.Attr<bool>("is_test");

      auto dims = paddle::framework::vectorize(input->dims());

      auto src_md = mkldnn::memory::desc(dims, platform::MKLDNNGetDataType<T>(),
                                         input->format());

      this->AcquireForwardPrimitiveDescriptor(
          is_test ? mkldnn::prop_kind::forward_inference
                  : mkldnn::prop_kind::forward_training,
          mkldnn::algorithm::lrn_across_channels, src_md, n, alpha, beta, k);
    }
696 697
  }

A
Adam 已提交
698 699
  LRNMKLDNNHandler(const std::vector<int64_t>& dims, const int n,
                   const float alpha, const float beta, const float k,
J
Jacek Czaja 已提交
700 701 702 703
                   const MKLDNNMemoryFormat fmt,
                   const MKLDNNMemoryFormat diff_fmt,
                   const platform::MKLDNNDeviceContext& dev_ctx,
                   platform::Place cpu_place, const std::string& unique_name)
704

J
Jacek Czaja 已提交
705 706
      : platform::MKLDNNHandlerT<T, mkldnn::lrn_forward, mkldnn::lrn_backward>(
            dev_ctx, dev_ctx.GetEngine(), cpu_place,
707
            platform::CreateKey(dev_ctx, dims, unique_name)) {
J
Jacek Czaja 已提交
708 709 710 711
    auto src_md =
        mkldnn::memory::desc(dims, platform::MKLDNNGetDataType<T>(), fmt);
    auto diff_md =
        mkldnn::memory::desc(dims, platform::MKLDNNGetDataType<T>(), diff_fmt);
712

J
Jacek Czaja 已提交
713
    this->AcquireBackwardPrimitiveDescriptor(
A
Adam 已提交
714 715
        mkldnn::algorithm::lrn_across_channels, src_md, diff_md, n, alpha, beta,
        k);
716 717
  }

J
Jacek Czaja 已提交
718 719 720
  std::shared_ptr<mkldnn::memory> AcquireWorkspaceMemory(
      framework::Tensor* workspace) {
    T* ptr = workspace->mutable_data<T>(
A
Adam 已提交
721 722 723
        this->place_, this->fwd_pd_->workspace_desc().get_size());
    return this->AcquireMemoryFromPrimitive(this->fwd_pd_->workspace_desc(),
                                            ptr, "@wrk_mem_p");
J
Jacek Czaja 已提交
724 725 726 727 728
  }

  std::shared_ptr<mkldnn::memory> AcquireBackwardWorkspaceMemory(
      const framework::Tensor* workspace) {
    const T* workspace_data = workspace->data<T>();
A
Adam 已提交
729 730 731
    return this->AcquireMemoryFromPrimitive(this->fwd_pd_->workspace_desc(),
                                            to_void_cast<T>(workspace_data),
                                            "@bwd-wrk_mem_p");
J
Jacek Czaja 已提交
732
  }
733 734
};

735 736 737
template <typename T>
class PoolingMKLDNNHandler : public MKLDNNHandlerT<T, mkldnn::pooling_forward,
                                                   mkldnn::pooling_backward> {
738
 public:
739 740 741 742 743
  PoolingMKLDNNHandler(const paddle::framework::ExecutionContext& ctx,
                       const MKLDNNDeviceContext& dev_ctx,
                       const mkldnn::engine mkldnn_engine,
                       platform::Place cpu_place, const Tensor* input,
                       Tensor* output, const std::string& unique_name)
744 745 746
      : platform::MKLDNNHandlerT<T, mkldnn::pooling_forward,
                                 mkldnn::pooling_backward>(
            dev_ctx, dev_ctx.GetEngine(), cpu_place,
747
            platform::CreateKey(dev_ctx, framework::vectorize(input->dims()),
748 749 750 751 752
                                framework::ToMKLDNNDataType(input->type()),
                                unique_name)) {
    if (!this->isCached()) {
      PADDLE_ENFORCE_EQ(input->layout(), DataLayout::kMKLDNN,
                        platform::errors::InvalidArgument(
G
GaoWei8 已提交
753
                            "Wrong layout set for Input tensor."));
754 755
      PADDLE_ENFORCE_NE(input->format(), MKLDNNMemoryFormat::undef,
                        platform::errors::InvalidArgument(
G
GaoWei8 已提交
756
                            "Wrong format set for Input tensor."));
757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773

      const std::string pooling_type = ctx.Attr<std::string>("pooling_type");

      std::vector<int> ksize_temp = ctx.Attr<std::vector<int>>("ksize");
      std::vector<int64_t> ksize(begin(ksize_temp), end(ksize_temp));

      std::vector<int> strides_temp = ctx.Attr<std::vector<int>>("strides");
      std::vector<int64_t> strides(begin(strides_temp), end(strides_temp));

      std::vector<int> paddings_temp = ctx.Attr<std::vector<int>>("paddings");
      std::vector<int64_t> paddings(begin(paddings_temp), end(paddings_temp));

      const bool global_pooling = ctx.Attr<bool>("global_pooling");
      const std::string padding_algorithm =
          ctx.Attr<std::string>("padding_algorithm");

      // Only 2D pooling is supported now
G
GaoWei8 已提交
774 775 776 777 778 779 780 781 782 783 784 785 786 787 788
      PADDLE_ENFORCE_EQ(
          ksize.size(), 2,
          platform::errors::InvalidArgument(
              "The ksize must be 2D, i.e. 2D pooling, but received %dD.",
              ksize.size()));
      PADDLE_ENFORCE_EQ(
          pooling_type == "max" || pooling_type == "avg", true,
          platform::errors::InvalidArgument(
              "The pooling_type must be 'max' or 'avg', but received %s.",
              pooling_type));
      PADDLE_ENFORCE_EQ(
          input->dims().size(), 4,
          platform::errors::InvalidArgument(
              "Input dim must be with 4, i.e. NCHW, but received %d.",
              input->dims().size()));
789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827

      const auto input_dims = input->dims();
      framework::DDim data_dims =
          framework::slice_ddim(input_dims, 2, input_dims.size());

      if (global_pooling) {
        operators::UpdateKsize(&ksize, data_dims);
      }

      operators::UpdatePadding(&paddings, global_pooling, 0, padding_algorithm,
                               data_dims, strides, ksize);

      const auto src_tz = paddle::framework::vectorize(input->dims());
      const auto dst_tz = paddle::framework::vectorize(output->dims());

      const auto is_test = ctx.Attr<bool>("is_test");

      const auto dt = framework::ToMKLDNNDataType(input->type());
      const auto fmt = input->format();

      const auto exclude_padding = ctx.Attr<bool>("exclusive");

      const auto src_md = mkldnn::memory::desc(src_tz, dt, fmt);
      /* create memory descriptor for pooling without specified format
       * ('any') which lets a primitive (pooling in this case) choose
       * the memory format preferred for best performance
       */

      const auto dst_md =
          platform::MKLDNNMemDesc(dst_tz, dt, MKLDNNMemoryFormat::any);

      auto mkldnn_paddings = ToMkldnnPadding(paddings);

      const bool ceil_mode = ctx.Attr<bool>("ceil_mode");

      if (ceil_mode) {
        CorrectOutputSize(src_tz, dst_tz, ksize, paddings, strides,
                          mkldnn_paddings[1]);
      }
828 829 830

      ComputeAdaptivePoolParameters(ctx, src_tz, ksize, strides);

831 832 833 834 835 836 837 838 839 840
      this->AcquireForwardPrimitiveDescriptor(
          is_test ? mkldnn::prop_kind::forward_inference
                  : mkldnn::prop_kind::forward_training,
          pooling_type == "max"
              ? mkldnn::algorithm::pooling_max
              : (exclude_padding
                     ? mkldnn::algorithm::pooling_avg_exclude_padding
                     : mkldnn::algorithm::pooling_avg_include_padding),
          src_md, dst_md, strides, ksize, mkldnn_paddings[0],
          mkldnn_paddings[1]);
841
    }
842 843 844
  }

  PoolingMKLDNNHandler(
A
Adam 已提交
845 846 847 848 849 850
      const std::vector<int64_t>& diff_dst_dims,
      const std::vector<int64_t>& diff_src_dims,
      const std::vector<int64_t>& ksize, const std::vector<int64_t>& strides,
      const std::vector<int64_t>& paddings, const std::string& pooling_type,
      bool ceil_mode, const MKLDNNMemoryFormat fmt,
      const MKLDNNMemoryFormat diff_dst_fmt, mkldnn::memory::data_type dt,
851
      const platform::MKLDNNDeviceContext& dev_ctx, platform::Place cpu_place,
852
      const std::string& unique_name, bool exclude_padding)
853 854 855
      : platform::MKLDNNHandlerT<T, mkldnn::pooling_forward,
                                 mkldnn::pooling_backward>(
            dev_ctx, dev_ctx.GetEngine(), cpu_place,
856
            platform::CreateKey(dev_ctx, diff_src_dims, dt, unique_name)) {
857 858 859 860 861 862
    auto diff_dst_md = mkldnn::memory::desc(
        diff_dst_dims, platform::MKLDNNGetDataType<T>(), diff_dst_fmt);
    auto diff_src_md =
        mkldnn::memory::desc(diff_src_dims, platform::MKLDNNGetDataType<T>(),
                             MKLDNNMemoryFormat::any);

863 864
    auto mkldnn_paddings = ToMkldnnPadding(paddings);

865
    this->AcquireBackwardPrimitiveDescriptor(
866 867 868 869 870
        pooling_type == "max"
            ? mkldnn::algorithm::pooling_max
            : (exclude_padding
                   ? mkldnn::algorithm::pooling_avg_exclude_padding
                   : mkldnn::algorithm::pooling_avg_include_padding),
871
        diff_src_md, diff_dst_md, strides, ksize, mkldnn_paddings[0],
A
Adam 已提交
872
        mkldnn_paddings[1]);
873 874 875
  }

  std::shared_ptr<mkldnn::memory> AcquireWorkspaceMemory(void) {
A
Adam 已提交
876
    mkldnn::memory::desc workspace_md = this->fwd_pd_->workspace_desc();
877 878 879
    // Pooling PD has to be passed to Grad op that
    // may be executed by diffrent thread, hence
    // for that one we use key that does not contain TID
880 881 882
    auto local_key = this->key_common_ + "@workspace";
    auto mem_p = std::static_pointer_cast<mkldnn::memory>(
        this->dev_ctx_.GetBlob(local_key));
883 884 885 886
    if (mem_p == nullptr) {
      static std::mutex acquire_barrier;
      std::lock_guard<std::mutex> block_threads_until_finish_this_job(
          acquire_barrier);
887 888
      mem_p = std::static_pointer_cast<mkldnn::memory>(
          this->dev_ctx_.GetBlob(local_key));
889
      if (mem_p == nullptr) {
A
Adam 已提交
890
        mem_p = std::make_shared<mkldnn::memory>(workspace_md, this->engine_);
891
        this->dev_ctx_.SetBlob(local_key, mem_p);
892 893 894 895 896
      }
    }
    return mem_p;
  }

897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917
  static void ComputeAdaptivePoolParameters(
      const paddle::framework::ExecutionContext& ctx,
      const std::vector<int64_t>& src_tz, std::vector<int64_t>& ksize,
      std::vector<int64_t>& strides) {
    if (ctx.Attr<bool>("adaptive")) {
      // (jczaja): oneDNN is supporting only unchangable in size pool window
      PADDLE_ENFORCE_EQ(
          src_tz[src_tz.size() - 1] % ksize[1], 0,
          platform::errors::Unimplemented(
              "Input dim must be divisible by corressponding ksize dim."));
      PADDLE_ENFORCE_EQ(
          src_tz[src_tz.size() - 2] % ksize[0], 0,
          platform::errors::Unimplemented(
              "Input dim must be divisible by corressponding ksize dim."));
      ksize[0] = src_tz[src_tz.size() - 2] / ksize[0];
      ksize[1] = src_tz[src_tz.size() - 1] / ksize[1];
      strides[0] = ksize[0];
      strides[1] = ksize[1];
    }
  }

918 919 920 921 922 923 924
 private:
  static inline int ComputeCeiledOutput(int input_size, int kernel_size,
                                        int padding, int stride) {
    return (input_size - kernel_size + 2 * padding) / stride + 1;
  }

  static inline void CorrectOutputSize(
A
Adam 已提交
925 926 927 928
      const std::vector<int64_t>& src_tz, const std::vector<int64_t>& dst_tz,
      const std::vector<int64_t>& kernel_size,
      const std::vector<int64_t>& paddings, const std::vector<int64_t>& strides,
      std::vector<int64_t>& right_bot_padding) {  // NOLINT
929 930 931 932
    for (size_t i = 0; i < right_bot_padding.size(); i++) {
      int desired_size = ComputeCeiledOutput(src_tz[i + 2], kernel_size[i],
                                             paddings[i], strides[i]);
      if (desired_size != dst_tz[i + 2]) {
J
Jacek Czaja 已提交
933
        right_bot_padding[i] += strides[i] - 1;
934 935 936 937 938
      }
    }
  }
};

939
template <typename T>
940 941
class TransposeMKLDNNHandler : public MKLDNNHandler {
 public:
A
Adam 已提交
942 943
  TransposeMKLDNNHandler(std::vector<int64_t>& dims,  // NOLINT
                         std::vector<int>& axis,      // NOLINT
944 945 946 947
                         const platform::MKLDNNDeviceContext& dev_ctx,
                         mkldnn::engine engine, const std::string& base_key)
      : platform::MKLDNNHandler(dev_ctx, engine, base_key),
        dims_(dims),
948 949 950 951
        axis_(axis),
        logical_axis_(dims.size(), 0) {}

  std::shared_ptr<mkldnn::memory> AcquireSrcMemory(
952
      const MKLDNNMemoryFormat& fmt, void* ptr) {
953 954 955 956 957 958 959 960 961
    auto local_key = key_ + "@user_src_mem_p";
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
      // Make memory descriptor using input format, unless it
      // cannot be trusted (nchw) then make up memory fmt manually
      for (size_t i = 0; i < logical_axis_.size(); ++i) {
        logical_axis_[i] = i;
      }
962

A
Adam 已提交
963
      auto src_md = fmt != MKLDNNMemoryFormat::nchw
964
                        ? platform::MKLDNNMemDesc(
965
                              dims_, platform::MKLDNNGetDataType<T>(), fmt)
966
                        : Axis2MemoryDesc(dims_, logical_axis_);
A
Adam 已提交
967
      mem_p = std::make_shared<mkldnn::memory>(src_md, engine_, ptr);
968 969 970 971 972 973
      dev_ctx_.SetBlob(local_key, mem_p);
    } else {
      mem_p->set_data_handle(ptr);
    }
    return mem_p;
  }
974 975 976 977 978 979 980

  std::shared_ptr<mkldnn::memory> AcquireDstMemory(framework::Tensor* output,
                                                   platform::Place place) {
    auto local_key = key_ + "@user_dst_mem_p";
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
A
Adam 已提交
981
      auto dst_md = Axis2MemoryDesc(dims_, axis_);
982

A
Adam 已提交
983
      auto dst_data = output->mutable_data<T>(place, dst_md.get_size());
984

A
Adam 已提交
985
      mem_p = std::make_shared<mkldnn::memory>(dst_md, engine_, dst_data);
986 987
      dev_ctx_.SetBlob(local_key, mem_p);
    } else {
988
      auto dst_data = output->mutable_data<T>(place);
989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008
      mem_p->set_data_handle(dst_data);
    }
    return mem_p;
  }

  std::shared_ptr<mkldnn::reorder> AcquireTranspose(
      std::shared_ptr<mkldnn::memory> dst_memory_p,
      std::shared_ptr<mkldnn::memory> src_memory_p) {
    auto prim_key = key_ + "@transpose_p";
    auto transpose_p =
        std::static_pointer_cast<mkldnn::reorder>(dev_ctx_.GetBlob(prim_key));
    if (transpose_p == nullptr) {
      transpose_p =
          std::make_shared<mkldnn::reorder>(*(src_memory_p), *(dst_memory_p));
      dev_ctx_.SetBlob(prim_key, transpose_p);
    }
    return transpose_p;
  }

 protected:
A
Adam 已提交
1009 1010 1011 1012
  mkldnn::memory::desc Axis2MemoryDesc(std::vector<int64_t>& nchw_tz,  // NOLINT
                                       std::vector<int>& axis          // NOLINT
                                       ) {
    size_t ndims = axis.size();
1013

A
Adam 已提交
1014
    std::vector<int64_t> strides(ndims);
1015
    unsigned int total_stride = 1;
A
Adam 已提交
1016 1017
    for (int i = ndims - 1; i >= 0; --i) {
      strides[axis[i]] = total_stride;
1018 1019
      total_stride *= nchw_tz[axis[i]];
    }
A
Adam 已提交
1020 1021 1022 1023
    mkldnn::memory::desc mem_d(nchw_tz, platform::MKLDNNGetDataType<T>(),
                               strides);

    return mem_d;
1024 1025 1026
  }

 private:
A
Adam 已提交
1027
  std::vector<int64_t> dims_;
1028
  std::vector<int> axis_;
1029
  std::vector<int> logical_axis_;
1030 1031
};

1032 1033
class ReorderMKLDNNHandler : public MKLDNNHandler {
 public:
A
Adam 已提交
1034
  ReorderMKLDNNHandler(std::vector<int64_t>& dims,  // NOLINT
1035 1036 1037 1038 1039 1040 1041 1042 1043 1044
                       framework::proto::VarType::Type vtype,
                       mkldnn::memory::data_type dtype,
                       const platform::MKLDNNDeviceContext& dev_ctx,
                       mkldnn::engine engine, const std::string& base_key)
      : platform::MKLDNNHandler(dev_ctx, engine, base_key),
        dims_(dims),
        vtype_(vtype),
        dtype_(dtype) {}

  std::shared_ptr<mkldnn::memory> AcquireSrcMemory(
1045
      const MKLDNNMemoryFormat& fmt, void* ptr) {
1046
    return this->AcquireMemory(dims_, dtype_, fmt, ptr, "@user_src_mem_p");
1047 1048 1049
  }

  std::shared_ptr<mkldnn::memory> AcquireDstMemory(
1050
      framework::Tensor* output, const MKLDNNMemoryFormat& fmt,
1051 1052 1053 1054 1055 1056
      platform::Place place) {
    auto local_key = key_ + "@user_dst_mem_p";
    auto mem_p =
        std::static_pointer_cast<mkldnn::memory>(dev_ctx_.GetBlob(local_key));
    if (mem_p == nullptr) {
      auto dst_md = platform::MKLDNNMemDesc(dims_, dtype_, fmt);
1057
      auto dst_data = output->mutable_data(place, vtype_, dst_md.get_size());
1058

A
Adam 已提交
1059
      mem_p = std::make_shared<mkldnn::memory>(dst_md, engine_, dst_data);
1060 1061
      dev_ctx_.SetBlob(local_key, mem_p);
    } else {
1062 1063 1064
      // Even if memory object exists , we may be using it for diffrent tensor
      auto dst_data =
          output->mutable_data(place, vtype_, mem_p->get_desc().get_size());
1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084
      mem_p->set_data_handle(dst_data);
    }
    return mem_p;
  }

  std::shared_ptr<mkldnn::reorder> AcquireReorder(
      std::shared_ptr<mkldnn::memory> dst_memory_p,
      std::shared_ptr<mkldnn::memory> src_memory_p) {
    auto prim_key = key_ + "@reorder_p";
    auto reorder_p =
        std::static_pointer_cast<mkldnn::reorder>(dev_ctx_.GetBlob(prim_key));
    if (reorder_p == nullptr) {
      reorder_p =
          std::make_shared<mkldnn::reorder>(*(src_memory_p), *(dst_memory_p));
      dev_ctx_.SetBlob(prim_key, reorder_p);
    }
    return reorder_p;
  }

 private:
A
Adam 已提交
1085
  std::vector<int64_t> dims_;
1086 1087 1088 1089
  framework::proto::VarType::Type vtype_;
  mkldnn::memory::data_type dtype_;
};

1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103
template <typename T>
struct convolutional_algorithm;

template <>
struct convolutional_algorithm<mkldnn::convolution_forward> {
  static constexpr mkldnn::algorithm T = mkldnn::algorithm::convolution_direct;
};

template <>
struct convolutional_algorithm<mkldnn::deconvolution_forward> {
  static constexpr mkldnn::algorithm T =
      mkldnn::algorithm::deconvolution_direct;
};

J
Jacek Czaja 已提交
1104 1105 1106
template <class forward_t, class backward_data_t, class backward_weights_t>
class ConvMKLDNNTemplateHandler : public MKLDNNHandler {
 public:
1107 1108 1109 1110
  ConvMKLDNNTemplateHandler(const platform::MKLDNNDeviceContext& dev_ctx,
                            mkldnn::engine engine, const std::string& base_key)
      : platform::MKLDNNHandler(dev_ctx, engine, base_key) {}

1111 1112 1113 1114 1115 1116 1117 1118 1119
  // TODO(jczaja): remove after conv int8 is adapted
  ConvMKLDNNTemplateHandler(
      std::shared_ptr<typename forward_t::primitive_desc> conv_pd,
      const platform::MKLDNNDeviceContext& dev_ctx, mkldnn::engine engine,
      const std::string& base_key)
      : platform::MKLDNNHandler(dev_ctx, engine, base_key) {
    conv_pd_ = conv_pd;
  }

J
Jacek Czaja 已提交
1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136
  ConvMKLDNNTemplateHandler(
      std::shared_ptr<typename forward_t::primitive_desc> conv_pd,
      std::shared_ptr<typename backward_data_t::primitive_desc>
          conv_bwd_data_pd,
      std::shared_ptr<typename backward_weights_t::primitive_desc>
          conv_bwd_weights_pd,
      const platform::MKLDNNDeviceContext& dev_ctx, mkldnn::engine engine,
      const std::string& base_key)
      : platform::MKLDNNHandler(dev_ctx, engine, base_key),
        conv_pd_(conv_pd),
        conv_bwd_weights_pd_(conv_bwd_weights_pd),
        conv_bwd_data_pd_(conv_bwd_data_pd) {
    // If we are in Grad operatgor then update a key with BWD suffix to
    // distinguish from FWD memory primitives
    key_ += "-BWD";
  }

A
Adam 已提交
1137
  size_t GetDstMemorySize() const { return conv_pd_->dst_desc().get_size(); }
J
Jacek Czaja 已提交
1138

1139
  MKLDNNMemoryFormat GetDstFormat() const {
A
Adam 已提交
1140
    return paddle::platform::GetMKLDNNFormat(conv_pd_->dst_desc());
J
Jacek Czaja 已提交
1141 1142 1143
  }

  size_t GetDiffWeightsMemorySize() const {
A
Adam 已提交
1144
    return conv_bwd_weights_pd_->diff_weights_desc().get_size();
J
Jacek Czaja 已提交
1145 1146 1147
  }

  size_t GetDiffSourceMemorySize() const {
A
Adam 已提交
1148
    return conv_bwd_data_pd_->diff_src_desc().get_size();
J
Jacek Czaja 已提交
1149 1150 1151 1152 1153
  }

  std::shared_ptr<mkldnn::memory> AcquireSrcMemoryFromWeightsPrimitive(
      const std::shared_ptr<mkldnn::memory> user_memory_p,
      std::vector<mkldnn::primitive>& pipeline) {  // NOLINT
A
Adam 已提交
1154 1155
    auto src_pd = conv_bwd_weights_pd_->src_desc();
    auto user_pd = user_memory_p->get_desc();
J
Jacek Czaja 已提交
1156 1157 1158 1159 1160 1161 1162
    return this->AcquireMemory(src_pd, user_pd, user_memory_p,
                               "@weights-src_mem_p", pipeline);
  }

  std::shared_ptr<mkldnn::memory> AcquireDiffDstMemoryFromWeightsPrimitive(
      const std::shared_ptr<mkldnn::memory> user_memory_p,
      std::vector<mkldnn::primitive>& pipeline) {  // NOLINT
A
Adam 已提交
1163 1164
    auto diff_dst_pd = conv_bwd_weights_pd_->diff_dst_desc();
    auto user_pd = user_memory_p->get_desc();
J
Jacek Czaja 已提交
1165 1166 1167 1168 1169 1170 1171
    return this->AcquireMemory(diff_dst_pd, user_pd, user_memory_p,
                               "@weights-diff_dst_mem_p", pipeline);
  }

  std::shared_ptr<mkldnn::memory> AcquireDiffWeightsMemoryFromWeightsPrimitive(
      void* ptr) {
    return this->AcquireMemoryFromPrimitive(
A
Adam 已提交
1172
        conv_bwd_weights_pd_->diff_weights_desc(), ptr, "@diff_weights_mem_p");
J
Jacek Czaja 已提交
1173 1174
  }

1175 1176 1177 1178 1179 1180
  std::shared_ptr<mkldnn::memory> AcquireDiffWeightsMemoryFromWeightsPrimitive(
      void) {
    return this->AcquireMemoryFromPrimitive(
        conv_bwd_weights_pd_->diff_weights_desc(), "@diff_weights_mem_p");
  }

J
Jacek Czaja 已提交
1181 1182 1183
  std::shared_ptr<mkldnn::memory> AcquireDiffDstMemoryFromDataPrimitive(
      const std::shared_ptr<mkldnn::memory> user_memory_p,
      std::vector<mkldnn::primitive>& pipeline) {  // NOLINT
A
Adam 已提交
1184 1185
    auto diff_dst_pd = conv_bwd_data_pd_->diff_dst_desc();
    auto user_pd = user_memory_p->get_desc();
J
Jacek Czaja 已提交
1186 1187 1188 1189 1190 1191 1192
    return this->AcquireMemory(diff_dst_pd, user_pd, user_memory_p,
                               "@data-diff_dst_mem_p", pipeline);
  }

  std::shared_ptr<mkldnn::memory> AcquireWeightsMemoryFromDataPrimitive(
      const std::shared_ptr<mkldnn::memory> user_weights_memory_p,
      std::vector<mkldnn::primitive>& pipeline) {  // NOLINT
A
Adam 已提交
1193 1194
    auto weights_pd = conv_bwd_data_pd_->weights_desc();
    auto user_pd = user_weights_memory_p->get_desc();
J
Jacek Czaja 已提交
1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214
    return this->AcquireMemory(weights_pd, user_pd, user_weights_memory_p,
                               "@data-weights_mem_p", pipeline);
  }

  std::shared_ptr<mkldnn::memory> AcquireResidualDataMemory(
      const mkldnn::memory::desc& md, void* ptr) {
    return this->AcquireMemory(md, ptr, "@user_residual_data_mem_p");
  }

  std::shared_ptr<mkldnn::memory> AcquireDstMemoryFromResidualDataMemory(
      const std::shared_ptr<mkldnn::memory>& user_residual_memory_p,
      void* dst_ptr,
      std::vector<mkldnn::primitive>& pipeline) {  // NOLINT
    return this->AcquireMemory(user_residual_memory_p,
                               this->AcquireDstMemoryFromPrimitive(dst_ptr),
                               "@residual_data_mem_p", pipeline);
  }

  std::shared_ptr<mkldnn::memory> AcquireDiffSrcMemoryFromDataPrimitive(
      void* ptr) {
A
Adam 已提交
1215 1216
    return this->AcquireMemoryFromPrimitive(conv_bwd_data_pd_->diff_src_desc(),
                                            ptr, "@diff_src_mem_p");
J
Jacek Czaja 已提交
1217 1218 1219
  }

  std::shared_ptr<mkldnn::memory> AcquireDstMemoryFromPrimitive(void* ptr) {
A
Adam 已提交
1220
    return this->AcquireMemoryFromPrimitive(conv_pd_->dst_desc(), ptr,
J
Jacek Czaja 已提交
1221 1222 1223 1224 1225 1226
                                            "@dst_mem_p");
  }

  std::shared_ptr<mkldnn::memory> AcquireSrcMemoryFromPrimitive(
      const std::shared_ptr<mkldnn::memory> user_memory_p,
      std::vector<mkldnn::primitive>& pipeline) {  // NOLINT
A
Adam 已提交
1227 1228
    auto src_pd = conv_pd_->src_desc();
    auto user_pd = user_memory_p->get_desc();
J
Jacek Czaja 已提交
1229 1230 1231 1232
    return this->AcquireMemory(src_pd, user_pd, user_memory_p, "@src_mem_p",
                               pipeline);
  }

A
Adam 已提交
1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243
  std::shared_ptr<mkldnn::memory> AcquireWeightsMemory(
      const mkldnn::memory::desc& md, void* ptr,
      user_function custom_func = {}) {
    return this->AcquireMemory(md, ptr, "@user_weights_mem_p", custom_func);
  }

  std::shared_ptr<mkldnn::memory> AcquireBiasMemory(
      const mkldnn::memory::desc& md, void* ptr) {
    return this->AcquireMemory(md, ptr, "@user_bias_mem_p");
  }

J
Jacek Czaja 已提交
1244 1245 1246
  std::shared_ptr<mkldnn::memory> AcquireWeightsMemoryFromPrimitive(
      const std::shared_ptr<mkldnn::memory> user_weights_memory_p,
      std::vector<mkldnn::primitive>& pipeline,  // NOLINT
1247 1248
      bool is_persistent = false, bool is_INT8 = false,
      std::vector<float> scale_data = {1.0f}, int mask = 0) {
A
Adam 已提交
1249 1250
    auto user_weights_pd = user_weights_memory_p->get_desc();
    auto weights_pd = conv_pd_->weights_desc();
1251 1252 1253
    return this->AcquireMemory(
        weights_pd, user_weights_pd, user_weights_memory_p, "@weights_mem_p",
        pipeline, is_persistent, is_INT8, scale_data, mask);
J
Jacek Czaja 已提交
1254 1255 1256 1257
  }

  std::shared_ptr<mkldnn::memory> AcquireBiasMemoryFromPrimitive(
      const std::shared_ptr<mkldnn::memory> user_bias_memory_p,
1258 1259 1260 1261
      std::vector<mkldnn::primitive>& pipeline,  // NOLINT
      bool is_persistent = false, bool is_INT8 = false,
      std::vector<float> scale_data = {1.0f},
      int mask = 0) {  // NOLINT
A
Adam 已提交
1262 1263
    auto user_bias_pd = user_bias_memory_p->get_desc();
    auto bias_pd = conv_pd_->bias_desc();
J
Jacek Czaja 已提交
1264
    return this->AcquireMemory(bias_pd, user_bias_pd, user_bias_memory_p,
1265 1266
                               "@bias_mem_p", pipeline, is_persistent, is_INT8,
                               scale_data, mask);
J
Jacek Czaja 已提交
1267 1268
  }

1269
  mkldnn::primitive_attr CreatePostOps(
1270 1271
      std::string fuse_activation, float fuse_alpha, float fuse_beta,
      bool fuse_residual_conn, const std::vector<float> output_shift_scale = {},
1272
      float sum_scale = 1.0f) const {
1273 1274
    mkldnn::primitive_attr conv_attr;
    mkldnn::post_ops post_operations;
1275 1276 1277 1278
    if (output_shift_scale.size() > 0) {
      int mask = output_shift_scale.size() > 1 ? 1 << 1 : 0;
      conv_attr.set_output_scales(mask, output_shift_scale);
    }
1279 1280 1281 1282 1283 1284
    // Fusion with Elementwise layer relies on adding a sum post-operation with
    // the scale parameter. It is assumed that when fuse_residual_connection is
    // true, the output tensor contains the data coming from residual
    // connection. The result of this post_op is:
    // Output = scale * Output + Conv_Out.
    if (fuse_residual_conn) {
1285
      post_operations.append_sum(sum_scale);
1286 1287 1288
    }
    // Fusion with ReLU layer is executed through the PostOps feature. Create a
    // PostOps object and configure it to execute an eltwise relu operation.
1289
    if (fuse_activation == "relu" || fuse_activation == "leaky_relu") {
1290 1291
      constexpr float scale = 1.0f;
      post_operations.append_eltwise(scale, mkldnn::algorithm::eltwise_relu,
1292
                                     fuse_alpha, fuse_beta);
1293
    } else if (fuse_activation == "relu6") {
1294 1295 1296
      constexpr float scale = 1.0f;
      post_operations.append_eltwise(scale,
                                     mkldnn::algorithm::eltwise_bounded_relu,
1297
                                     fuse_alpha, fuse_beta);
1298 1299 1300 1301
    } else if (fuse_activation == "swish") {
      constexpr float scale = 1.0f;
      post_operations.append_eltwise(scale, mkldnn::algorithm::eltwise_swish,
                                     fuse_alpha, fuse_beta);
1302
    }
1303 1304 1305 1306 1307 1308 1309 1310
    conv_attr.set_post_ops(post_operations);
    return conv_attr;
  }

  std::shared_ptr<typename forward_t::primitive_desc>
  AcquireConvolutionPrimitiveDescriptor(
      const mkldnn::memory::desc& src, const mkldnn::memory::desc& weights,
      boost::optional<const mkldnn::memory::desc&> bias,
A
Adam 已提交
1311
      const mkldnn::memory::desc& dst, const std::vector<int64_t>& strides,
1312
      const std::vector<int64_t>& dilations,
A
Adam 已提交
1313
      const std::vector<int64_t>& paddings, const mkldnn::engine& engine,
1314 1315
      const std::string& fuse_activation, float fuse_alpha, float fuse_beta,
      const bool fuse_residual_conn, mkldnn::prop_kind fwd_prop_kind,
1316 1317
      const std::vector<float> output_shift_scale = {},
      const float sum_scale = 1.0f) {
1318 1319 1320 1321
    // Conv PD has to be passed to Grad op that
    // may be exxecuted by diffrent thread, hence
    // for that one we use key that does not contain TID
    const std::string key_conv_pd = key_common_ + "@conv_pd";
1322

1323
    conv_pd_ = std::static_pointer_cast<typename forward_t::primitive_desc>(
1324 1325
        dev_ctx_.GetBlob(key_conv_pd));

1326 1327 1328 1329 1330 1331 1332 1333 1334
    if (conv_pd_ == nullptr) {
      static std::mutex acquire_barrier;
      std::lock_guard<std::mutex> block_threads_until_finish_this_job(
          acquire_barrier);

      conv_pd_ = std::static_pointer_cast<typename forward_t::primitive_desc>(
          dev_ctx_.GetBlob(key_conv_pd));
      if (conv_pd_ == nullptr) {
        mkldnn::memory::dims stride_dims = strides;
1335
        mkldnn::memory::dims dilations_dims = dilations;
1336
        auto mkldnn_paddings = ToMkldnnPadding(paddings);
1337 1338

        auto conv_desc =
A
Adam 已提交
1339 1340
            bias ? typename forward_t::desc(
                       fwd_prop_kind, convolutional_algorithm<forward_t>::T,
1341
                       src, weights, *bias, dst, stride_dims, dilations_dims,
A
Adam 已提交
1342 1343 1344
                       mkldnn_paddings[0], mkldnn_paddings[1])
                 : typename forward_t::desc(
                       fwd_prop_kind, convolutional_algorithm<forward_t>::T,
1345 1346
                       src, weights, dst, stride_dims, dilations_dims,
                       mkldnn_paddings[0], mkldnn_paddings[1]);
1347

1348
        mkldnn::primitive_attr conv_attr =
1349 1350
            CreatePostOps(fuse_activation, fuse_alpha, fuse_beta,
                          fuse_residual_conn, output_shift_scale, sum_scale);
1351 1352 1353 1354 1355 1356

        conv_pd_.reset(new typename forward_t::primitive_desc(
            conv_desc, conv_attr, engine));
        // Save conv_pd/src_memory/weights_memory for backward pass
        dev_ctx_.SetBlob(key_conv_pd, conv_pd_);
      }
1357 1358 1359 1360 1361
    }

    return conv_pd_;
  }

A
Adam 已提交
1362
  std::shared_ptr<forward_t> AcquireConvolution() {
J
Jacek Czaja 已提交
1363 1364 1365 1366
    auto prim_key = key_ + "@conv_p";
    auto conv_p =
        std::static_pointer_cast<forward_t>(dev_ctx_.GetBlob(prim_key));
    if (conv_p == nullptr) {
A
Adam 已提交
1367
      conv_p = std::make_shared<forward_t>(*conv_pd_);
J
Jacek Czaja 已提交
1368 1369 1370 1371 1372 1373

      dev_ctx_.SetBlob(prim_key, conv_p);
    }
    return conv_p;
  }

A
Adam 已提交
1374
  std::shared_ptr<backward_weights_t> AcquireConvolutionBackwardWeights() {
J
Jacek Czaja 已提交
1375 1376 1377 1378 1379
    auto prim_key = key_ + "@conv_bwd_weights_p";
    auto conv_bwd_weights_p = std::static_pointer_cast<backward_weights_t>(
        dev_ctx_.GetBlob(prim_key));
    if (conv_bwd_weights_p == nullptr) {
      // create backward conv primitive for weights
A
Adam 已提交
1380 1381
      conv_bwd_weights_p =
          std::make_shared<backward_weights_t>(*conv_bwd_weights_pd_);
J
Jacek Czaja 已提交
1382 1383 1384 1385 1386
      dev_ctx_.SetBlob(prim_key, conv_bwd_weights_p);
    }
    return conv_bwd_weights_p;
  }

A
Adam 已提交
1387
  std::shared_ptr<backward_data_t> AcquireConvolutionBackwardData() {
J
Jacek Czaja 已提交
1388 1389 1390 1391
    auto prim_key = key_ + "@conv_bwd_data_p";
    auto conv_bwd_data_p =
        std::static_pointer_cast<backward_data_t>(dev_ctx_.GetBlob(prim_key));
    if (conv_bwd_data_p == nullptr) {
A
Adam 已提交
1392
      conv_bwd_data_p = std::make_shared<backward_data_t>(*conv_bwd_data_pd_);
J
Jacek Czaja 已提交
1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413
      dev_ctx_.SetBlob(prim_key, conv_bwd_data_p);
    }
    return conv_bwd_data_p;
  }

 private:
  std::shared_ptr<typename forward_t::primitive_desc> conv_pd_;
  std::shared_ptr<typename backward_weights_t::primitive_desc>
      conv_bwd_weights_pd_;
  std::shared_ptr<typename backward_data_t::primitive_desc> conv_bwd_data_pd_;
};

using ConvMKLDNNHandler =
    ConvMKLDNNTemplateHandler<mkldnn::convolution_forward,
                              mkldnn::convolution_backward_data,
                              mkldnn::convolution_backward_weights>;

using ConvTransposeMKLDNNHandler =
    ConvMKLDNNTemplateHandler<mkldnn::deconvolution_forward,
                              mkldnn::deconvolution_backward_data,
                              mkldnn::deconvolution_backward_weights>;
1414

1415 1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433
template <typename T>
static std::shared_ptr<mkldnn::memory> SetDstMemory(
    const framework::ExecutionContext& ctx, framework::Tensor* output,
    const std::shared_ptr<ConvMKLDNNHandler>& handler) {
  T* output_data =
      output->mutable_data<T>(ctx.GetPlace(), handler->GetDstMemorySize());
  std::shared_ptr<mkldnn::memory> dst_memory_p =
      handler->AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));
  return dst_memory_p;
}

template <typename T>
static std::shared_ptr<mkldnn::memory> SetDstMemory(
    const framework::ExecutionContext& ctx, framework::Tensor* output,
    const framework::Tensor* residual_param,
    const mkldnn::memory::desc& user_residual_md,
    const std::shared_ptr<ConvMKLDNNHandler>& handler,
    std::vector<mkldnn::primitive>* pipeline) {
  const T* residual_param_data = residual_param->data<T>();
1434 1435 1436 1437
  PADDLE_ENFORCE_NOT_NULL(
      residual_param_data,
      platform::errors::PreconditionNotMet("Residual parameter is required for "
                                           "the DNNL conv+elementwise_add "
G
GaoWei8 已提交
1438
                                           "fusion, but now it is missing."));
1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458
  std::shared_ptr<mkldnn::memory> user_residual_memory_p =
      handler->AcquireResidualDataMemory(user_residual_md,
                                         to_void_cast<T>(residual_param_data));
  T* output_data = output->mutable_data<T>(ctx.GetPlace());
  std::shared_ptr<mkldnn::memory> dst_memory_p =
      handler->AcquireDstMemoryFromResidualDataMemory(
          user_residual_memory_p, to_void_cast<T>(output_data), *pipeline);
  return dst_memory_p;
}

template <typename T>
static void SetDstMemoryHandler(
    const framework::ExecutionContext& ctx, framework::Tensor* output,
    const std::shared_ptr<ConvMKLDNNHandler>& handler,
    std::shared_ptr<mkldnn::memory> dst_memory_p) {
  T* output_data =
      output->mutable_data<T>(ctx.GetPlace(), handler->GetDstMemorySize());
  dst_memory_p->set_data_handle(to_void_cast<T>(output_data));
}

1459 1460 1461
template <typename T>
static void SetDstMemoryQuantized(
    const framework::ExecutionContext& ctx, framework::Tensor* output,
A
Adam 已提交
1462 1463
    std::vector<int64_t> dst_tz, const mkldnn::engine& engine,
    std::shared_ptr<mkldnn::memory::desc>& dst_md,  // NOLINT
1464 1465
    std::shared_ptr<mkldnn::memory>& dst_memory,    // NOLINT
    MKLDNNMemoryFormat output_format) {
1466 1467
  T* output_data = output->mutable_data<T>(ctx.GetPlace());
  const size_t dst_dims = dst_tz.size();
1468
  MKLDNNMemoryFormat dst_fmt;
G
GaoWei8 已提交
1469 1470 1471 1472
  PADDLE_ENFORCE_LE(dst_dims, 5, platform::errors::InvalidArgument(
                                     "Dst memory for quantization can not have "
                                     "dims > 5. But received dst_dims is %d.",
                                     dst_dims));
1473
  dst_fmt = platform::MKLDNNFormatForSize(dst_dims, output_format);
1474

A
Adam 已提交
1475
  auto tmp_dst_md = platform::MKLDNNMemDesc(
1476
      {dst_tz}, paddle::framework::ToMKLDNNDataType(
1477
                    framework::DataTypeTrait<T>::DataType()),
1478
      dst_fmt);
A
Adam 已提交
1479 1480 1481
  dst_md.reset(new mkldnn::memory::desc(tmp_dst_md));
  dst_memory.reset(
      new mkldnn::memory(*dst_md, engine, to_void_cast<T>(output_data)));
1482 1483
}

J
Jacek Czaja 已提交
1484 1485
}  // namespace platform
}  // namespace paddle