conv_mkldnn_op.cc 43.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License. */

15
#include <unordered_map>
Y
Yu Yang 已提交
16 17
#include "paddle/fluid/framework/data_layout_transform.h"
#include "paddle/fluid/memory/malloc.h"
18
#include "paddle/fluid/operators/conv_op.h"
J
Jacek Czaja 已提交
19
#include "paddle/fluid/platform/mkldnn_reuse.h"
20 21 22 23

namespace paddle {
namespace operators {

24 25 26 27 28 29 30 31
using framework::DataLayout;
using mkldnn::memory;
using mkldnn::primitive;
using mkldnn::reorder;
using mkldnn::stream;
using platform::to_void_cast;
using platform::GetMKLDNNFormat;

Y
Yihua Xu 已提交
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71
inline void GetWeightsTz(std::vector<int>& weights_tz, int groups,  // NOLINT
                         bool is_conv3d) {
  if (groups > 1) {
    if (is_conv3d) {
      int output = weights_tz[0];
      int input = weights_tz[1];
      int dimension = weights_tz[2];
      int height = weights_tz[3];
      int width = weights_tz[4];
      weights_tz.resize(6);
      weights_tz[0] = groups;
      weights_tz[1] = output / groups;
      weights_tz[2] = input;
      weights_tz[3] = dimension;
      weights_tz[4] = height;
      weights_tz[5] = width;
    } else {
      int output = weights_tz[0];
      int input = weights_tz[1];
      int height = weights_tz[2];
      int width = weights_tz[3];
      weights_tz.resize(5);
      weights_tz[0] = groups;
      weights_tz[1] = output / groups;
      weights_tz[2] = input;
      weights_tz[3] = height;
      weights_tz[4] = width;
    }
  }
}

inline mkldnn::memory::format GetWeightsFormat(mkldnn::memory::format format,
                                               int groups, bool is_conv3d) {
  if (is_conv3d) {
    return (groups == 1) ? format : mkldnn::memory::format::goidhw;
  } else {
    return (groups == 1) ? format : mkldnn::memory::format::goihw;
  }
}

72
template <typename T, typename K>
73
class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
74 75 76 77
 public:
  void Compute(const paddle::framework::ExecutionContext& ctx) const override {
    PADDLE_ENFORCE(paddle::platform::is_cpu_place(ctx.GetPlace()),
                   "It must use CPUPlace.");
78 79 80 81 82 83 84 85
    bool is_INT8 =
        std::is_same<T, int8_t>::value || std::is_same<T, uint8_t>::value;
    if (!is_INT8) {
      ComputeFP32(ctx);
    } else {
      ComputeINT8(ctx);
    }
  }
86

87
  void ComputeFP32(const paddle::framework::ExecutionContext& ctx) const {
K
Krzysztof Binias 已提交
88 89
    const bool is_test = ctx.Attr<bool>("is_test");

90 91
    auto& dev_ctx =
        ctx.template device_context<paddle::platform::MKLDNNDeviceContext>();
92 93 94 95
    const auto& mkldnn_engine = dev_ctx.GetEngine();

    auto* input = ctx.Input<Tensor>("Input");
    auto* filter = ctx.Input<Tensor>("Filter");
96
    auto* bias = ctx.HasInput("Bias") ? ctx.Input<Tensor>("Bias") : nullptr;
97 98
    auto* output = ctx.Output<Tensor>("Output");

99 100 101 102 103 104
    PADDLE_ENFORCE(input->layout() == DataLayout::kMKLDNN &&
                       input->format() != memory::format::format_undef,
                   "Wrong layout/format set for Input tensor");
    PADDLE_ENFORCE(filter->layout() == DataLayout::kMKLDNN &&
                       filter->format() != memory::format::format_undef,
                   "Wrong layout/format set for Filter tensor");
105
    PADDLE_ENFORCE(input->dims().size() == 4 || input->dims().size() == 5,
Y
Yihua Xu 已提交
106
                   "Input must be with 4 or 5 dimensions, i.e. NCHW or NCDHW");
107 108
    PADDLE_ENFORCE(filter->dims().size() == 4 || filter->dims().size() == 5,
                   "Filter must be with 4 or 5 dimensions, i.e. OIHW or OIDHW");
109 110 111 112 113 114 115
    if (bias) {
      PADDLE_ENFORCE(bias->layout() == DataLayout::kMKLDNN &&
                         bias->format() != memory::format::format_undef,
                     "Wrong layout/format set for Bias tensor");
      PADDLE_ENFORCE(bias->dims().size() == 1,
                     "Bias must only have 1 dimension, i.e. X");
    }
116 117 118 119

    std::vector<int> strides = ctx.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = ctx.Attr<std::vector<int>>("paddings");
    std::vector<int> dilations = ctx.Attr<std::vector<int>>("dilations");
M
Michal Gallus 已提交
120
    bool fuse_relu = ctx.Attr<bool>("fuse_relu");
121
    bool fuse_residual_conn = ctx.Attr<bool>("fuse_residual_connection");
122 123
    int groups = ctx.Attr<int>("groups");

124
    bool is_conv3d = strides.size() == 3U;
125
    // TODO(tpatejko): add support for dilation
126
    PADDLE_ENFORCE(
127 128 129 130
        is_conv3d
            ? dilations.size() == 3 && dilations[0] == 1 && dilations[1] == 1 &&
                  dilations[2] == 1
            : dilations.size() == 2 && dilations[0] == 1 && dilations[1] == 1,
131 132 133 134 135 136 137 138
        "dilation in convolution is not implemented yet");

    const T* input_data = input->data<T>();
    const T* filter_data = filter->data<T>();

    std::vector<int> src_tz = paddle::framework::vectorize2int(input->dims());
    std::vector<int> weights_tz =
        paddle::framework::vectorize2int(filter->dims());
139
    int g = std::max(groups, 1);
Y
Yihua Xu 已提交
140
    GetWeightsTz(weights_tz, g, is_conv3d);
141 142
    std::vector<int> dst_tz = paddle::framework::vectorize2int(output->dims());

143
    // Get unique name for storing MKLDNN primitives
J
Jacek Czaja 已提交
144
    const std::string key = platform::ConvMKLDNNHandler::GetHash(
145
        src_tz, weights_tz, strides, paddings, dilations, groups,
146
        ctx.op().Input("Input") + ctx.op().Input("Filter"));
147 148 149 150
    const std::string key_conv_pd = key + "@conv_pd";

    std::vector<primitive> pipeline;

151 152
    auto src_format = input->format();
    mkldnn::memory::format weights_format =
Y
Yihua Xu 已提交
153
        GetWeightsFormat(filter->format(), g, is_conv3d);
154

155
    auto user_src_md = platform::MKLDNNMemDesc(
156
        {src_tz}, platform::MKLDNNGetDataType<T>(), src_format);
157
    auto user_weights_md = platform::MKLDNNMemDesc(
158
        {weights_tz}, platform::MKLDNNGetDataType<T>(), weights_format);
159 160 161 162 163

    /* create memory descriptor for convolution without specified format
     * ('any') which lets a primitive (convolution in this case) choose
     * the memory format preferred for best performance
     */
164 165 166 167
    std::string data_format = ctx.Attr<std::string>("data_format");
    auto chosen_memory_format =
        platform::data_format_to_memory_format(data_format);

168 169 170 171 172 173 174
    weights_format = mkldnn::memory::format::any;
    // Check the format for user's special output
    if (chosen_memory_format != mkldnn::memory::format::any) {
      if (is_conv3d) {
        chosen_memory_format =
            platform::MKLDNNFormatForSize(src_tz.size(), chosen_memory_format);
      }
175 176
    }

177
    auto src_md = platform::MKLDNNMemDesc(
178
        src_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
179
    auto weights_md = platform::MKLDNNMemDesc(
180
        weights_tz, platform::MKLDNNGetDataType<T>(), weights_format);
181 182
    std::vector<int> bias_tz;  // TODO(mgallus): avoid empty vector creation.
                               // Currently used whenever bias is != nullptr.
183
    auto dst_md = platform::MKLDNNMemDesc(
184
        dst_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
185 186

    // create a conv primitive descriptor and save it for usage in backward
187
    std::shared_ptr<mkldnn::convolution_forward::primitive_desc> conv_pd;
188 189
    auto fwd_prop_kind = is_test ? mkldnn::prop_kind::forward_inference
                                 : mkldnn::prop_kind::forward_training;
190 191 192 193
    if (bias) {
      bias_tz = paddle::framework::vectorize2int(bias->dims());
      auto bias_md = platform::MKLDNNMemDesc(
          bias_tz, platform::MKLDNNGetDataType<T>(), memory::format::x);
194 195 196
      conv_pd = ConvFwdPrimitiveDesc(
          src_md, weights_md, bias_md, dst_md, strides, paddings, mkldnn_engine,
          fuse_relu, fuse_residual_conn, fwd_prop_kind);
197
    } else {
198 199 200
      conv_pd = ConvFwdPrimitiveDesc(src_md, weights_md, dst_md, strides,
                                     paddings, mkldnn_engine, fuse_relu,
                                     fuse_residual_conn, fwd_prop_kind);
201
    }
202
    // Save conv_pd/src_memory/weights_memory for backward pass
203
    if (!is_test) dev_ctx.SetBlob(key_conv_pd, conv_pd);
204

J
Jacek Czaja 已提交
205
    platform::ConvMKLDNNHandler handler(conv_pd, dev_ctx, mkldnn_engine, key);
206

207 208 209 210 211 212
    // create mkldnn memory from input tensors (data/weights)
    auto user_src_memory_p =
        handler.AcquireSrcMemory(user_src_md, to_void_cast<T>(input_data));
    auto user_weights_memory_p = handler.AcquireWeightsMemory(
        user_weights_md, to_void_cast<T>(filter_data));

213 214 215 216 217
    // create reorder primitive if the input format is not the preferred one
    auto src_memory_p =
        handler.AcquireSrcMemoryFromPrimitive(user_src_memory_p, pipeline);
    auto weights_memory_p = handler.AcquireWeightsMemoryFromPrimitive(
        user_weights_memory_p, pipeline, is_test);
218 219

    std::shared_ptr<mkldnn::memory> dst_memory_p;
220

221
    if (fuse_residual_conn) {
222 223
      auto residual_param = ctx.Input<Tensor>("ResidualData");
      auto residual_param_data = residual_param->data<T>();
224

225 226 227 228 229 230
      PADDLE_ENFORCE(
          residual_param_data != nullptr,
          "Provide data if you want MKLDNN conv+elementwise_add fusion");
      PADDLE_ENFORCE_EQ(output->dims(), residual_param->dims(),
                        "Output and elementwise parameter need to have the "
                        "same dimension sizes");
231

232
      if (residual_param->format() != handler.GetDstFormat()) {
Y
Yu Yang 已提交
233 234 235
        auto output_data = output->mutable_data<T>(
            ctx.GetPlace(), ::paddle::memory::Allocator::kDefault,
            handler.GetDstMemorySize());
236 237 238 239 240 241 242 243 244
        auto residual_data_tz =
            paddle::framework::vectorize2int(residual_param->dims());
        auto residual_data_type =
            paddle::framework::ToMKLDNNDataType(residual_param->type());

        auto user_residual_md = platform::MKLDNNMemDesc(
            residual_data_tz, residual_data_type, residual_param->format());
        auto user_residual_memory_p = handler.AcquireResidualDataMemory(
            user_residual_md, to_void_cast<T>(residual_param_data));
245 246 247

        dst_memory_p = handler.AcquireDstMemoryFromResidualDataMemory(
            user_residual_memory_p, to_void_cast<T>(output_data), pipeline);
248 249
      } else {
        output->ShareDataWith(*residual_param);
250 251 252
        auto output_data = output->mutable_data<T>(ctx.GetPlace());
        dst_memory_p =
            handler.AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));
253
      }
254
    } else {
255 256 257
      auto output_data = output->mutable_data<T>(
          ctx.GetPlace(), paddle::memory::Allocator::kDefault,
          handler.GetDstMemorySize());
258 259
      dst_memory_p =
          handler.AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));
260
    }
261 262

    // create convolution op primitive
263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278
    std::shared_ptr<mkldnn::convolution_forward> conv_p;
    if (bias) {
      const T* bias_data = bias->data<T>();
      auto user_bias_md = platform::MKLDNNMemDesc(
          {bias_tz}, platform::MKLDNNGetDataType<T>(), memory::format::x);
      auto user_bias_memory_p =
          handler.AcquireBiasMemory(user_bias_md, to_void_cast<T>(bias_data));

      auto bias_memory_p =
          handler.AcquireBiasMemoryFromPrimitive(user_bias_memory_p, pipeline);
      conv_p = handler.AcquireConvolution(src_memory_p, weights_memory_p,
                                          bias_memory_p, dst_memory_p);
    } else {
      conv_p = handler.AcquireConvolution(src_memory_p, weights_memory_p,
                                          dst_memory_p);
    }
279 280

    // push primitive to stream and wait until it's executed
281
    pipeline.push_back(*conv_p);
282 283 284
    stream(stream::kind::eager).submit(pipeline).wait();

    output->set_layout(DataLayout::kMKLDNN);
285
    output->set_format(GetMKLDNNFormat(*dst_memory_p));
286
  }
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320
  void ComputeINT8(const paddle::framework::ExecutionContext& ctx) const {
    const bool is_test = ctx.Attr<bool>("is_test");

    auto& dev_ctx =
        ctx.template device_context<paddle::platform::MKLDNNDeviceContext>();
    const auto& mkldnn_engine = dev_ctx.GetEngine();

    auto* input = ctx.Input<Tensor>("Input");
    auto* filter = ctx.Input<Tensor>("Filter");
    auto* bias = ctx.HasInput("Bias") ? ctx.Input<Tensor>("Bias") : nullptr;
    auto* output = ctx.Output<Tensor>("Output");

    PADDLE_ENFORCE(input->layout() == DataLayout::kMKLDNN &&
                       input->format() != memory::format::format_undef,
                   "Wrong layout/format set for Input tensor");
    PADDLE_ENFORCE(filter->layout() == DataLayout::kMKLDNN &&
                       filter->format() != memory::format::format_undef,
                   "Wrong layout/format set for Filter tensor");
    PADDLE_ENFORCE(input->dims().size() == 4 || input->dims().size() == 5,
                   "Input must be with 4 or 5 dimensions, i.e. NCHW or NCDHW");
    PADDLE_ENFORCE(filter->dims().size() == 4 || filter->dims().size() == 5,
                   "Filter must be with 4 or 5 dimensions, i.e. OIHW or OIDHW");
    if (bias) {
      PADDLE_ENFORCE(bias->layout() == DataLayout::kMKLDNN &&
                         bias->format() != memory::format::format_undef,
                     "Wrong layout/format set for Bias tensor");
      PADDLE_ENFORCE(bias->dims().size() == 1,
                     "Bias must only have 1 dimension, i.e. X");
    }

    std::vector<int> strides = ctx.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = ctx.Attr<std::vector<int>>("paddings");
    std::vector<int> dilations = ctx.Attr<std::vector<int>>("dilations");
    int groups = ctx.Attr<int>("groups");
X
xiaolil1 已提交
321
    bool fuse_relu = ctx.Attr<bool>("fuse_relu");
X
xiaolil1 已提交
322
    bool fuse_residual_conn = ctx.Attr<bool>("fuse_residual_connection");
X
xiaolil1 已提交
323

324
    bool force_fp32_output = ctx.Attr<bool>("force_fp32_output");
X
xiaolil1 已提交
325 326 327 328
    if (fuse_residual_conn) {
      PADDLE_ENFORCE(force_fp32_output != true,
                     "residual fusion does not support force output with fp32");
    }
329 330 331 332 333 334 335 336 337

    bool is_conv3d = strides.size() == 3U;
    // TODO(tpatejko): add support for dilation
    PADDLE_ENFORCE(
        is_conv3d
            ? dilations.size() == 3 && dilations[0] == 1 && dilations[1] == 1 &&
                  dilations[2] == 1
            : dilations.size() == 2 && dilations[0] == 1 && dilations[1] == 1,
        "dilation in convolution is not implemented yet");
X
xiaolil1 已提交
338

339 340 341 342 343 344 345 346 347 348 349
    PADDLE_ENFORCE(is_conv3d != true, "int8 does not support conv3d currently");

    const T* input_data = input->data<T>();

    std::vector<int> src_tz = paddle::framework::vectorize2int(input->dims());
    std::vector<int> weights_tz =
        paddle::framework::vectorize2int(filter->dims());
    int g = std::max(groups, 1);
    GetWeightsTz(weights_tz, g, is_conv3d);
    std::vector<int> dst_tz = paddle::framework::vectorize2int(output->dims());

X
xiaolil1 已提交
350 351 352 353 354 355 356 357 358 359 360 361
    mkldnn::memory::data_type src_dt =
        paddle::framework::ToMKLDNNDataType(input->type());
    auto dst_dt = fuse_relu ? paddle::framework::ToMKLDNNDataType(
                                  framework::DataTypeTrait<uint8_t>::DataType)
                            : paddle::framework::ToMKLDNNDataType(
                                  framework::DataTypeTrait<int8_t>::DataType);

    if (force_fp32_output) {
      dst_dt = paddle::framework::ToMKLDNNDataType(
          framework::DataTypeTrait<float>::DataType);
    }

X
xiaolil1 已提交
362 363 364 365 366 367
    if (fuse_residual_conn) {
      auto residual = ctx.Input<Tensor>("ResidualData");
      auto residual_dt = paddle::framework::ToMKLDNNDataType(residual->type());
      if (dst_dt != residual_dt) dst_dt = residual_dt;
    }

368 369 370 371 372
    // Get unique name for storing MKLDNN primitives
    std::string key;
    key.reserve(MaxKeyLength);
    platform::ConvMKLDNNHandler::AppendKey(
        &key, src_tz, weights_tz, strides, paddings, dilations, groups, src_dt,
X
xiaolil1 已提交
373
        input->format(), fuse_relu, fuse_residual_conn,
374
        ctx.op().Input("Input") + ctx.op().Input("Filter"));
375 376
    const std::string key_conv_pd = key + "@conv_pd";

X
xiaolil1 已提交
377 378
    bool need_s8_to_u8 = false;

379 380 381 382 383 384 385 386 387 388 389 390 391 392
    std::shared_ptr<mkldnn::convolution_forward> conv_p = nullptr;
    std::shared_ptr<mkldnn::memory> src_memory_p = nullptr;
    std::shared_ptr<mkldnn::memory> user_src_memory_p = nullptr;
    std::shared_ptr<mkldnn::memory> dst_memory_p = nullptr;
    std::vector<primitive> pipeline;
    std::shared_ptr<mkldnn::convolution_forward::primitive_desc> conv_pd =
        nullptr;
    std::shared_ptr<platform::ConvMKLDNNHandler> handler = nullptr;

    auto prim_key = key + "@conv_p";
    auto dst_key = key + "@dst_mem_p";
    auto src_key = key + "@src_mem_p";
    auto user_src_key = key + "@user_src_mem_p";
    auto src_reorder_key = key + "@src_mem_preorder_p";
X
xiaolil1 已提交
393 394
    auto residual_reorder_key = key + "@residual_data_mem_preorder_p";

395 396
    conv_p = std::static_pointer_cast<mkldnn::convolution_forward>(
        dev_ctx.GetBlob(prim_key));
X
xiaolil1 已提交
397

398 399 400
    if (conv_p == nullptr || !is_test) {
      const K* filter_data = filter->data<K>();
      auto scale_in_data = ctx.Attr<float>("Scale_in");
X
xiaolil1 已提交
401
      auto scale_in_eltwise_data = ctx.Attr<float>("Scale_in_eltwise");
402 403 404
      auto scale_weights_data = ctx.Attr<std::vector<float>>("Scale_weights");
      auto scale_out_data =
          force_fp32_output ? 1.0f : ctx.Attr<float>("Scale_out");
X
xiaolil1 已提交
405 406
      float sum_scale =
          fuse_residual_conn ? scale_out_data / scale_in_eltwise_data : 1.0f;
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 437 438 439 440 441 442 443 444 445 446 447 448

      bool is_multi_channel = scale_weights_data.size() > 1;

      int count = is_multi_channel ? (g > 1 ? (weights_tz)[1] * (weights_tz)[0]
                                            : (weights_tz)[0])
                                   : 1;
      std::vector<float> output_shift_scale(count);
#pragma omp parallel for if (count > 1)
      for (int i = 0; i < count; i++) {
        if (scale_weights_data[i] == 0.0)
          output_shift_scale[i] =
              scale_out_data;  // weights data will contain 0
                               // in some models, then weights
                               // scale couldn't be calculated
        else
          output_shift_scale[i] =
              scale_out_data / (scale_in_data * scale_weights_data[i]);
      }

      auto user_src_md =
          platform::MKLDNNMemDesc({src_tz}, src_dt, input->format());
      auto user_weights_md = platform::MKLDNNMemDesc(
          {weights_tz}, platform::MKLDNNGetDataType<K>(),
          ((g) == 1) ? mkldnn::memory::format::oihw
                     : mkldnn::memory::format::goihw);

      /* create memory descriptor for convolution without specified format
      * ('any') which lets a primitive (convolution in this case) choose
      * the memory format preferred for best performance
      */
      std::string data_format = ctx.Attr<std::string>("data_format");
      auto chosen_memory_format =
          platform::data_format_to_memory_format(data_format);

      std::vector<int> bias_tz;

      auto src_md =
          platform::MKLDNNMemDesc(src_tz, src_dt, chosen_memory_format);
      auto weights_md = platform::MKLDNNMemDesc(
          weights_tz, memory::data_type::s8, chosen_memory_format);
      auto dst_md =
          platform::MKLDNNMemDesc(dst_tz, dst_dt, chosen_memory_format);
X
xiaolil1 已提交
449

450 451 452 453 454 455 456
      // create a conv primitive descriptor and save it for usage in backward
      if (bias) {
        bias_tz = paddle::framework::vectorize2int(bias->dims());
        auto bias_md = platform::MKLDNNMemDesc(bias_tz, memory::data_type::s32,
                                               memory::format::x);
        conv_pd = ConvFwdPrimitiveDesc(src_md, weights_md, bias_md, dst_md,
                                       strides, paddings, mkldnn_engine,
X
xiaolil1 已提交
457 458
                                       fuse_relu, fuse_residual_conn,
                                       output_shift_scale, sum_scale, is_test);
459
      } else {
X
xiaolil1 已提交
460 461 462 463
        conv_pd =
            ConvFwdPrimitiveDesc(src_md, weights_md, dst_md, strides, paddings,
                                 mkldnn_engine, fuse_relu, fuse_residual_conn,
                                 output_shift_scale, sum_scale, is_test);
464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487
      }
      // Save conv_pd/src_memory/weights_memory for backward pass
      dev_ctx.SetBlob(key_conv_pd, conv_pd);

      handler.reset(new platform::ConvMKLDNNHandler(conv_pd, dev_ctx,
                                                    mkldnn_engine, key));

      // create mkldnn memory from input tensors (data/weights)
      user_src_memory_p =
          handler->AcquireSrcMemory(user_src_md, to_void_cast<T>(input_data));
      auto user_weights_memory_p = handler->AcquireWeightsMemory(
          user_weights_md, to_void_cast<K>(filter_data));

      // create reorder primitive if the input format is not the preferred one
      src_memory_p =
          handler->AcquireSrcMemoryFromPrimitive(user_src_memory_p, pipeline);

      std::shared_ptr<mkldnn::memory> weights_memory_p;
      int mask_reorder =
          is_multi_channel ? ((g != 1) ? (1 << 1) + (1 << 0) : 1 << 0) : 0;
      weights_memory_p = handler->AcquireWeightsMemoryFromPrimitive(
          user_weights_memory_p, pipeline, is_test, true, scale_weights_data,
          mask_reorder);

X
xiaolil1 已提交
488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522
      if (fuse_residual_conn) {
        auto residual_param = ctx.Input<Tensor>("ResidualData");
        PADDLE_ENFORCE_EQ(output->dims(), residual_param->dims(),
                          "Output and elementwise parameter need to have the "
                          "same dimension sizes");
        auto residual_dt =
            paddle::framework::ToMKLDNNDataType(residual_param->type());
        if (residual_param->format() != handler->GetDstFormat()) {
          auto residual_data_tz =
              paddle::framework::vectorize2int(residual_param->dims());

          auto user_residual_md = platform::MKLDNNMemDesc(
              residual_data_tz, residual_dt, residual_param->format());

          if (residual_dt == mkldnn::memory::data_type::u8) {
            dst_memory_p = platform::SetDstMemory<uint8_t>(
                ctx, output, residual_param, user_residual_md, handler,
                &pipeline);
          } else {
            need_s8_to_u8 = fuse_relu;
            dst_memory_p = platform::SetDstMemory<int8_t>(
                ctx, output, residual_param, user_residual_md, handler,
                &pipeline);
          }
        } else {
          output->ShareDataWith(*residual_param);
          if (residual_dt == mkldnn::memory::data_type::u8) {
            dst_memory_p =
                platform::SetDstMemory<uint8_t>(ctx, output, handler);
          } else {
            need_s8_to_u8 = fuse_relu;
            dst_memory_p = platform::SetDstMemory<int8_t>(ctx, output, handler);
          }
        }
      } else if (!force_fp32_output) {
X
xiaolil1 已提交
523 524 525 526 527
        if (fuse_relu) {
          dst_memory_p = platform::SetDstMemory<uint8_t>(ctx, output, handler);
        } else {
          dst_memory_p = platform::SetDstMemory<int8_t>(ctx, output, handler);
        }
528 529 530 531 532 533 534
      } else {
        dst_memory_p = platform::SetDstMemory<float>(ctx, output, handler);
      }

      // create convolution op primitive
      auto scale_bias_key = key + "@scale_bias";
      if (bias) {
X
xiaolil1 已提交
535
        const K* bias_data = bias->data<K>();
536
        auto user_bias_md = platform::MKLDNNMemDesc(
X
xiaolil1 已提交
537
            {bias_tz}, platform::MKLDNNGetDataType<K>(), memory::format::x);
538
        auto user_bias_memory_p = handler->AcquireBiasMemory(
X
xiaolil1 已提交
539
            user_bias_md, to_void_cast<K>(bias_data));
540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584
        std::shared_ptr<mkldnn::memory> bias_memory_p;
        int mask_reorder = is_multi_channel ? 1 << 0 : 1;
        int count =
            is_multi_channel
                ? (g > 1 ? (weights_tz)[1] * (weights_tz)[0] : (weights_tz)[0])
                : 1;
        std::vector<float> scale_bias_data(count);
#pragma omp parallel for if (count > 1)
        for (int i = 0; i < count; i++) {
          scale_bias_data[i] = scale_in_data * scale_weights_data[i];
        }
        bias_memory_p = handler->AcquireBiasMemoryFromPrimitive(
            user_bias_memory_p, pipeline, is_test, true, scale_bias_data,
            mask_reorder);
        conv_p = handler->AcquireConvolution(src_memory_p, weights_memory_p,
                                             bias_memory_p, dst_memory_p);
      } else {
        conv_p = handler->AcquireConvolution(src_memory_p, weights_memory_p,
                                             dst_memory_p);
      }

      // push primitive to stream and wait until it's executed
      pipeline.push_back(*conv_p);
    } else {
      auto src_memory_reorder_p = std::static_pointer_cast<mkldnn::memory>(
          dev_ctx.GetBlob(src_reorder_key));
      src_memory_p =
          std::static_pointer_cast<mkldnn::memory>(dev_ctx.GetBlob(src_key));
      if (src_memory_reorder_p) {
        user_src_memory_p = std::static_pointer_cast<mkldnn::memory>(
            dev_ctx.GetBlob(user_src_key));
        user_src_memory_p->set_data_handle(to_void_cast<T>(input_data));
      } else if (src_memory_p) {
        src_memory_p->set_data_handle(to_void_cast<T>(input_data));
      }

      dst_memory_p =
          std::static_pointer_cast<mkldnn::memory>(dev_ctx.GetBlob(dst_key));
      conv_pd =
          std::static_pointer_cast<mkldnn::convolution_forward::primitive_desc>(
              dev_ctx.GetBlob(key_conv_pd));
      if (conv_pd) {
        handler.reset(new platform::ConvMKLDNNHandler(conv_pd, dev_ctx,
                                                      mkldnn_engine, key));
      }
X
xiaolil1 已提交
585 586 587 588 589 590 591 592 593 594 595 596 597 598

      if (fuse_residual_conn) {
        auto residual_param = ctx.Input<Tensor>("ResidualData");
        auto residual_dt =
            paddle::framework::ToMKLDNNDataType(residual_param->type());
        output->ShareDataWith(*residual_param);
        if (residual_dt == mkldnn::memory::data_type::u8) {
          platform::SetDstMemoryHandler<uint8_t>(ctx, output, handler,
                                                 &dst_memory_p);
        } else {
          platform::SetDstMemoryHandler<int8_t>(ctx, output, handler,
                                                &dst_memory_p);
        }
      } else if (!force_fp32_output) {
X
xiaolil1 已提交
599
        if (fuse_relu) {
X
xiaolil1 已提交
600 601
          platform::SetDstMemoryHandler<uint8_t>(ctx, output, handler,
                                                 &dst_memory_p);
X
xiaolil1 已提交
602
        } else {
X
xiaolil1 已提交
603 604
          platform::SetDstMemoryHandler<int8_t>(ctx, output, handler,
                                                &dst_memory_p);
X
xiaolil1 已提交
605
        }
606
      } else {
X
xiaolil1 已提交
607 608
        platform::SetDstMemoryHandler<float>(ctx, output, handler,
                                             &dst_memory_p);
609
      }
X
xiaolil1 已提交
610

611 612 613
      if (src_memory_reorder_p) {
        pipeline.push_back(*src_memory_reorder_p);
      }
X
xiaolil1 已提交
614 615 616 617 618 619 620

      auto residual_reorder_p = std::static_pointer_cast<mkldnn::memory>(
          dev_ctx.GetBlob(residual_reorder_key));
      if (residual_reorder_p) {
        pipeline.push_back(*residual_reorder_p);
      }

621 622 623 624 625
      pipeline.push_back(*conv_p);
    }
    // push primitive to stream and wait until it's executed
    stream(stream::kind::eager).submit(pipeline).wait();

X
xiaolil1 已提交
626 627 628 629
    if (need_s8_to_u8) {
      output->mutable_data<uint8_t>(ctx.GetPlace());
    }

630 631 632
    output->set_layout(DataLayout::kMKLDNN);
    output->set_format(GetMKLDNNFormat(*dst_memory_p));
  }
633

634
 private:
635
  mkldnn::primitive_attr CreatePostOps(bool fuse_relu,
636
                                       bool fuse_residual_conn) const {
M
Michal Gallus 已提交
637 638
    mkldnn::primitive_attr conv_attr;
    mkldnn::post_ops post_operations;
639
    // Fusion with Elementwise layer relies on adding a sum post-operation with
640 641 642 643 644
    // 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) {
645 646 647 648 649 650 651 652 653 654 655
      post_operations.append_sum(1.0f);
    }
    // Fusion with ReLU layer is executed through the PostOps feature. Create a
    // PostOps object and configure it to execute an eltwise relu operation.
    if (fuse_relu) {
      constexpr float scale = 1.0f;
      constexpr float negative_slope = 0.0f;
      constexpr float placeholder = 0.0f;
      post_operations.append_eltwise(scale, mkldnn::algorithm::eltwise_relu,
                                     negative_slope, placeholder);
    }
M
Michal Gallus 已提交
656 657 658 659
    conv_attr.set_post_ops(post_operations);
    return conv_attr;
  }

660
  mkldnn::primitive_attr CreatePostOps(
X
xiaolil1 已提交
661 662
      bool fuse_relu, bool fuse_residual_conn,
      const std::vector<float> output_shift_scale, float sum_scale) const {
663 664 665 666
    mkldnn::primitive_attr conv_attr;
    mkldnn::post_ops post_operations;
    int mask = output_shift_scale.size() > 1 ? 1 << 1 : 0;
    conv_attr.set_output_scales(mask, output_shift_scale);
X
xiaolil1 已提交
667 668 669
    if (fuse_residual_conn) {
      post_operations.append_sum(sum_scale);
    }
X
xiaolil1 已提交
670 671 672 673 674 675 676
    if (fuse_relu) {
      constexpr float scale = 1.0f;
      constexpr float negative_slope = 0.0f;
      constexpr float placeholder = 1.0f;  // beta
      post_operations.append_eltwise(scale, mkldnn::algorithm::eltwise_relu,
                                     negative_slope, placeholder);
    }
677 678 679 680
    conv_attr.set_post_ops(post_operations);
    return conv_attr;
  }

681 682 683 684
  std::unique_ptr<mkldnn::convolution_forward::primitive_desc>
  ConvFwdPrimitiveDesc(const memory::desc& src, const memory::desc& weights,
                       const memory::desc& dst, const std::vector<int>& strides,
                       const std::vector<int>& paddings,
685
                       const mkldnn::engine& engine, const bool fuse_relu,
686 687
                       const bool fuse_residual_conn,
                       mkldnn::prop_kind fwd_prop_kind) const {
688 689
    memory::dims stride_dims = strides;
    memory::dims padding_dims = paddings;
690

691
    auto conv_desc = mkldnn::convolution_forward::desc(
692 693
        fwd_prop_kind, mkldnn::convolution_direct, src, weights, dst,
        stride_dims, padding_dims, padding_dims, mkldnn::padding_kind::zero);
694

695 696
    mkldnn::primitive_attr conv_attr =
        CreatePostOps(fuse_relu, fuse_residual_conn);
M
Michal Gallus 已提交
697 698 699

    auto p_conv_pd = new mkldnn::convolution_forward::primitive_desc(
        conv_desc, conv_attr, engine);
700

701 702
    return std::unique_ptr<mkldnn::convolution_forward::primitive_desc>(
        p_conv_pd);
703
  }
704

705 706 707 708
  std::unique_ptr<mkldnn::convolution_forward::primitive_desc>
  ConvFwdPrimitiveDesc(const memory::desc& src, const memory::desc& weights,
                       const memory::desc& dst, const std::vector<int>& strides,
                       const std::vector<int>& paddings,
X
xiaolil1 已提交
709
                       const mkldnn::engine& engine, const bool fuse_relu,
X
xiaolil1 已提交
710
                       const bool fuse_residual_conn,
711
                       const std::vector<float> output_shift_scale,
X
xiaolil1 已提交
712
                       const float sum_scale, bool is_test) const {
713 714 715 716 717 718 719 720 721 722
    memory::dims stride_dims = {strides[0], strides[1]};
    memory::dims padding_dims = {paddings[0], paddings[1]};

    auto propagation = is_test ? mkldnn::prop_kind::forward_scoring
                               : mkldnn::prop_kind::forward_training;

    auto conv_desc = mkldnn::convolution_forward::desc(
        propagation, mkldnn::convolution_direct, src, weights, dst, stride_dims,
        padding_dims, padding_dims, mkldnn::padding_kind::zero);

X
xiaolil1 已提交
723 724
    mkldnn::primitive_attr conv_attr = CreatePostOps(
        fuse_relu, fuse_residual_conn, output_shift_scale, sum_scale);
725 726 727 728 729 730 731 732

    auto p_conv_pd = new mkldnn::convolution_forward::primitive_desc(
        conv_desc, conv_attr, engine);

    return std::unique_ptr<mkldnn::convolution_forward::primitive_desc>(
        p_conv_pd);
  }

733 734 735 736 737
  std::unique_ptr<mkldnn::convolution_forward::primitive_desc>
  ConvFwdPrimitiveDesc(const memory::desc& src, const memory::desc& weights,
                       const memory::desc& bias, const memory::desc& dst,
                       const std::vector<int>& strides,
                       const std::vector<int>& paddings,
738
                       const mkldnn::engine& engine, const bool fuse_relu,
739 740
                       const bool fuse_residual_conn,
                       mkldnn::prop_kind fwd_prop_kind) const {
741 742
    memory::dims stride_dims = strides;
    memory::dims padding_dims = paddings;
743 744

    auto conv_desc = mkldnn::convolution_forward::desc(
745 746
        fwd_prop_kind, mkldnn::convolution_direct, src, weights, bias, dst,
        stride_dims, padding_dims, padding_dims, mkldnn::padding_kind::zero);
747

748 749
    mkldnn::primitive_attr conv_attr =
        CreatePostOps(fuse_relu, fuse_residual_conn);
M
Michal Gallus 已提交
750 751 752

    auto p_conv_pd = new mkldnn::convolution_forward::primitive_desc(
        conv_desc, conv_attr, engine);
753 754 755 756

    return std::unique_ptr<mkldnn::convolution_forward::primitive_desc>(
        p_conv_pd);
  }
757 758 759 760 761 762

  std::unique_ptr<mkldnn::convolution_forward::primitive_desc>
  ConvFwdPrimitiveDesc(const memory::desc& src, const memory::desc& weights,
                       const memory::desc& bias, const memory::desc& dst,
                       const std::vector<int>& strides,
                       const std::vector<int>& paddings,
X
xiaolil1 已提交
763
                       const mkldnn::engine& engine, const bool fuse_relu,
X
xiaolil1 已提交
764
                       const bool fuse_residual_conn,
765
                       const std::vector<float> output_shift_scale,
X
xiaolil1 已提交
766
                       const float sum_scale, bool is_test) const {
767 768 769 770 771 772 773 774 775 776
    memory::dims stride_dims = {strides[0], strides[1]};
    memory::dims padding_dims = {paddings[0], paddings[1]};

    auto propagation = is_test ? mkldnn::prop_kind::forward_scoring
                               : mkldnn::prop_kind::forward_training;

    auto conv_desc = mkldnn::convolution_forward::desc(
        propagation, mkldnn::convolution_direct, src, weights, bias, dst,
        stride_dims, padding_dims, padding_dims, mkldnn::padding_kind::zero);

X
xiaolil1 已提交
777 778
    mkldnn::primitive_attr conv_attr = CreatePostOps(
        fuse_relu, fuse_residual_conn, output_shift_scale, sum_scale);
779 780 781 782 783 784 785

    auto p_conv_pd = new mkldnn::convolution_forward::primitive_desc(
        conv_desc, conv_attr, engine);

    return std::unique_ptr<mkldnn::convolution_forward::primitive_desc>(
        p_conv_pd);
  }
786 787 788
};

template <typename T>
789
class ConvMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
790 791 792 793 794
 public:
  void Compute(const paddle::framework::ExecutionContext& ctx) const override {
    PADDLE_ENFORCE(paddle::platform::is_cpu_place(ctx.GetPlace()),
                   "It must use CPUPlace.");

795 796
    auto& dev_ctx =
        ctx.template device_context<platform::MKLDNNDeviceContext>();
797 798 799 800 801 802 803 804 805
    const auto& mkldnn_engine = dev_ctx.GetEngine();

    const Tensor* input = ctx.Input<Tensor>("Input");
    const Tensor* filter = ctx.Input<Tensor>("Filter");
    const Tensor* output_grad =
        ctx.Input<Tensor>(framework::GradVarName("Output"));
    Tensor* input_grad = ctx.Output<Tensor>(framework::GradVarName("Input"));
    Tensor* filter_grad = ctx.Output<Tensor>(framework::GradVarName("Filter"));

806 807 808 809 810 811 812 813 814 815
    PADDLE_ENFORCE(input->layout() == DataLayout::kMKLDNN &&
                       input->format() != memory::format::format_undef,
                   "Wrong layout/format set for Input tensor");
    PADDLE_ENFORCE(filter->layout() == DataLayout::kMKLDNN &&
                       filter->format() != memory::format::format_undef,
                   "Wrong layout/format set for Filter tensor");
    PADDLE_ENFORCE(output_grad->layout() == DataLayout::kMKLDNN &&
                       output_grad->format() != memory::format::format_undef,
                   "Wrong layout/format set for output_grad tensor");

816 817 818 819
    PADDLE_ENFORCE(
        !ctx.Attr<bool>("is_test"),
        "is_test attribute should be set to False in training phase.");

820 821 822 823
    if (!input_grad && !filter_grad) return;

    std::vector<int> strides = ctx.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = ctx.Attr<std::vector<int>>("paddings");
824 825
    std::vector<int> dilations = ctx.Attr<std::vector<int>>("dilations");
    int groups = ctx.Attr<int>("groups");
826

827
    bool is_conv3d = strides.size() == 3U;
828 829 830 831 832 833 834 835 836
    const T* input_data = input->data<T>();
    const T* filter_data = filter->data<T>();
    const T* output_grad_data = output_grad->data<T>();
    T* input_grad_data = nullptr;
    T* filter_grad_data = nullptr;

    std::vector<int> src_tz = paddle::framework::vectorize2int(input->dims());
    std::vector<int> weights_tz =
        paddle::framework::vectorize2int(filter->dims());
837
    int g = std::max(groups, 1);
Y
Yihua Xu 已提交
838
    GetWeightsTz(weights_tz, g, is_conv3d);
839 840
    std::vector<int> dst_tz =
        paddle::framework::vectorize2int(output_grad->dims());
841

842 843
    auto src_format = input->format();
    mkldnn::memory::format weights_format =
Y
Yihua Xu 已提交
844
        GetWeightsFormat(filter->format(), g, is_conv3d);
845

846
    // Get an unique name from "argument" name of "input" and "Filter" variable
J
Jacek Czaja 已提交
847
    // as well as attributes of primitive to be created
848
    // This name will be used as key when saving info into device context
J
Jacek Czaja 已提交
849 850
    const std::string key = platform::ConvMKLDNNHandler::GetHash(
        src_tz, weights_tz, strides, paddings, dilations, groups,
851
        ctx.op().Input("Input") + ctx.op().Input("Filter"));
852 853

    const std::string key_conv_pd = key + "@conv_pd";
854
    std::vector<primitive> pipeline;
855

856 857
    // Create user memory descriptors
    auto user_src_md = platform::MKLDNNMemDesc(
858
        {src_tz}, platform::MKLDNNGetDataType<T>(), src_format);
859
    auto user_weights_md = platform::MKLDNNMemDesc(
860
        {weights_tz}, platform::MKLDNNGetDataType<T>(), weights_format);
861 862
    auto user_diff_dst_md = platform::MKLDNNMemDesc(
        {dst_tz}, platform::MKLDNNGetDataType<T>(), output_grad->format());
863 864 865 866 867

    /* create memory descriptor for conv backward without specified format
     * ('any') which lets a primitive (conv backward in this case) choose
     * the memory format preferred for best performance
     */
868 869 870 871
    std::string data_format = ctx.Attr<std::string>("data_format");
    auto chosen_memory_format =
        platform::data_format_to_memory_format(data_format);

872 873 874 875 876 877 878
    weights_format = mkldnn::memory::format::any;
    // Check the format for user's special output
    if (chosen_memory_format != mkldnn::memory::format::any) {
      if (is_conv3d) {
        chosen_memory_format =
            platform::MKLDNNFormatForSize(src_tz.size(), chosen_memory_format);
      }
879 880
    }

881
    auto src_md = platform::MKLDNNMemDesc(
882
        src_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
883
    auto diff_src_md = platform::MKLDNNMemDesc(
884
        src_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
885
    auto weights_md = platform::MKLDNNMemDesc(
886
        weights_tz, platform::MKLDNNGetDataType<T>(), weights_format);
887
    auto diff_weights_md = platform::MKLDNNMemDesc(
888
        weights_tz, platform::MKLDNNGetDataType<T>(), weights_format);
889
    auto diff_dst_md = platform::MKLDNNMemDesc(
890
        dst_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
891

892
    // Retrieve conv_pd from device context
893 894 895
    auto conv_pd =
        std::static_pointer_cast<mkldnn::convolution_forward::primitive_desc>(
            dev_ctx.GetBlob(key_conv_pd));
896 897 898
    PADDLE_ENFORCE(conv_pd != nullptr,
                   "Fail to find conv_pd in device context");

899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914
    // create backward convolution weights primitive descriptor
    auto conv_bwd_weights_desc = mkldnn::convolution_backward_weights::desc(
        mkldnn::convolution_direct, src_md, diff_weights_md, diff_dst_md,
        strides, paddings, paddings, mkldnn::padding_kind::zero);
    auto conv_bwd_weights_pd =
        std::make_shared<mkldnn::convolution_backward_weights::primitive_desc>(
            conv_bwd_weights_desc, mkldnn_engine, *conv_pd);

    // create backward convolution data primitive descriptor
    auto conv_bwd_data_desc = mkldnn::convolution_backward_data::desc(
        mkldnn::convolution_direct, diff_src_md, weights_md, diff_dst_md,
        strides, paddings, paddings, mkldnn::padding_kind::zero);
    auto conv_bwd_data_pd =
        std::make_shared<mkldnn::convolution_backward_data::primitive_desc>(
            conv_bwd_data_desc, mkldnn_engine, *conv_pd);

J
Jacek Czaja 已提交
915 916 917
    platform::ConvMKLDNNHandler handler(conv_pd, conv_bwd_data_pd,
                                        conv_bwd_weights_pd, dev_ctx,
                                        mkldnn_engine, key);
918 919 920 921 922 923 924 925 926

    // create mkldnn memory from input tensors (data/weights)
    auto user_src_memory_p =
        handler.AcquireSrcMemory(user_src_md, to_void_cast<T>(input_data));
    auto user_weights_memory_p = handler.AcquireWeightsMemory(
        user_weights_md, to_void_cast<T>(filter_data));
    auto user_diff_dst_memory_p = handler.AcquireDiffDstMemory(
        user_diff_dst_md, to_void_cast<T>(output_grad_data));

927 928
    // create backward conv primitive for weights
    if (filter_grad) {
929 930
      auto src_memory_p = handler.AcquireSrcMemoryFromWeightsPrimitive(
          user_src_memory_p, pipeline);
931

932 933 934 935
      auto diff_dst_memory_4filter_p =
          handler.AcquireDiffDstMemoryFromWeightsPrimitive(
              user_diff_dst_memory_p, pipeline);

936
      const size_t size = handler.GetDiffWeightsMemorySize();
Y
Yu Yang 已提交
937 938
      filter_grad_data = filter_grad->mutable_data<T>(
          ctx.GetPlace(), paddle::memory::Allocator::kDefault, size);
939

940 941 942 943 944 945 946 947 948
      auto diff_weights_memory_p =
          handler.AcquireDiffWeightsMemoryFromWeightsPrimitive(
              reinterpret_cast<void*>(filter_grad_data));

      auto conv_bwd_weights_p = handler.AcquireConvolutionBackwardWeights(
          src_memory_p, diff_dst_memory_4filter_p, diff_weights_memory_p);

      // push primitive to stream and wait until it's executed
      pipeline.push_back(*conv_bwd_weights_p);
949 950

      filter_grad->set_layout(DataLayout::kMKLDNN);
951
      filter_grad->set_format(GetMKLDNNFormat(*diff_weights_memory_p));
952 953 954
    }

    if (input_grad) {
955 956 957 958 959 960 961
      auto weights_memory_p = handler.AcquireWeightsMemoryFromDataPrimitive(
          user_weights_memory_p, pipeline);

      auto diff_dst_memory_4data_p =
          handler.AcquireDiffDstMemoryFromDataPrimitive(user_diff_dst_memory_p,
                                                        pipeline);

962
      const size_t size = handler.GetDiffSourceMemorySize();
Y
Yu Yang 已提交
963 964
      input_grad_data = input_grad->mutable_data<T>(
          ctx.GetPlace(), paddle::memory::Allocator::kDefault, size);
965

966 967 968 969 970 971 972
      auto diff_src_memory_p = handler.AcquireDiffSrcMemoryFromDataPrimitive(
          reinterpret_cast<void*>(input_grad_data));

      auto conv_bwd_data_p = handler.AcquireConvolutionBackwardData(
          diff_dst_memory_4data_p, weights_memory_p, diff_src_memory_p);

      pipeline.push_back(*conv_bwd_data_p);
973 974

      input_grad->set_layout(DataLayout::kMKLDNN);
975
      input_grad->set_format(GetMKLDNNFormat(*diff_src_memory_p));
976
    }
977
    stream(stream::kind::eager).submit(pipeline).wait();
X
xiaolil1 已提交
978
  }
979 980 981 982 983 984 985
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

X
Xin Pan 已提交
986 987 988
REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN,
                                    ::paddle::platform::CPUPlace, FP32,
                                    ops::kConvMKLDNNFP32,
989 990 991 992 993 994 995 996 997 998 999
                                    ops::ConvMKLDNNOpKernel<float, float>);

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN,
                                    ::paddle::platform::CPUPlace, U8,
                                    ops::kConvMKLDNNFP32,
                                    ops::ConvMKLDNNOpKernel<uint8_t, float>);

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN,
                                    ::paddle::platform::CPUPlace, S8,
                                    ops::kConvMKLDNNFP32,
                                    ops::ConvMKLDNNOpKernel<int8_t, float>);
X
Xin Pan 已提交
1000 1001 1002 1003 1004

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv2d_grad, MKLDNN,
                                    ::paddle::platform::CPUPlace, FP32,
                                    ops::kConvMKLDNNFP32,
                                    ops::ConvMKLDNNGradOpKernel<float>);
1005 1006 1007 1008

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv3d, MKLDNN,
                                    ::paddle::platform::CPUPlace, FP32,
                                    ops::kConvMKLDNNFP32,
1009
                                    ops::ConvMKLDNNOpKernel<float, float>);
1010 1011 1012 1013 1014

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv3d_grad, MKLDNN,
                                    ::paddle::platform::CPUPlace, FP32,
                                    ops::kConvMKLDNNFP32,
                                    ops::ConvMKLDNNGradOpKernel<float>);