text: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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.
   */
  
  /*!
   * \file mkldnn_batch_norm.cc
   * \brief
   * \author Tao Lv
  */
  
  #ifndef MXNET_OPERATOR_NN_MKLDNN_MKLDNN_BATCH_NORM_INL_H_
  #define MXNET_OPERATOR_NN_MKLDNN_MKLDNN_BATCH_NORM_INL_H_
  
  #if MXNET_USE_MKLDNN == 1
  #include <vector>
  #include <utility>
  #include <mkldnn.hpp>
  #include "../batch_norm-inl.h"
  #include "./mkldnn_ops-inl.h"
  #include "./mkldnn_base-inl.h"
  
  #define VARIANCE_TO_INVSTD(__var$,    __eps$)   (1.0/sqrt((__var$) + DType(__eps$)))
  #define INVSTD_TO_VARIANCE(__invstd$, __eps$)   ((1.0 / ((__invstd$) * (__invstd$))) - (__eps$))
  namespace mxnet {
  namespace op {
  
  typedef mkldnn::batch_normalization_forward::primitive_desc     t_bn_f_pdesc;
  typedef mkldnn::batch_normalization_forward::desc               t_bn_f_desc;
  typedef mkldnn::batch_normalization_backward::primitive_desc    t_bn_b_pdesc;
  typedef mkldnn::batch_normalization_backward::desc              t_bn_b_desc;
  
  using mkldnn::use_global_stats;
  using mkldnn::use_scale_shift;
  using mkldnn::forward_training;
  using mkldnn::forward_inference;
  
  inline static unsigned _GetFlags(const std::vector<NDArray> &in_data,
                                   const std::vector<NDArray> &aux_states,
                                   const BatchNormParam &param, bool is_train) {
    unsigned flags = 0U;
    if (in_data.size() == 3U) {
      flags |= use_scale_shift;
    }
  
    // aux_states[0]: inMean
    // aux_states[1]: inVariance
    if (aux_states.size() == 2U && !is_train) {
      flags |= use_global_stats;
    }
    return flags;
  }
  
  template <typename DType>
  inline static t_bn_f_pdesc _GetFwd(const mkldnn::memory &data_mem,
                                     bool is_train,
                                     DType eps,
                                     unsigned flags) {
    auto data_mpd   = data_mem.get_primitive_desc();
    auto data_md    = data_mpd.desc();
    auto engine     = CpuEngine::Get()->get_engine();
  
    if (is_train) {
      t_bn_f_desc bnFwd_desc(forward_training, data_md, eps, flags);
      return t_bn_f_pdesc(bnFwd_desc, engine);
    } else {
      t_bn_f_desc bnFwd_desc(forward_inference, data_md, eps, flags);
      return t_bn_f_pdesc(bnFwd_desc, engine);
    }
  }
  
  template <typename DType>
  inline static t_bn_b_pdesc _GetBwd(const mkldnn::memory &data_mem,
                                     const mkldnn::memory &diff_mem,
                                     DType eps,
                                     unsigned flags) {
    auto data_mpd   = data_mem.get_primitive_desc();
    auto data_md    = data_mpd.desc();
    auto diff_mpd   = diff_mem.get_primitive_desc();
    auto diff_md    = diff_mpd.desc();
    auto engine     = CpuEngine::Get()->get_engine();
  
    t_bn_b_desc  bnBwd_desc(mkldnn::prop_kind::backward, diff_md, data_md, eps, flags);
    return t_bn_b_pdesc(bnBwd_desc, engine, _GetFwd(data_mem, true, eps, flags));
  }
  
  typedef ParamOpSign<BatchNormParam> MKLDNNBNSignature;
  
  class MKLDNNBNForward {
    std::shared_ptr<const mkldnn::memory> data_m;
    std::shared_ptr<const mkldnn::memory> weight_m;
    std::shared_ptr<const mkldnn::memory> out_m;
    std::shared_ptr<const mkldnn::memory> mean_m;
    std::shared_ptr<const mkldnn::memory> var_m;
    std::shared_ptr<mkldnn::batch_normalization_forward> fwd;
    bool is_train;
    t_bn_f_pdesc pd;
  
   public:
    MKLDNNBNForward(const t_bn_f_pdesc &_pd, bool is_train): pd(_pd) {
      weight_m.reset(new mkldnn::memory(pd.weights_primitive_desc()));
      this->is_train = is_train;
    }
  
    const mkldnn::memory &GetWeight() const {
      return *weight_m;
    }
  
    const t_bn_f_pdesc &GetPd() const {
      return pd;
    }
  
    const mkldnn::memory &GetMean() const {
      return *mean_m;
    }
  
    const mkldnn::memory &GetVar() const {
      return *var_m;
    }
  
    void SetDataHandle(const NDArray &data, const NDArray &mean,
                       const NDArray &var, const mkldnn::memory &out) {
      auto _data = data.GetMKLDNNData();
      if (data_m) {
        data_m->set_data_handle(_data->get_data_handle());
      } else {
        data_m.reset(new mkldnn::memory(_data->get_primitive_desc(),
                                        _data->get_data_handle()));
      }
      if (out_m) {
        out_m->set_data_handle(out.get_data_handle());
      } else {
        out_m.reset(new mkldnn::memory(out.get_primitive_desc(),
                                       out.get_data_handle()));
      }
      auto mean_ptr = mean.data().dptr_;
      if (mean_m) {
        mean_m->set_data_handle(mean_ptr);
      } else {
        mean_m.reset(new mkldnn::memory(pd.mean_primitive_desc(),
                                        mean_ptr));
      }
      auto var_ptr = var.data().dptr_;
      if (var_m) {
        var_m->set_data_handle(var_ptr);
      } else {
        var_m.reset(new mkldnn::memory(pd.variance_primitive_desc(),
                                       var_ptr));
      }
  
      if (fwd == nullptr) {
        if (!is_train)
          fwd.reset(new mkldnn::batch_normalization_forward(
                  pd, *data_m, mkldnn::primitive::at(*mean_m),
                  mkldnn::primitive::at(*var_m), *weight_m, *out_m));
        else
          fwd.reset(new mkldnn::batch_normalization_forward(
                  pd, mkldnn::primitive::at(*data_m),
                  mkldnn::primitive::at(*weight_m), *out_m,
                  *mean_m, *var_m));
      }
    }
  
    const mkldnn::batch_normalization_forward &GetFwd() const {
      return *fwd;
    }
  };
  
  template<typename DType>
  static MKLDNNBNForward &GetBNForward(const BatchNormParam& param,
                                       const OpContext &ctx, const NDArray &in_data,
                                       unsigned flags) {
    static thread_local std::unordered_map<MKLDNNBNSignature, MKLDNNBNForward, OpHash> fwds;
    MKLDNNBNSignature key(param);
    key.AddSign(ctx.is_train);
    key.AddSign(in_data);
  
    auto it = fwds.find(key);
    if (it == fwds.end()) {
      auto fwd_pd = _GetFwd(*in_data.GetMKLDNNData(), ctx.is_train,
                            (DType) param.eps, flags);
      MKLDNNBNForward fwd(fwd_pd, ctx.is_train);
      auto ins_ret = fwds.insert(std::pair<MKLDNNBNSignature, MKLDNNBNForward>(
              key, fwd));
      CHECK(ins_ret.second);
      it = ins_ret.first;
    }
    return it->second;
  }
  
  template <typename DType>
  void MKLDNNBatchNormForward(const OpContext &ctx, const BatchNormParam &param,
                              const std::vector<NDArray>   &in_data,
                              const std::vector<OpReqType> &req,
                              const std::vector<NDArray>   &out_data,
                              const std::vector<NDArray>   &aux_states) {
    TmpMemMgr::Get()->Init(ctx.requested[batchnorm::kTempSpace]);
    unsigned flags      = _GetFlags(in_data, aux_states, param, ctx.is_train);
    const NDArray &data = in_data[batchnorm::kData];
  
    auto &fwd = GetBNForward<DType>(param, ctx, data, flags);
    const NDArray &out  = out_data[batchnorm::kOut];
  
    // for output memory
    auto out_mem = const_cast<NDArray &>(out).CreateMKLDNNData(fwd.GetPd().dst_primitive_desc());
  
    // mxnet will always use scale shift.
    // But if fix_gamma is true, then all scale elements will be set to 1.0f
    if (flags & use_scale_shift) {
      const NDArray &gamma    = in_data[batchnorm::kGamma];
      const NDArray &beta     = in_data[batchnorm::kBeta];
      CHECK_EQ(gamma.storage_type(), mxnet::kDefaultStorage);
      CHECK_EQ(beta.storage_type(), mxnet::kDefaultStorage);
  
      const mkldnn::memory &weight_mem = fwd.GetWeight();
      DType* weight_buf = reinterpret_cast<DType *>(weight_mem.get_data_handle());
  
      nnvm::dim_t channels_ = data.shape()[1];
      CHECK(weight_mem.get_primitive_desc().get_size() == channels_ * sizeof(DType) * 2);
      DType* weight_ptr = gamma.data().dptr<DType>();
      DType* bias_ptr = beta.data().dptr<DType>();
      if (!param.fix_gamma) {
        memcpy(weight_buf, weight_ptr, sizeof(weight_buf[0]) * channels_);
        memcpy(&weight_buf[channels_], bias_ptr, sizeof(weight_buf[0]) * channels_);
      } else if (IsBNWriting(req[batchnorm::kGamma])) {
        for (int i = 0; i < channels_; i++) {
          weight_buf[i] = (DType)1.0f;
          weight_ptr[i] = (DType)1.0f;
          weight_buf[channels_ + i] = bias_ptr[i];  // bias
        }
      } else {
        for (int i = 0; i < channels_; i++) {
          weight_buf[i] = (DType)1.0f;
          weight_buf[channels_ + i] = bias_ptr[i];  // bias
        }
      }
  
      if (!ctx.is_train) {
        DType* omean    = out_data[batchnorm::kMean].data().dptr<DType>();
        DType* ovar     = out_data[batchnorm::kVar].data().dptr<DType>();
        DType* inmean   = aux_states[batchnorm::kMovingMean].data().dptr<DType>();
        DType* invar    = aux_states[batchnorm::kMovingVar].data().dptr<DType>();
        // to align with origin implmentation: batch_norm.cc: L164
        for (int i = 0; i < channels_; i++) {
          omean[i] = inmean[i];
          ovar[i] = VARIANCE_TO_INVSTD(invar[i], param.eps);
        }
  
        fwd.SetDataHandle(data, aux_states[batchnorm::kMovingMean],
                          aux_states[batchnorm::kMovingVar],
                          *out_mem);
        MKLDNNStream::Get()->RegisterPrim(fwd.GetFwd());
        MKLDNNStream::Get()->Submit();
      } else {  // training
        const NDArray &outMean  = out_data[batchnorm::kMean];
        const NDArray &outVar   = out_data[batchnorm::kVar];
        DType* omean    = outMean.data().dptr<DType>();
        DType* ovar     = outVar.data().dptr<DType>();
  
        fwd.SetDataHandle(data, outMean, outVar, *out_mem);
        MKLDNNStream::Get()->RegisterPrim(fwd.GetFwd());
        MKLDNNStream::Get()->Submit();
        DType* mean_mem_ptr = reinterpret_cast<DType*>(fwd.GetMean().get_data_handle());
        DType* var_mem_ptr  = reinterpret_cast<DType*>(fwd.GetVar().get_data_handle());
        for (int i = 0; i < channels_; i++) {
          omean[i] = mean_mem_ptr[i];
          ovar[i]  = VARIANCE_TO_INVSTD(var_mem_ptr[i], param.eps);
        }
      }
    } else {  // no input gamma and beta
      LOG(FATAL) << "MKLDNN batch normalization: should not reach here ...";
    }
  }
  
  template <typename DType>
  void MKLDNNBatchNormBackward(const OpContext &ctx, const BatchNormParam &param,
                               const std::vector<NDArray>    &out_grad,
                               const std::vector<NDArray>    &in_data,
                               const std::vector<NDArray>    &out_data,
                               const std::vector<OpReqType>  &req,
                               const std::vector<NDArray>    &in_grad,
                               const std::vector<NDArray>    &aux_states) {
    TmpMemMgr::Get()->Init(ctx.requested[batchnorm::kTempSpace]);
    CHECK_EQ(out_grad.size(), 1U);
    CHECK_EQ(in_data.size(), 3U);
    CHECK_EQ(out_data.size(), 3U);
    CHECK_EQ(in_grad.size(), 3U);
    unsigned flags = _GetFlags(in_data, aux_states, param, ctx.is_train);
  
    const NDArray &data         = in_data[batchnorm::kData];
    const NDArray &diff         = out_grad[batchnorm::kOut];
    const NDArray &gradIn       = in_grad[batchnorm::kData];
    const NDArray &moving_mean  = aux_states[batchnorm::kMovingMean];
    const NDArray &moving_var   = aux_states[batchnorm::kMovingVar];
    const NDArray &out_mean     = out_data[batchnorm::kMean];
    const NDArray &out_var      = out_data[batchnorm::kVar];
  
    CHECK(out_mean.IsDefaultData());
    CHECK(out_var.IsDefaultData());
    CHECK(moving_mean.IsDefaultData());
    CHECK(moving_var.IsDefaultData());
  
    auto data_mem  = data.GetMKLDNNData();
    auto diff_mem  = diff.GetMKLDNNData();
    // MKLDNN batchnorm should run on special layouts. If one of them isn't, we
    // should reorder them.
    if (data.IsDefaultData())
      data_mem = data.GetMKLDNNDataReorder(diff_mem->get_primitive_desc());
    else if (diff.IsDefaultData())
      diff_mem = diff.GetMKLDNNDataReorder(data_mem->get_primitive_desc());
    auto bwd_pd = _GetBwd(*data_mem, *diff_mem, param.eps, flags);
    auto gradi_mem = const_cast<NDArray &>(gradIn).CreateMKLDNNData(data_mem->get_primitive_desc());
  
    if (flags & use_scale_shift) {
      const NDArray &gamma    = in_data[batchnorm::kGamma];
      const NDArray &beta     = in_data[batchnorm::kBeta];
      // TODO(tao): how to reuse this memory?
      std::shared_ptr<const mkldnn::memory> weight_mem(
                      new mkldnn::memory(bwd_pd.weights_primitive_desc()));
  
      DType* weight_buf = reinterpret_cast<DType *>(weight_mem->get_data_handle());
      nnvm::dim_t channels_ = data.shape()[1];
      for (int i = 0; i < channels_; i++) {
        if (!param.fix_gamma)
          weight_buf[i] = (gamma.data().dptr<DType>())[i];   // weight
        else
          weight_buf[i] = (DType)1.0f;
      }
  
      for (int i = 0; i < channels_; i++) {
        weight_buf[channels_ + i] = (beta.data().dptr<DType>())[i];  // bias
      }
  
      std::shared_ptr<const mkldnn::memory> gradw_mem(
                      new mkldnn::memory(bwd_pd.diff_weights_primitive_desc()));
      // training but no input mean and variance
      if (ctx.is_train && !param.use_global_stats) {
        DType* moving_mean_ptr  = reinterpret_cast<DType *>(moving_mean.data().dptr<DType>());
        DType* moving_var_ptr   = reinterpret_cast<DType *>(moving_var.data().dptr<DType>());
        DType* out_mean_ptr     = reinterpret_cast<DType *>(out_mean.data().dptr<DType>());
        DType* out_var_ptr      = reinterpret_cast<DType *>(out_var.data().dptr<DType>());
        mkldnn::memory var_mem(bwd_pd.variance_primitive_desc());
        DType *tmp_var_ptr = reinterpret_cast<DType *>(var_mem.get_data_handle());
  
        DType minus_mom = (1.0f - param.momentum);
        for (int i = 0; i < channels_; i++) {
          moving_mean_ptr[i] = moving_mean_ptr[i] * param.momentum +
                               out_mean_ptr[i] * minus_mom;
          float variance = INVSTD_TO_VARIANCE(out_var_ptr[i], param.eps);
          tmp_var_ptr[i] = variance;
          moving_var_ptr[i] = moving_var_ptr[i] * param.momentum +
                              variance * minus_mom;
        }
  
        std::shared_ptr<const mkldnn::memory> out_mean_mem(
                        new mkldnn::memory(bwd_pd.mean_primitive_desc(), out_mean_ptr));
        std::shared_ptr<const mkldnn::memory> out_var_mem(
                        new mkldnn::memory(bwd_pd.variance_primitive_desc(), out_var_ptr));
  
        auto bn_bwd = mkldnn::batch_normalization_backward(bwd_pd,
                                                           *data_mem,
                                                           mkldnn::primitive::at(*out_mean_mem),
                                                           mkldnn::primitive::at(var_mem),
                                                           *diff_mem,
                                                           *weight_mem,
                                                           *gradi_mem,
                                                           *gradw_mem);
  
        MKLDNNStream::Get()->RegisterPrim(bn_bwd);
        MKLDNNStream::Get()->Submit();
      } else {
        std::shared_ptr<const mkldnn::memory> imean_mem(
                        new mkldnn::memory(bwd_pd.mean_primitive_desc(),
                        moving_mean.data().dptr<DType>()));
        std::shared_ptr<const mkldnn::memory> ivar_mem(
                        new mkldnn::memory(bwd_pd.variance_primitive_desc(),
                        moving_var.data().dptr<DType>()));
        auto bn_bwd = mkldnn::batch_normalization_backward(bwd_pd,
                                                           *data_mem,
                                                           mkldnn::primitive::at(*imean_mem),
                                                           mkldnn::primitive::at(*ivar_mem),
                                                           *diff_mem,
                                                           *weight_mem,
                                                           *gradi_mem,
                                                           *gradw_mem);
  
        MKLDNNStream::Get()->RegisterPrim(bn_bwd);
        MKLDNNStream::Get()->Submit();
      }
  
      // copy data from gradw_mem to in_grad[1] and in_grad[2]
      DType* gw_buf = reinterpret_cast<DType *>(gradw_mem->get_data_handle());
      for (int i = 0; i < channels_; i++) {
        if (!param.fix_gamma)
          (in_grad[1].data().dptr<DType>())[i] = gw_buf[i];
        else
          (in_grad[1].data().dptr<DType>())[i] = 0.0f;
      }
  
      for (int i = 0; i < channels_; i++) {
        (in_grad[2].data().dptr<DType>())[i] = gw_buf[i + channels_];
      }
    } else {
      LOG(FATAL) << "MKLDNN batch normalization backward: should not reach here ...";
    }
  }
  }  // namespace op
  }  // namespace mxnet
  #endif  // MXNET_USE_MKLDNN
  #endif  // MXNET_OPERATOR_NN_MKLDNN_MKLDNN_BATCH_NORM_INL_H_
  