/*************************************************************************
 * 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 "group_gemm_descriptor.h"
#include <cstddef>
#include <cstdint>

namespace tmo {
namespace op_desc {

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

GroupGemmDesc::GroupGemmDesc(int num_expert,
                             int max_m,
                             int n,
                             int k,
                             cnnlDataType_t dtype,
                             bool idx_mode,
                             QuantMode quant_mode) {
  this->impl_ = instance.get();
  this->idx_mode_ = idx_mode;
  this->quant_mode_ = quant_mode;
  set(num_expert, max_m, n, k, dtype);
}

void GroupGemmDesc::set(int num_expert, int max_m, int n, int k, cnnlDataType_t dtype) {
  this->num_expert_ = num_expert;
  this->max_m_ = max_m;
  this->n_ = n;
  this->k_ = k;
  this->dtype_ = dtype;
  // no need to set compute_dtype, default is CNNL_DTYPE_FLOAT
  CNNL_CHECK_FATAL(cnnlSetGroupGemmDescAttr(impl_->group_gemm_desc_, CNNL_MATMUL_DESC_GROUP_M_MAX,
                                            &max_m, sizeof(int32_t)));
  CNNL_CHECK_FATAL(cnnlSetGroupGemmDescAttr(impl_->group_gemm_desc_, CNNL_MATMUL_DESC_GROUP_N_MAX,
                                            &n, sizeof(int32_t)));
  CNNL_CHECK_FATAL(cnnlSetGroupGemmDescAttr(impl_->group_gemm_desc_, CNNL_MATMUL_DESC_GROUP_K_MAX,
                                            &k, sizeof(int32_t)));
  int dim_nb = 1;
  int64_t dims[] = {num_expert};
  CNNL_CHECK_FATAL(cnnlSetTensorDescriptor_v2(impl_->group_host_tensor_, CNNL_LAYOUT_ARRAY,
                                              CNNL_DTYPE_INT32, dim_nb, dims));
  CNNL_CHECK_FATAL(cnnlSetTensorDescriptor_v2(impl_->group_device_tensor_, CNNL_LAYOUT_ARRAY,
                                              CNNL_DTYPE_INT32, dim_nb, dims));
  CNNL_CHECK_FATAL(cnnlSetTensorDescriptor_v2(impl_->scale_factor_tensor_, CNNL_LAYOUT_ARRAY,
                                              CNNL_DTYPE_FLOAT, dim_nb, dims));
  CNNL_CHECK_FATAL(
      cnnlSetTensorDescriptorPointerMode(impl_->group_host_tensor_, CNNL_POINTER_MODE_HOST));
  CNNL_CHECK_FATAL(
      cnnlSetTensorDescriptorPointerMode(impl_->group_device_tensor_, CNNL_POINTER_MODE_DEVICE));
  CNNL_CHECK_FATAL(
      cnnlSetTensorDescriptorPointerMode(impl_->scale_factor_tensor_, CNNL_POINTER_MODE_DEVICE));
}

void GroupGemmDesc::setPerRowColScaleBiasAct(void *a_scale, void *b_scale) {
  setPerRowColScaleBiasAct(a_scale, b_scale, nullptr, nullptr, CNNL_DTYPE_INVALID,
                           CNNL_ACTIVATION_IDENTITY, CNNL_PROPAGATE_NAN, 0.f, false, 0, 0, 0);
}

void GroupGemmDesc::setPerRowColScaleBiasAct(void *a_scale,
                                             void *b_scale,
                                             void *quant_flag,
                                             void *bias,
                                             cnnlDataType_t bias_dtype,
                                             int channels,
                                             int blk,
                                             int pos) {
  setPerRowColScaleBiasAct(a_scale, b_scale, quant_flag, bias, bias_dtype, CNNL_ACTIVATION_IDENTITY,
                           CNNL_PROPAGATE_NAN, 0.f, false, channels, blk, pos);
}

void GroupGemmDesc::setPerRowColScaleBiasAct(void *a_scale,
                                             void *b_scale,
                                             void *quant_flag,
                                             void *bias,
                                             cnnlDataType_t bias_dtype,
                                             cnnlActivationMode_t mode,
                                             cnnlNanPropagation_t nan_prop,
                                             float coef,
                                             bool has_active,
                                             int channels,
                                             int scale_block_size,
                                             int channel_pos) {
  bool is_a_quant = a_scale != nullptr;
  bool use_b_scale = false;
  bool quant_grouped = (scale_block_size > 0 && channels / scale_block_size > 0) ? true : false;
  if (a_scale) {
    CNNL_CHECK_FATAL(cnnlSetGroupGemmTensorDescriptor(a_scale_desc(), CNNL_DTYPE_FLOAT, a_scale,
                                                      CNNL_POINTER_MODE_DEVICE, nullptr));
  }
  if (!quant_grouped && b_scale) {
    use_b_scale = true;
    CNNL_CHECK_FATAL(cnnlSetGroupGemmTensorDescriptor(b_scale_desc(), CNNL_DTYPE_FLOAT, b_scale,
                                                      CNNL_POINTER_MODE_DEVICE, nullptr));
  }
  if (bias) {
    this->has_bias_ = true;
    CNNL_CHECK_FATAL(cnnlSetGroupGemmTensorDescriptor(bias_desc(), bias_dtype, bias,
                                                      CNNL_POINTER_MODE_DEVICE, nullptr));
  }
  this->has_active_ = has_active;
  if (has_active) {
    CNNL_CHECK_FATAL(cnnlSetActivationDescriptor(active_desc(), mode, nan_prop, coef));
  }

  if (quant_grouped) {
    const int dim_nb = 2;
    int group_num = channels / scale_block_size;
    int64_t dims[dim_nb] = {group_num, num_expert_ * n_};
    CNNL_CHECK_FATAL(cnnlSetTensorDescriptor_v2(b_scale_quant_grouped_desc(), CNNL_LAYOUT_ARRAY,
                                                CNNL_DTYPE_FLOAT, dim_nb, dims));

    CNNL_CHECK_FATAL(cnnlSetGroupGemmGroupwiseScale(
        group_gemm_desc(), CNNL_MATMUL_DESC_B_SCALE_POINTER, channels, scale_block_size,
        channel_pos, b_scale_quant_grouped_desc(), b_scale));
    if (quant_flag) {
      const int dim_qf = 2;
      int64_t dims[dim_qf] = {num_expert_, group_num};
      CNNL_CHECK_FATAL(
          cnnlSetTensorDescriptorPointerMode(quant_flag_desc(), CNNL_POINTER_MODE_HOST));
      CNNL_CHECK_FATAL(cnnlSetTensorDescriptor_v2(quant_flag_desc(), CNNL_LAYOUT_ARRAY,
                                                  CNNL_DTYPE_INT32, dim_qf, dims));
      CNNL_CHECK_FATAL(cnnlSetGroupGemmGroupMixedQuantBitFlag(
          group_gemm_desc(), CNNL_MATMUL_DESC_B_QUANT_FLAG_POINTER, channels, scale_block_size,
          quant_flag_desc(), quant_flag));
    }
  }

  CNNL_CHECK_FATAL(cnnlSetGroupGemmPerRowColScaleBiasAct(
      group_gemm_desc(), is_a_quant ? a_scale_desc() : nullptr,
      use_b_scale ? b_scale_desc() : nullptr, this->has_bias_ ? bias_desc() : nullptr,
      has_active ? active_desc() : nullptr));
}

void GroupGemmDesc::setInputOutputTensor(cnnlDataType_t a_dtype,
                                         cnnlDataType_t b_dtype,
                                         cnnlDataType_t d_dtype,
                                         cnnlDataType_t idx_dtype,
                                         void *a,
                                         void *b,
                                         void *c,
                                         void *d,
                                         void *idx,
                                         int64_t k,
                                         int64_t k_stride,
                                         int64_t total_m,
                                         bool has_c,
                                         int64_t *b_offset) {
  this->has_c_ = has_c;
  if (idx_mode()) {
    cnnlSetGroupGemmTensorGatherIdxDesc(a_desc(), a_dtype, idx_dtype, a, idx, k, k_stride, total_m);
  } else {
    CNNL_CHECK_FATAL(
        cnnlSetGroupGemmTensorDescriptor(a_desc(), a_dtype, a, CNNL_POINTER_MODE_DEVICE, nullptr));
  }
  auto b_ptr_mode = b_offset ? CNNL_POINTER_MODE_HOST : CNNL_POINTER_MODE_DEVICE;
  CNNL_CHECK_FATAL(cnnlSetGroupGemmTensorDescriptor(b_desc(), b_dtype, b, b_ptr_mode, b_offset));
  if (has_c) {
    CNNL_CHECK_FATAL(
        cnnlSetGroupGemmTensorDescriptor(c_desc(), d_dtype, c, CNNL_POINTER_MODE_DEVICE, nullptr));
  }
  CNNL_CHECK_FATAL(
      cnnlSetGroupGemmTensorDescriptor(d_desc(), d_dtype, d, CNNL_POINTER_MODE_DEVICE, nullptr));
}

}  // namespace op_desc
}  // namespace tmo
