/*************************************************************************
 * 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 "ffn_descriptor.h"

namespace tmo {
namespace op_desc {

namespace {
static OpDescPool<FeedForwardDescImpl> instance;
}  // end of anonymous namespace

FeedForwardDesc::FeedForwardDesc() : is_quanti_(false) {
  this->impl_ = instance.get();
}

FeedForwardDesc::FeedForwardDesc(
    const cnnlTransformerLayernormResidualStructure_t layernorm_residual_mode,
    const std::string &act_name,
    cnnlDataType_t compute_type,
    float eps,
    float alpha,
    float beta,
    float act_coef)
    : is_quanti_(false) {
  this->impl_ = instance.get();
  this->setFeedForwardDesc(layernorm_residual_mode, act_name, compute_type, eps, alpha, beta,
                           act_coef);
}

void FeedForwardDesc::setFeedForwardDesc(
    const cnnlTransformerLayernormResidualStructure_t layernorm_residual_mode,
    const std::string &act_name,
    cnnlDataType_t compute_type,
    float eps,
    float alpha,
    float beta,
    float act_coef) {
  // Now only using bfloat16 and half.
  // alpha must be 1.0f when NO_RESIDUAL mode.
  CNNL_CHECK_FATAL(cnnlSetTransformerFeedForwardDescriptor_v2(
      impl_->ffn_desc_, eps, alpha, beta, compute_type, layernorm_residual_mode));
  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, CNNL_ACTIVATION_FAST,
                                                  CNNL_NOT_PROPAGATE_NAN, act_coef,
                                                  // following parameters are not used
                                                  0, 0, 0, 0));
}

FeedForwardDescImpl::FeedForwardDescImpl() {
  CNNL_CHECK_FATAL(cnnlCreateTransformerFeedForwardDescriptor(&this->ffn_desc_));
  CNNL_CHECK_FATAL(cnnlCreateTransformerFeedForwardQuantizeDescriptor(&this->quant_desc_));
  CNNL_CHECK_FATAL(cnnlCreateActivationDescriptor(&this->act_desc_));
}

}  // namespace op_desc
}  // namespace tmo
