/*********************************************************
 *
 * 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.
 *
 ************************************************************/
#include "batchnorm.h"

#include <cctype>

namespace singa {

BatchNormHandle::BatchNormHandle(const float momentum, const Tensor& input) {
  factor = momentum;
  batchsize = input.shape(0);
  channels = input.shape(1);
  if (input.nDim() == 4u) {
    height = input.shape().at(2);
    width = input.shape().at(3);
    is_2d = false;
  } else if (input.nDim() == 2u) {
    height = 1;
    width = 1;
    is_2d = true;
  } else {
    LOG(FATAL) << "The dimension of input should either be 4D or 2D.";
  }

#ifdef USE_DNNL
  if (input.device()->lang() == kCpp) {
    use_dnnl = true;
    epsilon = 1e-5f;
    x_dims = dnnl::memory::dims(input.shape().begin(), input.shape().end());

    // support f32 only
    auto dtype_ = memory::data_type::f32;
    memory::format_tag format_tag_ = get_dnnl_format_tag(input);
    x_md = dnnl::memory::desc({x_dims}, dtype_, format_tag_);

    // add to
    bn_fwd_training_d = new dnnl::batch_normalization_forward::desc(
        dnnl::prop_kind::forward_training, x_md, epsilon,
        dnnl::normalization_flags::use_scale_shift);

    auto eng = input.device()->context(0)->dnnl_engine;
    bn_fwd_training_pd = new dnnl::batch_normalization_forward::primitive_desc(
        *bn_fwd_training_d, eng);
  }
#endif  // USE_DNNL
};

BatchNormHandle::~BatchNormHandle() {
#ifdef USE_DNNL
  if (use_dnnl) {
    delete (bn_fwd_training_d);
    delete (bn_fwd_training_pd);
  }
#endif  // USE_DNNL
}

#ifdef USE_DNNL

Tensor CpuBatchNormForwardInference(const BatchNormHandle& bnh, const Tensor& x,
                                    const Tensor& bnScale, const Tensor& bnBias,
                                    Tensor& running_mean, Tensor& running_var) {
  CHECK_EQ(x.device()->lang(), kCpp);
  Tensor y;
  y.ResetLike(x);

  Tensor w = get_bn_weight_from(bnScale, bnBias);

  y.device()->Exec(
      [y, w, x, &running_mean, &running_var, &bnh](Context* ctx) mutable {
        auto eng = ctx->dnnl_engine;
        using namespace dnnl;

        auto x_mem = memory(bnh.x_md, eng, x.block()->mutable_data());
        auto y_mem = memory(bnh.x_md, eng, y.block()->mutable_data());
        // indicates using scale&bias and running mean&var
        auto flags_ = normalization_flags::use_scale_shift |
                      normalization_flags::use_global_stats;

        auto bn_fwd_d = batch_normalization_forward::desc(
            prop_kind::forward_inference, bnh.x_md, bnh.epsilon, flags_);
        auto bn_fwd_pd =
            batch_normalization_forward::primitive_desc(bn_fwd_d, eng);
        auto m_mem = memory(bn_fwd_pd.mean_desc(), eng,
                            running_mean.block()->mutable_data());
        auto v_mem = memory(bn_fwd_pd.variance_desc(), eng,
                            running_var.block()->mutable_data());
        auto w_mem =
            memory(bn_fwd_pd.weights_desc(), eng, w.block()->mutable_data());

        // execution
        batch_normalization_forward(bn_fwd_pd).execute(
            ctx->dnnl_stream, {{DNNL_ARG_SRC, x_mem},
                               {DNNL_ARG_DST, y_mem},
                               {DNNL_ARG_SCALE_SHIFT, w_mem},
                               {DNNL_ARG_MEAN, m_mem},
                               {DNNL_ARG_VARIANCE, v_mem}});
        ctx->dnnl_stream.wait();
      },
      {x.block(), w.block(), running_mean.block(), running_var.block()},
      {y.block(), running_mean.block(), running_var.block()}, "CpuBatchNormForwardInference");

  return y;
}

const std::vector<Tensor> CpuBatchNormForwardTraining(
    const BatchNormHandle& bnh, const Tensor& x, const Tensor& bnScale,
    const Tensor& bnBias, Tensor& running_mean, Tensor& running_var) {
  CHECK_EQ(x.device()->lang(), kCpp);
  Tensor y;
  y.ResetLike(x);

  // mean and var for local batch
  Tensor mean;
  mean.ResetLike(running_mean);
  Tensor var;
  var.ResetLike(running_var);

  // combine scale and bias to construct weight tensor in required format for
  // backward
  Tensor w = get_bn_weight_from(bnScale, bnBias);

  y.device()->Exec(
      [y, mean, var, w, x, &running_mean, &running_var,
       &bnh](Context* ctx) mutable {
        auto eng = ctx->dnnl_engine;
        using namespace dnnl;

        auto x_mem = memory(bnh.x_md, eng, x.block()->mutable_data());
        auto y_mem = memory(bnh.x_md, eng, y.block()->mutable_data());
        auto m_mem = memory(bnh.bn_fwd_training_pd->mean_desc(), eng,
                            mean.block()->mutable_data());
        auto v_mem = memory(bnh.bn_fwd_training_pd->variance_desc(), eng,
                            var.block()->mutable_data());
        auto w_mem = memory(bnh.bn_fwd_training_pd->weights_desc(), eng,
                            w.block()->mutable_data());

        batch_normalization_forward(*bnh.bn_fwd_training_pd)
            .execute(ctx->dnnl_stream, {{DNNL_ARG_SRC, x_mem},
                                        {DNNL_ARG_DST, y_mem},
                                        {DNNL_ARG_SCALE_SHIFT, w_mem},
                                        {DNNL_ARG_MEAN, m_mem},
                                        {DNNL_ARG_VARIANCE, v_mem}});
        ctx->dnnl_stream.wait();

        // local implemented running mean as mkldnn does not support it yet:
        // https://github.com/intel/mkl-dnn/issues/371
        // https://github.com/intel/mkl-dnn/issues/517
        // https://arxiv.org/pdf/1502.03167.pdf
        auto s = x.shape();
        s[1] = 1;
        float p = Product(s);  // for unbiased variance
        running_mean = running_mean * (1 - bnh.factor) + mean * bnh.factor;
        running_var =
            running_var * (1 - bnh.factor) + var * (p / (p - 1)) * bnh.factor;
      },
      {x.block(), w.block(), running_mean.block(), running_var.block()},
      {y.block(), running_mean.block(), running_var.block(), mean.block(),
       var.block()}, "CpuBatchNormForwardTraining");

  return {y, mean, var};
}

const std::vector<Tensor> CpuBatchNormBackwardx(
    const BatchNormHandle& bnh, const Tensor& y, const Tensor& dy,
    const Tensor& x, const Tensor& bnScale, const Tensor& bnBias,
    const Tensor& mean, const Tensor& var) {
  CHECK_EQ(x.device()->lang(), kCpp);
  CHECK_EQ(y.device()->lang(), kCpp);
  CHECK_EQ(dy.device()->lang(), kCpp);
  CHECK_EQ(mean.device()->lang(), kCpp);
  CHECK_EQ(var.device()->lang(), kCpp);
  CHECK_EQ(bnScale.device()->lang(), kCpp);
  CHECK_EQ(bnBias.device()->lang(), kCpp);

  Tensor dx;
  dx.ResetLike(dy);

  // combine scale and bias to construct weight tensor in required format for
  // backward
  Tensor w = get_bn_weight_from(bnScale, bnBias);

  // Tensor dw(Shape{bnScale.Size(), 2});
  Tensor dw;
  dw.ResetLike(w);

  dx.device()->Exec(
      [w, dw, dx, dy, x, y, mean, var, &bnh](Context* ctx) mutable {
        auto eng = ctx->dnnl_engine;
        using namespace dnnl;

        auto x_mem = memory(bnh.x_md, eng, x.block()->mutable_data());
        auto dx_mem = memory(bnh.x_md, eng, dx.block()->mutable_data());
        auto y_mem = memory(bnh.x_md, eng, y.block()->mutable_data());
        auto dy_mem = memory(bnh.x_md, eng, dy.block()->mutable_data());

        auto m_mem = memory(bnh.bn_fwd_training_pd->mean_desc(), eng,
                            mean.block()->mutable_data());
        auto v_mem = memory(bnh.bn_fwd_training_pd->variance_desc(), eng,
                            var.block()->mutable_data());
        auto w_mem = memory(bnh.bn_fwd_training_pd->weights_desc(), eng,
                            w.block()->mutable_data());

        auto bn_bwd_d = batch_normalization_backward::desc(
            prop_kind::backward, bnh.x_md, bnh.x_md, bnh.epsilon,
            normalization_flags::use_scale_shift);
        auto bn_bwd_pd = batch_normalization_backward::primitive_desc(
            bn_bwd_d, eng, *bnh.bn_fwd_training_pd);

        auto dw_mem = memory(bn_bwd_pd.diff_weights_desc(), eng,
                             dw.block()->mutable_data());

        batch_normalization_backward(bn_bwd_pd).execute(
            ctx->dnnl_stream, {{DNNL_ARG_SRC, x_mem},
                               {DNNL_ARG_DIFF_SRC, dx_mem},
                               {DNNL_ARG_DIFF_DST, dy_mem},
                               {DNNL_ARG_MEAN, m_mem},
                               {DNNL_ARG_VARIANCE, v_mem},
                               {DNNL_ARG_DIFF_SCALE_SHIFT, dw_mem},
                               {DNNL_ARG_SCALE_SHIFT, w_mem}});
        ctx->dnnl_stream.wait();
      },
      {x.block(), dy.block(), mean.block(), var.block(), w.block(), y.block()},
      {dx.block(), dw.block()}, "CpuBatchNormBackwardx");

  singa::Tensor dbnScale(bnScale.shape());
  CopyDataToFrom(&dbnScale, dw, bnScale.Size(), 0, 0);
  singa::Tensor dbnBias(bnBias.shape());
  CopyDataToFrom(&dbnBias, dw, bnBias.Size(), 0, bnScale.Size());

  CHECK(dbnScale.nDim() == bnScale.nDim()) << "dbnScale ndim not match bnScale";
  CHECK(dbnBias.nDim() == bnBias.nDim()) << "dbnScale ndim not match bnScale";
  CHECK(dbnScale.shape()[0] == bnScale.shape()[0])
      << "dbnScale shape not match bnScale";
  CHECK(dbnBias.shape()[0] == bnBias.shape()[0])
      << "dbnBias shape not match bnBias";

  return {dx, dbnScale, dbnBias};
}

#endif  // USE_DNNL

#ifdef USE_CUDNN
CudnnBatchNormHandle::CudnnBatchNormHandle(const float momentum,
                                           const Tensor& input)
    : BatchNormHandle(momentum, input) {
  if (is_2d) {
    mode = CUDNN_BATCHNORM_PER_ACTIVATION;
  } else {
    mode = CUDNN_BATCHNORM_SPATIAL;
    if (const char* env_p = std::getenv("CUDNN_BATCHNORM_ALG")) {
      std::string alg = std::string(env_p);
      std::transform(alg.begin(), alg.end(), alg.begin(), toupper);
      if (alg == "CUDNN_BATCHNORM_SPATIAL_PERSISTENT")
        mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
      LOG(INFO) << " CUDNN_BATCHNORM_ALG: " << alg;
    }
  }
  DataType dtype = input.data_type();
  CUDNN_CHECK(cudnnCreateTensorDescriptor(&shape_desc));
  CUDNN_CHECK(cudnnCreateTensorDescriptor(&param_desc));
  CUDNN_CHECK(cudnnSetTensor4dDescriptor(shape_desc, CUDNN_TENSOR_NCHW,
                                         GetCudnnDataType(dtype), batchsize,
                                         channels, height, width));
  CUDNN_CHECK(cudnnSetTensor4dDescriptor(param_desc, CUDNN_TENSOR_NCHW,
                                         GetCudnnDataType(dtype), 1, channels,
                                         1, 1));
};

const std::vector<Tensor> GpuBatchNormForwardTraining(
    const CudnnBatchNormHandle& cbnh, const Tensor& x, const Tensor& bnScale,
    const Tensor& bnBias, Tensor& running_mean, Tensor& running_var) {
  CHECK_EQ(x.device()->lang(), kCuda);
  CHECK_EQ(bnScale.device()->lang(), kCuda);
  CHECK_EQ(bnBias.device()->lang(), kCuda);
  CHECK_EQ(running_mean.device()->lang(), kCuda);
  CHECK_EQ(running_var.device()->lang(), kCuda);

  Tensor mean;
  Tensor var;
  mean.ResetLike(running_mean);
  var.ResetLike(running_var);

  Shape shape = x.shape();

  Tensor input(x);  // for unification of 2d and 4d cases.
  if (cbnh.is_2d) input.Reshape(Shape{shape.at(0), shape.at(1), 1, 1});

  Tensor output;
  output.ResetLike(x);

  output.device()->Exec(
      [=, &bnScale, &bnBias, &running_mean, &running_var,
       &cbnh](Context* ctx) mutable {
        const float alpha = 1.0f, beta = 0.0f;
        double epsilon = CUDNN_BN_MIN_EPSILON;
        CUDNN_CHECK(cudnnBatchNormalizationForwardTraining(
            ctx->cudnn_handle, cbnh.mode, &alpha, &beta, cbnh.shape_desc,
            input.block()->data(), cbnh.shape_desc,
            output.block()->mutable_data(), cbnh.param_desc,
            bnScale.block()->data(), bnBias.block()->data(), cbnh.factor,
            running_mean.block()->mutable_data(),
            running_var.block()->mutable_data(), epsilon,
            mean.block()->mutable_data(), var.block()->mutable_data()));
      },
      {input.block(), bnScale.block(), bnBias.block(), running_mean.block(),
       running_var.block()},
      {output.block(), running_mean.block(), running_var.block(), mean.block(),
       var.block()}, "GpuBatchNormForwardTraining");
  if (cbnh.is_2d) output.Reshape(Shape{shape.at(0), shape.at(1)});
  return {output, mean, var};
}

Tensor GpuBatchNormForwardInference(const CudnnBatchNormHandle& cbnh,
                                    const Tensor& x, const Tensor& bnScale,
                                    const Tensor& bnBias,
                                    const Tensor& running_mean,
                                    const Tensor& running_var) {
  CHECK_EQ(x.device()->lang(), kCuda);
  CHECK_EQ(bnScale.device()->lang(), kCuda);
  CHECK_EQ(bnBias.device()->lang(), kCuda);
  CHECK_EQ(running_mean.device()->lang(), kCuda);
  CHECK_EQ(running_var.device()->lang(), kCuda);

  Shape shape = x.shape();

  Tensor input(x);  // for unification of 2d and 4d cases.
  if (cbnh.is_2d) input.Reshape(Shape{shape.at(0), shape.at(1), 1, 1});

  Tensor output;
  output.ResetLike(x);
  output.device()->Exec(
      [=, &bnScale, &bnBias, &running_mean, &running_var,
       &cbnh](Context* ctx) mutable {
        const float alpha = 1.0f, beta = 0.0f;
        double epsilon = CUDNN_BN_MIN_EPSILON;
        CUDNN_CHECK(cudnnBatchNormalizationForwardInference(
            ctx->cudnn_handle, cbnh.mode, &alpha, &beta, cbnh.shape_desc,
            input.block()->data(), cbnh.shape_desc,
            output.block()->mutable_data(), cbnh.param_desc,
            bnScale.block()->data(), bnBias.block()->data(),
            running_mean.block()->data(), running_var.block()->data(),
            epsilon));
      },
      {input.block(), bnScale.block(), bnBias.block(), running_mean.block(),
       running_var.block()},
      {output.block()}, "GpuBatchNormForwardInference");
  return output;
}

const std::vector<Tensor> GpuBatchNormBackward(
    const CudnnBatchNormHandle& cbnh, const Tensor& dy, const Tensor& x,
    const Tensor& bnScale, const Tensor& mean, const Tensor& var) {
  CHECK_EQ(dy.device()->lang(), kCuda);
  CHECK_EQ(x.device()->lang(), kCuda);
  CHECK_EQ(bnScale.device()->lang(), kCuda);
  CHECK_EQ(mean.device()->lang(), kCuda);
  CHECK_EQ(var.device()->lang(), kCuda);

  Tensor dx;
  dx.ResetLike(dy);

  Tensor dbnScale;
  dbnScale.ResetLike(bnScale);

  Tensor dbnBias;
  dbnBias.ResetLike(bnScale);

  dx.device()->Exec(
      [=, &bnScale, &cbnh](Context* ctx) mutable {
        const float alpha = 1.0f, beta = .0f;
        double epsilon = CUDNN_BN_MIN_EPSILON;
        CUDNN_CHECK(cudnnBatchNormalizationBackward(
            ctx->cudnn_handle, cbnh.mode, &alpha, &beta, &alpha, &beta,
            cbnh.shape_desc, x.block()->data(), cbnh.shape_desc,
            dy.block()->data(), cbnh.shape_desc, dx.block()->mutable_data(),
            cbnh.param_desc, bnScale.block()->data(),
            dbnScale.block()->mutable_data(), dbnBias.block()->mutable_data(),
            epsilon, mean.block()->data(), var.block()->data()));
      },
      {x.block(), dy.block(), bnScale.block(), mean.block(), var.block()},
      {dx.block(), dbnScale.block(), dbnBias.block()}, "GpuBatchNormBackward");

  if (cbnh.is_2d) dx.Reshape(Shape{dx.shape().at(0), dx.shape().at(1)});

  return {dx, dbnScale, dbnBias};
}

#endif  // USE_CUDNN
}  // namespace singa