/*************************************************************************
 * 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_ATTN_PROJ_DESCRIPTOR_H_
#define CSRC_OPS_OP_DESCRIPTOR_ATTN_PROJ_DESCRIPTOR_H_

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

namespace tmo {
namespace op_desc {

class TMO_HIDDEN AttnProjDescImpl {
 public:
  AttnProjDescImpl();
  // using clear function to avoid compilation errors.
  ~AttnProjDescImpl() { clear(); }

  DELETE_COPY_ASSIGN_CONSTRUCT(AttnProjDescImpl);

  cnnlTransformerAttnProjDescriptor_t attn_proj_desc_;
  cnnlTransformerAttnProjQuantizeDescriptor_t attn_proj_quant_desc_;

 private:
  inline void clear() {
    CNNL_CHECK_FATAL(cnnlDestroyTransformerAttnProjDescriptor(this->attn_proj_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyTransformerAttnProjQuantifyDescriptor(this->attn_proj_quant_desc_));
  }
};

class TMO_EXPORT AttnProjDesc {
 public:
  using impl_type_ = std::unique_ptr<AttnProjDescImpl, std::function<void(AttnProjDescImpl *)>>;
  AttnProjDesc();
  void setDesc(const cnnlTransformerLayernormResidualStructure_t layernorm_residual_mode,
               const cnnlDataType_t compute_dtype,
               const bool q_has_value,
               const bool k_has_value,
               const bool v_has_value,
               const bool has_bias,
               const bool is_pack_mode,
               const int packed_max_seq_len,
               const bool trans_out,
               const bool store_layernorm_result,
               const float alpha,
               const float beta,
               const float layernorm_eps,
               const bool is_quant = false);

  operator cnnlTransformerAttnProjDescriptor_t() { return this->impl_->attn_proj_desc_; }

  operator cnnlTransformerAttnProjDescriptor_t() const { return this->impl_->attn_proj_desc_; }

  operator cnnlTransformerAttnProjQuantizeDescriptor_t() {
    return is_quant_ ? this->impl_->attn_proj_quant_desc_ : nullptr;
  }

  operator cnnlTransformerAttnProjQuantizeDescriptor_t() const {
    return is_quant_ ? this->impl_->attn_proj_quant_desc_ : nullptr;
  }

 private:
  impl_type_ impl_;
  bool is_quant_ = false;
};

}  // namespace op_desc
}  // namespace tmo

#endif  // CSRC_OPS_OP_DESCRIPTOR_ATTN_PROJ_DESCRIPTOR_H_
