batch_norm_mkldnn_op.cc 22.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
/* 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. */

#include "mkldnn.hpp"
#include "paddle/fluid/operators/batch_norm_op.h"
#include "paddle/fluid/platform/mkldnn_helper.h"

namespace paddle {
namespace operators {

22 23 24 25 26 27
using batch_norm_bwd = mkldnn::batch_normalization_backward;
using batch_norm_fwd = mkldnn::batch_normalization_forward;
using mkldnn::memory;
using mkldnn::primitive;
using mkldnn::reorder;
using mkldnn::stream;
28 29
using paddle::platform::MKLDNNDeviceContext;
using paddle::platform::MKLDNNMemDesc;
30
using platform::to_void_cast;
31 32 33 34 35 36 37 38 39

namespace {
template <typename T>
struct bn_type_traits {
  using op_type = T;
  using op_desc = typename op_type::desc;
  using op_prim = typename op_type::primitive_desc;
};

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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155
class BatchNormMKLDNNHandler : public platform::MKLDNNHandler {
 public:
  BatchNormMKLDNNHandler(
      std::shared_ptr<batch_norm_fwd::primitive_desc> batch_norm_pd,
      const platform::MKLDNNDeviceContext &dev_ctx, mkldnn::engine engine,
      const std::string &base_key)
      : platform::MKLDNNHandler(dev_ctx, engine, base_key) {
    batch_norm_pd_ = batch_norm_pd;
  }

  std::shared_ptr<memory> AcquireScaleshiftMemoryFromPrimitive(void *ptr) {
    return this->AcquireMemoryFromPrimitive(
        batch_norm_pd_->weights_primitive_desc(), ptr, "@scaleshift_mem_p");
  }

  std::shared_ptr<memory> AcquireMeanMemoryFromPrimitive(void *ptr) {
    return this->AcquireMemoryFromPrimitive(
        batch_norm_pd_->mean_primitive_desc(), ptr, "@mean_mem_p");
  }

  std::shared_ptr<memory> AcquireVarianceMemoryFromPrimitive(void *ptr) {
    return this->AcquireMemoryFromPrimitive(
        batch_norm_pd_->variance_primitive_desc(), ptr, "@variance_mem_p");
  }

  std::shared_ptr<batch_norm_fwd> AcquireTestBatchNormFwd(
      std::shared_ptr<memory> src_memory,
      const mkldnn::primitive::at &mean_memory,
      const mkldnn::primitive::at &variance_memory,
      std::shared_ptr<memory> scaleshift_memory,
      std::shared_ptr<memory> dst_memory) {
    auto prim_key = key_ + "@batch_norm_p";
    auto batch_norm_p =
        std::static_pointer_cast<batch_norm_fwd>(dev_ctx_.GetBlob(prim_key));
    PADDLE_ENFORCE(
        (batch_norm_p != nullptr) || (is_reusing_ == false),
        "Fail to find batch norm primitive for test in device context");
    if (batch_norm_p == nullptr) {
      batch_norm_p = std::make_shared<batch_norm_fwd>(
          *batch_norm_pd_, *src_memory, mean_memory, variance_memory,
          *scaleshift_memory, *dst_memory);

      dev_ctx_.SetBlob(prim_key, batch_norm_p);
    } else {
      is_reusing_ = true;
    }
    return batch_norm_p;
  }

  std::shared_ptr<batch_norm_fwd> AcquireTrainingBatchNormFwd(
      std::shared_ptr<memory> src_memory,
      std::shared_ptr<memory> scaleshift_memory,
      std::shared_ptr<memory> dst_memory, std::shared_ptr<memory> mean_memory,
      std::shared_ptr<memory> variance_memory) {
    auto prim_key = key_ + "@batch_norm_p";
    auto batch_norm_p =
        std::static_pointer_cast<batch_norm_fwd>(dev_ctx_.GetBlob(prim_key));
    PADDLE_ENFORCE(
        (batch_norm_p != nullptr) || (is_reusing_ == false),
        "Fail to find batch norm primitive for training in device context");
    if (batch_norm_p == nullptr) {
      batch_norm_p = std::make_shared<batch_norm_fwd>(
          *batch_norm_pd_, *src_memory, *scaleshift_memory, *dst_memory,
          *mean_memory, *variance_memory);

      dev_ctx_.SetBlob(prim_key, batch_norm_p);
    } else {
      is_reusing_ = true;
    }
    return batch_norm_p;
  }
  //
  static std::string GetHash(const memory::dims &input_dims, float epsilon,
                             unsigned flag, bool is_test, memory::format format,
                             const std::string &suffix) {
    auto dims2str = [](const memory::dims &operand_dims) {
      std::string dstr = "";
      for (size_t i = 0; i < operand_dims.size(); ++i) {
        dstr += std::to_string(operand_dims[i]) + "-";
      }
      return dstr;
    };
    return dims2str(input_dims) + std::to_string(epsilon) +
           std::to_string(flag) + std::to_string(is_test) +
           std::to_string(format) + suffix;
  }

 private:
  std::shared_ptr<batch_norm_fwd::primitive_desc> batch_norm_pd_;
};

std::string gethash(const memory::dims &input_dims, float epsilon,
                    unsigned flag, bool is_test, memory::format format) {
  auto dims2str = [](const memory::dims &operand_dims) {
    std::string dstr = "";
    for (size_t i = 0; i < operand_dims.size(); ++i) {
      dstr += std::to_string(operand_dims[i]) + "-";
    }
    return dstr;
  };
  return dims2str(input_dims) + std::to_string(epsilon) + std::to_string(flag) +
         std::to_string(is_test) + std::to_string(format);
}

std::shared_ptr<memory> UpdateMemoryData(
    const platform::MKLDNNDeviceContext &dev_ctx, const std::string &key,
    void *new_ptr) {
  auto mem = std::static_pointer_cast<memory>(dev_ctx.GetBlob(key));
  PADDLE_ENFORCE(
      mem != nullptr,
      (std::string("Fail to find memory in device context [key: ") + key + "]")
          .c_str());
  mem->set_data_handle(new_ptr);
  return mem;
}

156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175
template <typename T, typename Container>
void copy_to_weights(T scale_begin, T scale_end, T shift_begin, T shift_end,
                     Container *c) {
  auto it = std::begin(*c);

  std::copy(scale_begin, scale_end, std::inserter(*c, it));
  std::copy(
      shift_begin, shift_end,
      std::inserter(*c, std::next(it, std::distance(scale_begin, scale_end))));
}

}  // namespace

template <typename T>
class BatchNormMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext &ctx) const override {
    const float epsilon = ctx.Attr<float>("epsilon");
    const float momentum = ctx.Attr<float>("momentum");
    const bool is_test = ctx.Attr<bool>("is_test");
176
    const bool fuse_with_relu = ctx.Attr<bool>("fuse_with_relu");
177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193

    const auto *x = ctx.Input<Tensor>("X");
    const auto *mean = ctx.Input<Tensor>("Mean");
    const auto *variance = ctx.Input<Tensor>("Variance");

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

    auto *y = ctx.Output<Tensor>("Y");
    auto *mean_out = ctx.Output<Tensor>("MeanOut");
    auto *variance_out = ctx.Output<Tensor>("VarianceOut");
    auto *batch_mean = ctx.Output<Tensor>("SavedMean");
    auto *batch_variance = ctx.Output<Tensor>("SavedVariance");

    const auto *scale = ctx.Input<Tensor>("Scale");
    const auto *shift = ctx.Input<Tensor>("Bias");

194 195 196 197 198 199 200 201 202 203 204 205
    PADDLE_ENFORCE(x->layout() == DataLayout::kMKLDNN &&
                       x->format() != memory::format::format_undef,
                   "Wrong layout/format set for Input x tensor");

    const T *x_data = x->data<T>();
    const T *mean_data = mean->data<T>();
    const T *variance_data = variance->data<T>();
    T *y_data = y->mutable_data<T>(ctx.GetPlace());
    T *mean_out_data = mean_out->mutable_data<T>(ctx.GetPlace());
    T *variance_out_data = variance_out->mutable_data<T>(ctx.GetPlace());
    T *batch_mean_data = nullptr;
    T *batch_variance_data = nullptr;
206 207

    if (!is_test) {
208 209
      batch_mean_data = batch_mean->mutable_data<T>(ctx.GetPlace());
      batch_variance_data = batch_variance->mutable_data<T>(ctx.GetPlace());
210 211 212 213 214
    }

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

215 216 217 218
    auto src_tz = paddle::framework::vectorize2int(x->dims());
    auto scale_tz = paddle::framework::vectorize2int(scale->dims());
    PADDLE_ENFORCE(scale_tz.size() == 1, "Dims of scale tensor is NOT 1");
    const unsigned int ic = scale_tz[0];
219

220 221 222 223 224 225 226 227
    // MKLDNN requires a single piece of memory for scale and shift/bias data
    const size_t scaleshift_size = 2 * ic;
    std::vector<T> scaleshift_data;
    scaleshift_data.reserve(scaleshift_size);

    copy_to_weights(scale->data<T>(), scale->data<T>() + ic, shift->data<T>(),
                    shift->data<T>() + ic, &scaleshift_data);

228 229
    unsigned flags = mkldnn::use_scale_shift;
    if (is_test) flags |= mkldnn::use_global_stats;
230
    if (fuse_with_relu) flags |= mkldnn::fuse_bn_relu;
231

232
    // create mkldnn memory from input x tensor
233 234
    mkldnn::memory::format input_format =
        platform::MKLDNNFormatForSize(src_tz.size(), x->format());
235

236 237 238 239 240 241 242 243
    // keys for backward pass
    const std::string key = BatchNormMKLDNNHandler::GetHash(
        src_tz, epsilon, flags, is_test, input_format,
        ctx.op().Output("SavedMean"));
    const std::string key_batch_norm_fwd_pd = key + "@bn_fwd_pd";

    auto user_src_md = platform::MKLDNNMemDesc(
        {src_tz}, platform::MKLDNNGetDataType<T>(), input_format);
244 245

    // create primitive descriptor for batch norm forward
246
    using bn_fwd_types = bn_type_traits<mkldnn::batch_normalization_forward>;
247 248 249 250 251
    auto batch_norm_fwd_desc =
        bn_fwd_types::op_desc{propagation, user_src_md, epsilon, flags};
    auto batch_norm_fwd_pd = std::make_shared<batch_norm_fwd::primitive_desc>(
        batch_norm_fwd_desc, mkldnn_engine);
    // Save conv_pd/src_memory/weights_memory for backward pass
252
    dev_ctx.SetBlob(key_batch_norm_fwd_pd, batch_norm_fwd_pd);
253

254 255
    BatchNormMKLDNNHandler handler(batch_norm_fwd_pd, dev_ctx, mkldnn_engine,
                                   key);
256

257 258
    auto src_memory =
        handler.AcquireSrcMemory(user_src_md, to_void_cast(x_data));
259

260
    // crate mkldnn memory for weights(scale/shift)
261 262
    auto scaleshift_memory =
        handler.AcquireScaleshiftMemoryFromPrimitive(scaleshift_data.data());
263

264
    // create mkldnn memory for output y tensor
265 266
    auto dst_memory = handler.AcquireDstMemory(
        batch_norm_fwd_pd->dst_primitive_desc().desc(), y_data);
267

268
    std::shared_ptr<batch_norm_fwd> batch_norm_p;
269 270
    if (is_test) {
      // create mkldnn memory for stats (as input)
271 272 273 274 275 276 277 278 279
      std::shared_ptr<memory> mean_memory =
          handler.AcquireMeanMemoryFromPrimitive(to_void_cast(mean_data));
      std::shared_ptr<memory> variance_memory =
          handler.AcquireVarianceMemoryFromPrimitive(
              to_void_cast(variance_data));

      batch_norm_p = handler.AcquireTestBatchNormFwd(
          src_memory, (const mkldnn::primitive::at &)*mean_memory,
          (const mkldnn::primitive::at &)*variance_memory, scaleshift_memory,
280
          dst_memory);
281
    } else {
282
      // create mkldnn memory for stats (as output)
283 284 285 286 287 288 289 290
      std::shared_ptr<memory> mean_memory =
          handler.AcquireMeanMemoryFromPrimitive(batch_mean_data);
      std::shared_ptr<memory> variance_memory =
          handler.AcquireVarianceMemoryFromPrimitive(batch_variance_data);

      batch_norm_p = handler.AcquireTrainingBatchNormFwd(
          src_memory, scaleshift_memory, dst_memory, mean_memory,
          variance_memory);
291 292
    }

293 294 295 296 297 298 299
    y->set_layout(DataLayout::kMKLDNN);
    y->set_format(platform::GetMKLDNNFormat(*dst_memory));

    std::vector<mkldnn::primitive> pipeline;
    pipeline.push_back(*batch_norm_p);
    mkldnn::stream(mkldnn::stream::kind::eager).submit(pipeline).wait();

300
    if (!is_test) {
301 302 303 304 305 306 307 308 309
      // mkldnn only compute stats for current batch
      // so we need compute momentum stats via Eigen lib
      EigenVectorArrayMap<T> batch_mean_e(batch_mean_data, ic);
      EigenVectorArrayMap<T> batch_variance_e(batch_variance_data, ic);
      ConstEigenVectorArrayMap<T> mean_e(mean_data, ic);
      ConstEigenVectorArrayMap<T> variance_e{variance_data, ic};

      EigenVectorArrayMap<T> running_mean_e(mean_out_data, ic);
      EigenVectorArrayMap<T> running_variance_e(variance_out_data, ic);
310 311

      auto one_minus_momentum = 1. - momentum;
312 313 314
      running_mean_e = mean_e * momentum + batch_mean_e * one_minus_momentum;
      running_variance_e =
          variance_e * momentum + batch_variance_e * one_minus_momentum;
315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338
    }
  }
};

template <typename T>
class BatchNormMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
 public:
  void Compute(const paddle::framework::ExecutionContext &ctx) const override {
    auto &dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
    auto mkldnn_engine = dev_ctx.GetEngine();

    const float epsilon = ctx.Attr<float>("epsilon");

    const auto *x = ctx.Input<Tensor>("X");
    const auto *scale = ctx.Input<Tensor>("Scale");
    const auto *shift = ctx.Input<Tensor>("Bias");
    const auto *batch_mean = ctx.Input<Tensor>("SavedMean");
    const auto *batch_variance = ctx.Input<Tensor>("SavedVariance");

    const auto *diff_y = ctx.Input<Tensor>(framework::GradVarName("Y"));
    auto *diff_x = ctx.Output<Tensor>(framework::GradVarName("X"));
    auto *diff_scale = ctx.Output<Tensor>(framework::GradVarName("Scale"));
    auto *diff_shift = ctx.Output<Tensor>(framework::GradVarName("Bias"));

339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362
    PADDLE_ENFORCE(diff_y->layout() == DataLayout::kMKLDNN &&
                       diff_y->format() != memory::format::format_undef,
                   "Wrong layout/format set for Input diff_y tensor");

    const T *x_data = x->data<T>();
    const T *diff_y_data = diff_y->data<T>();
    const T *batch_mean_data = batch_mean->data<T>();
    const T *batch_variance_data = batch_variance->data<T>();
    const T *scale_data = scale->data<T>();
    const T *shift_data = shift->data<T>();
    T *diff_x_data = diff_x->mutable_data<T>(ctx.GetPlace());
    T *diff_scale_data = diff_scale->mutable_data<T>(ctx.GetPlace());
    T *diff_shift_data = diff_shift->mutable_data<T>(ctx.GetPlace());

    auto src_tz = paddle::framework::vectorize2int(x->dims());
    auto diff_src_tz = src_tz;
    auto dst_tz = src_tz;
    auto diff_dst_tz = dst_tz;
    auto scale_tz = paddle::framework::vectorize2int(scale->dims());
    PADDLE_ENFORCE(scale_tz.size() == 1, "Dims of scale tensor is NOT 1");

    const unsigned int ic = scale_tz[0];

    using bn_bwd_types = bn_type_traits<mkldnn::batch_normalization_backward>;
363

364 365 366 367 368
    mkldnn::memory::format dst_format =
        platform::MKLDNNFormatForSize(src_tz.size(), diff_y->format());

    mkldnn::memory::format input_format =
        platform::MKLDNNFormatForSize(src_tz.size(), x->format());
369

370 371 372 373 374 375 376
    unsigned flags = mkldnn::use_scale_shift;

    // keys from forward pass
    const std::string key = BatchNormMKLDNNHandler::GetHash(
        src_tz, epsilon, flags, false, input_format,
        ctx.op().Input("SavedMean"));
    const std::string key_batch_norm_fwd_pd = key + "@bn_fwd_pd";
377

378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396
    // keys for primitives reuse
    const std::string key_with_hash =
        key + gethash(src_tz, epsilon, flags, false, input_format);
    const std::string key_batch_norm_bwd_p =
        key_with_hash + "@batch_norm_bwd_p";
    const std::string key_batch_norm_src_mem_p =
        key_with_hash + "@batch_norm_bwd_src_mem_p";
    const std::string key_batch_norm_mean_mem_p =
        key_with_hash + "@batch_norm_bwd_mean_mem_p";
    const std::string key_batch_norm_variance_mem_p =
        key_with_hash + "@batch_norm_bwd_variance_mem_p";
    const std::string key_batch_norm_scaleshift_mem_p =
        key_with_hash + "@batch_norm_bwd_scaleshift_mem_p";
    const std::string key_batch_norm_diff_scaleshift_mem_p =
        key_with_hash + "@batch_norm_bwd_diff_scaleshift_mem_p";
    const std::string key_batch_norm_diff_src_mem_p =
        key_with_hash + "@batch_norm_bwd_diff_src_mem_p";
    const std::string key_batch_norm_diff_dst_mem_p =
        key_with_hash + "@batch_norm_bwd_diff_dst_mem_p";
397

398 399
    primitive reorder_diff_dst;
    bool is_diff_dst_reordered = false;
400 401 402
    auto user_diff_dst_memory = memory(
        {{{diff_dst_tz}, memory::data_type::f32, dst_format}, mkldnn_engine},
        to_void_cast(diff_y_data));
403

404
    // MKLDNN requires a single piece of memory for scale and shift/bias data
405 406 407 408
    const size_t scaleshift_size = 2 * ic;

    std::vector<T> scaleshift_data;
    scaleshift_data.reserve(scaleshift_size);
409 410
    copy_to_weights(scale_data, scale_data + ic, shift_data, shift_data + ic,
                    &scaleshift_data);
411 412 413

    std::vector<T> diff_scaleshift_data;
    diff_scaleshift_data.reserve(scaleshift_size);
414

415 416 417 418 419
    auto batch_norm_fwd_pd =
        std::static_pointer_cast<batch_norm_fwd::primitive_desc>(
            dev_ctx.GetBlob(key_batch_norm_fwd_pd));
    PADDLE_ENFORCE(batch_norm_fwd_pd != nullptr,
                   "Fail to find batch_norm_fwd_pd in device context");
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 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 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
    auto batch_norm_bwd_p = std::static_pointer_cast<batch_norm_bwd>(
        dev_ctx.GetBlob(key_batch_norm_bwd_p));

    if (batch_norm_bwd_p == nullptr) {
      auto src_memory = std::shared_ptr<memory>(new memory(
          {{{src_tz}, memory::data_type::f32, input_format}, mkldnn_engine},
          to_void_cast(x_data)));

      // for diff_dst, try to use same format as dst in forward pass
      auto diff_dst_pd = batch_norm_fwd_pd.get()->dst_primitive_desc();
      auto diff_dst_md = diff_dst_pd.desc();

      // create primitive descriptor for batch norm backward
      auto batch_norm_bwd_desc = bn_bwd_types::op_desc{
          mkldnn::prop_kind::backward, diff_dst_md,
          src_memory->get_primitive_desc().desc(), epsilon, flags};
      auto batch_norm_bwd_pd = bn_bwd_types::op_prim{
          batch_norm_bwd_desc, mkldnn_engine, *batch_norm_fwd_pd};

      // reorder user_diff_dst if it's not in preferred format
      auto diff_dst_memory = std::make_shared<memory>(user_diff_dst_memory);
      if (diff_dst_pd != user_diff_dst_memory.get_primitive_desc()) {
        diff_dst_memory = std::make_shared<memory>(diff_dst_pd);
        reorder_diff_dst = reorder(user_diff_dst_memory, *diff_dst_memory);
        is_diff_dst_reordered = true;
      }

      // create mkldnn memory for input tensors (src/mean/variance)
      auto mean_memory =
          std::make_shared<memory>(batch_norm_bwd_pd.mean_primitive_desc(),
                                   to_void_cast(batch_mean_data));
      auto variance_memory =
          std::make_shared<memory>(batch_norm_bwd_pd.variance_primitive_desc(),
                                   to_void_cast(batch_variance_data));

      // create mkldnn memory for input tensors (scale/shift)
      auto scaleshift_memory = std::make_shared<memory>(
          batch_norm_bwd_pd.weights_primitive_desc(), scaleshift_data.data());

      // create mkldnn memory for output diff weights (combined scale/shift)
      auto diff_scaleshift_memory = std::make_shared<memory>(
          batch_norm_bwd_pd.diff_weights_primitive_desc(),
          diff_scaleshift_data.data());

      // here assume diff_src is in the same format of src
      auto diff_src_memory = std::make_shared<memory>(
          src_memory->get_primitive_desc(), diff_x_data);

      // finally create batch_norm backward primitive
      batch_norm_bwd_p = std::make_shared<batch_norm_bwd>(
          batch_norm_bwd_pd, *src_memory, *mean_memory, *variance_memory,
          *diff_dst_memory, *scaleshift_memory, *diff_src_memory,
          *diff_scaleshift_memory);

      dev_ctx.SetBlob(key_batch_norm_bwd_p, batch_norm_bwd_p);
      dev_ctx.SetBlob(key_batch_norm_src_mem_p, src_memory);
      dev_ctx.SetBlob(key_batch_norm_mean_mem_p, mean_memory);
      dev_ctx.SetBlob(key_batch_norm_variance_mem_p, variance_memory);
      dev_ctx.SetBlob(key_batch_norm_scaleshift_mem_p, scaleshift_memory);
      dev_ctx.SetBlob(key_batch_norm_diff_scaleshift_mem_p,
                      diff_scaleshift_memory);
      dev_ctx.SetBlob(key_batch_norm_diff_src_mem_p, diff_src_memory);
      dev_ctx.SetBlob(key_batch_norm_diff_dst_mem_p, diff_dst_memory);

      // set layout/format of output tensors
      diff_x->set_layout(DataLayout::kMKLDNN);
      diff_x->set_format((memory::format)diff_src_memory->get_primitive_desc()
                             .desc()
                             .data.format);
    } else {
      // primitives already exist
      UpdateMemoryData(dev_ctx, key_batch_norm_src_mem_p, to_void_cast(x_data));
      UpdateMemoryData(dev_ctx, key_batch_norm_mean_mem_p,
                       to_void_cast(batch_mean_data));
      UpdateMemoryData(dev_ctx, key_batch_norm_variance_mem_p,
                       to_void_cast(batch_variance_data));
      UpdateMemoryData(dev_ctx, key_batch_norm_scaleshift_mem_p,
                       scaleshift_data.data());
      UpdateMemoryData(dev_ctx, key_batch_norm_diff_scaleshift_mem_p,
                       diff_scaleshift_data.data());
      auto diff_src_memory = UpdateMemoryData(
          dev_ctx, key_batch_norm_diff_src_mem_p, to_void_cast(diff_x_data));
      auto diff_dst_memory = UpdateMemoryData(
          dev_ctx, key_batch_norm_diff_dst_mem_p, to_void_cast(diff_y_data));

      // reorder user_diff_dst if it's not in preferred format
      if (diff_dst_memory->get_primitive_desc() !=
          user_diff_dst_memory.get_primitive_desc()) {
        reorder_diff_dst = reorder(user_diff_dst_memory, *diff_dst_memory);
        is_diff_dst_reordered = true;
      }

      // set layout/format of output tensors
      diff_x->set_layout(DataLayout::kMKLDNN);
      diff_x->set_format((memory::format)diff_src_memory->get_primitive_desc()
                             .desc()
                             .data.format);
    }
519 520 521 522

    // execute optional reorder and batch_norm backward primitive
    std::vector<primitive> pipeline;
    if (is_diff_dst_reordered) pipeline.push_back(reorder_diff_dst);
523
    pipeline.push_back(*batch_norm_bwd_p);
524 525 526 527
    stream(stream::kind::eager).submit(pipeline).wait();

    // copy back diff sacle/shift to output tensors (diff scale/shift)
    diff_scaleshift_data.resize(scaleshift_size);
528
    auto it = std::begin(diff_scaleshift_data);
529
    std::copy(it, std::next(it, ic), diff_scale_data);
530
    std::copy(std::next(it, ic), std::end(diff_scaleshift_data),
531
              diff_shift_data);
532 533 534 535 536 537
  }
};
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
538
REGISTER_OP_KERNEL(batch_norm, MKLDNN, ::paddle::platform::CPUPlace,
539
                   ops::BatchNormMKLDNNOpKernel<float>);
540
REGISTER_OP_KERNEL(batch_norm_grad, MKLDNN, ::paddle::platform::CPUPlace,
541
                   ops::BatchNormMKLDNNGradOpKernel<float>);