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

namespace tmo {
namespace op_desc {

namespace {
static OpDescPool<MatMulDescImpl> matmul_instance;
}

MatMulDesc::MatMulDesc(const cnnlTensorDescriptor_t &bias_tensor,
                       const void *bias_data,
                       size_t &workspace_size,
                       int32_t lda,
                       int32_t ldb,
                       int32_t ldc,
                       const std::string &act_name,
                       bool use_beta,
                       bool trans_a,
                       bool trans_b,
                       bool fast_act,
                       bool approximate) {
  impl_ = matmul_instance.get();
  auto matmul_desc_ = impl_->matmul_desc_;
  auto act_desc_ = impl_->act_desc_;
  workspace_size = max_workspace_size_;

  int32_t matmul_trans_a = int(trans_a);
  int32_t matmul_trans_b = int(trans_b);
  int32_t matmul_use_tf32 = 0;
  int32_t matmul_use_beta = int(use_beta);
  CNNL_CHECK_FATAL(cnnlSetMatMulExDescAttr(matmul_desc_, CNNL_MATMUL_EX_DESC_TRANSA,
                                           &(matmul_trans_a), sizeof(int32_t)));
  CNNL_CHECK_FATAL(cnnlSetMatMulExDescAttr(matmul_desc_, CNNL_MATMUL_EX_DESC_TRANSB,
                                           &(matmul_trans_b), sizeof(int32_t)));
  CNNL_CHECK_FATAL(cnnlSetMatMulExDescAttr(matmul_desc_, CNNL_MATMUL_EX_ALLOW_TF32,
                                           &(matmul_use_tf32), sizeof(int32_t)));
  CNNL_CHECK_FATAL(cnnlSetMatMulExDescAttr(matmul_desc_, CNNL_MATMUL_EX_USE_BETA,
                                           &(matmul_use_beta), sizeof(int32_t)));
  CNNL_CHECK_FATAL(
      cnnlSetMatMulExDescAttr(matmul_desc_, CNNL_MATMUL_EX_DESC_LDA, &(lda), sizeof(int32_t)));
  CNNL_CHECK_FATAL(
      cnnlSetMatMulExDescAttr(matmul_desc_, CNNL_MATMUL_EX_DESC_LDB, &(ldb), sizeof(int32_t)));
  CNNL_CHECK_FATAL(
      cnnlSetMatMulExDescAttr(matmul_desc_, CNNL_MATMUL_EX_DESC_LDC, &(ldc), sizeof(int32_t)));

  CNNL_CHECK_FATAL(cnnlSetMatMulExDescAttr(matmul_desc_, CNNL_MATMUL_EX_PREF_MAX_WORKSPACE_BYTES,
                                           &(max_workspace_size_), sizeof(int64_t)));
  if (act_name == "none") {
    if (bias_data == nullptr) {
      cnnlMatMulEpilogueType_t fuse_type = CNNL_MATMUL_EPI_NONE;
      CNNL_CHECK_FATAL(cnnlSetMatMulExDescAttr(
          matmul_desc_, cnnlMatMulExDescAttribute_t::CNNL_MATMUL_EX_DESC_EPILOGUE_TYPE,
          (void *)&fuse_type, sizeof(fuse_type)));

    } else {
      cnnlMatMulEpilogueType_t fuse_type = CNNL_MATMUL_EPI_BIAS;
      CNNL_CHECK_FATAL(cnnlSetMatMulExDescAttr(
          matmul_desc_, cnnlMatMulExDescAttribute_t::CNNL_MATMUL_EX_DESC_EPILOGUE_TYPE,
          (void *)&fuse_type, sizeof(fuse_type)));
      CNNL_CHECK_FATAL(cnnlSetMatMulExBias(matmul_desc_, bias_tensor, (void *)bias_data));
    }
  } else {
    float act_coef = 0.0f;
    cnnlActivationMode_t act_mode = strToActivationMode(act_name);
    if (act_name == "silu") {
      act_coef = 1.0f;
      act_mode = CNNL_ACTIVATION_SILU;
    }
    cnnlActivationPreference_t act_pref = CNNL_ACTIVATION_HIGH_PRECISION;
    if (fast_act) {
      act_pref = CNNL_ACTIVATION_FAST;
    }
    CNNL_CHECK_FATAL(cnnlSetActivationDescriptor_v6(
        act_desc_, act_mode, act_pref, CNNL_NOT_PROPAGATE_NAN, act_coef, 0, 0, 0, 0, approximate));
    // fuse info
    cnnlMatMulEpilogueType_t fuse_type = CNNL_MATMUL_EPI_BIAS_SCALE_BN_ACTIVATION;
    CNNL_CHECK_FATAL(cnnlSetMatMulExDescAttr(
        matmul_desc_, cnnlMatMulExDescAttribute_t::CNNL_MATMUL_EX_DESC_EPILOGUE_TYPE, &(fuse_type),
        sizeof(fuse_type)));
    CNNL_CHECK_FATAL(cnnlSetMatMulExBiasScaleBNActive(matmul_desc_, bias_tensor, bias_data, nullptr,
                                                      nullptr, nullptr, nullptr, nullptr, nullptr,
                                                      0, 0, 0, 0, 0, act_desc_));
  }
}

MatMulDescImpl::MatMulDescImpl() {
  CNNL_CHECK_FATAL(cnnlCreateMatMulExDescriptor(&this->matmul_desc_));
  CNNL_CHECK_FATAL(cnnlCreateActivationDescriptor(&this->act_desc_));
}

}  // namespace op_desc
}  // namespace tmo
