// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <unordered_map>
#include <vector>
#include "paddle/phi/backends/xpu/enforce_xpu.h"
#include "paddle/phi/backends/xpu/xpu_header.h"
#include "paddle/phi/backends/xpu/xpu_info.h"
#include "paddle/phi/kernels/funcs/blas/blas.h"

#ifdef PADDLE_WITH_XPU_XRE5
#include "xblas/cublasLt.h"
namespace xblas = baidu::xpu::xblas;
#endif

namespace phi {

using XPUTypeFP16 = typename XPUTypeTrait<phi::float16>::Type;
using XPUTypeBF16 = typename XPUTypeTrait<phi::bfloat16>::Type;

enum XPUFCCalcType {
  FC_INT16 = 0,
  FC_INT32,
  FC_FLOAT,
  FC_INT32_WITH_LL,
  FC_TF32,
  FC_FLOAT16,
};

using XPUFCCalcTypeMap = std::vector<std::pair<const char*, XPUFCCalcType>>;

inline XPUFCCalcType GetFCCalcTypeFromEnv(const XPUFCCalcTypeMap& env_map,
                                          XPUFCCalcType default_calc_type) {
  for (auto [env_name, calc_type] : env_map) {
    if (std::getenv(env_name) != nullptr) {
      return calc_type;
    }
  }
  return default_calc_type;
}

template <typename T>
inline XPUFCCalcType FCCalcType() {
  // FLOAT32
  XPUFCCalcTypeMap calc_type_map = {
      {"XPU_PADDLE_FC_FLOAT", XPUFCCalcType::FC_FLOAT},
      {"XPU_PADDLE_FC_LOCAL_INT16", XPUFCCalcType::FC_FLOAT},
      {"XPU_PADDLE_FC_TF32", XPUFCCalcType::FC_TF32},
      {"XPU_PADDLE_FC_INT16", XPUFCCalcType::FC_INT16},
      {"XPU_PADDLE_FC_INT32", XPUFCCalcType::FC_INT32},
      {"XPU_PADDLE_FC_INT32_WITH_LL", XPUFCCalcType::FC_INT32_WITH_LL},
  };
#ifdef PADDLE_WITH_XPU_XRE5
  // Use full float32 accumulation by default for better precision, matching
  // the GPU default (FLAGS_cublas_allow_tf32=false). Users who need TF32
  // performance can set env var XPU_PADDLE_FC_TF32.
  auto default_calc_type = XPUFCCalcType::FC_FLOAT;
#else
  auto default_calc_type = XPUFCCalcType::FC_INT16;
#endif
  return GetFCCalcTypeFromEnv(calc_type_map, default_calc_type);
}

template <>
inline XPUFCCalcType FCCalcType<XPUTypeFP16>() {
  XPUFCCalcTypeMap calc_type_map = {
      {"XPU_PADDLE_FC_FLOAT16", XPUFCCalcType::FC_FLOAT16},
      {"XPU_PADDLE_FC_INT16", XPUFCCalcType::FC_INT16},
      {"XPU_PADDLE_FC_FLOAT", XPUFCCalcType::FC_FLOAT},
      {"XPU_PADDLE_FC_LOCAL_INT16", XPUFCCalcType::FC_FLOAT}};
#ifdef PADDLE_WITH_XPU_XRE5
  auto default_calc_type = XPUFCCalcType::FC_FLOAT16;
#else
  auto default_calc_type = XPUFCCalcType::FC_INT16;
#endif
  return GetFCCalcTypeFromEnv(calc_type_map, default_calc_type);
}

template <>
inline XPUFCCalcType FCCalcType<XPUTypeBF16>() {
  XPUFCCalcTypeMap calc_type_map = {
      // TF32 is the default, do not need to be listed here.
      {"XPU_PADDLE_FC_FLOAT", XPUFCCalcType::FC_FLOAT},
      {"XPU_PADDLE_FC_LOCAL_INT16", XPUFCCalcType::FC_FLOAT}};
  auto default_calc_type = XPUFCCalcType::FC_TF32;
  return GetFCCalcTypeFromEnv(calc_type_map, default_calc_type);
}

struct XpuFcInfo {
  int64_t bs;
  int64_t m;
  int64_t n;
  int64_t k;
  bool trans_x;
  bool trans_y;
  int64_t stride_x;
  int64_t stride_y;
  int64_t stride_out;
  float* max_x;
  float* max_y;
  float* max_out;
  const float* bias;
  bool is_x_need_broadcast;
  bool is_y_need_broadcast;
  const float* scale_x;
  const float* scale_y;
  int scale_x_mode;
  int scale_y_mode;

  XpuFcInfo()
      : bs(0),
        m(0),
        n(0),
        k(0),
        trans_x(false),
        trans_y(false),
        stride_x(0),
        stride_y(0),
        stride_out(0),
        max_x(nullptr),
        max_y(nullptr),
        max_out(nullptr),
        bias(nullptr),
        is_x_need_broadcast(false),
        is_y_need_broadcast(false),
        scale_x(nullptr),
        scale_y(nullptr),
        scale_x_mode(0),
        scale_y_mode(0) {}
  void InitFcInfo(int64_t bs,
                  int64_t m,
                  int64_t n,
                  int64_t k,
                  bool trans_x,
                  bool trans_y,
                  float* max_x,
                  float* max_y,
                  float* max_out) {
    this->bs = bs;
    this->m = m;
    this->n = n;
    this->k = k;
    this->trans_x = trans_x;
    this->trans_y = trans_y;
    this->max_x = max_x;
    this->max_y = max_y;
    this->max_out = max_out;

    if (this->bs <= 1) {
      this->stride_x = trans_x ? m : k;
      this->stride_y = trans_y ? k : n;
      this->stride_out = n;
    } else {
      this->stride_x = m * k;
      this->stride_y = k * n;
      this->stride_out = m * n;
    }
  }
};

static std::ostream& operator<<(std::ostream& os, const XpuFcInfo& fc_inf) {
  os << "fc_inf[ bs, m, n, k, trans_x, trans_y, stride_x, stride_y, "
        "stride_out] = "
     << "[" << fc_inf.bs << ", " << fc_inf.m << ", " << fc_inf.n << ", "
     << fc_inf.k << ", " << fc_inf.trans_x << ", " << fc_inf.trans_y << ", "
     << fc_inf.stride_x << ", " << fc_inf.stride_y << ", " << fc_inf.stride_out;
  return os;
}

static void GetFCInfo(const DDim& x_dims,
                      const DDim& y_dims,
                      bool trans_x,
                      bool trans_y,
                      XpuFcInfo* info) {
  DDim new_x_dims = (x_dims.size() > 1) ? x_dims : make_ddim({1, x_dims[0]});
  DDim new_y_dims = (y_dims.size() > 1) ? y_dims : make_ddim({y_dims[0], 1});

  auto mat_dim_a = funcs::CreateMatrixDescriptor(new_x_dims, 0, trans_x);
  auto mat_dim_b = funcs::CreateMatrixDescriptor(new_y_dims, 0, trans_y);

  if (x_dims.size() >= 3 && y_dims.size() <= 2) {
    if (!trans_x || mat_dim_a.batch_size_ == 1) {
      mat_dim_a.height_ *= mat_dim_a.batch_size_;
      mat_dim_a.batch_size_ = 0;
    } else {
      info->is_y_need_broadcast = true;
    }
  }

  if (y_dims.size() >= 3 && x_dims.size() <= 2) {
    info->is_x_need_broadcast = (mat_dim_b.batch_size_ > 1);
  }
  PADDLE_ENFORCE_EQ(mat_dim_a.width_,
                    mat_dim_b.height_,
                    common::errors::InvalidArgument(
                        "Shape mistake in matmul_op xdims = %s ydims = %s "
                        "x_trans = %d y_trans = %d",
                        x_dims.to_str(),
                        y_dims.to_str(),
                        mat_dim_a.trans_,
                        mat_dim_b.trans_));

  if (mat_dim_a.batch_size_ == 0 && mat_dim_b.batch_size_ == 1) {
    mat_dim_a.batch_size_ = mat_dim_b.batch_size_ = 0;
  }
  if (mat_dim_a.batch_size_ == 1 && mat_dim_b.batch_size_ == 0) {
    mat_dim_a.batch_size_ = mat_dim_b.batch_size_ = 0;
  }

  info->m = mat_dim_a.height_;
  info->n = mat_dim_b.width_;
  info->k = mat_dim_a.width_;
  info->bs = std::max(mat_dim_a.batch_size_, mat_dim_b.batch_size_);
  info->trans_x = trans_x;
  info->trans_y = trans_y;

  if (info->bs <= 1) {
    info->stride_x = trans_x ? info->m : info->k;
    info->stride_y = trans_y ? info->k : info->n;
    info->stride_out = info->n;
  } else {
    info->stride_x = info->m * info->k;
    info->stride_y = info->k * info->n;
    info->stride_out = info->m * info->n;
  }
}

template <typename XPUType, typename FCT>
static void xblas_fc_wrapper(xpu::Context* xpu_ctx,
                             const XPUType* x,
                             const XPUType* w,
                             XPUType* y,
                             int64_t m,
                             int64_t n,
                             int64_t k,
                             bool x_trans,
                             bool w_trans,
                             const float* x_maxptr,
                             const float* w_maxptr,
                             float* y_maxptr,
                             int64_t ldx,
                             int64_t ldw,
                             int64_t ldy,
                             float alpha,
                             float beta,
                             const float* bias,
                             const xpu::Activation_t& act,
                             const float* scale_x,
                             const float* scale_w,
                             int scale_x_mode,
                             int scale_w_mode) {
  int r = 0;
  xpu::ctx_guard RAII_GUARD(xpu_ctx);
  if (x_trans && std::getenv("XPU_PADDLE_FC_TRANS_A") != nullptr &&
      std::is_same<float, XPUType>::value) {
    XPUType* l3_addr = nullptr;
    l3_addr = RAII_GUARD.alloc_l3_or_gm<XPUType>(m * k);
    PADDLE_ENFORCE_XDNN_NOT_NULL(l3_addr);

    std::vector<int64_t> shape = {k, m};
    std::vector<int64_t> axis = {1, 0};
    r = xpu::transpose<XPUType>(xpu_ctx, x, l3_addr, shape, axis);
    PADDLE_ENFORCE_XDNN_SUCCESS(r, "transpose");
#ifdef PADDLE_WITH_XPU_XRE5
    r = xblas::fc_fusion<XPUType, XPUType, XPUType, FCT>(xpu_ctx,
                                                         l3_addr,
                                                         w,
                                                         y,
                                                         m,
                                                         n,
                                                         k,
                                                         false,
                                                         w_trans,
                                                         x_maxptr,
                                                         w_maxptr,
                                                         y_maxptr,
                                                         k,
                                                         ldw,
                                                         ldy,
                                                         alpha,
                                                         beta,
                                                         bias,
                                                         act,
                                                         scale_x,
                                                         scale_w,
                                                         scale_x_mode,
                                                         scale_w_mode);
    PADDLE_ENFORCE_XBLAS_SUCCESS(r, "xblas_fc_fusion");
#else
    r = xpu::fc_fusion<XPUType, XPUType, XPUType, FCT>(xpu_ctx,
                                                       l3_addr,
                                                       w,
                                                       y,
                                                       m,
                                                       n,
                                                       k,
                                                       false,
                                                       w_trans,
                                                       x_maxptr,
                                                       w_maxptr,
                                                       y_maxptr,
                                                       k,
                                                       ldw,
                                                       ldy,
                                                       alpha,
                                                       beta,
                                                       bias,
                                                       act);
    PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_fc_fusion");
#endif
  } else {
#ifdef PADDLE_WITH_XPU_XRE5
    if constexpr (std::is_same<XPUTypeBF16, XPUType>::value) {
      if (std::getenv("XPU_PADDLE_FC_BFLOAT16_XTE") != nullptr) {
        const int MAXPTR_N = xpu_ctx->max_ptr_size();
        int64_t x_len = m * k;
        XPUTypeFP16* x_fp16 = nullptr;
        x_fp16 = RAII_GUARD.alloc_l3_or_gm<XPUTypeFP16>(x_len);
        PADDLE_ENFORCE_XDNN_NOT_NULL(x_fp16);
        int64_t w_len = k * n;
        XPUTypeFP16* w_fp16 = nullptr;
        w_fp16 = RAII_GUARD.alloc_l3_or_gm<XPUTypeFP16>(w_len);
        PADDLE_ENFORCE_XDNN_NOT_NULL(w_fp16);

        float* xte_scale_x = nullptr;
        float* xte_scale_w = nullptr;
        xte_scale_x = RAII_GUARD.alloc_l3_or_gm<float>(1);
        PADDLE_ENFORCE_XDNN_NOT_NULL(xte_scale_x);
        xte_scale_w = RAII_GUARD.alloc_l3_or_gm<float>(1);
        PADDLE_ENFORCE_XDNN_NOT_NULL(xte_scale_w);

        float* xte_x_maxptr = nullptr;
        float* xte_w_maxptr = nullptr;
        if (x_maxptr == nullptr) {
          xte_x_maxptr = RAII_GUARD.alloc_l3_or_gm<float>(MAXPTR_N);
          PADDLE_ENFORCE_XDNN_NOT_NULL(xte_x_maxptr);
          int r = xpu::findmax(xpu_ctx, x, xte_x_maxptr, x_len);
          PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_findmax");
          r = xpu::cast_te(
              xpu_ctx, x, xte_x_maxptr, x_fp16, xte_scale_x, x_len);
          PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_cast_te");
        } else {
          r = xpu::cast_te(xpu_ctx, x, x_maxptr, x_fp16, xte_scale_x, x_len);
          PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_cast_te");
        }
        if (w_maxptr == nullptr) {
          xte_w_maxptr = RAII_GUARD.alloc_l3_or_gm<float>(MAXPTR_N);
          PADDLE_ENFORCE_XDNN_NOT_NULL(xte_w_maxptr);
          r = xpu::findmax(xpu_ctx, w, xte_w_maxptr, w_len);
          PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_findmax");
          r = xpu::cast_te(
              xpu_ctx, w, xte_w_maxptr, w_fp16, xte_scale_w, w_len);
          PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_cast_te");
        } else {
          r = xpu::cast_te(xpu_ctx, w, w_maxptr, w_fp16, xte_scale_w, w_len);
          PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_cast_te");
        }

        r = xblas::
            fc_fusion<XPUTypeFP16, XPUTypeFP16, XPUTypeBF16, XPUTypeFP16>(
                xpu_ctx,
                x_fp16,
                w_fp16,
                y,
                m,
                n,
                k,
                x_trans,
                w_trans,
                x_maxptr ? x_maxptr : xte_x_maxptr,
                w_maxptr ? w_maxptr : xte_w_maxptr,
                y_maxptr,
                ldx,
                ldw,
                ldy,
                alpha,
                beta,
                bias,
                act,
                xte_scale_x,
                xte_scale_w,
                scale_x_mode,
                scale_w_mode);

        PADDLE_ENFORCE_XBLAS_SUCCESS(r, "xblas_fc_fusion");
        return;
      }
    }
    r = xblas::fc_fusion<XPUType, XPUType, XPUType, FCT>(xpu_ctx,
                                                         x,
                                                         w,
                                                         y,
                                                         m,
                                                         n,
                                                         k,
                                                         x_trans,
                                                         w_trans,
                                                         x_maxptr,
                                                         w_maxptr,
                                                         y_maxptr,
                                                         ldx,
                                                         ldw,
                                                         ldy,
                                                         alpha,
                                                         beta,
                                                         bias,
                                                         act,
                                                         scale_x,
                                                         scale_w,
                                                         scale_x_mode,
                                                         scale_w_mode);

    PADDLE_ENFORCE_XBLAS_SUCCESS(r, "xblas_fc_fusion");
#else
    r = xpu::fc_fusion<XPUType, XPUType, XPUType, FCT>(xpu_ctx,
                                                       x,
                                                       w,
                                                       y,
                                                       m,
                                                       n,
                                                       k,
                                                       x_trans,
                                                       w_trans,
                                                       x_maxptr,
                                                       w_maxptr,
                                                       y_maxptr,
                                                       ldx,
                                                       ldw,
                                                       ldy,
                                                       alpha,
                                                       beta,
                                                       bias,
                                                       act);
    PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_fc_fusion");
#endif
  }
}

#define DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUType, FCT)          \
  template <>                                                       \
  void xblas_fc_wrapper<XPUType, FCT>(xpu::Context * xpu_ctx,       \
                                      const XPUType* x,             \
                                      const XPUType* w,             \
                                      XPUType* y,                   \
                                      int64_t m,                    \
                                      int64_t n,                    \
                                      int64_t k,                    \
                                      bool x_trans,                 \
                                      bool w_trans,                 \
                                      const float* x_maxptr,        \
                                      const float* w_maxptr,        \
                                      float* y_maxptr,              \
                                      int64_t ldx,                  \
                                      int64_t ldw,                  \
                                      int64_t ldy,                  \
                                      float alpha,                  \
                                      float beta,                   \
                                      const float* bias,            \
                                      const xpu::Activation_t& act, \
                                      const float* scale_x,         \
                                      const float* scale_w,         \
                                      int scale_x_mode,             \
                                      int scale_w_mode) {           \
    int r = xpu::Error_t::INVALID_PARAM;                            \
    PADDLE_ENFORCE_XDNN_SUCCESS(r, "xblas_fc_wrapper");             \
  }

DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeBF16, int_with_ll_t)
DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeBF16, int16_t)
DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeBF16, int32_t)
DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeBF16, XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeFP16, int32_t)
DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeFP16, tfloat32)

#ifndef PADDLE_WITH_XPU_XRE5
// disable in kl2
DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(float, XPUTypeFP16)
#endif

template <typename XPUType, typename FCT, typename TGEMM_OUT>
static void xblas_fc_batch_wrapper(xpu::Context* xpu_ctx,
                                   int64_t bs,
                                   bool trans_x,
                                   bool trans_w,
                                   int64_t m,
                                   int64_t n,
                                   int64_t k,
                                   float alpha,
                                   const XPUType* x,
                                   int64_t stride_x,
                                   const XPUType* w,
                                   int64_t stride_w,
                                   float beta,
                                   XPUType* y,
                                   int64_t stride_y,
                                   const float* x_maxptr,
                                   const float* w_maxptr) {
#ifdef PADDLE_WITH_XPU_XRE5
  int r = xblas::fc_batched<XPUType, XPUType, XPUType, FCT, TGEMM_OUT, 0>(
      xpu_ctx,
      bs,
      trans_x,
      trans_w,
      m,
      n,
      k,
      alpha,
      reinterpret_cast<const XPUType*>(x),
      stride_x,
      reinterpret_cast<const XPUType*>(w),
      stride_w,
      beta,
      reinterpret_cast<XPUType*>(y),
      stride_y,
      x_maxptr,
      w_maxptr);
  PADDLE_ENFORCE_XBLAS_SUCCESS(r, "xblas_fc_batched");
#else
  int r = xpu::fc_batched<XPUType, XPUType, XPUType, FCT>(
      xpu_ctx,
      bs,
      trans_x,
      trans_w,
      m,
      n,
      k,
      alpha,
      reinterpret_cast<const XPUType*>(x),
      stride_x,
      reinterpret_cast<const XPUType*>(w),
      stride_w,
      beta,
      reinterpret_cast<XPUType*>(y),
      stride_y,
      x_maxptr,
      w_maxptr);
  PADDLE_ENFORCE_XDNN_SUCCESS(r, "xdnn_fc_batched");
#endif
}

#define DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUType, FCT, TGEMM_OUT) \
  template <>                                                               \
  void xblas_fc_batch_wrapper<XPUType, FCT, TGEMM_OUT>(                     \
      xpu::Context * xpu_ctx,                                               \
      int64_t bs,                                                           \
      bool trans_x,                                                         \
      bool trans_w,                                                         \
      int64_t m,                                                            \
      int64_t n,                                                            \
      int64_t k,                                                            \
      float alpha,                                                          \
      const XPUType* x,                                                     \
      int64_t stride_x,                                                     \
      const XPUType* w,                                                     \
      int64_t stride_w,                                                     \
      float beta,                                                           \
      XPUType* y,                                                           \
      int64_t stride_y,                                                     \
      const float* x_maxptr,                                                \
      const float* w_maxptr) {                                              \
    int r = xpu::Error_t::INVALID_PARAM;                                    \
    PADDLE_ENFORCE_XDNN_SUCCESS(r, "xblas_fc_batched");                     \
  }

DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16,
                                           int_with_ll_t,
                                           XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16, tfloat32, XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, float, XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16,
                                           XPUTypeFP16,
                                           XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, int32_t, XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, int16_t, XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16, int32_t, XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, int_with_ll_t, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, XPUTypeFP16, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16, tfloat32, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, int32_t, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, int16_t, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16, int32_t, float)

#ifndef PADDLE_WITH_XPU_XRE5
// disable in kl2
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, tfloat32, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, float, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(float, XPUTypeFP16, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(float, XPUTypeFP16, XPUTypeFP16)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16, XPUTypeFP16, float)
DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16,
                                           XPUTypeFP16,
                                           XPUTypeFP16)
#endif

template <typename T>
static void MatMulXPUFunction(
    xpu::Context* xpu_ctx,
    const T* x,
    const T* y,
    T* out,
    const XpuFcInfo& fcinfo,
    float alpha,
    float beta = 0.f,
    bool is_grad = false,
    xpu::Activation_t act = xpu::Activation_t::LINEAR) {
  using XPUType = typename XPUTypeTrait<T>::Type;
  int fc_calc_type = FCCalcType<XPUType>();

  decltype(&xblas_fc_wrapper<XPUType, int16_t>) xblas_fc_api_list[6] = {
      &xblas_fc_wrapper<XPUType, int16_t>,
      &xblas_fc_wrapper<XPUType, int32_t>,
      &xblas_fc_wrapper<XPUType, float>,
      &xblas_fc_wrapper<XPUType, int_with_ll_t>,
      &xblas_fc_wrapper<XPUType, tfloat32>,
      &xblas_fc_wrapper<XPUType, XPUTypeFP16>,
  };

  decltype(&xblas_fc_batch_wrapper<XPUType, int16_t, float>)
      xblas_fc_batch_api_list[6] = {
          &xblas_fc_batch_wrapper<XPUType, int16_t, float>,
          &xblas_fc_batch_wrapper<XPUType, int32_t, float>,
          &xblas_fc_batch_wrapper<XPUType, float, float>,
          &xblas_fc_batch_wrapper<XPUType, int_with_ll_t, float>,
          &xblas_fc_batch_wrapper<XPUType, tfloat32, float>,
          &xblas_fc_batch_wrapper<XPUType, XPUTypeFP16, float>,
      };

  auto xblas_fc_api = xblas_fc_api_list[fc_calc_type];

  if (std::getenv("XPU_PADDLE_FC_GRAD_LOCAL") != nullptr) {
    if (is_grad) {
      xblas_fc_api = xblas_fc_api_list[2];
    }
  }
  auto xblas_fc_batch_api = xblas_fc_batch_api_list[fc_calc_type];

  if (fc_calc_type == XPUFCCalcType::FC_FLOAT16 &&
      std::getenv("XPU_PADDLE_FC_FLOAT16") != nullptr) {
    xblas_fc_batch_api =
        &xblas_fc_batch_wrapper<XPUType, XPUTypeFP16, XPUTypeFP16>;
  }
  int64_t m = fcinfo.m;
  int64_t n = fcinfo.n;
  int64_t k = fcinfo.k;
  int64_t batch_size = fcinfo.bs;
  int64_t ldx = fcinfo.stride_x;
  int64_t ldy = fcinfo.stride_y;
  int64_t ldout = fcinfo.stride_out;
  bool trans_x = fcinfo.trans_x;
  bool trans_y = fcinfo.trans_y;
  float* max_x = fcinfo.max_x;
  float* max_y = fcinfo.max_y;
  float* max_out = fcinfo.max_out;
  bool is_x_need_broadcast = fcinfo.is_x_need_broadcast;
  bool is_y_need_broadcast = fcinfo.is_y_need_broadcast;
  const float* bias = fcinfo.bias;
  const float* scale_x = fcinfo.scale_x;
  const float* scale_y = fcinfo.scale_y;
  int scale_x_mode = fcinfo.scale_x_mode;
  int scale_y_mode = fcinfo.scale_y_mode;

  xpu::ctx_guard RAII_GUARD(xpu_ctx);
  if (batch_size <= 1) {
    xblas_fc_api(xpu_ctx,
                 reinterpret_cast<const XPUType*>(x),
                 reinterpret_cast<const XPUType*>(y),
                 reinterpret_cast<XPUType*>(out),
                 m,
                 n,
                 k,
                 trans_x,
                 trans_y,
                 max_x,
                 max_y,
                 max_out,
                 ldx,
                 ldy,
                 ldout,
                 alpha,
                 beta,
                 bias,
                 act,
                 scale_x,
                 scale_y,
                 scale_x_mode,
                 scale_y_mode);
  } else {
    const XPUType* x_data = reinterpret_cast<const XPUType*>(x);
    if (is_x_need_broadcast) {
      XPUType* x_broadcast_data = nullptr;
      x_broadcast_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(batch_size * m * k);
      PADDLE_ENFORCE_XDNN_NOT_NULL(x_broadcast_data);
      std::vector<int64_t> x_shape = {1, m, k};
      std::vector<int64_t> new_x_shape = {batch_size, m, k};
      int r = xpu::broadcast<XPUType>(
          xpu_ctx, x_data, x_broadcast_data, x_shape, new_x_shape);
      PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast");
      x_data = x_broadcast_data;
    }
    const XPUType* y_data = reinterpret_cast<const XPUType*>(y);
    if (is_y_need_broadcast) {
      XPUType* y_broadcast_data = nullptr;
      y_broadcast_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(batch_size * k * n);
      PADDLE_ENFORCE_XDNN_NOT_NULL(y_broadcast_data);
      std::vector<int64_t> y_shape = {1, k, n};
      std::vector<int64_t> new_y_shape = {batch_size, k, n};
      int r = xpu::broadcast<XPUType>(
          xpu_ctx, y_data, y_broadcast_data, y_shape, new_y_shape);
      PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast");
      y_data = y_broadcast_data;
    }
    // batch matmul
    xblas_fc_batch_api(xpu_ctx,                          // Context* xpu_ctx,
                       batch_size,                       // int64_t batch_size,
                       trans_x,                          // bool x_trans,
                       trans_y,                          // bool w_trans,
                       m,                                // int64_t m,
                       n,                                // int64_t n,
                       k,                                // int64_t k,
                       alpha,                            // float alpha,
                       x_data,                           // const TX* x,
                       ldx,                              // int64_t stride_a,
                       y_data,                           // const TW* w,
                       ldy,                              // int64_t stride_b,
                       beta,                             // float beta,
                       reinterpret_cast<XPUType*>(out),  // TY* y,
                       ldout,                            // int64_t stride_c,
                       max_x,   // const float* x_maxptr,
                       max_y);  // const float* w_maxptr
  }
}

template <typename T>
static std::tuple<XpuFcInfo, XpuFcInfo, const T*, const T*, const T*, const T*>
MatmulGradFcInfo(xpu::Context* xpu_ctx,
                 xpu::ctx_guard* RAII_GUARD,
                 const XpuFcInfo& dout_shape,
                 bool trans_x,
                 bool trans_y,
                 const T* x,
                 const T* y,
                 const T* dout) {
  XpuFcInfo dx_shape, dy_shape;
  const T* dx_a = NULL;
  const T* dx_b = NULL;
  const T* dy_a = NULL;
  const T* dy_b = NULL;
  bool copy_to_l3 = false;
  float* max_dout = NULL;
  int maxptr_size = xpu_ctx->max_ptr_size();
  uint64_t l3_size = uint64_t(xpu_ctx->_l3_mgr.get_size());
  int64_t bs = (dout_shape.bs <= 1) ? (1) : (dout_shape.bs);
  int64_t dx_size = bs * dout_shape.m * dout_shape.k;
  int64_t dy_size = bs * dout_shape.k * dout_shape.n;
  int64_t dout_size = bs * dout_shape.m * dout_shape.n;
  if (trans_x && trans_y) {
    copy_to_l3 = l3_size >= (dout_size * 2 + dy_size) * sizeof(T);
  } else if (trans_x) {
    copy_to_l3 = l3_size >= dout_size * sizeof(T);
  } else if (trans_y) {
    copy_to_l3 = l3_size >= dout_size * 2 * sizeof(T);
  } else {
    copy_to_l3 = l3_size >= (dout_size + dx_size) * sizeof(T);
  }

  const T* dout_new = dout;
  int r = 0;
  if (copy_to_l3) {
    T* dout_l3 = RAII_GUARD->alloc_l3<T>(dout_size);
    PADDLE_ENFORCE_XDNN_NOT_NULL(dout_l3);
    if ((dout_shape.bs > 1) || ((dout_shape.bs <= 1) &&
                                (FCCalcType<T>() == XPUFCCalcType::FC_FLOAT))) {
      r = xpu::copy(xpu_ctx, dout, dout_l3, dout_size);
      PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy");
      dout_new = dout_l3;
    } else {
      max_dout = RAII_GUARD->alloc_l3_or_gm<float>(maxptr_size);
      PADDLE_ENFORCE_XDNN_NOT_NULL(max_dout);

      r = xpu::findmax_copy_fusion(xpu_ctx, dout, max_dout, dout_l3, dout_size);
      PADDLE_ENFORCE_XDNN_SUCCESS(r, "findmax_copy_fusion");
      dout_new = dout_l3;
    }
  } else if (((dout_shape.bs <= 1) &&
              (FCCalcType<T>() != XPUFCCalcType::FC_FLOAT))) {
    max_dout = RAII_GUARD->alloc_l3_or_gm<float>(maxptr_size);
    PADDLE_ENFORCE_XDNN_NOT_NULL(max_dout);
    r = xpu::findmax_copy_fusion(
        xpu_ctx, dout, max_dout, reinterpret_cast<T*>(NULL), dout_size);
    PADDLE_ENFORCE_XDNN_SUCCESS(r, "findmax_copy_fusion");
  }

  if (trans_x && trans_y) {
    // dx = T(y) * T(dout)
    dx_shape.InitFcInfo(dout_shape.bs,
                        dout_shape.k,
                        dout_shape.m,
                        dout_shape.n,
                        true,
                        true,
                        nullptr,
                        max_dout,
                        nullptr);
    dx_a = y, dx_b = dout_new;
    dx_shape.is_x_need_broadcast = dout_shape.is_y_need_broadcast;
    // dy = T(dout) * T(x)
    dy_shape.InitFcInfo(dout_shape.bs,
                        dout_shape.n,
                        dout_shape.k,
                        dout_shape.m,
                        true,
                        true,
                        max_dout,
                        nullptr,
                        nullptr);
    dy_a = dout_new, dy_b = x;
    dy_shape.is_y_need_broadcast = dout_shape.is_x_need_broadcast;
  } else if (trans_x) {
    // dx = y * T(dout)
    dx_shape.InitFcInfo(dout_shape.bs,
                        dout_shape.k,
                        dout_shape.m,
                        dout_shape.n,
                        false,
                        true,
                        nullptr,
                        max_dout,
                        nullptr);
    dx_a = y, dx_b = dout_new;
    dx_shape.is_x_need_broadcast = dout_shape.is_y_need_broadcast;
    // dy = x * dout
    dy_shape.InitFcInfo(dout_shape.bs,
                        dout_shape.k,
                        dout_shape.n,
                        dout_shape.m,
                        false,
                        false,
                        nullptr,
                        max_dout,
                        nullptr);
    dy_shape.is_x_need_broadcast = dout_shape.is_x_need_broadcast;
    dy_a = x, dy_b = dout_new;
  } else if (trans_y) {
    // dx = dout * y
    dx_shape.InitFcInfo(dout_shape.bs,
                        dout_shape.m,
                        dout_shape.k,
                        dout_shape.n,
                        false,
                        false,
                        max_dout,
                        nullptr,
                        nullptr);
    dx_a = dout_new, dx_b = y;
    dx_shape.is_y_need_broadcast = dout_shape.is_y_need_broadcast;
    // dy =  T(dout) * x
    dy_shape.InitFcInfo(dout_shape.bs,
                        dout_shape.n,
                        dout_shape.k,
                        dout_shape.m,
                        true,
                        false,
                        max_dout,
                        nullptr,
                        nullptr);
    dy_a = dout_new, dy_b = x;
    dy_shape.is_y_need_broadcast = dout_shape.is_x_need_broadcast;
  } else {
    // dx = dout * T(y)
    dx_shape.InitFcInfo(dout_shape.bs,
                        dout_shape.m,
                        dout_shape.k,
                        dout_shape.n,
                        false,
                        true,
                        max_dout,
                        nullptr,
                        nullptr);
    dx_a = dout_new, dx_b = y;
    dx_shape.is_y_need_broadcast = dout_shape.is_y_need_broadcast;
    // dy = T(x) * dout
    dy_shape.InitFcInfo(dout_shape.bs,
                        dout_shape.k,
                        dout_shape.n,
                        dout_shape.m,
                        true,
                        false,
                        nullptr,
                        max_dout,
                        nullptr);
    dy_shape.is_x_need_broadcast = dout_shape.is_x_need_broadcast;
    dy_a = x, dy_b = dout_new;
  }
  std::tuple<XpuFcInfo, XpuFcInfo, const T*, const T*, const T*, const T*>
      result = std::make_tuple(dx_shape, dy_shape, dx_a, dx_b, dy_a, dy_b);

  return result;
}

}  // namespace phi