MKLDNNBase.h 2.5 KB
Newer Older
1
/* Copyright (c) 2017 PaddlePaddle Authors. All Rights Reserved.
T
tensor-tang 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21

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

#include "mkldnn.hpp"

namespace paddle {

typedef enum {
T
tensor-tang 已提交
22 23
  MKLDNN_BASE = 1,   // basical info of MKLDNN
  MKLDNN_TESTS = 1,  // gtest info of MKLDNN
24 25
  MKLDNN_FMTS = 2,   // format info of MKLDNN
  MKLDNN_SIZES = 3,  // size info of MKLDNN
T
tensor-tang 已提交
26 27
  MKLDNN_ALL = 4,    // show all info of MKLDNN
} MKLDNN_LOG_LEVEL;
T
tensor-tang 已提交
28 29 30 31 32

/**
 * @brief MKLDNN CPU engine.
 *
 */
33
class CPUEngine {
W
Wu Yi 已提交
34
 public:
35
  static CPUEngine& Instance() {
T
tensor-tang 已提交
36
    // Thread-safe in C++11.
37
    static CPUEngine myInstance;
T
tensor-tang 已提交
38 39 40 41
    return myInstance;
  }

  // Disallow copy or move
42 43 44 45
  CPUEngine(const CPUEngine&) = delete;             // Copy constructor
  CPUEngine(CPUEngine&&) = delete;                  // Move constructor
  CPUEngine& operator=(const CPUEngine&) = delete;  // Copy assignment
  CPUEngine& operator=(CPUEngine&&) = delete;       // Move assignment
T
tensor-tang 已提交
46 47 48

  mkldnn::engine& getEngine() { return cpuEngine_; }

W
Wu Yi 已提交
49
 protected:
50 51 52
  CPUEngine() : cpuEngine_(mkldnn::engine::cpu, 0) {}
  //    CPUEngine() : cpuEngine_(mkldnn::engine::cpu_lazy, 0) {}
  ~CPUEngine() {}
T
tensor-tang 已提交
53

W
Wu Yi 已提交
54
 private:
T
tensor-tang 已提交
55 56 57 58 59 60 61
  mkldnn::engine cpuEngine_;
};

/**
 * @brief MKLDNN Stream.
 *
 */
62
class MKLDNNStream {
W
Wu Yi 已提交
63
 public:
64
  MKLDNNStream() : ready_(false) { resetState(); }
T
tensor-tang 已提交
65

66
  virtual ~MKLDNNStream() {}
T
tensor-tang 已提交
67 68 69 70

  /**
   * @brief Submit stream
   * @param prims The primitives vector
T
tensor-tang 已提交
71
   * @param block Waiting for the stream to complete
T
tensor-tang 已提交
72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87
   */
  void submit(std::vector<mkldnn::primitive>& prims, bool block = true) {
    resetState();
    stream_->submit(prims).wait(block);
    ready_ = false;
  }

  /**
   * @brief Reset the mkldnn stream
   */
  void resetState() {
    if (ready_) {
      return;
    }
    // TODO(TJ): change me when mkldnn have method to reset this state
    // stream_.reset(new mkldnn::stream(mkldnn::stream::kind::lazy));
T
tensor-tang 已提交
88
    stream_.reset(new mkldnn::stream(mkldnn::stream::kind::eager));
T
tensor-tang 已提交
89 90 91
    ready_ = true;
  }

W
Wu Yi 已提交
92
 private:
T
tensor-tang 已提交
93 94 95 96 97
  bool ready_;
  std::shared_ptr<mkldnn::stream> stream_;
};

}  // namespace paddle