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

namespace tmo {
namespace op_desc {

namespace {
static OpDescPool<AttnProjDescImpl> attn_proj_instance;
}

AttnProjDesc::AttnProjDesc() {
  impl_ = attn_proj_instance.get();
}

void AttnProjDesc::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) {
  CNNL_CHECK_FATAL(cnnlSetTransformerAttnProjDescriptor(
      impl_->attn_proj_desc_, layernorm_residual_mode, nullptr, /*activation desc*/
      compute_dtype, q_has_value, k_has_value, v_has_value, has_bias, is_pack_mode,
      packed_max_seq_len, trans_out, store_layernorm_result, alpha, beta, layernorm_eps));
  const int allow_tf32 = 0; /*Not allow*/
  CNNL_CHECK_FATAL(
      cnnlSetTransformerAttnProjDescriptorAllowTF32(impl_->attn_proj_desc_, allow_tf32));
  this->is_quant_ = is_quant;
}

AttnProjDescImpl::AttnProjDescImpl() {
  CNNL_CHECK_FATAL(cnnlCreateTransformerAttnProjDescriptor(&this->attn_proj_desc_));
  CNNL_CHECK_FATAL(cnnlCreateTransformerAttnProjQuantifyDescriptor(&this->attn_proj_quant_desc_));
}

}  // namespace op_desc
}  // namespace tmo
