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

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 <vector>
#include "Layer.h"
19
#include "MKLDNNBase.h"
T
tensor-tang 已提交
20
#include "mkldnn.hpp"
T
tensor-tang 已提交
21
#include "paddle/math/MKLDNNMatrix.h"
T
tensor-tang 已提交
22

T
tensor-tang 已提交
23 24
DECLARE_bool(use_mkldnn);

T
tensor-tang 已提交
25 26
namespace paddle {

27 28
class MKLDNNLayer;
typedef std::shared_ptr<MKLDNNLayer> MKLDNNLayerPtr;
T
tensor-tang 已提交
29 30

/**
31
 * @brief Base class of MKLDNNlayer.
T
tensor-tang 已提交
32 33
 *
 */
34
class MKLDNNLayer : public Layer {
T
tensor-tang 已提交
35 36 37 38 39 40 41 42
protected:
  // batch size
  int bs_;
  // input image channel, height and width
  int ic_, ih_, iw_;
  // output image channel, height and width
  int oc_, oh_, ow_;

T
tensor-tang 已提交
43 44 45
  // backward also need reset after reset forward handle
  bool needResetBwd_;

T
tensor-tang 已提交
46 47
  // mkldnn engine, stream and primivtives
  mkldnn::engine engine_;
48
  std::shared_ptr<MKLDNNStream> stream_;
T
tensor-tang 已提交
49
  std::shared_ptr<mkldnn::primitive> fwd_;
T
tensor-tang 已提交
50 51
  std::shared_ptr<mkldnn::primitive> bwdWgt_;
  std::shared_ptr<mkldnn::primitive> bwdData_;
T
tensor-tang 已提交
52 53 54
  std::vector<mkldnn::primitive> pipelineFwd_;
  std::vector<mkldnn::primitive> pipelineBwd_;

T
tensor-tang 已提交
55
  // MKLDNNMatrixPtr
T
tensor-tang 已提交
56
  MKLDNNMatrixPtr inVal_;
T
tensor-tang 已提交
57
  MKLDNNMatrixPtr inGrad_;
T
tensor-tang 已提交
58
  MKLDNNMatrixPtr outVal_;
T
tensor-tang 已提交
59
  MKLDNNMatrixPtr outGrad_;
T
tensor-tang 已提交
60
  MKLDNNMatrixPtr wgtVal_;
T
tensor-tang 已提交
61
  MKLDNNMatrixPtr wgtGrad_;
T
tensor-tang 已提交
62
  MKLDNNMatrixPtr biasVal_;
T
tensor-tang 已提交
63
  MKLDNNMatrixPtr biasGrad_;
T
tensor-tang 已提交
64

T
tensor-tang 已提交
65
public:
66
  explicit MKLDNNLayer(const LayerConfig& config)
T
tensor-tang 已提交
67 68 69 70 71 72 73 74
      : Layer(config),
        bs_(0),
        ic_(0),
        ih_(0),
        iw_(0),
        oc_(0),
        oh_(0),
        ow_(0),
T
tensor-tang 已提交
75
        needResetBwd_(true),
T
tensor-tang 已提交
76
        engine_(mkldnn::engine::cpu, 0),
T
tensor-tang 已提交
77 78 79 80
        stream_(nullptr),
        fwd_(nullptr),
        bwdWgt_(nullptr),
        bwdData_(nullptr) {}
T
tensor-tang 已提交
81

82
  ~MKLDNNLayer() {}
T
tensor-tang 已提交
83

T
tensor-tang 已提交
84 85
  virtual bool init(const LayerMap& layerMap,
                    const ParameterMap& parameterMap) {
T
tensor-tang 已提交
86 87 88
    CHECK(FLAGS_use_mkldnn) << "MkldnnLayers only support use_mkldnn."
                            << "Please set WITH_MKLDNN=ON "
                            << "and set use_mkldnn=True";
T
refine  
tensor-tang 已提交
89
    CHECK(!useGpu_) << "Do not support GPU yet";
T
tensor-tang 已提交
90 91 92 93 94

    // set device id before Layer::init
    setDevice(MKLDNN_DEVICE);
    // change param device to MKLDNN device
    setParamsDevice(MKLDNN_DEVICE, parameterMap);
T
tensor-tang 已提交
95 96 97 98
    if (!Layer::init(layerMap, parameterMap)) {
      return false;
    }

99 100
    stream_.reset(new MKLDNNStream());
    engine_ = CPUEngine::Instance().getEngine();
T
tensor-tang 已提交
101 102
    return true;
  }
T
tensor-tang 已提交
103 104 105 106 107

  /**
   * convert weight from paddle format to mkldnn format
   * weight_ will be override
   */
108
  virtual void convertWeightsFromPaddle() {}
T
tensor-tang 已提交
109 110 111 112 113

  /**
   * convert mkldnn weight to paddle format
   * weight_ will be override
   */
114
  virtual void convertWeightsToPaddle() {}
T
tensor-tang 已提交
115

T
tensor-tang 已提交
116 117 118 119 120 121 122 123
  /**
   * print info about sizes
   */
  virtual void printSizeInfo() {
    VLOG(MKLDNN_SIZES) << getName() << ": bs: " << bs_ << ", ic: " << ic_
                       << ", ih: " << ih_ << ", iw: " << iw_ << ", oc: " << oc_
                       << ", oh: " << oh_ << ", ow: " << ow_;
  }
T
tensor-tang 已提交
124

125 126 127 128 129 130 131 132
  /**
   * Print the mkldnn memory format flow of value
   */
  virtual void printValueFormatFlow() {
    if (inVal_ && outVal_) {
      VLOG(MKLDNN_FMTS) << "value format flow --- " << inVal_->getFormat()
                        << " >>> " << outVal_->getFormat();
    }
T
tensor-tang 已提交
133
  }
T
tensor-tang 已提交
134

135 136 137 138 139 140 141 142
  /**
   * Print the mkldnn memory format flow of grad
   */
  virtual void printGradFormatFlow() {
    if (inGrad_ && outGrad_) {
      VLOG(MKLDNN_FMTS) << "grad format flow --- " << inGrad_->getFormat()
                        << " <<< " << outGrad_->getFormat();
    }
T
tensor-tang 已提交
143 144 145
  }

protected:
146
  /**
T
refine  
tensor-tang 已提交
147
   * copy image size and sequence info to other device
T
rename  
tensor-tang 已提交
148 149
   * @note: can not directly use Layer::copyOutputToOtherDevice since here only
   *        copy base info and do not copy data value
150
   */
T
refine  
tensor-tang 已提交
151
  void copyOutputInfoToOtherDevice() {
T
tensor-tang 已提交
152
    int cnt = 0;
153
    for (size_t i = 0; i < outputOtherDevice_.size(); i++) {
T
refine  
tensor-tang 已提交
154 155 156 157 158 159 160
      outputOtherDevice_[i].setFrameHeight(output_.getFrameHeight());
      outputOtherDevice_[i].setFrameWidth(output_.getFrameWidth());
      outputOtherDevice_[i].sequenceStartPositions =
          output_.sequenceStartPositions;
      outputOtherDevice_[i].subSequenceStartPositions =
          output_.subSequenceStartPositions;
      outputOtherDevice_[i].cpuSequenceDims = output_.cpuSequenceDims;
T
tensor-tang 已提交
161 162 163 164 165 166
      if (outputOtherDevice_[i].deviceId == CPU_DEVICE) {
        ++cnt;
      }
    }
    if (cnt > 1) {
      LOG(WARNING) << "should not have more than one CPU devie";
167 168 169 170
    }
  }

  /**
T
rename  
tensor-tang 已提交
171
   * If input only has MKLDNN device.
T
refine  
tensor-tang 已提交
172
   * Otherwise, only support the previous layer using CPU device.
173
   */
T
rename  
tensor-tang 已提交
174
  bool inputIsOnlyMKLDNN(int index = 0) {
175 176 177 178 179 180 181 182 183 184
    int prevDevice = getPrev(index)->getDeviceId();
    if (prevDevice == MKLDNN_DEVICE) {
      return true;
    } else {
      // do not support GPU yet
      CHECK_EQ(prevDevice, CPU_DEVICE) << "Only support CPU yet";
      return false;
    }
  }

T
refine  
tensor-tang 已提交
185 186 187 188
  /**
   * If output only has MKLDNN device.
   * Otherwise, other devices should only using CPU device.
   */
T
rename  
tensor-tang 已提交
189
  bool outputIsOnlyMKLDNN() {
T
refine  
tensor-tang 已提交
190 191 192 193 194 195 196
    for (size_t i = 0; i < outputOtherDevice_.size(); i++) {
      CHECK_EQ(outputOtherDevice_[i].deviceId, CPU_DEVICE)
          << "Only support other device is CPU yet";
    }
    return outputOtherDevice_.size() == 0;
  }

197 198 199 200
  /**
   * Sync input value data
   */
  void syncInputValue() {
T
rename  
tensor-tang 已提交
201
    if (inputIsOnlyMKLDNN()) {
202 203 204 205 206
      return;
    }
    real* iData = getInputValue(0, CPU_DEVICE)->getData();
    // update input data
    // since it might be changed if this is after data layer
207
    inVal_->setData(iData);
208 209 210 211 212 213
  }

  /**
   * Sync output grad data
   */
  void syncOutputGrad() {
T
rename  
tensor-tang 已提交
214
    if (outputIsOnlyMKLDNN()) {
215 216 217 218 219
      return;
    }

    // update diff
    real* oDiff = getOutput(CPU_DEVICE).grad->getData();
220
    outGrad_->setData(oDiff);
221 222
  }

T
tensor-tang 已提交
223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
  /**
   * Set deviceId of this layer.
   */
  void setDevice(int id) { deviceId_ = id; }

  /**
   * Set deviceId of the params used in this layer.
   */
  void setParamsDevice(int id, const ParameterMap& parameterMap) {
    for (auto& inputConfig : config_.inputs()) {
      if (inputConfig.has_input_parameter_name()) {
        ParameterPtr parameter;
        std::string name = inputConfig.input_parameter_name();
        CHECK(mapGet(name, parameterMap, &parameter))
            << "Cannot find input parameter " << name << " for layer "
            << getName();
        parameter->setDevice(id);
      }
    }
    if (config_.has_bias_parameter_name()) {
      ParameterPtr parameter;
      std::string name = config_.bias_parameter_name();
      CHECK(mapGet(name, parameterMap, &parameter))
          << "Cannot find bias parameter " << name << " for layer "
          << getName();
      parameter->setDevice(id);
    }
T
tensor-tang 已提交
250
  }
T
tensor-tang 已提交
251 252 253
};

}  // namespace paddle