MKLDNNLayer.h 12.9 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"
22
#include "paddle/utils/Stat.h"
T
tensor-tang 已提交
23

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

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

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

/**
32
 * @brief Base class of MKLDNNlayer.
T
tensor-tang 已提交
33 34
 *
 */
35
class MKLDNNLayer : public Layer {
T
tensor-tang 已提交
36
protected:
37 38
  // input value element count
  size_t inputElemenCnt_;
T
tensor-tang 已提交
39 40 41 42 43 44 45
  // 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 已提交
46 47 48
  // backward also need reset after reset forward handle
  bool needResetBwd_;

49 50 51
  // is output only mkldnn
  bool outputOnlyMKLDNN_;

T
tensor-tang 已提交
52 53
  // mkldnn engine, stream and primivtives
  mkldnn::engine engine_;
54
  std::shared_ptr<MKLDNNStream> stream_;
T
tensor-tang 已提交
55
  std::shared_ptr<mkldnn::primitive> fwd_;
T
tensor-tang 已提交
56 57
  std::shared_ptr<mkldnn::primitive> bwdWgt_;
  std::shared_ptr<mkldnn::primitive> bwdData_;
T
tensor-tang 已提交
58 59 60
  std::vector<mkldnn::primitive> pipelineFwd_;
  std::vector<mkldnn::primitive> pipelineBwd_;

T
tensor-tang 已提交
61 62 63 64 65 66 67 68 69
  /* Value and grad are seperated as internal and external buffers.
   * Each MKLDNNLayer must init or reset internal buffer at least,
   * and the external buffer format is always nchw of nc(when h==w==1),
   * which is the same format as paddle.
   * The output_.value and output_.grad always save the external data,
   * when mixed with cpu device.
   * When all layers are mkldnn layers, they could save internal data.
   */
  // below MKLDNNMatrix buffers are all internal buffers
T
tensor-tang 已提交
70
  MKLDNNMatrixPtr inVal_;
T
tensor-tang 已提交
71
  MKLDNNMatrixPtr inGrad_;
T
tensor-tang 已提交
72
  MKLDNNMatrixPtr outVal_;
T
tensor-tang 已提交
73
  MKLDNNMatrixPtr outGrad_;
74 75 76 77 78 79 80 81 82 83 84 85
  // below are external value and grad
  MKLDNNMatrixPtr extInVal_;
  MKLDNNMatrixPtr extInGrad_;
  MKLDNNMatrixPtr extOutVal_;
  MKLDNNMatrixPtr extOutGrad_;
  // convert handle between external and internal buffers
  std::shared_ptr<mkldnn::reorder> cvtInVal_;
  std::shared_ptr<mkldnn::reorder> cvtInGrad_;
  std::shared_ptr<mkldnn::reorder> cvtOutVal_;
  std::shared_ptr<mkldnn::reorder> cvtOutGrad_;

  // weight and bias are always internal buffers
T
tensor-tang 已提交
86
  MKLDNNMatrixPtr wgtVal_;
T
tensor-tang 已提交
87
  MKLDNNMatrixPtr wgtGrad_;
T
tensor-tang 已提交
88
  MKLDNNMatrixPtr biasVal_;
T
tensor-tang 已提交
89
  MKLDNNMatrixPtr biasGrad_;
T
tensor-tang 已提交
90

T
tensor-tang 已提交
91 92
  // merge grad primitive
  std::shared_ptr<mkldnn::primitive> mergeGrad_;
93
  std::vector<mkldnn::primitive> pipelineMergeGrad_;
T
tensor-tang 已提交
94 95
  // tmp input argument to save input grad, only used to merge grad
  Argument tmpInArg_;
96 97 98 99 100
  // since mkldnn sum do not support different formats:
  // can refer to https://github.com/01org/mkl-dnn/issues/134
  // so need create reorder manually and save tmp MKLDNNMatrix
  MKLDNNMatrixPtr tmpOutGrad_;
  std::shared_ptr<mkldnn::primitive> tmpCvt_;
T
tensor-tang 已提交
101

T
tensor-tang 已提交
102
public:
103
  explicit MKLDNNLayer(const LayerConfig& config)
T
tensor-tang 已提交
104
      : Layer(config),
105
        inputElemenCnt_(0),
T
tensor-tang 已提交
106 107 108 109 110 111 112
        bs_(0),
        ic_(0),
        ih_(0),
        iw_(0),
        oc_(0),
        oh_(0),
        ow_(0),
T
tensor-tang 已提交
113
        needResetBwd_(true),
114
        outputOnlyMKLDNN_(false),
T
tensor-tang 已提交
115
        engine_(mkldnn::engine::cpu, 0),
T
tensor-tang 已提交
116 117 118 119
        stream_(nullptr),
        fwd_(nullptr),
        bwdWgt_(nullptr),
        bwdData_(nullptr) {}
T
tensor-tang 已提交
120

121
  ~MKLDNNLayer() {}
T
tensor-tang 已提交
122

T
tensor-tang 已提交
123
  virtual bool init(const LayerMap& layerMap, const ParameterMap& parameterMap);
T
tensor-tang 已提交
124 125
  virtual void forward(PassType passType);
  virtual void backward(const UpdateCallback& callback);
126 127 128 129

  /**
   * reshape the input image sizes
   * and reset output image and buffer size
130
   * output channel can not be changed
131
   */
132 133
  virtual void reshape(
      int& bs, int& ic, int& ih, int& iw, int oc, int& oh, int& ow) = 0;
134 135

  /**
136
   * reset the mkldnn forward primitve and memories
137 138
   * only would be called when input size changes
   */
139 140 141 142 143
  virtual void resetFwd(std::vector<mkldnn::primitive>& pipeline,
                        MKLDNNMatrixPtr& in,
                        MKLDNNMatrixPtr& wgt,
                        MKLDNNMatrixPtr& bias,
                        MKLDNNMatrixPtr& out) = 0;
144 145

  /**
146
   * reset the mkldnn backward primitve and memories
147 148
   * only would be called when needed
   */
149 150 151 152 153
  virtual void resetBwd(std::vector<mkldnn::primitive>& pipeline,
                        MKLDNNMatrixPtr& in,
                        MKLDNNMatrixPtr& wgt,
                        MKLDNNMatrixPtr& bias,
                        MKLDNNMatrixPtr& out) = 0;
154 155 156 157 158 159

  /**
   * Update weights and biases if necessary.
   */
  virtual void updateWeights(const UpdateCallback& callback) {}

T
tensor-tang 已提交
160 161 162 163
  /**
   * convert weight from paddle format to mkldnn format
   * weight_ will be override
   */
164
  virtual void convertWeightsFromPaddle() {}
T
tensor-tang 已提交
165 166 167 168 169

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

172
  /**
173
   * add this interface as public for unit test
174
   */
175 176 177 178 179 180
  void addOutputArgument(int deviceId) { Layer::addOutputArgument(deviceId); }

protected:
  /**
   * reshape the input image sizes and input batchsize
   */
T
tensor-tang 已提交
181
  void reshapeInput(int& batchsize, int& height, int& width);
182 183 184 185

  /**
   * reshape output image sizes
   */
T
tensor-tang 已提交
186
  void reshapeOutput(size_t height, size_t width);
187

T
tensor-tang 已提交
188
  /**
189 190 191 192 193
   * reset MKLDNNMatrix from Matrix and internal primitive desc.
   * reset nullptr if matrix or primitive desc is empty
   */
  void resetWithMatrix(MKLDNNMatrixPtr& dnn,
                       const MatrixPtr& mat,
T
tensor-tang 已提交
194
                       mkldnn::memory::primitive_desc pd);
195 196 197 198 199 200 201

  /**
   * reset input value from input MKLDNNMatrix and internal primitive desc.
   * reset both internal and external buffer and create reorder if necessary.
   */
  void resetInValue(
      MKLDNNMatrixPtr& in,
202 203
      const std::shared_ptr<mkldnn::memory::primitive_desc>& intPD = nullptr,
      size_t inputIdx = 0);
204 205 206 207 208 209

  /**
   * reset output value from internal primitive desc.
   * reset both internal and external buffer and create reorder if necessary.
   */
  void resetOutValue(MKLDNNMatrixPtr& out,
T
tensor-tang 已提交
210
                     mkldnn::memory::primitive_desc intPD);
211 212 213 214 215

  /**
   * reset input grad from internal primitive desc.
   * reset both internal and external buffer and create reorder if necessary.
   */
216 217 218
  void resetInGrad(MKLDNNMatrixPtr& in,
                   mkldnn::memory::primitive_desc intPD,
                   size_t inputIdx = 0);
219 220 221 222 223

  /**
   * reset output grad from internal primitive desc.
   * merge grad if necessary.
   * reset both internal and external buffer and create reorder if necessary.
T
tensor-tang 已提交
224
   * note: about merge grad, when this layer has several outputs,
T
tensor-tang 已提交
225 226
   *       it could not be mixed with cpu device,
   *       since it can not get memory desc from cpu device.
T
tensor-tang 已提交
227
   */
T
tensor-tang 已提交
228
  void resetOutGrad(MKLDNNMatrixPtr& out, mkldnn::memory::primitive_desc intPD);
229 230 231

  /**
   * reset the merge grad primitive if necessary.
T
tensor-tang 已提交
232
   * note: do not support the grads mixed with cpu device,
233 234
   *       since it can not get memory desc from cpu device.
   */
T
tensor-tang 已提交
235 236 237 238 239 240 241
  void resetMergeGrad(MKLDNNMatrixPtr& out);

protected:
  /**
   * Set deviceId of this layer.
   */
  void setDevice(int id) { deviceId_ = id; }
242

T
tensor-tang 已提交
243 244 245 246 247 248 249 250 251 252
  /**
   * check the format is nchw or nc,
   * which is supported by Paddle default memory layout
   */
  bool isPaddleFormat(mkldnn::memory::format fmt) {
    if (fmt == mkldnn::memory::format::nchw ||
        fmt == mkldnn::memory::format::nc) {
      return true;
    } else {
      return false;
253
    }
T
tensor-tang 已提交
254 255 256 257 258 259 260 261 262 263 264 265 266
  }

  /**
   * If input only has MKLDNN device.
   * Otherwise, only support the previous layer using CPU device.
   */
  bool inputIsOnlyMKLDNN(int index = 0) {
    int prevDevice = getPrev(index)->getDeviceId();
    if (prevDevice == MKLDNN_DEVICE) {
      return true;
    } else {
      CHECK_EQ(prevDevice, CPU_DEVICE) << "Only support CPU yet";
      return false;
267
    }
T
tensor-tang 已提交
268
  }
269

T
tensor-tang 已提交
270 271 272 273 274 275 276 277 278 279 280
  /**
   * If output only has MKLDNN device.
   * Otherwise, other devices should only using CPU device.
   */
  bool outputIsOnlyMKLDNN() {
    for (size_t i = 0; i < outputOtherDevice_.size(); i++) {
      CHECK_EQ(outputOtherDevice_[i].deviceId, CPU_DEVICE)
          << "Only support other device is CPU yet";
    }
    outputOnlyMKLDNN_ = outputOtherDevice_.size() == 0;
    return outputOnlyMKLDNN_;
T
tensor-tang 已提交
281 282
  }

T
tensor-tang 已提交
283 284 285 286 287 288 289 290
  /**
   * 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 已提交
291

292
  /**
293
   * print the mkldnn memory format of value
294
   */
295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312
  virtual void printValueFormat() {
    if (extInVal_) {
      VLOG(MKLDNN_FMTS) << extInVal_->getFormat() << " >>> ";
    }
    if (inVal_) {
      VLOG(MKLDNN_FMTS) << inVal_->getFormat() << " >>>";
    }
    if (outVal_) {
      VLOG(MKLDNN_FMTS) << outVal_->getFormat() << " >>> ";
    }
    if (extOutVal_) {
      VLOG(MKLDNN_FMTS) << extOutVal_->getFormat();
    }
    if (wgtVal_) {
      VLOG(MKLDNN_FMTS) << "Weight value format: " << wgtVal_->getFormat();
    }
    if (biasVal_) {
      VLOG(MKLDNN_FMTS) << "Bias value format: " << biasVal_->getFormat();
313
    }
T
tensor-tang 已提交
314
  }
T
tensor-tang 已提交
315

316
  /**
317
   * print the mkldnn memory format of grad
318
   */
319
  virtual void printGradFormat() {
T
tensor-tang 已提交
320 321
    if (extOutGrad_) {
      VLOG(MKLDNN_FMTS) << extOutGrad_->getFormat();
322 323 324 325
    }
    if (outGrad_) {
      VLOG(MKLDNN_FMTS) << outGrad_->getFormat() << " <<< ";
    }
T
tensor-tang 已提交
326 327 328 329 330
    if (inGrad_) {
      VLOG(MKLDNN_FMTS) << inGrad_->getFormat() << " <<<";
    }
    if (extInGrad_) {
      VLOG(MKLDNN_FMTS) << extInGrad_->getFormat() << " <<< ";
331 332 333 334 335 336
    }
    if (wgtGrad_) {
      VLOG(MKLDNN_FMTS) << "Weight grad format: " << wgtGrad_->getFormat();
    }
    if (biasGrad_) {
      VLOG(MKLDNN_FMTS) << "Bias grad format: " << biasGrad_->getFormat();
337
    }
T
tensor-tang 已提交
338 339
  }

340
private:
341 342 343 344
  /**
   * clear all grad
   */
  void clearGrads() {
T
tensor-tang 已提交
345 346 347
    if (output_.grad) {
      output_.grad->zeroMem();
    }
348
    for (size_t i = 0; i < outputOtherDevice_.size(); i++) {
T
tensor-tang 已提交
349 350 351
      if (outputOtherDevice_[i].grad) {
        outputOtherDevice_[i].grad->zeroMem();
      }
352 353 354
    }
  }

T
tensor-tang 已提交
355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376
  /**
   * 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 已提交
377
  }
378

T
tensor-tang 已提交
379 380 381 382 383 384 385 386 387 388
  /**
   * Set output map of prev layers.
   */
  void setOutputMap() {
    outputMap_.clear();
    for (size_t i = 0; i < inputLayers_.size(); ++i) {
      inputLayers_[i]->setOutput(getName(), &tmpInArg_);
    }
  }

389 390 391 392 393 394 395 396 397 398 399 400 401
  /**
   * if have cpu device, share value and grad data with output_
   */
  void shareCPUDevice() {
    if (outputIsOnlyMKLDNN()) {
      return;
    }
    for (size_t i = 0; i < outputOtherDevice_.size(); i++) {
      outputOtherDevice_[i].value = output_.value;
      outputOtherDevice_[i].grad = output_.grad;
    }
  }

402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436
  /**
   * Check the cpu device number of outputOtherDevice_.
   * should have only one at most.
   */
  void checkCPUOutputsNumber(int max = 1) {
    int cnt = 0;
    for (size_t i = 0; i < outputOtherDevice_.size(); i++) {
      if (outputOtherDevice_[i].deviceId == CPU_DEVICE) {
        ++cnt;
      }
    }
    CHECK_LE(cnt, max) << "too much CPU devies";
  }

  /**
   * copy SeqInfo from input layer to this output and other output devices.
   * @note: do not use getInput(0) since it used this deviceId_,
   *        use "inputLayers_[0]->getOutput()" instead.
   */
  void copySeqInfoToOutputs() {
    if (inputLayers_.empty() || !needSequenceInfo_) {
      return;
    }
    const Argument& input = inputLayers_[0]->getOutput();
    output_.sequenceStartPositions = input.sequenceStartPositions;
    output_.subSequenceStartPositions = input.subSequenceStartPositions;
    output_.cpuSequenceDims = input.cpuSequenceDims;
    for (size_t i = 0; i < outputOtherDevice_.size(); i++) {
      outputOtherDevice_[i].sequenceStartPositions =
          output_.sequenceStartPositions;
      outputOtherDevice_[i].subSequenceStartPositions =
          output_.subSequenceStartPositions;
      outputOtherDevice_[i].cpuSequenceDims = output_.cpuSequenceDims;
    }
  }
T
tensor-tang 已提交
437 438 439
};

}  // namespace paddle