/* Copyright (c) 2021 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 "glog/logging.h"

#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/autotune/cache_base.h"
#include "paddle/phi/kernels/cast_kernel.h"
#include "paddle/phi/kernels/contiguous_kernel.h"
#include "paddle/phi/kernels/funcs/blas/blas.h"
#ifdef PADDLE_WITH_HIP
#include "paddle/phi/kernels/funcs/blas/blaslt_impl.hip.h"
#else
#include "paddle/phi/kernels/funcs/blas/blaslt_impl.cu.h"
#endif
#include "paddle/phi/kernels/funcs/complex_functors.h"
#include "paddle/phi/kernels/scale_kernel.h"
#if defined(PADDLE_WITH_CUDA)
#include "paddle/phi/kernels/funcs/cublaslt.h"
#include "paddle/phi/kernels/gpu/cuda_gemm_kernel.h"
#include "paddle/phi/kernels/transpose_kernel.h"
#elif defined(PADDLE_WITH_HIP)
#include "paddle/phi/kernels/funcs/hipblaslt.h"
#endif
#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 11060
#include "paddle/phi/kernels/autotune/auto_tune_base.h"
#endif
#include "paddle/phi/kernels/full_kernel.h"

COMMON_DECLARE_bool(cuda_core_int8_gemm);
COMMON_DECLARE_bool(use_legacy_gemm);
COMMON_DECLARE_bool(use_accuracy_compatible_kernel);

namespace phi {

static void GetBroadcastFromDims(const int x_ndim,
                                 const std::int64_t* x_dims,
                                 const int y_ndim,
                                 const std::int64_t* y_dims,
                                 std::int64_t* x_bd_dims,
                                 std::int64_t* y_bd_dims,
                                 std::int64_t* out_bd_dims) {
  const int ndim = (std::max)(x_ndim, y_ndim);
  std::fill(x_bd_dims, x_bd_dims + ndim - x_ndim, 1);
  std::fill(y_bd_dims, y_bd_dims + ndim - y_ndim, 1);
  std::copy(x_dims, x_dims + x_ndim, x_bd_dims + ndim - x_ndim);
  std::copy(y_dims, y_dims + y_ndim, y_bd_dims + ndim - y_ndim);

  for (int i = 0; i < ndim; ++i) {
    PADDLE_ENFORCE_EQ(
        x_bd_dims[i] == y_bd_dims[i] || x_bd_dims[i] <= 1 || y_bd_dims[i] <= 1,
        true,
        common::errors::InvalidArgument(
            "Input(X) and Input(Y) has error dim. "
            "X_broadcast's shape[%s] must be equal to Y_broadcast's shape[%s], "
            "or X_broadcast's shape[%s] <= 1, or Y_broadcast's shape[%s] <= 1, "
            "but received X_broadcast's shape[%s] = [%s]"
            "received Y_broadcast's shape[%s] = [%s].",
            i,
            i,
            i,
            i,
            i,
            x_bd_dims[i],
            i,
            y_bd_dims[i]));
    if (x_bd_dims[i] == 0 || y_bd_dims[i] == 0) {
      out_bd_dims[i] = 0;
    } else {
      out_bd_dims[i] = (std::max)(x_bd_dims[i], y_bd_dims[i]);
    }
  }
}

static int64_t GetIndexMessage(const int n,
                               const int64_t* dims,
                               const int64_t* index) {
  int64_t sum = 0;
  for (int i = 0; i < n; ++i) {
    if (dims[i] > 1) {
      sum = sum * dims[i] + index[i];
    }
  }
  return sum;
}

static void IndexIncreaseFromDims(const int ndim,
                                  const int64_t* dims,
                                  int64_t* index) {
  for (int i = ndim - 1; i >= 0; --i) {
    ++index[i];
    if (index[i] >= dims[i]) {
      index[i] -= dims[i];
    } else {
      break;
    }
  }
}

// The general implementation with blas.
template <typename Context, typename T>
void MatMulFunctionImplWithBlas(
    const Context& dev_ctx,
    const DenseTensor& X,
    const DenseTensor& Y,
    const std::vector<std::int64_t>& x_dims,
    const std::vector<std::int64_t>& y_dims,
    DenseTensor* Out,
    bool trans_x,
    bool trans_y,
    bool flag = false,
    funcs::MatmulPlanner* matmul_planner UNUSED = nullptr) {
  const int x_ndim = x_dims.size();
  const int y_ndim = y_dims.size();

  // Get data ptr
  const T* x_data = X.data<T>();
  const T* y_data = Y.data<T>();

  auto blas = funcs::GetBlas<Context, T>(dev_ctx);

  if (x_ndim == 1 && y_ndim == 1) {
    const int64_t M = X.numel();
    const int64_t N = Y.numel();
    PADDLE_ENFORCE_EQ(
        M,
        N,
        common::errors::InvalidArgument(
            "X's numbers must be equal to Y's numbers, "
            "when X/Y's dims =1. But received X has [%d] elements, "
            "received Y has [%d] elements.",
            M,
            N));
    VLOG(3) << "MatMul's case 1";
    Out->Resize({});
    dev_ctx.template Alloc<T>(Out);
    if (FLAGS_use_legacy_gemm) {
      blas.GEMM(CblasNoTrans,
                CblasTrans,
                1,
                1,
                M,
                static_cast<T>(1),
                y_data,
                x_data,
                static_cast<T>(flag),
                dev_ctx.template Alloc<T>(Out));
      return;
    } else {
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
      if (std::is_same<Context, GPUContext>::value) {
        blas.CUDOT(M, X.data<T>(), 1, Y.data<T>(), 1, Out->data<T>());
      } else  // NOLINT
#endif
      {
        blas.GEMM(CblasNoTrans,
                  CblasTrans,
                  1,
                  1,
                  M,
                  static_cast<T>(1),
                  y_data,
                  x_data,
                  static_cast<T>(flag),
                  dev_ctx.template Alloc<T>(Out));
      }
      return;
    }
  }

  if (x_ndim == 1) {
    const int64_t N = X.numel();
    if (trans_y) {
      PADDLE_ENFORCE_EQ(
          y_dims[y_ndim - 1],
          N,
          common::errors::InvalidArgument("Input(Y) has error dim. "
                                          "Y'dims[%d] must be equal to %d, "
                                          "but received Y'dims[%d] is %d.",
                                          y_ndim - 1,
                                          N,
                                          y_ndim - 1,
                                          y_dims[y_ndim - 1]));
    } else {
      PADDLE_ENFORCE_EQ(
          y_dims[y_ndim - 2],
          N,
          common::errors::InvalidArgument("Input(Y) has error dim. "
                                          "Y'dims[%d] must be equal to %d, "
                                          "but received Y'dims[%d] is %d.",
                                          y_ndim - 2,
                                          N,
                                          y_ndim - 2,
                                          y_dims[y_ndim - 2]));
    }
    std::vector<std::int64_t> out_dims(y_ndim - 1);
    if (trans_y) {
      std::copy_n(y_dims.cbegin(), y_ndim - 1, out_dims.begin());
    } else {
      std::copy_n(y_dims.cbegin(), y_ndim - 2, out_dims.begin());
      out_dims.back() = y_dims.back();
    }
    Out->ResizeAndAllocate(make_ddim(out_dims));
    dev_ctx.template Alloc<T>(Out);
    if (trans_y) {
      const int64_t M = Y.numel() / N;
      VLOG(3) << "MatMul's case 2";
      blas.GEMV(false,
                M,
                N,
                static_cast<T>(1),
                y_data,
                x_data,
                static_cast<T>(flag),
                dev_ctx.template Alloc<T>(Out));
    } else {
      const int64_t M = y_dims[y_ndim - 1];
      const int64_t batch_size = Y.numel() / (M * N);
      if (batch_size == 1) {
        VLOG(3) << "MatMul's case 3";
        blas.GEMV(true,
                  N,
                  M,
                  static_cast<T>(1),
                  y_data,
                  x_data,
                  static_cast<T>(flag),
                  dev_ctx.template Alloc<T>(Out));
      } else {
        VLOG(3) << "MatMul's case 4";
        blas.BatchedGEMM(CblasTrans,
                         CblasNoTrans,
                         M,
                         1,
                         N,
                         static_cast<T>(1),
                         y_data,
                         x_data,
                         static_cast<T>(flag),
                         dev_ctx.template Alloc<T>(Out),
                         batch_size,
                         M * N,
                         0);
      }
    }
    return;
  }

  if (y_ndim == 1) {
    const int64_t N = Y.numel();
    if (trans_x) {
      PADDLE_ENFORCE_EQ(
          x_dims[x_ndim - 2],
          N,
          common::errors::InvalidArgument("Input(X) has error dim. "
                                          "X'dims[%d] must be equal to %d. "
                                          "But received X'dims[%d] is %d",
                                          x_ndim - 2,
                                          N,
                                          x_ndim - 2,
                                          x_dims[x_ndim - 2]));
    } else {
      PADDLE_ENFORCE_EQ(
          x_dims[x_ndim - 1],
          N,
          common::errors::InvalidArgument("Input(X) has error dim. "
                                          "X'dims[%d] must be equal to %d. "
                                          "But received X'dims[%d] is %d",
                                          x_ndim - 1,
                                          N,
                                          x_ndim - 1,
                                          x_dims[x_ndim - 1]));
    }
    std::vector<std::int64_t> out_dims(x_ndim - 1);
    if (trans_x) {
      std::copy_n(x_dims.cbegin(), x_ndim - 2, out_dims.begin());
      out_dims.back() = x_dims.back();
    } else {
      std::copy_n(x_dims.cbegin(), x_ndim - 1, out_dims.begin());
    }
    Out->ResizeAndAllocate(make_ddim(out_dims));
    dev_ctx.template Alloc<T>(Out);

    if (trans_x) {
      const int64_t M = x_dims[x_ndim - 1];
      const int64_t batch_size = X.numel() / (M * N);
      if (batch_size == 1) {
        VLOG(3) << "MatMul's case 5";
        blas.GEMV(true,
                  N,
                  M,
                  static_cast<T>(1),
                  x_data,
                  y_data,
                  static_cast<T>(flag),
                  dev_ctx.template Alloc<T>(Out));
      } else {
        VLOG(3) << "MatMul's case 6";
        blas.BatchedGEMM(CblasTrans,
                         CblasNoTrans,
                         M,
                         1,
                         N,
                         static_cast<T>(1),
                         x_data,
                         y_data,
                         static_cast<T>(flag),
                         dev_ctx.template Alloc<T>(Out),
                         batch_size,
                         M * N,
                         0);
      }
    } else {
      const int64_t M = X.numel() / N;
      VLOG(3) << "MatMul's case 7";
      blas.GEMV(false,
                M,
                N,
                static_cast<T>(1),
                x_data,
                y_data,
                static_cast<T>(flag),
                dev_ctx.template Alloc<T>(Out));
    }
    return;
  }

  const int64_t M = trans_x ? x_dims[x_ndim - 1] : x_dims[x_ndim - 2];
  const int64_t K = trans_x ? x_dims[x_ndim - 2] : x_dims[x_ndim - 1];
  if (trans_y) {
    PADDLE_ENFORCE_EQ(
        y_dims[y_ndim - 1],
        K,
        common::errors::InvalidArgument("Input(Y) has error dim. "
                                        "Y'dims[%d] must be equal to %d, "
                                        "but received Y'dims[%d] is %d.",
                                        y_ndim - 1,
                                        K,
                                        y_ndim - 1,
                                        y_dims[y_ndim - 1]));
  } else {
    PADDLE_ENFORCE_EQ(
        y_dims[y_ndim - 2],
        K,
        common::errors::InvalidArgument("Input(Y) has error dim. "
                                        "Y'dims[%d] must be equal to %d, "
                                        "but received Y'dims[%d] is %d.",
                                        y_ndim - 2,
                                        K,
                                        y_ndim - 2,
                                        y_dims[y_ndim - 2]));
  }
  const int64_t N = trans_y ? y_dims[y_ndim - 2] : y_dims[y_ndim - 1];
  const int64_t ndim = (std::max)(x_ndim, y_ndim);
  std::vector<std::int64_t> x_broadcast_dims(ndim);
  std::vector<std::int64_t> y_broadcast_dims(ndim);
  std::vector<std::int64_t> out_broadcast_dims(ndim);

  GetBroadcastFromDims(x_ndim - 2,
                       x_dims.data(),
                       y_ndim - 2,
                       y_dims.data(),
                       x_broadcast_dims.data(),
                       y_broadcast_dims.data(),
                       out_broadcast_dims.data());
  out_broadcast_dims[ndim - 2] = M;
  out_broadcast_dims[ndim - 1] = N;

  Out->ResizeAndAllocate(make_ddim(out_broadcast_dims));
  dev_ctx.template Alloc<T>(Out);

  const int batch_dim = ndim - 2;
  // broadcast message
  const bool is_broadcast_dims =
      !std::equal(x_broadcast_dims.cbegin(),
                  x_broadcast_dims.cbegin() + batch_dim,
                  y_broadcast_dims.cbegin());

  const std::int64_t x_batch_size =
      std::accumulate(x_broadcast_dims.cbegin(),
                      x_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  const std::int64_t y_batch_size =
      std::accumulate(y_broadcast_dims.cbegin(),
                      y_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  const std::int64_t out_batch_size =
      std::accumulate(out_broadcast_dims.cbegin(),
                      out_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  if (out_batch_size == 0) return;
  if (x_batch_size == 1 && y_batch_size == 1) {
    VLOG(3) << "MatMul's case 8";
    blas.GEMM(trans_x ? CblasTrans : CblasNoTrans,
              trans_y ? CblasTrans : CblasNoTrans,
              M,
              N,
              K,
              static_cast<T>(1),
              x_data,
              y_data,
              static_cast<T>(flag),
              dev_ctx.template Alloc<T>(Out));
  } else if (x_batch_size == 1) {
    if (M == 1 && trans_y) {
      VLOG(3) << "MatMul's case 9";
      PADDLE_ENFORCE_LE_INT_MAX(y_batch_size * N, "GEMV M");
      blas.GEMV(false,
                static_cast<int>(y_batch_size * N),
                K,
                static_cast<T>(1),
                y_data,
                x_data,
                static_cast<T>(flag),
                dev_ctx.template Alloc<T>(Out));
    } else {
      VLOG(3) << "MatMul's case 10";
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
      // NOTE(Pan Zhaowu): This change may cause performance degrade on current
      // low-efficiency implementation of batched transpose. So enable with flag
      // "use_accuracy_compatible_kernel"
      if (!FLAGS_use_legacy_gemm && FLAGS_use_accuracy_compatible_kernel) {
        // x_batch_size == 1 && M != 1 || !transy
        DenseTensor processedY = trans_y ? Y : TransposeLast2Dim<T>(dev_ctx, Y);
        DenseTensor processedX = X;
        blas.GEMM(CblasNoTrans,
                  trans_x ? CblasNoTrans : CblasTrans,
                  Y.numel() / K,
                  M,
                  K,
                  static_cast<T>(1),
                  processedY.data<T>(),
                  processedX.data<T>(),
                  static_cast<T>(flag),
                  dev_ctx.template Alloc<T>(Out));
        // The actual layout is (B, N, M), need to reshape and
        // transpose to (B, M, N), this requires batched transpose kernel
        // to be implemented in high efficiency.
        const auto out_original_shape = Out->dims();
        std::vector<int64_t> actual_dim = vectorize(processedY.dims());
        actual_dim[actual_dim.size() - 1] =
            out_original_shape[out_original_shape.size() - 2];
        Out->Resize(actual_dim);
        DenseTensor transposedOut = TransposeLast2Dim<T>(dev_ctx, *Out);
        *Out = transposedOut;
        Out->Resize(out_original_shape);
      } else  // NOLINT
#endif
      {  // NOLINT
        blas.BatchedGEMM(trans_x ? CblasTrans : CblasNoTrans,
                         trans_y ? CblasTrans : CblasNoTrans,
                         M,
                         N,
                         K,
                         static_cast<T>(1),
                         x_data,
                         y_data,
                         static_cast<T>(flag),
                         dev_ctx.template Alloc<T>(Out),
                         out_batch_size,
                         0,
                         K * N);
      }
    }
  } else if (y_batch_size == 1) {
    if (!trans_x) {
      VLOG(3) << "MatMul's case 11";
      PADDLE_ENFORCE_LE_INT_MAX(x_batch_size * M, "GEMM M");
      blas.GEMM(CblasNoTrans,
                trans_y ? CblasTrans : CblasNoTrans,
                x_batch_size * M,
                N,
                K,
                static_cast<T>(1),
                x_data,
                y_data,
                static_cast<T>(flag),
                dev_ctx.template Alloc<T>(Out));
    } else {
      VLOG(3) << "MatMul's case 12";
      blas.BatchedGEMM(CblasTrans,
                       trans_y ? CblasTrans : CblasNoTrans,
                       M,
                       N,
                       K,
                       static_cast<T>(1),
                       x_data,
                       y_data,
                       static_cast<T>(flag),
                       dev_ctx.template Alloc<T>(Out),
                       out_batch_size,
                       M * K,
                       0);
    }
  } else if (!is_broadcast_dims) {
    VLOG(3) << "MatMul's case 13";
    blas.BatchedGEMM(trans_x ? CblasTrans : CblasNoTrans,
                     trans_y ? CblasTrans : CblasNoTrans,
                     M,
                     N,
                     K,
                     static_cast<T>(1),
                     x_data,
                     y_data,
                     static_cast<T>(flag),
                     dev_ctx.template Alloc<T>(Out),
                     out_batch_size,
                     M * K,
                     K * N);
  } else {
    // in the case, can't use stridedgemm
    PADDLE_ENFORCE_LE_INT_MAX(out_batch_size, "out_batch_size");
    std::vector<const T*> x_ptr(out_batch_size);
    std::vector<const T*> y_ptr(out_batch_size);
    std::vector<T*> out_ptr(out_batch_size);
    std::vector<std::int64_t> index(batch_dim, 0);
    for (std::int64_t i = 0; i < out_batch_size; ++i) {
      // using the index to get offset
      const std::int64_t x_index =
          GetIndexMessage(batch_dim, x_broadcast_dims.data(), index.data());
      const std::int64_t y_index =
          GetIndexMessage(batch_dim, y_broadcast_dims.data(), index.data());

      x_ptr[i] = x_data + x_index * M * K;
      y_ptr[i] = y_data + y_index * K * N;
      out_ptr[i] = dev_ctx.template Alloc<T>(Out) + i * M * N;
      IndexIncreaseFromDims(batch_dim, out_broadcast_dims.data(), index.data());
    }
    VLOG(3) << "MatMul's case 14";
    blas.BatchedGEMM(trans_x ? CblasTrans : CblasNoTrans,
                     trans_y ? CblasTrans : CblasNoTrans,
                     M,
                     N,
                     K,
                     static_cast<T>(1),
                     x_ptr.data(),
                     y_ptr.data(),
                     static_cast<T>(flag),
                     out_ptr.data(),
                     static_cast<int>(out_batch_size));
  }
}

#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 11060
// This is almost a copy from MatMulFunctionImplWithBlas,
// compare cublas with cublasLt kernels when Matmul autotune is on
template <typename Context, typename T>
void MatMulFunctionImplWithCublasLt(
    const Context& dev_ctx,
    const DenseTensor& X,
    const DenseTensor& Y,
    const std::vector<std::int64_t>& x_dims,
    const std::vector<std::int64_t>& y_dims,
    DenseTensor* Out,
    bool trans_x,
    bool trans_y,
    bool flag = false,
    funcs::MatmulPlanner* matmul_planner = nullptr) {
  const int x_ndim = x_dims.size();
  const int y_ndim = y_dims.size();
  const T* x_data = X.data<T>();
  const T* y_data = Y.data<T>();
  using blaslt = funcs::MatmulWithCublasLt<T>;

  if (x_ndim == 1 && y_ndim == 1) {
    const int64_t M = X.numel();
    const int64_t N = Y.numel();
    PADDLE_ENFORCE_EQ(
        M,
        N,
        common::errors::InvalidArgument(
            "X's numbers must be equal to Y's numbers, "
            "when X/Y's dims =1. But received X has [%d] elements, "
            "received Y has [%d] elements",
            M,
            N));

    // MatMul's case 0  =>  vector * vector
    Out->Resize({});
    dev_ctx.template Alloc<T>(Out);
    VLOG(3) << "MatMul with blaslt case 1";
    blaslt::Run(dev_ctx,
                y_data,
                x_data,
                dev_ctx.template Alloc<T>(Out),
                1,
                1,
                M,
                false,
                true,
                matmul_planner);
    return;
  }

  if (x_ndim == 1) {
    const int64_t N = X.numel();
    if (trans_y) {
      PADDLE_ENFORCE_EQ(
          y_dims[y_ndim - 1],
          N,
          common::errors::InvalidArgument("Input(Y) has error dim. "
                                          "Y'dims[%d] must be equal to %d, "
                                          "but received Y'dims[%d] is %d.",
                                          y_ndim - 1,
                                          N,
                                          y_ndim - 1,
                                          y_dims[y_ndim - 1]));
    } else {
      PADDLE_ENFORCE_EQ(
          y_dims[y_ndim - 2],
          N,
          common::errors::InvalidArgument("Input(Y) has error dim. "
                                          "Y'dims[%d] must be equal to %d, "
                                          "but received Y'dims[%d] is %d.",
                                          y_ndim - 2,
                                          N,
                                          y_ndim - 2,
                                          y_dims[y_ndim - 2]));
    }
    std::vector<std::int64_t> out_dims(y_ndim - 1);
    if (trans_y) {
      std::copy_n(y_dims.cbegin(), y_ndim - 1, out_dims.begin());
    } else {
      std::copy_n(y_dims.cbegin(), y_ndim - 2, out_dims.begin());
      out_dims.back() = y_dims.back();
    }
    Out->ResizeAndAllocate(make_ddim(out_dims));
    dev_ctx.template Alloc<T>(Out);
    if (trans_y) {
      const int64_t M = Y.numel() / N;
      VLOG(3) << "MatMul with blaslt 2";
      blaslt::Run(dev_ctx,
                  y_data,
                  x_data,
                  dev_ctx.template Alloc<T>(Out),
                  M,
                  1,
                  N,
                  false,
                  false,
                  matmul_planner);
    } else {
      const int64_t M = y_dims[y_ndim - 1];
      const int64_t batch_size = Y.numel() / (M * N);
      if (batch_size == 1) {
        VLOG(3) << "MatMul with blaslt 3";
        blaslt::Run(dev_ctx,
                    y_data,
                    x_data,
                    dev_ctx.template Alloc<T>(Out),
                    M,
                    1,
                    N,
                    true,
                    false,
                    matmul_planner);
      } else {
        VLOG(3) << "MatMul with blaslt 4";
        PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size");
        blaslt::RunWithBatch(dev_ctx,
                             y_data,
                             x_data,
                             dev_ctx.template Alloc<T>(Out),
                             M,
                             1,
                             N,
                             true,
                             false,
                             static_cast<int>(batch_size),
                             M * N,
                             0,
                             M,
                             matmul_planner);
      }
    }
    return;
  }

  if (y_ndim == 1) {
    const int64_t N = Y.numel();
    if (trans_x) {
      PADDLE_ENFORCE_EQ(
          x_dims[x_ndim - 2],
          N,
          common::errors::InvalidArgument("Input(X) has error dim."
                                          "X'dims[%d] must be equal to %d. "
                                          "But received X'dims[%d] is %d",
                                          x_ndim - 2,
                                          N,
                                          x_ndim - 2,
                                          x_dims[x_ndim - 2]));
    } else {
      PADDLE_ENFORCE_EQ(
          x_dims[x_ndim - 1],
          N,
          common::errors::InvalidArgument("Input(X) has error dim."
                                          "X'dims[%d] must be equal to %d, "
                                          "but received X'dims[%d] is %d",
                                          x_ndim - 1,
                                          N,
                                          x_ndim - 1,
                                          x_dims[x_ndim - 1]));
    }
    std::vector<std::int64_t> out_dims(x_ndim - 1);
    if (trans_x) {
      std::copy_n(x_dims.cbegin(), x_ndim - 2, out_dims.begin());
      out_dims.back() = x_dims.back();
    } else {
      std::copy_n(x_dims.cbegin(), x_ndim - 1, out_dims.begin());
    }
    Out->ResizeAndAllocate(make_ddim(out_dims));
    dev_ctx.template Alloc<T>(Out);

    if (trans_x) {
      const int64_t M = x_dims[x_ndim - 1];
      const int64_t batch_size = X.numel() / (M * N);
      if (batch_size == 1) {
        VLOG(3) << "MatMul with blaslt 5";
        blaslt::Run(dev_ctx,
                    x_data,
                    y_data,
                    dev_ctx.template Alloc<T>(Out),
                    M,
                    1,
                    N,
                    true,
                    false,
                    matmul_planner);
      } else {
        VLOG(3) << "MatMul with blaslt 6";
        PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size");
        blaslt::RunWithBatch(dev_ctx,
                             x_data,
                             y_data,
                             dev_ctx.template Alloc<T>(Out),
                             M,
                             1,
                             N,
                             true,
                             false,
                             static_cast<int>(batch_size),
                             M * N,
                             0,
                             M,
                             matmul_planner);
      }
    } else {
      const int64_t M = X.numel() / N;
      VLOG(3) << "MatMul with blaslt 7";
      blaslt::Run(dev_ctx,
                  x_data,
                  y_data,
                  dev_ctx.template Alloc<T>(Out),
                  M,
                  1,
                  N,
                  false,
                  false,
                  matmul_planner);
    }
    return;
  }

  const int64_t M = trans_x ? x_dims[x_ndim - 1] : x_dims[x_ndim - 2];
  const int64_t K = trans_x ? x_dims[x_ndim - 2] : x_dims[x_ndim - 1];
  if (trans_y) {
    PADDLE_ENFORCE_EQ(
        y_dims[y_ndim - 1],
        K,
        common::errors::InvalidArgument("Input(Y) has error dim. "
                                        "Y'dims[%d] must be equal to %d, "
                                        "but received Y'dims[%d] is %d.",
                                        y_ndim - 1,
                                        K,
                                        y_ndim - 1,
                                        y_dims[y_ndim - 1]));
  } else {
    PADDLE_ENFORCE_EQ(
        y_dims[y_ndim - 2],
        K,
        common::errors::InvalidArgument("Input(Y) has error dim. "
                                        "Y'dims[%d] must be equal to %d, "
                                        "but received Y'dims[%d] is %d.",
                                        y_ndim - 2,
                                        K,
                                        y_ndim - 2,
                                        y_dims[y_ndim - 2]));
  }
  const int64_t N = trans_y ? y_dims[y_ndim - 2] : y_dims[y_ndim - 1];
  const int ndim = (std::max)(x_ndim, y_ndim);
  std::vector<std::int64_t> x_broadcast_dims(ndim);
  std::vector<std::int64_t> y_broadcast_dims(ndim);
  std::vector<std::int64_t> out_broadcast_dims(ndim);

  GetBroadcastFromDims(x_ndim - 2,
                       x_dims.data(),
                       y_ndim - 2,
                       y_dims.data(),
                       x_broadcast_dims.data(),
                       y_broadcast_dims.data(),
                       out_broadcast_dims.data());
  out_broadcast_dims[ndim - 2] = M;
  out_broadcast_dims[ndim - 1] = N;

  Out->ResizeAndAllocate(make_ddim(out_broadcast_dims));
  dev_ctx.template Alloc<T>(Out);

  const int batch_dim = ndim - 2;
  // broadcast message
  const bool is_broadcast_dims =
      !std::equal(x_broadcast_dims.cbegin(),
                  x_broadcast_dims.cbegin() + batch_dim,
                  y_broadcast_dims.cbegin());

  const std::int64_t x_batch_size =
      std::accumulate(x_broadcast_dims.cbegin(),
                      x_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  const std::int64_t y_batch_size =
      std::accumulate(y_broadcast_dims.cbegin(),
                      y_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  const std::int64_t out_batch_size =
      std::accumulate(out_broadcast_dims.cbegin(),
                      out_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  if (out_batch_size == 0) return;
  PADDLE_ENFORCE_LE_INT_MAX(out_batch_size, "out_batch_size");
  if (x_batch_size == 1 && y_batch_size == 1) {
    VLOG(3) << "MatMul with blaslt 8";
    blaslt::Run(dev_ctx,
                x_data,
                y_data,
                dev_ctx.template Alloc<T>(Out),
                M,
                N,
                K,
                trans_x,
                trans_y,
                matmul_planner);
  } else if (x_batch_size == 1) {
    if (M == 1 && trans_y) {
      VLOG(3) << "MatMul with blaslt 9";
      PADDLE_ENFORCE_LE_INT_MAX(y_batch_size * N, "GEMV M");
      blaslt::Run(dev_ctx,
                  y_data,
                  x_data,
                  dev_ctx.template Alloc<T>(Out),
                  y_batch_size * N,
                  1,
                  K,
                  false,
                  false,
                  matmul_planner);
    } else {
      VLOG(3) << "MatMul with blaslt 10";
      blaslt::RunWithBatch(dev_ctx,
                           x_data,
                           y_data,
                           dev_ctx.template Alloc<T>(Out),
                           M,
                           N,
                           K,
                           trans_x,
                           trans_y,
                           static_cast<int>(out_batch_size),
                           0,
                           K * N,
                           M * N,
                           matmul_planner);
    }
  } else if (y_batch_size == 1) {
    if (!trans_x) {
      VLOG(3) << "MatMul with blaslt 11";
      PADDLE_ENFORCE_LE_INT_MAX(x_batch_size * M, "GEMM M");
      blaslt::Run(dev_ctx,
                  x_data,
                  y_data,
                  dev_ctx.template Alloc<T>(Out),
                  x_batch_size * M,
                  N,
                  K,
                  false,
                  trans_y,
                  matmul_planner);
    } else {
      VLOG(3) << "MatMul with blaslt 12";
      blaslt::RunWithBatch(dev_ctx,
                           x_data,
                           y_data,
                           dev_ctx.template Alloc<T>(Out),
                           M,
                           N,
                           K,
                           true,
                           trans_y,
                           static_cast<int>(out_batch_size),
                           M * K,
                           0,
                           M * N,
                           matmul_planner);
    }
  } else if (!is_broadcast_dims) {
    VLOG(3) << "MatMul with blaslt 13";
    blaslt::RunWithBatch(dev_ctx,
                         x_data,
                         y_data,
                         dev_ctx.template Alloc<T>(Out),
                         M,
                         N,
                         K,
                         trans_x,
                         trans_y,
                         static_cast<int>(out_batch_size),
                         M * K,
                         K * N,
                         M * N,
                         matmul_planner);
  } else {
    // in the case, can't use stridedgemm
    std::vector<const T*> x_ptr(out_batch_size);
    std::vector<const T*> y_ptr(out_batch_size);
    std::vector<T*> out_ptr(out_batch_size);
    std::vector<std::int64_t> index(batch_dim, 0);
    for (std::int64_t i = 0; i < out_batch_size; ++i) {
      // using the index to get offset
      const std::int64_t x_index =
          GetIndexMessage(batch_dim, x_broadcast_dims.data(), index.data());
      const std::int64_t y_index =
          GetIndexMessage(batch_dim, y_broadcast_dims.data(), index.data());

      x_ptr[i] = x_data + x_index * M * K;
      y_ptr[i] = y_data + y_index * K * N;
      out_ptr[i] = dev_ctx.template Alloc<T>(Out) + i * M * N;
      IndexIncreaseFromDims(batch_dim, out_broadcast_dims.data(), index.data());
    }
    VLOG(3) << "MatMul with blaslt 14";
    blaslt::RunWithBatch(dev_ctx,
                         x_ptr.data(),
                         y_ptr.data(),
                         out_ptr.data(),
                         M,
                         N,
                         K,
                         trans_x,
                         trans_y,
                         static_cast<int>(out_batch_size),
                         matmul_planner);
  }
}
#endif

template <typename Context, typename T>
struct MatMulDispatcher {
  void operator()(const Context& dev_ctx,
                  const DenseTensor& x,
                  const DenseTensor& y,
                  const std::vector<std::int64_t>& x_dims,
                  const std::vector<std::int64_t>& y_dims,
                  DenseTensor* out,
                  bool trans_x,
                  bool trans_y,
                  bool flag = false) {
    MatMulFunctionImplWithBlas<Context, T>(
        dev_ctx, x, y, x_dims, y_dims, out, trans_x, trans_y, flag);
  }
};

#ifdef PADDLE_WITH_CUDA
template <typename T>
struct MatMulDispatcher<GPUContext, T> {
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor& x,
                  const DenseTensor& y,
                  const std::vector<std::int64_t>& x_dims,
                  const std::vector<std::int64_t>& y_dims,
                  DenseTensor* out,
                  bool trans_x,
                  bool trans_y,
                  bool flag = false) {
#if CUDA_VERSION >= 11060
    auto* tuner =
        autotune::MakeMatmulTuner<T>(MatMulFunctionImplWithBlas<GPUContext, T>);
    tuner->AddCallBack(MatMulFunctionImplWithCublasLt<GPUContext, T>);
    funcs::MatmulPlanner matmul_planner(x_dims,
                                        y_dims,
                                        trans_x,
                                        trans_y,
                                        CppTypeToDataType<T>::Type(),
                                        funcs::MatmulFusedType::kMatmul,
                                        /* bias_data */ nullptr,
                                        /* reserve_data */ nullptr,
                                        /* use_addto */ flag,
                                        /* no_exchange */ true);
    tuner->Run(dev_ctx,
               matmul_planner.GetKey(),
               dev_ctx,
               x,
               y,
               x_dims,
               y_dims,
               out,
               trans_x,
               trans_y,
               flag,
               &matmul_planner);
#else
    MatMulFunctionImplWithBlas<GPUContext, T>(
        dev_ctx, x, y, x_dims, y_dims, out, trans_x, trans_y, flag);
#endif
  }
};

#endif  // PADDLE_WITH_CUDA

template <typename Context, typename T>
void MatMulFunction(const Context& dev_ctx,
                    const DenseTensor& x,
                    const DenseTensor& y,
                    const std::vector<std::int64_t>& x_dims,
                    const std::vector<std::int64_t>& y_dims,
                    DenseTensor* out,
                    bool trans_x,
                    bool trans_y,
                    bool flag = false) {
  MatMulDispatcher<Context, T>()(
      dev_ctx, x, y, x_dims, y_dims, out, trans_x, trans_y, flag);
}

template <typename Context>
bool MatMulInt8Function(const Context& dev_ctx,
                        const DenseTensor& x,
                        const DenseTensor& y,
                        const std::vector<std::int64_t>& x_dims,
                        const std::vector<std::int64_t>& y_dims,
                        DenseTensor* out,
                        bool trans_x,
                        bool trans_y) {
  return false;
}

#ifdef PADDLE_WITH_CUDA
template <>
bool inline MatMulInt8Function(const GPUContext& dev_ctx,
                               const DenseTensor& x,
                               const DenseTensor& y,
                               const std::vector<std::int64_t>& x_dims,
                               const std::vector<std::int64_t>& y_dims,
                               DenseTensor* out,
                               bool trans_x,
                               bool trans_y) {
  if (x.dtype() != DataType::INT8 || y.dtype() != DataType::INT8) {
    return false;
  }
#if CUDA_VERSION >= 11060
  const int64_t x_ndim = x_dims.size();
  const int64_t y_ndim = y_dims.size();
  const int8_t* x_data = x.data<int8_t>();
  const int8_t* y_data = y.data<int8_t>();
  using blaslt = funcs::MatmulWithCublasLt<int8_t, int32_t>;

  funcs::MatmulPlanner matmul_planner(x_dims,
                                      y_dims,
                                      trans_x,
                                      trans_y,
                                      CppTypeToDataType<int8_t>::Type(),
                                      funcs::MatmulFusedType::kMatmul,
                                      /* bias_data */ nullptr,
                                      /* reserve_data */ nullptr,
                                      /* use_addto */ false,
                                      /* no_exchange */ true);

  if (x_ndim == 1 && y_ndim == 1) {
    const int64_t M = x.numel();
    const int64_t N = y.numel();
    PADDLE_ENFORCE_EQ(
        M,
        N,
        common::errors::InvalidArgument(
            "X's numbers must be equal to Y's numbers, "
            "when X/Y's dims =1. But received X has [%d] elements, "
            "received Y has [%d] elements",
            M,
            N));
    if (!(M % 4 == 0)) {
      return false;
    }

    out->Resize({});
    dev_ctx.template Alloc<int32_t>(out);
    blaslt::Run(dev_ctx,
                y_data,
                x_data,
                dev_ctx.template Alloc<int32_t>(out),
                1,
                1,
                M,
                false,
                true,
                &matmul_planner);
    return true;
  }
  if (x_ndim == 1) {
    const int64_t N = x.numel();
    if (trans_y) {
      PADDLE_ENFORCE_EQ(
          y_dims[y_ndim - 1],
          N,
          common::errors::InvalidArgument("Input(Y) has error dim. "
                                          "Y'dims[%d] must be equal to %d, "
                                          "but received Y'dims[%d] is %d.",
                                          y_ndim - 1,
                                          N,
                                          y_ndim - 1,
                                          y_dims[y_ndim - 1]));
      if (!(N % 4 == 0)) {
        return false;
      }
    } else {
      PADDLE_ENFORCE_EQ(
          y_dims[y_ndim - 2],
          N,
          common::errors::InvalidArgument("Input(Y) has error dim. "
                                          "Y'dims[%d] must be equal to %d, "
                                          "but received Y'dims[%d] is %d.",
                                          y_ndim - 2,
                                          N,
                                          y_ndim - 2,
                                          y_dims[y_ndim - 2]));
      const int64_t M = y.numel() / N;
      if (!(M == 1 || M % 4 == 0)) {
        return false;
      }
    }
    std::vector<std::int64_t> out_dims(y_ndim - 1);
    if (trans_y) {
      std::copy_n(y_dims.cbegin(), y_ndim - 1, out_dims.begin());
    } else {
      std::copy_n(y_dims.cbegin(), y_ndim - 2, out_dims.begin());
      out_dims.back() = y_dims.back();
    }
    out->ResizeAndAllocate(make_ddim(out_dims));
    dev_ctx.template Alloc<int32_t>(out);
    if (trans_y) {
      const int64_t M = y.numel() / N;
      blaslt::Run(dev_ctx,
                  y_data,
                  x_data,
                  dev_ctx.template Alloc<int32_t>(out),
                  M,
                  1,
                  N,
                  false,
                  false,
                  &matmul_planner);
    } else {
      const int64_t M = y_dims[y_ndim - 1];
      const int64_t batch_size = y.numel() / (M * N);
      if (batch_size == 1) {
        blaslt::Run(dev_ctx,
                    y_data,
                    x_data,
                    dev_ctx.template Alloc<int32_t>(out),
                    M,
                    1,
                    N,
                    true,
                    false,
                    &matmul_planner);
      } else {
        PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size");
        blaslt::RunWithBatch(dev_ctx,
                             y_data,
                             x_data,
                             dev_ctx.template Alloc<int32_t>(out),
                             M,
                             1,
                             N,
                             true,
                             false,
                             static_cast<int>(batch_size),
                             M * N,
                             0,
                             M,
                             &matmul_planner);
      }
    }
    return true;
  }

  if (y_ndim == 1) {
    const int64_t N = y.numel();
    if (trans_x) {
      PADDLE_ENFORCE_EQ(
          x_dims[x_ndim - 2],
          N,
          common::errors::InvalidArgument("Input(X) has error dim."
                                          "X'dims[%d] must be equal to %d, "
                                          "but received X'dims[%d] is %d",
                                          x_ndim - 2,
                                          N,
                                          x_ndim - 2,
                                          x_dims[x_ndim - 2]));
      const int64_t M = x.numel() / N;
      if (!((M == 1 || M % 4 == 0))) {
        return false;
      }
    } else {
      PADDLE_ENFORCE_EQ(
          x_dims[x_ndim - 1],
          N,
          common::errors::InvalidArgument("Input(X) has error dim."
                                          "X'dims[%d] must be equal to %d, "
                                          "but received X'dims[%d] is %d",
                                          x_ndim - 1,
                                          N,
                                          x_ndim - 1,
                                          x_dims[x_ndim - 1]));
      if (N % 4 != 0) {
        return false;
      }
    }
    std::vector<std::int64_t> out_dims(x_ndim - 1);
    if (trans_x) {
      std::copy_n(x_dims.cbegin(), x_ndim - 2, out_dims.begin());
      out_dims.back() = x_dims.back();
    } else {
      std::copy_n(x_dims.cbegin(), x_ndim - 1, out_dims.begin());
    }
    out->ResizeAndAllocate(make_ddim(out_dims));
    dev_ctx.template Alloc<int32_t>(out);

    if (trans_x) {
      const int64_t M = x_dims[x_ndim - 1];
      const int64_t batch_size = x.numel() / (M * N);
      if (batch_size == 1) {
        blaslt::Run(dev_ctx,
                    x_data,
                    y_data,
                    dev_ctx.template Alloc<int32_t>(out),
                    M,
                    1,
                    N,
                    true,
                    false,
                    &matmul_planner);
      } else {
        PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size");
        blaslt::RunWithBatch(dev_ctx,
                             x_data,
                             y_data,
                             dev_ctx.template Alloc<int32_t>(out),
                             M,
                             1,
                             N,
                             true,
                             false,
                             static_cast<int>(batch_size),
                             M * N,
                             0,
                             M,
                             &matmul_planner);
      }
    } else {
      const int64_t M = x.numel() / N;
      blaslt::Run(dev_ctx,
                  x_data,
                  y_data,
                  dev_ctx.template Alloc<int32_t>(out),
                  M,
                  1,
                  N,
                  false,
                  false,
                  &matmul_planner);
    }
    return true;
  }

  const int64_t M = trans_x ? x_dims[x_ndim - 1] : x_dims[x_ndim - 2];
  const int64_t K = trans_x ? x_dims[x_ndim - 2] : x_dims[x_ndim - 1];
  if (trans_y) {
    PADDLE_ENFORCE_EQ(
        y_dims[y_ndim - 1],
        K,
        common::errors::InvalidArgument("Input(Y) has error dim. "
                                        "Y'dims[%d] must be equal to %d, "
                                        "but received Y'dims[%d] is %d.",
                                        y_ndim - 1,
                                        K,
                                        y_ndim - 1,
                                        y_dims[y_ndim - 1]));
  } else {
    PADDLE_ENFORCE_EQ(
        y_dims[y_ndim - 2],
        K,
        common::errors::InvalidArgument("Input(Y) has error dim. "
                                        "Y'dims[%d] must be equal to %d, "
                                        "but received Y'dims[%d] is %d.",
                                        y_ndim - 2,
                                        K,
                                        y_ndim - 2,
                                        y_dims[y_ndim - 2]));
  }
  const int64_t N = trans_y ? y_dims[y_ndim - 2] : y_dims[y_ndim - 1];
  const int ndim = (std::max)(x_ndim, y_ndim);
  std::vector<std::int64_t> x_broadcast_dims(ndim);
  std::vector<std::int64_t> y_broadcast_dims(ndim);
  std::vector<std::int64_t> out_broadcast_dims(ndim);
  GetBroadcastFromDims(x_ndim - 2,
                       x_dims.data(),
                       y_ndim - 2,
                       y_dims.data(),
                       x_broadcast_dims.data(),
                       y_broadcast_dims.data(),
                       out_broadcast_dims.data());
  out_broadcast_dims[ndim - 2] = M;
  out_broadcast_dims[ndim - 1] = N;

  out->ResizeAndAllocate(make_ddim(out_broadcast_dims));
  dev_ctx.template Alloc<int32_t>(out);

  const int batch_dim = ndim - 2;
  // broadcast message
  const bool is_broadcast_dims =
      !std::equal(x_broadcast_dims.cbegin(),
                  x_broadcast_dims.cbegin() + batch_dim,
                  y_broadcast_dims.cbegin());

  const std::int64_t x_batch_size =
      std::accumulate(x_broadcast_dims.cbegin(),
                      x_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  const std::int64_t y_batch_size =
      std::accumulate(y_broadcast_dims.cbegin(),
                      y_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  const std::int64_t out_batch_size =
      std::accumulate(out_broadcast_dims.cbegin(),
                      out_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  if (out_batch_size == 0) return true;

  PADDLE_ENFORCE_LE_INT_MAX(out_batch_size, "out_batch_size");
  if (x_batch_size == 1 && M == 1 && trans_y) {
    if (!(K % 4 == 0)) {
      return false;
    }
  } else if (!trans_x && !trans_y) {
    if (!(N % 4 == 0 || N == 1) || !(K % 4 == 0) || (M == 1 && N == 1)) {
      return false;
    }
  } else if (!trans_x && trans_y) {
    if (!(K % 4 == 0)) {
      return false;
    }
  } else if (trans_x && !trans_y) {
    if (!(M % 4 == 0 || M == 1) || !(N % 4 == 0 || N == 1)) {
      return false;
    }
  } else {
    if (!(M % 4 == 0 || M == 1) || !(K % 4 == 0)) {
      return false;
    }
  }
  if (x_batch_size == 1 && y_batch_size == 1) {
    blaslt::Run(dev_ctx,
                x_data,
                y_data,
                dev_ctx.template Alloc<int32_t>(out),
                M,
                N,
                K,
                trans_x,
                trans_y,
                &matmul_planner);
  } else if (x_batch_size == 1) {
    if (M == 1 && trans_y) {
      blaslt::Run(dev_ctx,
                  y_data,
                  x_data,
                  dev_ctx.template Alloc<int32_t>(out),
                  y_batch_size * N,
                  1,
                  K,
                  false,
                  false,
                  &matmul_planner);
    } else {
      blaslt::RunWithBatch(dev_ctx,
                           x_data,
                           y_data,
                           dev_ctx.template Alloc<int32_t>(out),
                           M,
                           N,
                           K,
                           trans_x,
                           trans_y,
                           static_cast<int>(out_batch_size),
                           0,
                           K * N,
                           M * N,
                           &matmul_planner);
    }
  } else if (y_batch_size == 1) {
    if (!trans_x) {
      blaslt::Run(dev_ctx,
                  x_data,
                  y_data,
                  dev_ctx.template Alloc<int32_t>(out),
                  x_batch_size * M,
                  N,
                  K,
                  false,
                  trans_y,
                  &matmul_planner);
    } else {
      blaslt::RunWithBatch(dev_ctx,
                           x_data,
                           y_data,
                           dev_ctx.template Alloc<int32_t>(out),
                           M,
                           N,
                           K,
                           true,
                           trans_y,
                           static_cast<int>(out_batch_size),
                           M * K,
                           0,
                           M * N,
                           &matmul_planner);
    }
  } else if (!is_broadcast_dims) {
    blaslt::RunWithBatch(dev_ctx,
                         x_data,
                         y_data,
                         dev_ctx.template Alloc<int32_t>(out),
                         M,
                         N,
                         K,
                         trans_x,
                         trans_y,
                         static_cast<int>(out_batch_size),
                         M * K,
                         K * N,
                         M * N,
                         &matmul_planner);
  } else {
    // in the case, can't use stridedgemm
    std::vector<const int8_t*> x_ptr(out_batch_size);
    std::vector<const int8_t*> y_ptr(out_batch_size);
    std::vector<int32_t*> out_ptr(out_batch_size);
    std::vector<std::int64_t> index(batch_dim, 0);
    for (std::int64_t i = 0; i < out_batch_size; ++i) {
      // using the index to get offset
      const std::int64_t x_index =
          GetIndexMessage(batch_dim, x_broadcast_dims.data(), index.data());
      const std::int64_t y_index =
          GetIndexMessage(batch_dim, y_broadcast_dims.data(), index.data());

      x_ptr[i] = x_data + x_index * M * K;
      y_ptr[i] = y_data + y_index * K * N;
      out_ptr[i] = dev_ctx.template Alloc<int32_t>(out) + i * M * N;
      IndexIncreaseFromDims(batch_dim, out_broadcast_dims.data(), index.data());
    }
    blaslt::RunWithBatch(dev_ctx,
                         x_ptr.data(),
                         y_ptr.data(),
                         out_ptr.data(),
                         M,
                         N,
                         K,
                         trans_x,
                         trans_y,
                         static_cast<int>(out_batch_size),
                         &matmul_planner);
  }
  return true;
#else
  return false;
#endif
}
#endif

#ifdef PADDLE_WITH_HIP
template <>
bool inline MatMulInt8Function(const GPUContext& dev_ctx,
                               const DenseTensor& x,
                               const DenseTensor& y,
                               const std::vector<std::int64_t>& x_dims,
                               const std::vector<std::int64_t>& y_dims,
                               DenseTensor* out,
                               bool trans_x,
                               bool trans_y) {
  if (x.dtype() != DataType::INT8 || y.dtype() != DataType::INT8) {
    return false;
  }
  const int x_ndim = x_dims.size();
  const int y_ndim = y_dims.size();

  // Get data ptr
  const int8_t* x_data = x.data<int8_t>();
  const int8_t* y_data = y.data<int8_t>();

  if (x_ndim == 1 && y_ndim == 1) {
    const int64_t M = x.numel();
    const int64_t N = y.numel();
    PADDLE_ENFORCE_EQ(
        M,
        N,
        common::errors::InvalidArgument(
            "X's numbers must be equal to Y's numbers, "
            "when X/Y's dims =1. But received X has [%d] elements, "
            "received Y has [%d] elements.",
            M,
            N));
    PADDLE_ENFORCE_LE_INT_MAX(M, "M");
    VLOG(3) << "MatMul's case 1";
    out->Resize({});
    dev_ctx.template Alloc<int32_t>(out);
    funcs::Int8GEMM(dev_ctx,
                    CblasNoTrans,
                    CblasTrans,
                    1,
                    1,
                    static_cast<int>(M),
                    static_cast<int32_t>(1),
                    y_data,
                    x_data,
                    static_cast<int32_t>(0),
                    dev_ctx.template Alloc<int32_t>(out));
    return true;
  }

  if (x_ndim == 1) {
    const int64_t N = x.numel();
    if (trans_y) {
      PADDLE_ENFORCE_EQ(
          y_dims[y_ndim - 1],
          N,
          common::errors::InvalidArgument("Input(Y) has error dim. "
                                          "Y'dims[%d] must be equal to %d, "
                                          "but received Y'dims[%d] is %d.",
                                          y_ndim - 1,
                                          N,
                                          y_ndim - 1,
                                          y_dims[y_ndim - 1]));
    } else {
      PADDLE_ENFORCE_EQ(
          y_dims[y_ndim - 2],
          N,
          common::errors::InvalidArgument("Input(Y) has error dim. "
                                          "Y'dims[%d] must be equal to %d, "
                                          "but received Y'dims[%d] is %d.",
                                          y_ndim - 2,
                                          N,
                                          y_ndim - 2,
                                          y_dims[y_ndim - 2]));
    }
    PADDLE_ENFORCE_LE_INT_MAX(N, "N");
    std::vector<std::int64_t> out_dims(y_ndim - 1);
    if (trans_y) {
      std::copy_n(y_dims.cbegin(), y_ndim - 1, out_dims.begin());
    } else {
      std::copy_n(y_dims.cbegin(), y_ndim - 2, out_dims.begin());
      out_dims.back() = y_dims.back();
    }
    out->ResizeAndAllocate(make_ddim(out_dims));
    dev_ctx.template Alloc<int32_t>(out);
    if (trans_y) {
      const int64_t M = y.numel() / N;
      PADDLE_ENFORCE_LE_INT_MAX(M, "M");
      VLOG(3) << "MatMul's case 2";
      funcs::Int8GEMV(dev_ctx,
                      false,
                      static_cast<int>(M),
                      static_cast<int>(N),
                      static_cast<int32_t>(1),
                      y_data,
                      x_data,
                      static_cast<int32_t>(0),
                      dev_ctx.template Alloc<int32_t>(out));
    } else {
      const int64_t M = y_dims[y_ndim - 1];
      PADDLE_ENFORCE_LE_INT_MAX(M, "M");
      const int64_t batch_size = y.numel() / (M * N);
      if (batch_size == 1) {
        VLOG(3) << "MatMul's case 3";
        funcs::Int8GEMV(dev_ctx,
                        true,
                        static_cast<int>(N),
                        static_cast<int>(M),
                        static_cast<int32_t>(1),
                        y_data,
                        x_data,
                        static_cast<int32_t>(0),
                        dev_ctx.template Alloc<int32_t>(out));
      } else {
        VLOG(3) << "MatMul's case 4";
        PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size");
        funcs::Int8BatchedGEMM(dev_ctx,
                               CblasTrans,
                               CblasNoTrans,
                               static_cast<int>(M),
                               1,
                               static_cast<int>(N),
                               static_cast<int32_t>(1),
                               y_data,
                               x_data,
                               static_cast<int32_t>(0),
                               dev_ctx.template Alloc<int32_t>(out),
                               static_cast<int>(batch_size),
                               M * N,
                               0);
      }
    }
    return true;
  }

  if (y_ndim == 1) {
    const int64_t N = y.numel();
    if (trans_x) {
      PADDLE_ENFORCE_EQ(
          x_dims[x_ndim - 2],
          N,
          common::errors::InvalidArgument("Input(X) has error dim. "
                                          "X'dims[%d] must be equal to %d, "
                                          "But received X'dims[%d] is %d",
                                          x_ndim - 2,
                                          N,
                                          x_ndim - 2,
                                          x_dims[x_ndim - 2]));
    } else {
      PADDLE_ENFORCE_EQ(
          x_dims[x_ndim - 1],
          N,
          common::errors::InvalidArgument("Input(X) has error dim. "
                                          "X'dims[%d] must be equal to %d, "
                                          "But received X'dims[%d] is %d",
                                          x_ndim - 1,
                                          N,
                                          x_ndim - 1,
                                          x_dims[x_ndim - 1]));
    }
    PADDLE_ENFORCE_LE_INT_MAX(N, "N");
    std::vector<std::int64_t> out_dims(x_ndim - 1);
    if (trans_x) {
      std::copy_n(x_dims.cbegin(), x_ndim - 2, out_dims.begin());
      out_dims.back() = x_dims.back();
    } else {
      std::copy_n(x_dims.cbegin(), x_ndim - 1, out_dims.begin());
    }
    out->ResizeAndAllocate(make_ddim(out_dims));
    dev_ctx.template Alloc<int32_t>(out);

    if (trans_x) {
      const int64_t M = x_dims[x_ndim - 1];
      PADDLE_ENFORCE_LE_INT_MAX(M, "M");
      const int64_t batch_size = x.numel() / (M * N);
      if (batch_size == 1) {
        VLOG(3) << "MatMul's case 5";
        funcs::Int8GEMV(dev_ctx,
                        true,
                        static_cast<int>(N),
                        static_cast<int>(M),
                        static_cast<int32_t>(1),
                        x_data,
                        y_data,
                        static_cast<int32_t>(0),
                        dev_ctx.template Alloc<int32_t>(out));
      } else {
        VLOG(3) << "MatMul's case 6";
        PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size");
        funcs::Int8BatchedGEMM(dev_ctx,
                               CblasTrans,
                               CblasNoTrans,
                               static_cast<int>(M),
                               1,
                               static_cast<int>(N),
                               static_cast<int32_t>(1),
                               x_data,
                               y_data,
                               static_cast<int32_t>(0),
                               dev_ctx.template Alloc<int32_t>(out),
                               static_cast<int>(batch_size),
                               M * N,
                               0);
      }
    } else {
      const int64_t M = x.numel() / N;
      PADDLE_ENFORCE_LE_INT_MAX(M, "M");
      VLOG(3) << "MatMul's case 7";
      funcs::Int8GEMV(dev_ctx,
                      false,
                      static_cast<int>(M),
                      static_cast<int>(N),
                      static_cast<int32_t>(1),
                      x_data,
                      y_data,
                      static_cast<int32_t>(0),
                      dev_ctx.template Alloc<int32_t>(out));
    }
    return true;
  }

  const int64_t M = trans_x ? x_dims[x_ndim - 1] : x_dims[x_ndim - 2];
  const int64_t K = trans_x ? x_dims[x_ndim - 2] : x_dims[x_ndim - 1];
  if (trans_y) {
    PADDLE_ENFORCE_EQ(
        y_dims[y_ndim - 1],
        K,
        common::errors::InvalidArgument("Input(Y) has error dim. "
                                        "Y'dims[%d] must be equal to %d, "
                                        "but received Y'dims[%d] is %d.",
                                        y_ndim - 1,
                                        K,
                                        y_ndim - 1,
                                        y_dims[y_ndim - 1]));
  } else {
    PADDLE_ENFORCE_EQ(
        y_dims[y_ndim - 2],
        K,
        common::errors::InvalidArgument("Input(Y) has error dim. "
                                        "Y'dims[%d] must be equal to %d, "
                                        "but received Y'dims[%d] is %d.",
                                        y_ndim - 2,
                                        K,
                                        y_ndim - 2,
                                        y_dims[y_ndim - 2]));
  }
  const int64_t N = trans_y ? y_dims[y_ndim - 2] : y_dims[y_ndim - 1];
  PADDLE_ENFORCE_LE_INT_MAX(M, "M");
  PADDLE_ENFORCE_LE_INT_MAX(N, "N");
  PADDLE_ENFORCE_LE_INT_MAX(K, "K");
  const int ndim = (std::max)(x_ndim, y_ndim);
  std::vector<std::int64_t> x_broadcast_dims(ndim);
  std::vector<std::int64_t> y_broadcast_dims(ndim);
  std::vector<std::int64_t> out_broadcast_dims(ndim);

  GetBroadcastFromDims(x_ndim - 2,
                       x_dims.data(),
                       y_ndim - 2,
                       y_dims.data(),
                       x_broadcast_dims.data(),
                       y_broadcast_dims.data(),
                       out_broadcast_dims.data());
  out_broadcast_dims[ndim - 2] = M;
  out_broadcast_dims[ndim - 1] = N;

  out->ResizeAndAllocate(make_ddim(out_broadcast_dims));
  dev_ctx.template Alloc<int32_t>(out);

  const int batch_dim = ndim - 2;
  // broadcast message
  const bool is_broadcast_dims =
      !std::equal(x_broadcast_dims.cbegin(),
                  x_broadcast_dims.cbegin() + batch_dim,
                  y_broadcast_dims.cbegin());

  const std::int64_t x_batch_size =
      std::accumulate(x_broadcast_dims.cbegin(),
                      x_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  const std::int64_t y_batch_size =
      std::accumulate(y_broadcast_dims.cbegin(),
                      y_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  const std::int64_t out_batch_size =
      std::accumulate(out_broadcast_dims.cbegin(),
                      out_broadcast_dims.cbegin() + batch_dim,
                      1LL,
                      std::multiplies<std::int64_t>());
  if (out_batch_size == 0) return true;
  PADDLE_ENFORCE_LE_INT_MAX(out_batch_size, "out_batch_size");
  if (x_batch_size == 1 && y_batch_size == 1) {
    VLOG(3) << "MatMul's case 8";
    funcs::Int8GEMM(dev_ctx,
                    trans_x ? CblasTrans : CblasNoTrans,
                    trans_y ? CblasTrans : CblasNoTrans,
                    static_cast<int>(M),
                    static_cast<int>(N),
                    static_cast<int>(K),
                    static_cast<int32_t>(1),
                    x_data,
                    y_data,
                    static_cast<int32_t>(0),
                    dev_ctx.template Alloc<int32_t>(out));
  } else if (x_batch_size == 1) {
    if (M == 1 && trans_y) {
      VLOG(3) << "MatMul's case 9";
      PADDLE_ENFORCE_LE_INT_MAX(y_batch_size * N, "GEMV M");
      funcs::Int8GEMV(dev_ctx,
                      false,
                      static_cast<int>(y_batch_size * N),
                      static_cast<int>(K),
                      static_cast<int32_t>(1),
                      y_data,
                      x_data,
                      static_cast<int32_t>(0),
                      dev_ctx.template Alloc<int32_t>(out));
    } else {
      VLOG(3) << "MatMul's case 10";
      funcs::Int8BatchedGEMM(dev_ctx,
                             trans_x ? CblasTrans : CblasNoTrans,
                             trans_y ? CblasTrans : CblasNoTrans,
                             static_cast<int>(M),
                             static_cast<int>(N),
                             static_cast<int>(K),
                             static_cast<int32_t>(1),
                             x_data,
                             y_data,
                             static_cast<int32_t>(0),
                             dev_ctx.template Alloc<int32_t>(out),
                             static_cast<int>(out_batch_size),
                             0,
                             K * N);
    }
  } else if (y_batch_size == 1) {
    if (!trans_x) {
      VLOG(3) << "MatMul's case 11";
      PADDLE_ENFORCE_LE_INT_MAX(x_batch_size * M, "GEMM M");
      funcs::Int8GEMM(dev_ctx,
                      CblasNoTrans,
                      trans_y ? CblasTrans : CblasNoTrans,
                      static_cast<int>(x_batch_size * M),
                      static_cast<int>(N),
                      static_cast<int>(K),
                      static_cast<int32_t>(1),
                      x_data,
                      y_data,
                      static_cast<int32_t>(0),
                      dev_ctx.template Alloc<int32_t>(out));
    } else {
      VLOG(3) << "MatMul's case 12";
      funcs::Int8BatchedGEMM(dev_ctx,
                             CblasTrans,
                             trans_y ? CblasTrans : CblasNoTrans,
                             static_cast<int>(M),
                             static_cast<int>(N),
                             static_cast<int>(K),
                             static_cast<int32_t>(1),
                             x_data,
                             y_data,
                             static_cast<int32_t>(0),
                             dev_ctx.template Alloc<int32_t>(out),
                             static_cast<int>(out_batch_size),
                             M * K,
                             0);
    }
  } else if (!is_broadcast_dims) {
    VLOG(3) << "MatMul's case 13";
    funcs::Int8BatchedGEMM(dev_ctx,
                           trans_x ? CblasTrans : CblasNoTrans,
                           trans_y ? CblasTrans : CblasNoTrans,
                           static_cast<int>(M),
                           static_cast<int>(N),
                           static_cast<int>(K),
                           static_cast<int32_t>(1),
                           x_data,
                           y_data,
                           static_cast<int32_t>(0),
                           dev_ctx.template Alloc<int32_t>(out),
                           static_cast<int>(out_batch_size),
                           M * K,
                           K * N);
  } else {
    // in the case, can't use stridedgemm
    std::vector<const int8_t*> x_ptr(out_batch_size);
    std::vector<const int8_t*> y_ptr(out_batch_size);
    std::vector<int32_t*> out_ptr(out_batch_size);
    std::vector<std::int64_t> index(batch_dim, 0);
    for (std::int64_t i = 0; i < out_batch_size; ++i) {
      // using the index to get offset
      const std::int64_t x_index =
          GetIndexMessage(batch_dim, x_broadcast_dims.data(), index.data());
      const std::int64_t y_index =
          GetIndexMessage(batch_dim, y_broadcast_dims.data(), index.data());

      x_ptr[i] = x_data + x_index * M * K;
      y_ptr[i] = y_data + y_index * K * N;
      out_ptr[i] = dev_ctx.template Alloc<int32_t>(out) + i * M * N;
      IndexIncreaseFromDims(batch_dim, out_broadcast_dims.data(), index.data());
    }
    VLOG(3) << "MatMul's case 14";
    funcs::Int8BatchedGEMM(dev_ctx,
                           trans_x ? CblasTrans : CblasNoTrans,
                           trans_y ? CblasTrans : CblasNoTrans,
                           static_cast<int>(M),
                           static_cast<int>(N),
                           static_cast<int>(K),
                           static_cast<int32_t>(1),
                           x_ptr.data(),
                           y_ptr.data(),
                           static_cast<int32_t>(0),
                           out_ptr.data(),
                           static_cast<int>(out_batch_size));
  }
  return true;
}
#endif

template <typename Context, typename T>
typename std::enable_if<std::is_integral<T>::value>::type
MatmulJudgeDtypeKernel(const Context& dev_ctx,
                       const DenseTensor& x,
                       const DenseTensor& y,
                       const std::vector<std::int64_t>& x_dims,
                       const std::vector<std::int64_t>& y_dims,
                       DenseTensor* out,
                       bool transpose_x,
                       bool transpose_y) {
#if defined(PADDLE_WITH_CUDA)
  if constexpr (std::is_same<Context, GPUContext>::value &&
                std::is_same<T, int8_t>::value) {
    if (x.dtype() == DataType::INT8 && x_dims[0] <= 4 && y_dims.size() == 2 &&
        y_dims[0] % 16 == 0 && y_dims[1] % 16 == 0 &&
        FLAGS_cuda_core_int8_gemm && dev_ctx.GetComputeCapability() >= 70 &&
        transpose_y) {
      CudaGemm<T, Context>(dev_ctx, x, y, out);
      return;
    }
  }
#endif
  bool try_matmul_int8 = MatMulInt8Function<Context>(
      dev_ctx, x, y, x_dims, y_dims, out, transpose_x, transpose_y);
  if (try_matmul_int8) {
    return;
  }
  auto x_tmp = Cast<T, Context>(dev_ctx, x, DataType::FLOAT32);
  auto y_tmp = Cast<T, Context>(dev_ctx, y, DataType::FLOAT32);
  DenseTensor out_tmp;
  MatMulFunction<Context, float>(dev_ctx,
                                 x_tmp,
                                 y_tmp,
                                 x_dims,
                                 y_dims,
                                 &out_tmp,
                                 transpose_x,
                                 transpose_y);
  if (x.dtype() == DataType::INT8) {
    CastKernel<float>(dev_ctx, out_tmp, DataType::INT32, out);
    return;
  }
  CastKernel<float>(dev_ctx, out_tmp, x.dtype(), out);
}

#if defined(PADDLE_WITH_CUDA)
#if CUDA_VERSION >= 12010
template <typename Context>
typename std::enable_if<std::is_same<Context, GPUContext>::value>::type
DispatchMatmulFP8Kernel(const Context& dev_ctx,
                        const DenseTensor& x,
                        const DenseTensor& y,
                        const std::vector<std::int64_t>& x_dims,
                        const std::vector<std::int64_t>& y_dims,
                        DenseTensor* out,
                        bool transpose_x,
                        bool transpose_y) {
  if (x.dtype() != DataType::FLOAT8_E4M3FN ||
      y.dtype() != DataType::FLOAT8_E4M3FN) {
    PADDLE_THROW(common::errors::InvalidArgument(
        "float8 matmul needs input x and y be float8_e4m3fn"));
  }

  PADDLE_ENFORCE_EQ(
      x_dims.size(),
      2,
      common::errors::InvalidArgument(
          "Mat x for matmul fp8 just support 2-dim tensor, but got %d",
          x_dims.size()));
  PADDLE_ENFORCE_EQ(
      y_dims.size(),
      2,
      common::errors::InvalidArgument(
          "Mat y for matmul fp8 just support 2-dim tensor, but got %d",
          y_dims.size()));
  PADDLE_ENFORCE_EQ(x_dims[1],
                    y_dims[0],
                    common::errors::InvalidArgument(
                        "x_dims[1] needs to equal to y_dims[0], but "
                        "got x_dims[1] = %d, y_dims[0] = %d",
                        x_dims[1],
                        y_dims[0]));

  PADDLE_ENFORCE_EQ(
      x_dims[1] % 16,
      0,
      common::errors::InvalidArgument(
          "fp8 matmul need x_dims[1] % 16 = 0, got x_dims[1] = %d", x_dims[1]));
  PADDLE_ENFORCE_EQ(
      y_dims[0] % 16,
      0,
      common::errors::InvalidArgument(
          "fp8 matmul need y_dims[0] % 16 = 0, got y_dims[0] = %d", y_dims[0]));

  DenseTensor workspace;
  workspace.Resize({30 * 1024 * 1024});
  dev_ctx.template Alloc<int8_t>(&workspace);
  dev_ctx.template Alloc<float16>(out);

  CublasLtMatmulFP8<float16>(dev_ctx, x, y, &workspace, out);
}

template <typename Context>
typename std::enable_if<std::is_same<Context, CPUContext>::value>::type
DispatchMatmulFP8Kernel(const Context& dev_ctx,
                        const DenseTensor& x,
                        const DenseTensor& y,
                        const std::vector<std::int64_t>& x_dims,
                        const std::vector<std::int64_t>& y_dims,
                        DenseTensor* out,
                        bool transpose_x,
                        bool transpose_y) {}

template <typename Context, typename T>
typename std::enable_if<std::is_same<T, float8_e4m3fn>::value>::type
DispatchMatmulKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const DenseTensor& y,
                     const std::vector<std::int64_t>& x_dims,
                     const std::vector<std::int64_t>& y_dims,
                     DenseTensor* out,
                     bool transpose_x,
                     bool transpose_y) {
  DispatchMatmulFP8Kernel<Context>(
      dev_ctx, x, y, x_dims, y_dims, out, transpose_x, transpose_y);
}
#endif
#endif

template <typename Context, typename T>
typename std::enable_if<!std::is_same<T, float8_e4m3fn>::value>::type
DispatchMatmulKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const DenseTensor& y,
                     const std::vector<std::int64_t>& x_dims,
                     const std::vector<std::int64_t>& y_dims,
                     DenseTensor* out,
                     bool transpose_x,
                     bool transpose_y) {
  MatMulFunction<Context, T>(
      dev_ctx, x, y, x_dims, y_dims, out, transpose_x, transpose_y);
}

template <typename Context, typename T>
typename std::enable_if<!std::is_integral<T>::value>::type
MatmulJudgeDtypeKernel(const Context& dev_ctx,
                       const DenseTensor& x,
                       const DenseTensor& y,
                       const std::vector<std::int64_t>& x_dims,
                       const std::vector<std::int64_t>& y_dims,
                       DenseTensor* out,
                       bool transpose_x,
                       bool transpose_y) {
  DispatchMatmulKernel<Context, T>(
      dev_ctx, x, y, x_dims, y_dims, out, transpose_x, transpose_y);
}

template <typename T, typename Context>
void MmOutDtypeKernel(const Context& dev_ctx,
                      const DenseTensor& x,
                      const DenseTensor& y,
                      DataType out_dtype,
                      DenseTensor* out) {
  PADDLE_ENFORCE_EQ(
      out_dtype,
      DataType::FLOAT32,
      common::errors::InvalidArgument(
          "The out_dtype of paddle.mm currently only supports float32."));
  PADDLE_ENFORCE_EQ(
      x.dtype(),
      DataType::BFLOAT16,
      common::errors::InvalidArgument(
          "The out_dtype of paddle.mm currently only supports bfloat16 "
          "Input(X)."));
  PADDLE_ENFORCE_EQ(
      y.dtype(),
      DataType::BFLOAT16,
      common::errors::InvalidArgument(
          "The out_dtype of paddle.mm currently only supports bfloat16 "
          "Input(Y)."));
  const std::vector<std::int64_t> x_dims = vectorize(x.dims());
  const std::vector<std::int64_t> y_dims = vectorize(y.dims());
  PADDLE_ENFORCE_EQ(
      x_dims.size(),
      2UL,
      common::errors::InvalidArgument(
          "The out_dtype of paddle.mm currently only supports 2-D Input(X)."));
  PADDLE_ENFORCE_EQ(
      y_dims.size(),
      2UL,
      common::errors::InvalidArgument(
          "The out_dtype of paddle.mm currently only supports 2-D Input(Y)."));
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP)
  if constexpr (std::is_same<Context, phi::GPUContext>::value &&
                std::is_same<T, phi::bfloat16>::value) {
    const int64_t M = x_dims[0];
    const int64_t K = x_dims[1];
    const int64_t N = y_dims[1];
    PADDLE_ENFORCE_EQ(
        K,
        y_dims[0],
        common::errors::InvalidArgument(
            "Input(X)'s width must equal Input(Y)'s height, but received %d "
            "and %d.",
            K,
            y_dims[0]));
    if (x.numel() == 0 || y.numel() == 0) {
      Full<float, Context>(dev_ctx, out->dims(), 0, out);
      return;
    }
    DenseTensor x_contiguous;
    DenseTensor y_contiguous;
    const DenseTensor* x_ptr = &x;
    const DenseTensor* y_ptr = &y;
    if (!x.meta().is_contiguous()) {
      ContiguousKernel<T, Context>(dev_ctx, x, &x_contiguous);
      x_ptr = &x_contiguous;
    }
    if (!y.meta().is_contiguous()) {
      ContiguousKernel<T, Context>(dev_ctx, y, &y_contiguous);
      y_ptr = &y_contiguous;
    }
    dev_ctx.template Alloc<float>(out);
    funcs::Blas<Context> blas(dev_ctx);
    blas.GEMM(CblasNoTrans,
              CblasNoTrans,
              M,
              N,
              K,
              1.0f,
              x_ptr->data<phi::bfloat16>(),
              y_ptr->data<phi::bfloat16>(),
              0.0f,
              out->data<float>());
  } else {
    PADDLE_THROW(common::errors::Unimplemented(
        "The out_dtype of paddle.mm currently only supports CUDA bfloat16 "
        "inputs."));
  }
#else
  PADDLE_THROW(common::errors::Unimplemented(
      "The out_dtype of paddle.mm currently only supports CUDA."));
#endif
}

template <typename T, typename Context>
void MatmulKernel(const Context& dev_ctx,
                  const DenseTensor& x,
                  const DenseTensor& y,
                  bool transpose_x,
                  bool transpose_y,
                  DenseTensor* out) {
  if (x.numel() == 0 || y.numel() == 0) {
    // input shape [1, 1, 5, 0], [1, 1, 0, 5], result shape is [1, 1, 5, 5]
    Full<T, Context>(dev_ctx, out->dims(), 0, out);
    return;
  }
  PADDLE_ENFORCE_GE(
      common::product(x.dims()),
      0,
      common::errors::InvalidArgument(
          "The dims of Input(X) should be greater than or equal to 0."));
  PADDLE_ENFORCE_GE(
      common::product(y.dims()),
      0,
      common::errors::InvalidArgument(
          "The dims of Input(Y) should be greater than or equal to 0."));
  const std::vector<std::int64_t> x_dims = vectorize(x.dims());
  const std::vector<std::int64_t> y_dims = vectorize(y.dims());
  MatmulJudgeDtypeKernel<Context, T>(
      dev_ctx, x, y, x_dims, y_dims, out, transpose_x, transpose_y);
}

template <typename T, typename Context>
void MatmulWithFlattenKernelImpl(const Context& dev_ctx,
                                 const DenseTensor& x,
                                 const DenseTensor& y,
                                 int x_num_col_dims,
                                 int y_num_col_dims,
                                 DenseTensor* out) {
  const DenseTensor x_matrix =
      x.dims().size() > 2 ? ReshapeToMatrix(x, x_num_col_dims) : x;
  const DenseTensor y_matrix =
      y.dims().size() > 2 ? ReshapeToMatrix(y, y_num_col_dims) : y;

  dev_ctx.template Alloc<T>(out);
  auto z_dim = out->dims();
  if (z_dim.size() != 2) {
    out->Resize({x_matrix.dims()[0], y_matrix.dims()[1]});
  }

  auto blas = funcs::GetBlas<Context, T>(dev_ctx);

  blas.MatMul(x_matrix, y_matrix, out);
  if (z_dim.size() != 2) {
    out->Resize(z_dim);
  }
}

#ifdef PADDLE_WITH_CUDA

template <typename Context>
void MatmulWithFlattenKernelInt8Impl(const Context& dev_ctx,
                                     const DenseTensor& x,
                                     const DenseTensor& y,
                                     int x_num_col_dims,
                                     int y_num_col_dims,
                                     DenseTensor* out) {
  PADDLE_ENFORCE_EQ(
      x.dtype(),
      DataType::INT8,
      common::errors::InvalidArgument(
          "The type of input(x) used in int8 mul must be (%s) "
          "does not match the "
          "type of data (%s) currently contained in the container.",
          CppTypeToDataType<int8_t>::Type(),
          x.dtype()));
  PADDLE_ENFORCE_EQ(
      y.dtype(),
      DataType::INT8,
      common::errors::InvalidArgument(
          "The type of input(y) used in int8 mul must be (%s) "
          "does not match the "
          "type of data (%s) currently contained in the container.",
          CppTypeToDataType<int8_t>::Type(),
          y.dtype()));

  const DenseTensor x_matrix =
      x.dims().size() > 2 ? ReshapeToMatrix(x, x_num_col_dims) : x;
  const DenseTensor y_matrix =
      y.dims().size() > 2 ? ReshapeToMatrix(y, y_num_col_dims) : y;

  PADDLE_ENFORCE_EQ(
      x_matrix.dims()[1],
      y_matrix.dims()[0],
      common::errors::InvalidArgument(
          "X's numbers of columns must be equal to Y's numbers of rows."
          "But received X has [%d] columns,"
          "received Y has [%d] rows",
          x_matrix.dims()[1],
          y_matrix.dims()[0]));

  PADDLE_ENFORCE_EQ((y_matrix.dims()[1] % 4 == 0 || y_matrix.dims()[1] == 1),
                    true,
                    common::errors::InvalidArgument(
                        "The dimension size N used in int8 mul must be 1 "
                        "or a multiple of 4 does not match the size (%d)"
                        "currently contained in the container.",
                        y_matrix.dims()[1]));
  PADDLE_ENFORCE_EQ((x_matrix.dims()[1] % 4 == 0),
                    true,
                    common::errors::InvalidArgument(
                        "The dimension size K used in int8 mul must be a "
                        "multiple of 4 does not match the size (%d) currently "
                        "contained in the container.",
                        x_matrix.dims()[1]));

  dev_ctx.template Alloc<int32_t>(out);
  auto z_dim = out->dims();
  if (z_dim.size() != 2) {
    out->Resize({x_matrix.dims()[0], y_matrix.dims()[1]});
  }

#if CUDA_VERSION >= 11060
  using blaslt = funcs::MatmulWithCublasLt<int8_t, int32_t>;

  const int8_t* x_data = x_matrix.data<int8_t>();
  const int8_t* y_data = y_matrix.data<int8_t>();

  std::vector<std::int64_t> x_dims = {x_matrix.dims()[0], x_matrix.dims()[1]};
  std::vector<std::int64_t> y_dims = {y_matrix.dims()[0], y_matrix.dims()[1]};
  funcs::MatmulPlanner matmul_planner(x_dims,
                                      y_dims,
                                      false,
                                      false,
                                      CppTypeToDataType<int8_t>::Type(),
                                      funcs::MatmulFusedType::kMatmul,
                                      /* bias_data */ nullptr,
                                      /* reserve_data */ nullptr,
                                      /* use_addto */ false,
                                      /* no_exchange */ true);

  blaslt::Run(dev_ctx,
              x_data,
              y_data,
              dev_ctx.template Alloc<int32_t>(out),
              x_matrix.dims()[0],
              y_matrix.dims()[1],
              x_matrix.dims()[1],
              false,
              false,
              &matmul_planner);

  if (z_dim.size() != 2) {
    out->Resize(z_dim);
  }
#endif
}
#endif

#ifdef PADDLE_WITH_CUDA
template <typename Context>
typename std::enable_if<std::is_same<Context, GPUContext>::value, void>::type
DispatchMatmulWithFlattenInt8Kernel(const GPUContext& dev_ctx,
                                    const DenseTensor& x,
                                    const DenseTensor& y,
                                    int x_num_col_dims,
                                    int y_num_col_dims,
                                    DenseTensor* out) {
  MatmulWithFlattenKernelInt8Impl<Context>(
      dev_ctx, x, y, x_num_col_dims, y_num_col_dims, out);
}
#endif

template <typename Context>
typename std::enable_if<std::is_same<Context, CPUContext>::value, void>::type
DispatchMatmulWithFlattenInt8Kernel(const CPUContext& dev_ctx,
                                    const DenseTensor& x,
                                    const DenseTensor& y,
                                    int x_num_col_dims,
                                    int y_num_col_dims,
                                    DenseTensor* out) {
  PADDLE_THROW(common::errors::Unimplemented(
      "MatmulWithFlatten with CPU is NOT implemented "
      "yet."));
}

template <typename T, typename Context>
typename std::enable_if<std::is_same<T, int8_t>::value, void>::type
DispatchMatmulFlattenKernel(const Context& dev_ctx,
                            const DenseTensor& x,
                            const DenseTensor& y,
                            int x_num_col_dims,
                            int y_num_col_dims,
                            DenseTensor* out) {
  DispatchMatmulWithFlattenInt8Kernel<Context>(
      dev_ctx, x, y, x_num_col_dims, y_num_col_dims, out);
}

template <typename T, typename Context>
typename std::enable_if<!std::is_same<T, int8_t>::value, void>::type
DispatchMatmulFlattenKernel(const Context& dev_ctx,
                            const DenseTensor& x,
                            const DenseTensor& y,
                            int x_num_col_dims,
                            int y_num_col_dims,
                            DenseTensor* out) {
  MatmulWithFlattenKernelImpl<T, Context>(
      dev_ctx, x, y, x_num_col_dims, y_num_col_dims, out);
}

template <typename T, typename Context>
void MatmulWithFlattenKernel(const Context& dev_ctx,
                             const DenseTensor& x,
                             const DenseTensor& y,
                             int x_num_col_dims,
                             int y_num_col_dims,
                             DenseTensor* out) {
  DispatchMatmulFlattenKernel<T, Context>(
      dev_ctx, x, y, x_num_col_dims, y_num_col_dims, out);
}

template <typename T, typename Context>
void LegacyMatmulKernel(const Context& dev_ctx,
                        const DenseTensor& x,
                        const DenseTensor& y,
                        bool transpose_x,
                        bool transpose_y,
                        float alpha,
                        DenseTensor* out) {
  MatmulKernel<T, Context>(dev_ctx, x, y, transpose_x, transpose_y, out);
  if (std::fabs(alpha - 1.f) > 1e-6f) {
    ScaleKernel<T, Context>(
        dev_ctx, *out, Scalar(alpha), Scalar(0), false, out);
  }
}
}  // namespace phi