/*************************************************************************
 * 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_GROUP_GEMM_DESCRIPTOR_H_
#define CSRC_OPS_OP_DESCRIPTOR_GROUP_GEMM_DESCRIPTOR_H_

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

namespace tmo {
namespace op_desc {

struct TMO_HIDDEN GroupGemmDescImpl {
  GroupGemmDescImpl() {
    CNNL_CHECK_FATAL(cnnlCreateGroupGemmDescriptor(&group_gemm_desc_));
    CNNL_CHECK_FATAL(
        cnnlCreateGroupGemmTensorDescriptor(&a_idx_desc_, CNNL_GROUP_GEMM_ADDRESSING_GATHER_IDX));
    CNNL_CHECK_FATAL(
        cnnlCreateGroupGemmTensorDescriptor(&a_offset_desc_, CNNL_GROUP_GEMM_ADDRESSING_OFFSETS));
    CNNL_CHECK_FATAL(
        cnnlCreateGroupGemmTensorDescriptor(&b_desc_, CNNL_GROUP_GEMM_ADDRESSING_OFFSETS));
    CNNL_CHECK_FATAL(
        cnnlCreateGroupGemmTensorDescriptor(&c_desc_, CNNL_GROUP_GEMM_ADDRESSING_OFFSETS));
    CNNL_CHECK_FATAL(
        cnnlCreateGroupGemmTensorDescriptor(&d_desc_, CNNL_GROUP_GEMM_ADDRESSING_OFFSETS));
    CNNL_CHECK_FATAL(
        cnnlCreateGroupGemmTensorDescriptor(&a_scale_desc_, CNNL_GROUP_GEMM_ADDRESSING_OFFSETS));
    CNNL_CHECK_FATAL(
        cnnlCreateGroupGemmTensorDescriptor(&b_scale_desc_, CNNL_GROUP_GEMM_ADDRESSING_OFFSETS));
    CNNL_CHECK_FATAL(
        cnnlCreateGroupGemmTensorDescriptor(&bias_desc_, CNNL_GROUP_GEMM_ADDRESSING_OFFSETS));
    CNNL_CHECK_FATAL(cnnlCreateTensorDescriptor(&b_scale_quant_grouped_desc_));
    CNNL_CHECK_FATAL(cnnlCreateTensorDescriptor(&group_host_tensor_));
    CNNL_CHECK_FATAL(cnnlCreateTensorDescriptor(&group_device_tensor_));
    CNNL_CHECK_FATAL(cnnlCreateTensorDescriptor(&scale_factor_tensor_));
    CNNL_CHECK_FATAL(cnnlCreateTensorDescriptor(&quant_flag_tensor_));
    CNNL_CHECK_FATAL(cnnlCreateActivationDescriptor(&active_desc_));
  }
  // using clear function to avoid compilation errors.
  ~GroupGemmDescImpl() { clear(); }
  DELETE_COPY_ASSIGN_CONSTRUCT(GroupGemmDescImpl);
  // variable
  cnnlGroupGemmDescriptor_t group_gemm_desc_;
  cnnlGroupGemmTensorDescriptor_t a_idx_desc_;
  cnnlGroupGemmTensorDescriptor_t a_offset_desc_;
  cnnlGroupGemmTensorDescriptor_t b_desc_;
  cnnlGroupGemmTensorDescriptor_t c_desc_;
  cnnlGroupGemmTensorDescriptor_t d_desc_;
  cnnlGroupGemmTensorDescriptor_t a_scale_desc_;
  cnnlGroupGemmTensorDescriptor_t b_scale_desc_;
  cnnlGroupGemmTensorDescriptor_t bias_desc_;
  cnnlTensorDescriptor_t b_scale_quant_grouped_desc_;
  cnnlTensorDescriptor_t group_host_tensor_;
  cnnlTensorDescriptor_t group_device_tensor_;
  cnnlTensorDescriptor_t scale_factor_tensor_;
  cnnlTensorDescriptor_t quant_flag_tensor_;
  cnnlActivationDescriptor_t active_desc_;

 private:
  inline void clear() {
    CNNL_CHECK_FATAL(
        cnnlSetTensorDescriptorPointerMode(group_host_tensor_, CNNL_POINTER_MODE_DEVICE));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmDescriptor(group_gemm_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmTensorDescriptor(a_idx_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmTensorDescriptor(a_offset_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmTensorDescriptor(b_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmTensorDescriptor(c_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmTensorDescriptor(d_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmTensorDescriptor(a_scale_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmTensorDescriptor(b_scale_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyGroupGemmTensorDescriptor(bias_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyTensorDescriptor(b_scale_quant_grouped_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyTensorDescriptor(group_host_tensor_));
    CNNL_CHECK_FATAL(cnnlDestroyTensorDescriptor(group_device_tensor_));
    CNNL_CHECK_FATAL(cnnlDestroyTensorDescriptor(scale_factor_tensor_));
    CNNL_CHECK_FATAL(cnnlDestroyTensorDescriptor(quant_flag_tensor_));
    CNNL_CHECK_FATAL(cnnlDestroyActivationDescriptor(active_desc_));
  }
};

class TMO_EXPORT GroupGemmDesc {
 public:
  enum class QuantMode { noQuant, W8, W4, W4W8 };
  using impl_type = std::unique_ptr<GroupGemmDescImpl, std::function<void(GroupGemmDescImpl *)>>;
  GroupGemmDesc() = delete;
  GroupGemmDesc(int num_expert,
                int max_m,
                int n,
                int k,
                cnnlDataType_t dtype,
                bool idx_mode,
                QuantMode quant_mode);
  void set(int num_expert, int max_m, int n, int k, cnnlDataType_t dtype);
  void setPerRowColScaleBiasAct(void *a_scale, void *b_scale);
  void setPerRowColScaleBiasAct(void *a_scale,
                                void *b_scale,
                                void *quant_flag,
                                void *bias,
                                cnnlDataType_t bias_dtype,
                                int channels,
                                int blk,
                                int pos);
  void setPerRowColScaleBiasAct(void *a_scale,
                                void *b_scale,
                                void *quant_flag,               // for w4/w8 quantize
                                void *bias,                     // nullptr
                                cnnlDataType_t bias_dtype,      // CNNL_DTYPE_INVALID
                                cnnlActivationMode_t mode,      // CNNL_ACTIVATION_SILU
                                cnnlNanPropagation_t nan_prop,  // CNNL_PROPAGATE_NAN
                                float coef,                     // 0.f
                                bool has_active,                // false
                                int channels,
                                int scale_block_size,
                                int channel_pos);
  void 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 = nullptr);
  cnnlGroupGemmDescriptor_t group_gemm_desc() { return impl_->group_gemm_desc_; }
  cnnlGroupGemmTensorDescriptor_t a_desc() {
    return idx_mode_ ? impl_->a_idx_desc_ : impl_->a_offset_desc_;
  }
  cnnlGroupGemmTensorDescriptor_t b_desc() { return impl_->b_desc_; }
  cnnlGroupGemmTensorDescriptor_t c_desc() { return impl_->c_desc_; }
  cnnlGroupGemmTensorDescriptor_t d_desc() { return impl_->d_desc_; }
  cnnlGroupGemmTensorDescriptor_t a_scale_desc() { return impl_->a_scale_desc_; }
  cnnlGroupGemmTensorDescriptor_t b_scale_desc() { return impl_->b_scale_desc_; }
  cnnlGroupGemmTensorDescriptor_t bias_desc() { return impl_->bias_desc_; }
  cnnlTensorDescriptor_t b_scale_quant_grouped_desc() { return impl_->b_scale_quant_grouped_desc_; }
  cnnlTensorDescriptor_t group_host_tensor() { return impl_->group_host_tensor_; }
  cnnlTensorDescriptor_t group_device_tensor() { return impl_->group_device_tensor_; }
  cnnlTensorDescriptor_t scale_factor_tensor() { return impl_->scale_factor_tensor_; }
  cnnlTensorDescriptor_t quant_flag_desc() { return impl_->quant_flag_tensor_; }
  cnnlActivationDescriptor_t active_desc() { return impl_->active_desc_; }
  bool idx_mode() { return idx_mode_; }
  bool has_c() { return has_c_; }

 private:
  impl_type impl_;
  bool idx_mode_ = false;
  bool has_c_ = false;
  bool has_bias_ = false;
  bool has_active_ = false;
  int num_expert_;
  int max_m_;
  int n_;
  int k_;
  cnnlDataType_t dtype_;
  QuantMode quant_mode_ = QuantMode::noQuant;
};

}  // namespace op_desc
}  // namespace tmo

#endif  // CSRC_OPS_OP_DESCRIPTOR_GROUP_GEMM_DESCRIPTOR_H_
