/*************************************************************************
 * 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.
 *************************************************************************/

#ifndef CSRC_OPS_OP_DESCRIPTOR_QUANT_MATMUL_DESCRIPTOR_H_
#define CSRC_OPS_OP_DESCRIPTOR_QUANT_MATMUL_DESCRIPTOR_H_

#include <functional>
#include <memory>
#include "base_descriptor.h"
#include "common/utils.h"

namespace tmo {
namespace op_desc {

class TMO_HIDDEN QuantMatmulDescImpl {
 public:
  QuantMatmulDescImpl();
  // using clear function to avoid compilation errors.
  ~QuantMatmulDescImpl() { clear(); }
  DELETE_COPY_ASSIGN_CONSTRUCT(QuantMatmulDescImpl);
  // variable
  cnnlLLMQuantMatmulDescriptor_t op_desc_;
  cnnlQuantizeDescriptor_t a_quant_desc_;
  cnnlQuantizeDescriptor_t b_quant_desc_;
  cnnlQuantizeDescriptor_t c_quant_desc_;
  cnnlQuantizeDescriptor_t d_quant_desc_;
  cnnlActivationDescriptor_t act_desc_;

 private:
  inline void clear() {
    CNNL_CHECK_FATAL(cnnlDestroyLLMQuantMatmulDescriptor(this->op_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyQuantizeDescriptor(this->a_quant_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyQuantizeDescriptor(this->b_quant_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyQuantizeDescriptor(this->c_quant_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyQuantizeDescriptor(this->d_quant_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyActivationDescriptor(this->act_desc_));
  }
};

class TMO_EXPORT QuantMatmulDesc {
 public:
  using impl_type_ =
      std::unique_ptr<QuantMatmulDescImpl, std::function<void(QuantMatmulDescImpl *)>>;
  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);

  // set descriptor
  void 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);

  operator cnnlLLMQuantMatmulDescriptor_t() { return this->impl_->op_desc_; }
  operator cnnlLLMQuantMatmulDescriptor_t() const { return this->impl_->op_desc_; }

 private:
  impl_type_ impl_;
};

}  // namespace op_desc
}  // namespace tmo

#endif  // CSRC_OPS_OP_DESCRIPTOR_QUANT_MATMUL_DESCRIPTOR_H_
