/*
 * 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 dnnl_fc.cc
 * \brief DNNL (Quantized) FullyConnected operator based on subgraph
 * \author Ciyong Chen
 */

#if MXNET_USE_ONEDNN == 1

#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include "operator/nn/dnnl/dnnl_act-inl.h"
#include "operator/nn/dnnl/dnnl_base-inl.h"
#include "operator/nn/dnnl/dnnl_fully_connected-inl.h"
#include "operator/quantization/quantization_utils.h"
#include "operator/tensor/matrix_op-inl.h"
#include "operator/subgraph/common.h"
#include "dnnl_common.h"
#include "dnnl_fc-inl.h"

namespace mxnet {
namespace op {

class SgDNNLFCOp {
 public:
  explicit SgDNNLFCOp(const nnvm::NodeAttrs& attrs)
      : subgraph_sym_(*attrs.subgraphs[0]),
        attrs(attrs),
        full_param_(nnvm::get<DNNLFCFullParam>(attrs.parsed)) {}

  void Forward(const OpContext& ctx,
               const std::vector<NDArray>& inputs,
               const std::vector<OpReqType>& req,
               const std::vector<NDArray>& outputs);

  void Backward(const OpContext& ctx,
                const std::vector<NDArray>& inputs,
                const std::vector<OpReqType>& req,
                const std::vector<NDArray>& outputs) {
    LOG(FATAL) << "Not implemented: subgraph oneDNN fully connected only supports "
                  "inference computation.";
  }

 private:
  enum { kDataMin = 0, kDataMax, kWeightMin, kWeightMax, kBiasMin, kBiasMax, kSumMin, kSumMax };
  const size_t MIN_MAX_COUNT = 8;
  const float u8_to_s8_scale = 0.5;

  NDArray PrepareOutputWithSum(const NDArray& sum_input, const NDArray& output);
  bool CheckInitializationConditions(const std::vector<NDArray>& inputs,
                                     const std::vector<float>& min_max_vec,
                                     bool is_channel_wise);
  bool PrepareQuantization(const OpContext& ctx,
                           const std::vector<NDArray>& in_data,
                           const NDArray& output,
                           const std::vector<float>& min_max_vec);
  dnnl::memory::desc CreateOutputMemoryDesc(const mxnet::TShape& oshape, int out_dtype);
  void GetCachedWeightsAndBias(const NDArray& weight,
                               bool support_channelwise_scale,
                               bool has_bias);
  nnvm::Symbol subgraph_sym_;
  nnvm::NodeAttrs attrs;
  DNNLFCFullParam full_param_;
  bool initialized_{false};
  bool reorder_data_{false};
  bool inplace_{false};
  dnnl_args_map_t args_;
  std::shared_ptr<DNNLFullyConnectedForward> fwd_;
  std::shared_ptr<dnnl::memory> cached_data_mem_;
  std::shared_ptr<dnnl::memory> cached_out_mem_;
  NDArray cached_weight_;
  NDArray cached_bias_;
  float cached_data_min_;
  float cached_data_max_;
  float cached_weight_min_;
  float cached_weight_max_;
  float cached_sum_min_;
  float cached_sum_max_;
  float cached_bias_min_;
  float cached_bias_max_;
  size_t weight_ver_;
  size_t bias_ver_;
  float cached_output_min_;
  float cached_output_max_;
  float data_scale_{0.0f};
  std::vector<float> weight_scales_;
};

void SgDNNLFCOp::Forward(const OpContext& ctx,
                         const std::vector<NDArray>& in_data,
                         const std::vector<OpReqType>& req,
                         const std::vector<NDArray>& out_data) {
  const FCInputIndex idx(full_param_);
  CHECK_EQ(in_data.size(), idx.GetTotal());
  const int out_index     = 0;
  const int out_min_index = 1;
  const int out_max_index = 2;

  const auto& default_param = full_param_.default_param;
  const auto& dnnl_param    = full_param_.dnnl_param;

  NDArray output;
  if (dnnl_param.with_sum) {
    output = PrepareOutputWithSum(in_data[idx.sum], out_data[out_index]);
  } else {
    output = out_data[out_index];
  }

  if (!dnnl_param.quantized) {
    // dispatch to float version
    DNNLFCForwardImpl(full_param_, ctx, in_data, req, {output});
    return;
  }

  const bool has_bias = !default_param.no_bias;
  const bool channel_wise =
      dnnl_param.channel_wise_quantize.has_value() && dnnl_param.channel_wise_quantize.value();

  std::vector<float> min_max_vec(MIN_MAX_COUNT);
  min_max_vec[kDataMin]   = 0.0f;
  min_max_vec[kDataMax]   = 0.0f;
  min_max_vec[kWeightMin] = 0.0f;
  min_max_vec[kWeightMax] = 0.0f;
  min_max_vec[kBiasMin]   = 0.0f;
  min_max_vec[kBiasMax]   = 0.0f;

  min_max_vec[kSumMin]  = idx.sum_min ? in_data[idx.sum_min].data().dptr<float>()[0] : 0.0f;
  min_max_vec[kSumMax]  = idx.sum_max ? in_data[idx.sum_max].data().dptr<float>()[0] : 0.0f;
  NDArray data          = in_data[idx.data];
  const NDArray& weight = in_data[idx.weight];

  if (!channel_wise) {
    min_max_vec[kWeightMin] = in_data[idx.weight_min].data().dptr<float>()[0];
    min_max_vec[kWeightMax] = in_data[idx.weight_max].data().dptr<float>()[0];
    if (has_bias) {
      min_max_vec[kBiasMin] = in_data[idx.bias_min].data().dptr<float>()[0];
      min_max_vec[kBiasMax] = in_data[idx.bias_max].data().dptr<float>()[0];
    }
  }
  min_max_vec[kDataMin] = in_data[idx.data_min].data().dptr<float>()[0];
  min_max_vec[kDataMax] = in_data[idx.data_max].data().dptr<float>()[0];

  initialized_ = CheckInitializationConditions(in_data, min_max_vec, channel_wise);

  if (!initialized_) {
    const auto engine  = CpuEngine::Get()->get_engine();
    cached_data_min_   = min_max_vec[kDataMin];
    cached_data_max_   = min_max_vec[kDataMax];
    cached_weight_min_ = min_max_vec[kWeightMin];
    cached_weight_max_ = min_max_vec[kWeightMax];
    weight_ver_        = weight.version();
    cached_weight_     = weight;
    cached_sum_min_    = min_max_vec[kSumMin];
    cached_sum_max_    = min_max_vec[kSumMax];
    if (has_bias) {
      cached_bias_min_ = min_max_vec[kBiasMin];
      cached_bias_max_ = min_max_vec[kBiasMax];
      bias_ver_        = in_data[idx.bias].version();
      cached_bias_     = in_data[idx.bias];
    } else {
      cached_bias_ = NDArray();
    }
    const mxnet::TShape ishape = data.shape();
    const auto data_ndim       = ishape.ndim();
    if (data.IsDNNLData()) {
      reorder_data_ = true;
      data          = data.Reorder2Default();
    }
    if (data_ndim != 2) {
      if (!default_param.flatten) {
        data =
            data.DNNLDataReshape(Shape2(ishape.ProdShape(0, data_ndim - 1), ishape[data_ndim - 1]));
      } else {
        data = data.DNNLDataReshape(Shape2(ishape[0], ishape.ProdShape(1, data_ndim)));
      }
    }

    // create cached out_md
    dnnl::memory::desc out_md = CreateOutputMemoryDesc(output.shape(), output.dtype());
    cached_out_mem_           = std::make_shared<dnnl::memory>(out_md, engine);

    bool support_channelwise_scale = PrepareQuantization(ctx, in_data, output, min_max_vec);

    fwd_.reset(new DNNLFullyConnectedForward(full_param_,
                                             ctx.is_train,
                                             data,
                                             cached_weight_,
                                             (has_bias ? &cached_bias_ : nullptr),
                                             out_md));
    GetCachedWeightsAndBias(weight, support_channelwise_scale, has_bias);

    const auto data_mem = static_cast<const dnnl::memory*>(data.GetDNNLData());
    cached_data_mem_    = std::make_shared<dnnl::memory>(data_mem->get_desc(), engine);

    args_[DNNL_ARG_SRC]     = *cached_data_mem_;
    args_[DNNL_ARG_WEIGHTS] = *static_cast<const dnnl::memory*>(cached_weight_.GetDNNLData());
    if (has_bias)
      args_[DNNL_ARG_BIAS] = *static_cast<const dnnl::memory*>(cached_bias_.GetDNNLData());
    args_[DNNL_ARG_DST] = *cached_out_mem_;
    initialized_        = true;
  }

  if (reorder_data_) {
    data = data.Reorder2Default();
  }

  cached_data_mem_->set_data_handle(reinterpret_cast<void*>(data.data().dptr_));
  cached_out_mem_->set_data_handle(reinterpret_cast<void*>(output.data().dptr_));

  DNNLStream::Get()->RegisterPrimArgs(fwd_->GetFwd(), args_);
  DNNLStream::Get()->Submit();

  if (!dnnl_param.enabled_float_output.has_value()) {
    float* output_min_ptr = out_data[out_min_index].data().dptr<float>();
    float* output_max_ptr = out_data[out_max_index].data().dptr<float>();

    *output_min_ptr = cached_output_min_;
    *output_max_ptr = cached_output_max_;
  }
}

NDArray SgDNNLFCOp::PrepareOutputWithSum(const NDArray& sum_input, const NDArray& output) {
  if (!initialized_) {
    // TODO(zhennan): Currently, dnnl fallback mechanism will break inplace option,
    // which make check (req[out_index] == kWriteInplace) useless.
    auto in_dnnl_mem  = static_cast<const dnnl::memory*>(sum_input.GetDNNLData());
    auto out_dnnl_mem = static_cast<const dnnl::memory*>(output.GetDNNLData());
    if (in_dnnl_mem->get_data_handle() == out_dnnl_mem->get_data_handle() &&
        sum_input.dtype() == output.dtype()) {
      inplace_ = true;
    }
  }
  if (inplace_) {
    return sum_input;
  } else {
    // Not in place: copy sum_input into output.
    auto in_dnnl_mem  = static_cast<const dnnl::memory*>(sum_input.GetDNNLData());
    auto out_dnnl_mem = static_cast<const dnnl::memory*>(output.GetDNNLData());
    if (output.dtype() == mshadow::kInt32) {
      auto mem_desc           = in_dnnl_mem->get_desc();
      auto this_dtype         = get_dnnl_type(mshadow::kInt32);
      mem_desc.data.data_type = static_cast<dnnl_data_type_t>(this_dtype);
      dnnl_mem_ptr tmp_mem(new dnnl::memory(
          mem_desc, CpuEngine::Get()->get_engine(), out_dnnl_mem->get_data_handle()));
      DNNLStream::Get()->RegisterMem(tmp_mem);
      DNNLStream::Get()->RegisterPrimArgs(dnnl::reorder(*in_dnnl_mem, *tmp_mem),
                                          {{DNNL_ARG_FROM, *in_dnnl_mem}, {DNNL_ARG_TO, *tmp_mem}});
      return NDArray(tmp_mem);
    } else if (sum_input.dtype() == mshadow::kUint8 && output.dtype() == mshadow::kInt8) {
      auto sum_mem_desc           = in_dnnl_mem->get_desc();
      auto out_dtype              = get_dnnl_type(mshadow::kInt8);
      sum_mem_desc.data.data_type = static_cast<dnnl_data_type_t>(out_dtype);
      dnnl_mem_ptr tmp_mem(new dnnl::memory(
          sum_mem_desc, CpuEngine::Get()->get_engine(), out_dnnl_mem->get_data_handle()));
      DNNLStream::Get()->RegisterMem(tmp_mem);
      std::vector<float> reorder_scale = {u8_to_s8_scale};
      dnnl::primitive_attr reorder_attr;
      reorder_attr.set_output_scales(0, reorder_scale);
      const auto reorder_pd = dnnl::reorder::primitive_desc(CpuEngine::Get()->get_engine(),
                                                            in_dnnl_mem->get_desc(),
                                                            CpuEngine::Get()->get_engine(),
                                                            sum_mem_desc,
                                                            reorder_attr);
      DNNLStream::Get()->RegisterPrimArgs(dnnl::reorder(reorder_pd),
                                          {{DNNL_ARG_FROM, *in_dnnl_mem}, {DNNL_ARG_TO, *tmp_mem}});
      return NDArray(tmp_mem);
    } else {
      dnnl_mem_ptr tmp_mem(new dnnl::memory(in_dnnl_mem->get_desc(),
                                            CpuEngine::Get()->get_engine(),
                                            out_dnnl_mem->get_data_handle()));
      DNNLStream::Get()->RegisterMem(tmp_mem);
      DNNLMemoryCopy(*in_dnnl_mem, tmp_mem.get());
      return NDArray(tmp_mem);
    }
  }
}

bool SgDNNLFCOp::CheckInitializationConditions(const std::vector<NDArray>& inputs,
                                               const std::vector<float>& min_max_vec,
                                               bool is_channel_wise) {
  if (initialized_ && dmlc::GetEnv("MXNET_ONEDNN_QFC_DYNAMIC_PARAMS", 0)) {
    bool has_bias = !full_param_.default_param.no_bias;
    if (cached_data_min_ != min_max_vec[kDataMin] || cached_data_max_ != min_max_vec[kDataMax] ||
        cached_sum_min_ != min_max_vec[kSumMin] || cached_sum_max_ != min_max_vec[kSumMax]) {
      return false;
    }

    if (is_channel_wise) {
      if (weight_ver_ != inputs[fullc::kWeight].version() ||
          (has_bias && (bias_ver_ != inputs[fullc::kBias].version()))) {
        return false;
      }
    } else {
      if (cached_weight_min_ != min_max_vec[kWeightMin] ||
          cached_weight_max_ != min_max_vec[kWeightMax] ||
          (has_bias && (cached_bias_min_ != min_max_vec[kBiasMin] ||
                        cached_bias_max_ != min_max_vec[kBiasMax]))) {
        return false;
      }
    }
    return true;
  }
  return false;
}

dnnl::memory::desc SgDNNLFCOp::CreateOutputMemoryDesc(const mxnet::TShape& oshape, int out_dtype) {
  auto default_param = full_param_.default_param;
  dnnl::memory::dims out_dims(2);
  if (oshape.ndim() == 2) {
    out_dims[0] = static_cast<index_t>(oshape[0]);
    out_dims[1] = static_cast<index_t>(oshape[1]);
  } else {
    if (!default_param.flatten) {
      out_dims[0] = static_cast<index_t>(oshape.ProdShape(0, oshape.ndim() - 1));
      out_dims[1] = static_cast<index_t>(oshape[oshape.ndim() - 1]);
    } else {
      out_dims[0] = static_cast<index_t>(oshape[0]);
      out_dims[1] = static_cast<index_t>(oshape.ProdShape(1, oshape.ndim()));
    }
  }
  dnnl::memory::desc out_md =
      dnnl::memory::desc(out_dims,
                         get_dnnl_type(out_dtype),
                         static_cast<dnnl::memory::format_tag>(GetDefaultFormat(2)));
  return out_md;
}

bool SgDNNLFCOp::PrepareQuantization(const OpContext& ctx,
                                     const std::vector<NDArray>& in_data,
                                     const NDArray& output,
                                     const std::vector<float>& min_max_vec) {
  const auto nthreads = engine::OpenMP::Get()->GetRecommendedOMPThreadCount();
  const FCInputIndex idx(full_param_);
  bool support_channelwise_scale = false;
  auto dnnl_param                = full_param_.dnnl_param;
  bool has_bias                  = !full_param_.default_param.no_bias;
  const NDArray& data            = in_data[fullc::kData];
  const NDArray& weight          = in_data[fullc::kWeight];
  const bool channel_wise =
      dnnl_param.channel_wise_quantize.has_value() && dnnl_param.channel_wise_quantize.value();

  CHECK(data.dtype() == mshadow::kInt8 || data.dtype() == mshadow::kUint8);
  data_scale_ = GetQuantizeScale(data.dtype(), cached_data_min_, cached_data_max_);

  bool fuse_requantize = false;
  // Channelwise scaling is only supported when fusion is enabled (requantize or dequantize).
  if (dnnl_param.min_calib_range.has_value() && dnnl_param.max_calib_range.has_value()) {
    cached_output_min_        = dnnl_param.min_calib_range.value();
    cached_output_max_        = dnnl_param.max_calib_range.value();
    support_channelwise_scale = true;
    fuse_requantize           = true;
  }
  if (dnnl_param.enabled_float_output.has_value()) {
    support_channelwise_scale = true;
  }
  // channel_wise  support_channelwise_scale  result
  // True          True                       True
  // True          False                      Error
  // False         True/False                 False
  if (channel_wise && !support_channelwise_scale) {
    LOG(FATAL) << "Currently, channel-wise quantization requires fuse requantize or dequantize."
               << " Please make sure the `min_calib_range` and `max_calib_range` are set when only"
               << " fuse requantize (outputs of FullyConnected are collected during calibration "
                  "phase),"
               << " or the env var of `MXNET_DISABLE_ONEDNN_QFC_FLOAT_OUTPUT` and "
               << " `MXNET_DISABLE_ONEDNN_QFC_FUSE_ALL` are not set to true (default is false)";
  }
  support_channelwise_scale = support_channelwise_scale && channel_wise;

  if (support_channelwise_scale) {
    MSHADOW_REAL_TYPE_SWITCH(cached_weight_.dtype(), DType, {
      weight_scales_ = GetWeightScales<DType>(cached_weight_,
                                              has_bias ? &cached_bias_ : nullptr,
                                              data_scale_,
                                              support_channelwise_scale);
    });
  } else {
    weight_scales_.resize(1);
    weight_scales_[0] =
        GetQuantizeScale(cached_weight_.dtype(), cached_weight_min_, cached_weight_max_);
    if (has_bias) {
      if (cached_bias_.dtype() == mshadow::kInt8) {
        float bias_scale = GetQuantizeScale(mshadow::kInt8, cached_bias_min_, cached_bias_max_);

        float bias_int32_rescale = data_scale_ * weight_scales_[0] / bias_scale;
        // TODO(zhennan): dnnl has bug to handle INT_MAX in bias, so set
        // the maximum value of bias to INT_MAX / 2.
        float bias_max_rescale =
            MaxValue<int32_t>() / 2 / MaxAbs(cached_bias_min_, cached_bias_max_) / bias_scale;
        if (bias_int32_rescale > bias_max_rescale) {
          // avoid overflow on bias
          bias_int32_rescale   = bias_max_rescale;
          float weight_rescale = bias_int32_rescale * bias_scale / data_scale_ / weight_scales_[0];
          int8_t* weight_ptr   = weight.data().dptr<int8_t>();
          size_t weight_size   = weight.shape().Size();
#pragma omp parallel for num_threads(nthreads)
          for (index_t i = 0; i < static_cast<index_t>(weight_size); ++i) {
            weight_ptr[i] = std::round(weight_ptr[i] * weight_rescale);
          }
          weight_scales_[0] *= weight_rescale;
        }
        NDArray bias = in_data[fullc::kBias];
        cached_bias_ =
            NDArray(bias.storage_type(), bias.shape(), bias.ctx(), true, mshadow::kInt32);
        int8_t* bias_ptr            = bias.data().dptr<int8_t>();
        int32_t* quantized_bias_ptr = cached_bias_.data().dptr<int32_t>();
        size_t bias_size            = bias.shape().Size();

#pragma omp parallel for num_threads(nthreads)
        for (index_t i = 0; i < static_cast<index_t>(bias_size); ++i) {
          quantized_bias_ptr[i] = std::round(bias_ptr[i] * bias_int32_rescale);
        }
      }
    }
  }

  size_t num_channel = cached_weight_.shape()[0];
  float out_scale    = 1.0f;
  if (fuse_requantize || dnnl_param.enabled_float_output.has_value()) {
    float tmp_scale_ = 1.0f;
    if (fuse_requantize) {
      if (dnnl_param.with_eltwise) {
        tmp_scale_ = 1.0 / data_scale_;
        full_param_.eltwise_param.scale =
            GetQuantizeScale(output.dtype(), cached_output_min_, cached_output_max_);
      } else {
        out_scale  = GetQuantizeScale(output.dtype(), cached_output_min_, cached_output_max_);
        tmp_scale_ = out_scale / data_scale_;
      }
    } else {
      tmp_scale_ = 1.0 / data_scale_;
    }

    if (support_channelwise_scale) {
      full_param_.output_scales.resize(num_channel);
#pragma omp parallel for num_threads(nthreads)
      for (index_t i = 0; i < static_cast<index_t>(num_channel); ++i) {
        full_param_.output_scales[i] = tmp_scale_ / weight_scales_[i];
      }
    } else {
      full_param_.output_scales.resize(1);
      full_param_.output_scales[0] = tmp_scale_ / weight_scales_[0];
    }
  } else {
    Stream<cpu>* s = ctx.get_stream<cpu>();
    if (data.dtype() == mshadow::kInt8) {
      mxnet_op::Kernel<QuantizationRangeForS8S8MultiplicationStruct, cpu>::Launch(
          s,
          1,
          &cached_output_min_,
          &cached_output_max_,
          &(min_max_vec[kDataMin]),
          &(min_max_vec[kDataMax]),
          &(min_max_vec[kWeightMin]),
          &(min_max_vec[kWeightMax]));
    } else {
      mxnet_op::Kernel<QuantizationRangeForS8U8MultiplicationStruct, cpu>::Launch(
          s,
          1,
          &cached_output_min_,
          &cached_output_max_,
          &(min_max_vec[kDataMin]),
          &(min_max_vec[kDataMax]),
          &(min_max_vec[kWeightMin]),
          &(min_max_vec[kWeightMax]));
    }
    full_param_.output_scales.resize(0);
    out_scale = data_scale_ * weight_scales_[0];
  }

  if (dnnl_param.with_sum && !dnnl_param.enabled_float_output.has_value()) {
    float sum_in_scale =
        GetQuantizeScale(in_data[idx.sum].dtype(), cached_sum_min_, cached_sum_max_);
    full_param_.sum_scale = out_scale / sum_in_scale;
    if (in_data[idx.sum].dtype() == mshadow::kUint8 && output.dtype() == mshadow::kInt8) {
      // In this case, reorder with scale 0.5 is used on in_data[idx.sum] to
      // scale it to s8 range, so sum_scale has to be rescaled as well
      full_param_.sum_scale *= 1.0 / u8_to_s8_scale;
    }
  }
  return support_channelwise_scale;
}

void SgDNNLFCOp::GetCachedWeightsAndBias(const NDArray& weight,
                                         bool support_channelwise_scale,
                                         bool has_bias) {
#if DMLC_CXX11_THREAD_LOCAL
  static thread_local std::unordered_map<DNNLFullyconSignature, std::pair<NDArray, NDArray>, OpHash>
      fcWeightsAndBias;
#else
  static MX_THREAD_LOCAL
      std::unordered_map<DNNLFullyconSignature, std::pair<NDArray, NDArray>, OpHash>
          fcWeightsAndBias;
#endif
  static const bool use_cache = !(dmlc::GetEnv("MXNET_ONEDNN_DISABLE_FC_CACHE", 0));
  const bool has_id           = attrs.dict.count("__identifier__");
  bool read_from_cache        = false;

  DNNLFullyconSignature key(full_param_);
  if (use_cache && has_id) {
    key.AddSign(fwd_->fwd_pd.weights_desc());
    key.AddSign(attrs.dict["__identifier__"]);
    key.AddSign(attrs.name);

    auto it = fcWeightsAndBias.find(key);
    if (it != fcWeightsAndBias.end()) {
      cached_weight_  = it->second.first;
      cached_bias_    = it->second.second;
      read_from_cache = true;
      common::LogOnce(
          "oneDNN optimized version of FullyConnected for inference is being used. Weights and "
          "bias are cached and can not be dynamically changed during runtime. To disable caching "
          "mechanism use MXNET_ONEDNN_DISABLE_FC_CACHE=1.");
    }
  }

  if (!read_from_cache) {
    // convert weight and bias to the format that oneDNN requires
    if (!full_param_.dnnl_param.quantized || support_channelwise_scale) {
      dnnl::memory::desc bias_md;
      if (has_bias)
        bias_md = fwd_->fwd_pd.bias_desc();
      ConvertWeightBias2DNNL(&cached_weight_,
                             &cached_bias_,
                             has_bias,
                             fwd_->fwd_pd.weights_desc(),
                             has_bias ? &bias_md : nullptr,
                             1,
                             data_scale_,
                             weight_scales_,
                             false);
    } else {
      const auto def_weight_mem = weight.GetDNNLData();
      if (def_weight_mem->get_desc() != fwd_->fwd_pd.weights_desc()) {
        auto weight_desc       = fwd_->fwd_pd.weights_desc();
        cached_weight_         = NDArray(&weight_desc);
        auto cached_weight_mem = cached_weight_.GetDNNLData();
        std::unordered_map<int, dnnl::memory> args(
            {{DNNL_ARG_FROM, *def_weight_mem}, {DNNL_ARG_TO, *cached_weight_mem}});
        DNNLStream::Get()->RegisterPrimArgs(dnnl::reorder(*def_weight_mem, *cached_weight_mem),
                                            args);
      }
    }
    if (use_cache && has_id)
      AddToCache(&fcWeightsAndBias, key, {cached_weight_, cached_bias_});
  }
}

static void SgDNNLFCParamParser(nnvm::NodeAttrs* attrs) {
  // For backward compatible, with_relu->with_eltwise
  auto legacy = attrs->dict.find("with_relu");
  if (legacy != attrs->dict.end()) {
    attrs->dict["with_eltwise"] = attrs->dict["with_relu"];
    attrs->dict.erase(legacy);
  }

  DNNLFCFullParam full_param;
  try {
    full_param.dnnl_param.Init(attrs->dict);
  } catch (const dmlc::ParamError& e) {
    std::ostringstream os;
    os << e.what();
    os << ", in operator " << attrs->op->name << "("
       << "name=\"" << attrs->name << "\"";
    for (const auto& k : attrs->dict) {
      os << ", " << k.first << "=\"" << k.second << "\"";
    }
    os << ")";
    throw dmlc::ParamError(os.str());
  }
  auto subgraph_sym = attrs->subgraphs[0];
  DFSVisit(subgraph_sym->outputs, [&](const nnvm::ObjectPtr& node) {
    if (node->is_variable())
      return;
    auto& op_name = node->op()->name;
    if (op_name == "FullyConnected") {
      full_param.default_param = nnvm::get<FullyConnectedParam>(node->attrs.parsed);
    } else if (SupportDNNLFCEltwiseFusion(op_name)) {
      if (op_name == "Activation") {
        const ActivationParam act_param = nnvm::get<ActivationParam>(node->attrs.parsed);
        full_param.eltwise_param.alg    = GetDNNLActAlgo(act_param);
      } else if (op_name == "LeakyReLU") {
        const auto act_param           = nnvm::get<LeakyReLUParam>(node->attrs.parsed);
        full_param.eltwise_param.alpha = act_param.slope;
        full_param.eltwise_param.alg   = GetDNNLActAlgo(act_param);
      } else if (op_name == "clip") {
        const ClipParam clip_param     = nnvm::get<ClipParam>(node->attrs.parsed);
        full_param.eltwise_param.alg   = dnnl::algorithm::eltwise_bounded_relu;
        full_param.eltwise_param.alpha = clip_param.a_max;
      } else {
        full_param.eltwise_param.alg = GetDNNLEltwiseAlgo(op_name);
      }
    }
  });
  attrs->parsed = std::move(full_param);
}

static std::vector<std::string> SgDNNLFCListInputNames(const NodeAttrs& attrs) {
  auto const& full_param               = nnvm::get<DNNLFCFullParam>(attrs.parsed);
  auto const& dnnl_param               = full_param.dnnl_param;
  std::vector<std::string> input_names = DefaultSubgraphOpListInputs(attrs);
  if (dnnl_param.quantized) {
    const bool channel_wise =
        dnnl_param.channel_wise_quantize.has_value() && dnnl_param.channel_wise_quantize;
    input_names.emplace_back("data_min");
    input_names.emplace_back("data_max");
    if (!channel_wise) {
      input_names.emplace_back("weight_min");
      input_names.emplace_back("weight_max");
      if (!full_param.default_param.no_bias) {
        input_names.emplace_back("bias_min");
        input_names.emplace_back("bias_max");
      }
    }
    if (dnnl_param.with_sum && !dnnl_param.enabled_float_output.has_value()) {
      input_names.emplace_back("sum_min");
      input_names.emplace_back("sum_max");
    }
  }
  return input_names;
}

static std::vector<std::string> SgDNNLFCListOutputNames(const NodeAttrs& attrs) {
  auto const& full_param = nnvm::get<DNNLFCFullParam>(attrs.parsed);
  if (full_param.dnnl_param.quantized) {
    if (full_param.dnnl_param.enabled_float_output.has_value())
      return std::vector<std::string>{"output"};
    else
      return std::vector<std::string>{"output", "output_min", "output_max"};
  } else {
    return std::vector<std::string>{"output"};
  }
}

template <typename T>
static inline void FillBaseInputOutputInfo(const DNNLFCFullParam& param,
                                           std::vector<T>* base_in_attrs,
                                           std::vector<T>* base_out_attrs,
                                           std::vector<T>* in_attrs,
                                           std::vector<T>* out_attrs) {
  auto base_num_inputs = FCInputIndex(param).GetBase();

  base_out_attrs->push_back(out_attrs->at(0));
  for (int i = 0; i < base_num_inputs; ++i) {
    base_in_attrs->push_back(in_attrs->at(i));
  }
}

static bool SgDNNLFCInferShape(const nnvm::NodeAttrs& attrs,
                               mxnet::ShapeVector* in_shapes,
                               mxnet::ShapeVector* out_shapes) {
  auto const& full_param = nnvm::get<DNNLFCFullParam>(attrs.parsed);
  if (full_param.dnnl_param.quantized) {
    mxnet::ShapeVector base_in_shapes;
    mxnet::ShapeVector base_out_shapes;
    FillBaseInputOutputInfo(full_param, &base_in_shapes, &base_out_shapes, in_shapes, out_shapes);
    bool ret = DefaultSubgraphOpShape(attrs, &base_in_shapes, &base_out_shapes);

    for (size_t i = 0; i < in_shapes->size(); ++i) {
      if (i < base_in_shapes.size())
        in_shapes->at(i) = base_in_shapes[i];
      else
        SHAPE_ASSIGN_CHECK(*in_shapes, i, Shape1(1));
    }

    out_shapes->at(0) = base_out_shapes[0];
    if (!full_param.dnnl_param.enabled_float_output.has_value()) {
      SHAPE_ASSIGN_CHECK(*out_shapes, 1, Shape1(1));
      SHAPE_ASSIGN_CHECK(*out_shapes, 2, Shape1(1));
    }
    return ret;
  } else {
    return DefaultSubgraphOpShape(attrs, in_shapes, out_shapes);
  }
}

static bool SgDNNLFCInferType(const nnvm::NodeAttrs& attrs,
                              std::vector<int>* in_types,
                              std::vector<int>* out_types) {
  auto const& full_param = nnvm::get<DNNLFCFullParam>(attrs.parsed);
  if (full_param.dnnl_param.quantized) {
    const bool channel_wise = full_param.dnnl_param.channel_wise_quantize.has_value() &&
                              full_param.dnnl_param.channel_wise_quantize;
    const FCInputIndex idx(full_param);

    if (in_types->at(idx.data) == mshadow::kBfloat16) {
      return false;
    }

    CHECK(in_types->at(idx.data) == mshadow::kInt8 || in_types->at(idx.data) == mshadow::kUint8)
        << "QuantizedFullyConnected  data input only supports int8/uint8, while "
        << in_types->at(idx.data) << " is given.";
    if (channel_wise) {
      TYPE_ASSIGN_CHECK(*in_types, idx.weight, mshadow::kFloat32);
      if (idx.IsBiasExist()) {
        TYPE_ASSIGN_CHECK(*in_types, idx.bias, mshadow::kFloat32);
      }
    } else {
      TYPE_ASSIGN_CHECK(*in_types, idx.weight, mshadow::kInt8);
      if (idx.IsBiasExist()) {
        if (in_types->at(idx.bias) == -1) {
          TYPE_ASSIGN_CHECK(*in_types, idx.bias, mshadow::kInt32);
        } else {
          CHECK(in_types->at(idx.bias) == mshadow::kInt8 ||
                in_types->at(idx.bias) == mshadow::kInt32)
              << "QuantizedFullyConnected bias input only supports int8/int32, while "
              << in_types->at(idx.bias) << " is given.";
        }
      }
    }
    if (idx.IsSumExist()) {
      if (full_param.dnnl_param.enabled_float_output.has_value()) {
        TYPE_ASSIGN_CHECK(*in_types, idx.sum, full_param.dnnl_param.enabled_float_output.value());
      } else {
        CHECK(in_types->at(idx.sum) == mshadow::kInt8 || in_types->at(idx.sum) == mshadow::kUint8)
            << "QuantizedFullyConnected sum input only supports int8/uint8, while "
            << in_types->at(idx.sum) << " is given.";
      }
    }
    for (size_t i = idx.data_min; i < in_types->size(); ++i) {
      TYPE_ASSIGN_CHECK(*in_types, i, mshadow::kFloat32);
    }

    if (full_param.dnnl_param.enabled_float_output.has_value()) {
      TYPE_ASSIGN_CHECK(*out_types, 0, full_param.dnnl_param.enabled_float_output.value());
    } else {
      if (full_param.dnnl_param.min_calib_range.has_value() &&
          full_param.dnnl_param.max_calib_range.has_value()) {
        if (IsOutputUint8(full_param) &&
            (!idx.IsSumExist() || in_types->at(idx.sum) == mshadow::kUint8)) {
          TYPE_ASSIGN_CHECK(*out_types, 0, mshadow::kUint8);
        } else {
          TYPE_ASSIGN_CHECK(*out_types, 0, mshadow::kInt8);
        }
      } else {
        TYPE_ASSIGN_CHECK(*out_types, 0, mshadow::kInt32);
      }
      TYPE_ASSIGN_CHECK(*out_types, 1, mshadow::kFloat32);
      TYPE_ASSIGN_CHECK(*out_types, 2, mshadow::kFloat32);
    }
    return true;
  } else {
    bool result = DefaultSubgraphOpType(attrs, in_types, out_types);
    if (full_param.dnnl_param.enabled_float_output.has_value()) {
      (*out_types)[0] = full_param.dnnl_param.enabled_float_output.value();
    }
    return result;
  }
}

static bool SgDNNLFCStorageType(const nnvm::NodeAttrs& attrs,
                                const int dev_mask,
                                DispatchMode* dispatch_mode,
                                std::vector<int>* in_attrs,
                                std::vector<int>* out_attrs) {
  auto const& full_param = nnvm::get<DNNLFCFullParam>(attrs.parsed);
  if (full_param.dnnl_param.quantized) {
    std::vector<int> base_in_attrs;
    std::vector<int> base_out_attrs;
    FillBaseInputOutputInfo(full_param, &base_in_attrs, &base_out_attrs, in_attrs, out_attrs);
    bool ret = DefaultSubgraphOpStorageType(
        attrs, dev_mask, dispatch_mode, &base_in_attrs, &base_out_attrs);

    for (size_t i = 0; i < in_attrs->size(); ++i) {
      if (i < base_in_attrs.size())
        in_attrs->at(i) = base_in_attrs[i];
      else
        type_assign(&in_attrs->at(i), mxnet::kDefaultStorage);
    }

    out_attrs->at(0) = base_out_attrs[0];
    if (!full_param.dnnl_param.enabled_float_output.has_value()) {
      type_assign(&out_attrs->at(1), mxnet::kDefaultStorage);
      type_assign(&out_attrs->at(2), mxnet::kDefaultStorage);
    }
    return ret;
  } else {
    return DefaultSubgraphOpStorageType(attrs, dev_mask, dispatch_mode, in_attrs, out_attrs);
  }
}

std::vector<std::pair<int, int>> SgDNNLFCInplaceOption(const NodeAttrs& attrs) {
  auto const& param = nnvm::get<DNNLFCFullParam>(attrs.parsed);
  if (param.dnnl_param.with_sum) {
    return std::vector<std::pair<int, int>>{{FCInputIndex(param).sum, 0}};
  } else {
    return std::vector<std::pair<int, int>>();
  }
}

static OpStatePtr CreateSgDNNLFCState(const nnvm::NodeAttrs& attrs,
                                      Context ctx,
                                      const mxnet::ShapeVector& in_shapes,
                                      const std::vector<int>& in_types) {
  return OpStatePtr::Create<SgDNNLFCOp>(attrs);
}

static void SgDNNLFCForward(const OpStatePtr& state_pointer,
                            const OpContext& ctx,
                            const std::vector<NDArray>& inputs,
                            const std::vector<OpReqType>& req,
                            const std::vector<NDArray>& outputs) {
  SgDNNLFCOp& op = state_pointer.get_state<SgDNNLFCOp>();
  op.Forward(ctx, inputs, req, outputs);
}

nnvm::ObjectPtr SgDNNLFCQuantizedOp(const NodeAttrs& attrs) {
  nnvm::ObjectPtr node          = nnvm::Node::Create();
  node->attrs.op                = Op::Get("_sg_onednn_fully_connected");
  node->attrs.name              = "quantized_" + attrs.name;
  node->attrs.dict              = attrs.dict;
  node->attrs.dict["quantized"] = "True";
  node->attrs.subgraphs.reserve(attrs.subgraphs.size());
  for (auto sub : attrs.subgraphs) {
    node->attrs.subgraphs.push_back(sub);
  }
  node->op()->attr_parser(&(node->attrs));
  return node;
}

static bool SgDNNLAvoidFCQuantizeInput(const NodeAttrs& attrs,
                                       const size_t index_to_check,
                                       const std::string quantize_granularity) {
  auto const& full_param = nnvm::get<DNNLFCFullParam>(attrs.parsed);
  std::unordered_set<size_t> avoid_indexes;
  FCInputIndex idx(full_param);
  if (quantize_granularity == "channel-wise") {
    avoid_indexes.insert(fullc::kWeight);  // weight
    if (!full_param.default_param.no_bias) {
      avoid_indexes.insert(fullc::kBias);  // bias
    }
  }
  if (idx.IsSumInputFloat()) {
    avoid_indexes.insert(idx.sum);
  }
  return avoid_indexes.count(index_to_check);
}

NNVM_REGISTER_OP(_sg_onednn_fully_connected)
    .add_alias("_sg_mkldnn_fully_connected")
    .describe(R"code(_sg_onednn_fully_connected)code" ADD_FILELINE)
    .set_num_inputs([](const NodeAttrs& attrs) {
      auto const& full_param = nnvm::get<DNNLFCFullParam>(attrs.parsed);
      return FCInputIndex(full_param).GetTotal();
    })
    .set_num_outputs([](const NodeAttrs& attrs) {
      auto const& full_param = nnvm::get<DNNLFCFullParam>(attrs.parsed);
      if (full_param.dnnl_param.quantized &&
          !full_param.dnnl_param.enabled_float_output.has_value()) {
        return 3;
      }
      return 1;
    })
    .set_attr_parser(SgDNNLFCParamParser)
    .set_attr<nnvm::FListInputNames>("FListInputNames", SgDNNLFCListInputNames)
    .set_attr<nnvm::FListOutputNames>("FListOutputNames", SgDNNLFCListOutputNames)
    .set_attr<mxnet::FInferShape>("FInferShape", SgDNNLFCInferShape)
    .set_attr<nnvm::FInferType>("FInferType", SgDNNLFCInferType)
    .set_attr<FInferStorageType>("FInferStorageType", SgDNNLFCStorageType)
    .set_attr<FCreateOpState>("FCreateOpState", CreateSgDNNLFCState)
    .set_attr<FStatefulComputeEx>("FStatefulComputeEx<cpu>", SgDNNLFCForward)
    .set_attr<bool>("TIsDNNL", true)
    // TODO(Xinyu): a temp solution to enable GluonCV INT8 flow,
    // will be reverted after the improvement of CachedOP is done.
    .set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
    .set_attr<FResourceRequest>("FResourceRequest",
                                [](const NodeAttrs& n) {
                                  return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
                                })
    .set_attr<nnvm::FMutateInputs>("FMutateInputs", DefaultSubgraphOpMutableInputs)
    .set_attr<std::string>("key_var_num_args", "num_args")
    .set_attr<nnvm::FInplaceOption>("FInplaceOption", SgDNNLFCInplaceOption)
    .set_attr<FQuantizable>("FQuantizable",
                            [](const NodeAttrs& attrs) { return QuantizeType::kMust; })
    .set_attr<FQuantizedOp>("FQuantizedOp", SgDNNLFCQuantizedOp)
    .set_attr<FNeedRequantize>("FNeedRequantize", [](const NodeAttrs& attrs) { return true; })
    .set_attr<FAvoidQuantizeInput>("FAvoidQuantizeInput", SgDNNLAvoidFCQuantizeInput);

}  // namespace op
}  // namespace mxnet

#endif  // if MXNET_USE_ONEDNN == 1