/*************************************************************************
 * 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_FFN_DESCRIPTOR_H_
#define CSRC_OPS_OP_DESCRIPTOR_FFN_DESCRIPTOR_H_

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

namespace tmo {
namespace op_desc {

/**
 * Note [FeedForwardDesc]
 * ~~~~~~~~~~~~~~~~
 * FeedForwardDesc is class to create a Transformer FeedForward op descriptor.
 *
 * ffn_inner_size is intermediate_size;
 * act_name mean act type, and will convert act type when set op desc;
 * act_coef is only using when act_name is CNNL_ACTIVATION_CLIPPED_RELU,
 * CNNL_ACTIVATION_ELU, CNNL_ACTIVATION_ELU_V2, CNNL_ACTIVATION_LEAKYRELU,
 * CNNL_ACTIVATION_TF_LEAKYRELU, CNNL_ACTIVATION_CAFFE_RELU6.
 *
 */

struct TMO_HIDDEN FeedForwardDescImpl {
  FeedForwardDescImpl();
  // using clear function to avoid compilation errors.
  ~FeedForwardDescImpl() { clear(); }
  DELETE_COPY_ASSIGN_CONSTRUCT(FeedForwardDescImpl);
  // variable
  cnnlTransformerFeedForwardDescriptor_t ffn_desc_;
  cnnlTransformerFeedForwardQuantizeDescriptor_t quant_desc_;
  cnnlActivationDescriptor_t act_desc_;

 private:
  inline void clear() {
    CNNL_CHECK_FATAL(cnnlDestroyTransformerFeedForwardDescriptor(this->ffn_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyTransformerFeedForwardQuantizeDescriptor(this->quant_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyActivationDescriptor(this->act_desc_));
  }
};

class TMO_EXPORT FeedForwardDesc {
 public:
  using impl_type =
      std::unique_ptr<FeedForwardDescImpl, std::function<void(FeedForwardDescImpl *)>>;
  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 = 0);

  void 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 = 0);

  // TODO(SG): quanti op desc is not support now, will support later.
  // void setQuantiDesc();

  // TODO(SG): using std::string to pass layer_norm type? It's not clear now.
  // Add layer_norm fused later.
  // void setLayerNormInfo();

  CLASS_CAST_TYPE_OPERATOR_DEFINE(cnnlTransformerFeedForwardDescriptor_t, impl_->ffn_desc_);
  CLASS_CAST_TYPE_OPERATOR_DEFINE(cnnlActivationDescriptor_t, impl_->act_desc_);

  operator cnnlTransformerFeedForwardQuantizeDescriptor_t() const {
    return is_quanti_ == false
               ? nullptr
               : const_cast<cnnlTransformerFeedForwardQuantizeDescriptor_t>(impl_->quant_desc_);
  }
  operator cnnlTransformerFeedForwardQuantizeDescriptor_t() {
    return is_quanti_ == false ? nullptr : impl_->quant_desc_;
  }

 private:
  impl_type impl_;
  bool is_quanti_{false};
};

}  // namespace op_desc
}  // namespace tmo

#endif  // CSRC_OPS_OP_DESCRIPTOR_FFN_DESCRIPTOR_H_
