/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/

#include "quant_matmul_descriptor.h"

namespace tmo {
namespace op_desc {

namespace {
static OpDescPool<QuantMatmulDescImpl> instance;
}

QuantMatmulDesc::QuantMatmulDesc(const std::string &quant_algo,
                                 const std::string &a_quant_layout,
                                 const std::string &b_quant_layout,
                                 unsigned quant_bit_size,
                                 cnnlDataType_t compute_dtype,
                                 const std::string &act_name,
                                 bool use_hp_active,
                                 float act_coef,
                                 bool trans_a,
                                 bool trans_b) {
  impl_ = instance.get();
  this->setQuantMatmulDesc(quant_algo, a_quant_layout, b_quant_layout, quant_bit_size,
                           compute_dtype, act_name, use_hp_active, act_coef, trans_a, trans_b);
}

void QuantMatmulDesc::setQuantMatmulDesc(const std::string &quant_algo,
                                         const std::string &a_quant_layout,
                                         const std::string &b_quant_layout,
                                         unsigned quant_bit_size,
                                         cnnlDataType_t compute_dtype,
                                         const std::string &act_name,
                                         bool use_hp_active,
                                         float act_coef,
                                         bool trans_a,
                                         bool trans_b) {
  cnnlQuantizeDescriptor_t no_quant;
  CNNL_CHECK_FATAL(cnnlCreateQuantizeDescriptor(&no_quant));
  // set quantize algo and layout
  cnnlLLMQuantAlgo_t cnnl_quant_algo = strToQuantizeAlgo(quant_algo);
  CNNL_CHECK_FATAL(cnnlSetQuantizeDescriptorBehaviour(
      impl_->a_quant_desc_, strToQuantizeLayout(a_quant_layout), quant_bit_size, 0, false));
  CNNL_CHECK_FATAL(cnnlSetQuantizeDescriptorBehaviour(
      impl_->b_quant_desc_, strToQuantizeLayout(b_quant_layout), quant_bit_size, 0, false));
  // set activation
  act_coef = act_name == "silu" ? 1.0f : act_coef;
  cnnlActivationMode_t act_mode = strToActivationMode(act_name);
  CNNL_CHECK_FATAL(cnnlSetActivationDescriptor_v5(
      impl_->act_desc_, act_mode,
      use_hp_active ? CNNL_ACTIVATION_HIGH_PRECISION : CNNL_ACTIVATION_FAST, CNNL_NOT_PROPAGATE_NAN,
      act_coef, 0, 0, 0, 0));
  // set into op_desc
  CNNL_CHECK_FATAL(cnnlSetLLMQuantMatmulDescriptor(
      impl_->op_desc_, impl_->a_quant_desc_, impl_->b_quant_desc_, no_quant, no_quant,
      impl_->act_desc_, compute_dtype, cnnl_quant_algo, trans_a, trans_b));
  CNNL_CHECK_FATAL(cnnlDestroyQuantizeDescriptor(no_quant));
}

QuantMatmulDescImpl::QuantMatmulDescImpl() {
  CNNL_CHECK_FATAL(cnnlCreateLLMQuantMatmulDescriptor(&this->op_desc_));
  CNNL_CHECK_FATAL(cnnlCreateQuantizeDescriptor(&this->a_quant_desc_));
  CNNL_CHECK_FATAL(cnnlCreateQuantizeDescriptor(&this->b_quant_desc_));
  CNNL_CHECK_FATAL(cnnlCreateQuantizeDescriptor(&this->c_quant_desc_));
  CNNL_CHECK_FATAL(cnnlCreateQuantizeDescriptor(&this->d_quant_desc_));
  CNNL_CHECK_FATAL(cnnlCreateActivationDescriptor(&this->act_desc_));
}

}  // namespace op_desc
}  // namespace tmo
