PaddlePaddle
diff --git a/‎paddle/fluid/operators/matmul_op.cc
Copy file name to clipboard
+95-47Lines changed: 95 additions & 47 deletions b/‎paddle/fluid/operators/matmul_op.cc
Copy file name to clipboard
+95-47Lines changed: 95 additions & 47 deletions
diff --git a/‎paddle/phi/api/yaml/legacy_ops.yaml
Copy file name to clipboard
-8Lines changed: 0 additions & 8 deletions b/‎paddle/phi/api/yaml/legacy_ops.yaml
Copy file name to clipboard
-8Lines changed: 0 additions & 8 deletions
diff --git a/‎paddle/phi/backends/gpu/cuda/cuda_helper.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/phi/backends/gpu/cuda/cuda_helper.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/binary.cc
Copy file name to clipboard
+8-70Lines changed: 8 additions & 70 deletions b/‎paddle/phi/infermeta/binary.cc
Copy file name to clipboard
+8-70Lines changed: 8 additions & 70 deletions
diff --git a/‎paddle/phi/infermeta/binary.h
Copy file name to clipboard
-6Lines changed: 0 additions & 6 deletions b/‎paddle/phi/infermeta/binary.h
Copy file name to clipboard
-6Lines changed: 0 additions & 6 deletions
diff --git a/‎paddle/phi/kernels/autotune/auto_tune_base.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/phi/kernels/autotune/auto_tune_base.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
@@ -16,6 +16,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/impl/matmul_kernel_impl.h"
 
 namespace paddle {
 namespace operators {
@@ -53,61 +54,90 @@ static framework::DDim ColumnMatrixFromVector(const framework::DDim &y_dim) {
   return phi::make_ddim({y_dim[0], 1});
 }
 
-template <typename DeviceContext, typename T>
-class MatMulKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext &context) const override {
-    auto &x = GET_DATA_SAFELY(
-        context.Input<phi::DenseTensor>("X"), "Input", "X", "MatMul");
-    auto &y = GET_DATA_SAFELY(
-        context.Input<phi::DenseTensor>("Y"), "Input", "Y", "MatMul");
-    auto *out = context.Output<phi::DenseTensor>("Out");
-
-    auto &dev_ctx = context.template device_context<DeviceContext>();
-    dev_ctx.template Alloc<T>(out, out->numel() * sizeof(T));
-
-    auto blas = phi::funcs::GetBlas<DeviceContext, T>(dev_ctx);
-    auto mat_dim_a = phi::funcs::CreateMatrixDescriptor(
-        RowMatrixFromVector(x.dims()), 0, context.Attr<bool>("transpose_X"));
-    auto mat_dim_b = phi::funcs::CreateMatrixDescriptor(
-        ColumnMatrixFromVector(y.dims()), 0, context.Attr<bool>("transpose_Y"));
-    auto scale = static_cast<T>(context.Attr<float>("alpha"));
+#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 11060
+template <typename T, typename DeviceContext>
+typename std::enable_if<std::is_integral<T>::value, void>::type
+ComputeMatmulImpl(const framework::ExecutionContext &context) {
+  auto &dev_ctx = context.template device_context<DeviceContext>();
+
+  auto &x = GET_DATA_SAFELY(
+      context.Input<phi::DenseTensor>("X"), "Input", "X", "MatMul");
+  auto &y = GET_DATA_SAFELY(
+      context.Input<phi::DenseTensor>("Y"), "Input", "Y", "MatMul");
+  auto *out = context.Output<phi::DenseTensor>("Out");
+
+  dev_ctx.template Alloc<T>(out, out->numel() * sizeof(T));
+
+  phi::MatmulKernel<T>(dev_ctx,
+                       x,
+                       y,
+                       context.Attr<bool>("transpose_X"),
+                       context.Attr<bool>("transpose_Y"),
+                       out);
+}
+#endif
 
-    int head_number = 1;
+template <typename T, typename DeviceContext>
+typename std::enable_if<!std::is_integral<T>::value, void>::type
+ComputeMatmulImpl(const framework::ExecutionContext &context) {
+  auto &x = GET_DATA_SAFELY(
+      context.Input<phi::DenseTensor>("X"), "Input", "X", "MatMul");
+  auto &y = GET_DATA_SAFELY(
+      context.Input<phi::DenseTensor>("Y"), "Input", "Y", "MatMul");
+  auto *out = context.Output<phi::DenseTensor>("Out");
+
+  auto &dev_ctx = context.template device_context<DeviceContext>();
+  dev_ctx.template Alloc<T>(out, out->numel() * sizeof(T));
+
+  auto blas = phi::funcs::GetBlas<DeviceContext, T>(dev_ctx);
+  auto mat_dim_a = phi::funcs::CreateMatrixDescriptor(
+      RowMatrixFromVector(x.dims()), 0, context.Attr<bool>("transpose_X"));
+  auto mat_dim_b = phi::funcs::CreateMatrixDescriptor(
+      ColumnMatrixFromVector(y.dims()), 0, context.Attr<bool>("transpose_Y"));
+  auto scale = static_cast<T>(context.Attr<float>("alpha"));
+
+  int head_number = 1;
 #if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) && \
     !defined(PADDLE_WITH_HIP)
-    head_number = context.Attr<int>("head_number");
+  head_number = context.Attr<int>("head_number");
 #endif
 
-    const auto &x_dims = x.dims();
-    const auto &y_dims = y.dims();
-    if (head_number <= 1 && x_dims.size() == 3 && y_dims.size() <= 2) {
-      // the transpose_X must be false, if is true, the transpose cost much time
-      if (!context.Attr<bool>("transpose_X")) {
-        mat_dim_a.height_ *= mat_dim_a.batch_size_;
-        mat_dim_a.batch_size_ = 0;
-      }
+  const auto &x_dims = x.dims();
+  const auto &y_dims = y.dims();
+  if (head_number <= 1 && x_dims.size() == 3 && y_dims.size() <= 2) {
+    // the transpose_X must be false, if is true, the transpose cost much time
+    if (!context.Attr<bool>("transpose_X")) {
+      mat_dim_a.height_ *= mat_dim_a.batch_size_;
+      mat_dim_a.batch_size_ = 0;
     }
+  }
 #if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) && \
     !defined(PADDLE_WITH_HIP)
-    bool split_vertical_y = (mat_dim_a.width_ != mat_dim_b.height_);
-
-    if (head_number > 1) {
-      blas.MatMulWithHead(x,
-                          mat_dim_a,
-                          y,
-                          mat_dim_b,
-                          scale,
-                          head_number,
-                          out,
-                          T(0),
-                          split_vertical_y);
-    } else {
-      blas.MatMul(x, mat_dim_a, y, mat_dim_b, scale, out, T(0));
-    }
-#else
+  bool split_vertical_y = (mat_dim_a.width_ != mat_dim_b.height_);
+
+  if (head_number > 1) {
+    blas.MatMulWithHead(x,
+                        mat_dim_a,
+                        y,
+                        mat_dim_b,
+                        scale,
+                        head_number,
+                        out,
+                        T(0),
+                        split_vertical_y);
+  } else {
     blas.MatMul(x, mat_dim_a, y, mat_dim_b, scale, out, T(0));
+  }
+#else
+  blas.MatMul(x, mat_dim_a, y, mat_dim_b, scale, out, T(0));
 #endif
+}
+
+template <typename DeviceContext, typename T>
+class MatMulKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext &context) const override {
+    ComputeMatmulImpl<T, DeviceContext>(context);
   }
 };
 
@@ -926,12 +956,31 @@ REGISTER_OP_CPU_KERNEL(matmul_grad_grad,
                        ops::MatMulDoubleGradKernel<phi::CPUContext, float>,
                        ops::MatMulDoubleGradKernel<phi::CPUContext, double>);
 
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#if defined(PADDLE_WITH_HIP)
 REGISTER_OP_CUDA_KERNEL(
     matmul,
     ops::MatMulKernel<phi::GPUContext, float>,
     ops::MatMulKernel<phi::GPUContext, double>,
     ops::MatMulKernel<phi::GPUContext, paddle::platform::float16>);
+#endif
+
+#if defined(PADDLE_WITH_CUDA)
+#if CUDA_VERSION >= 11060
+REGISTER_OP_CUDA_KERNEL(
+    matmul,
+    ops::MatMulKernel<phi::GPUContext, int8_t>,
+    ops::MatMulKernel<phi::GPUContext, float>,
+    ops::MatMulKernel<phi::GPUContext, double>,
+    ops::MatMulKernel<phi::GPUContext, paddle::platform::float16>);
+#else
+REGISTER_OP_CUDA_KERNEL(
+    matmul,
+    ops::MatMulKernel<phi::GPUContext, float>,
+    ops::MatMulKernel<phi::GPUContext, double>,
+    ops::MatMulKernel<phi::GPUContext, paddle::platform::float16>);
+#endif
+#endif
+
 REGISTER_OP_CUDA_KERNEL(
     matmul_grad,
     ops::MatMulGradKernel<phi::GPUContext, float>,
@@ -940,7 +989,6 @@ REGISTER_OP_CUDA_KERNEL(
 REGISTER_OP_CUDA_KERNEL(matmul_grad_grad,
                         ops::MatMulDoubleGradKernel<phi::GPUContext, float>,
                         ops::MatMulDoubleGradKernel<phi::GPUContext, double>);
-#endif
 
 REGISTER_OP_VERSION(matmul).AddCheckpoint(
     R"ROC(Register matmul for adding the attribute of
 
@@ -685,14 +685,6 @@
     func : matmul
   backward : matmul_grad
 
-- op : matmul_int8
-  args : (Tensor x, Tensor y, bool transpose_x = false, bool transpose_y = false)
-  output : Tensor
-  infer_meta :
-    func : MatmulInt8InferMeta
-  kernel :
-    func : matmul_int8
-
 - op : matrix_rank
   args : (Tensor x, float tol, bool use_default_tol=true, bool hermitian=false)
   output : Tensor(out)
 
@@ -88,6 +88,12 @@ cudaDataType_t ToCudaDataType() {
 #if CUDA_VERSION >= 11000
   } else if (std::is_same<T, phi::dtype::bfloat16>::value) {
     return CUDA_R_16BF;
+#endif
+#if CUDA_VERSION >= 11060
+  } else if (std::is_same<T, int8_t>::value) {
+    return CUDA_R_8I;
+  } else if (std::is_same<T, int32_t>::value) {
+    return CUDA_R_32I;
 #endif
   } else {
     PADDLE_THROW(phi::errors::InvalidArgument(
 
@@ -2169,77 +2169,11 @@ void MatmulInferMeta(const MetaTensor& x,
   auto ddim_out = phi::make_ddim(new_dims);
 
   out->set_dims(ddim_out);
-  out->set_dtype(x.dtype());
-  out->set_layout(x.layout());
-}
-
-void MatmulInt8InferMeta(const MetaTensor& x,
-                         const MetaTensor& y,
-                         bool trans_x,
-                         bool trans_y,
-                         MetaTensor* out) {
-  std::vector<int64_t> dims_x = phi::vectorize(x.dims());
-  std::vector<int64_t> dims_y = phi::vectorize(y.dims());
-  auto ndims_x = dims_x.size();
-  auto ndims_y = dims_y.size();
-  PADDLE_ENFORCE_GT(ndims_x,
-                    0UL,
-                    phi::errors::InvalidArgument(
-                        "The Input(x) dims size must be greater than 0,"
-                        " but reviced dims size is 0. "));
-  PADDLE_ENFORCE_GT(ndims_y,
-                    0UL,
-                    phi::errors::InvalidArgument(
-                        "The Input(y) dims size must be greater than 0,"
-                        " but reviced dims size is 0. "));
-
-  bool x_broadcasted = false, y_broadcasted = false;
-  if (ndims_x == 1) {
-    dims_x.insert(dims_x.begin(), 1);
-    ndims_x = 2;
-    x_broadcasted = true;
-  }
-
-  if (ndims_y == 1) {
-    dims_y.push_back(1);
-    ndims_y = 2;
-    y_broadcasted = true;
-  }
-
-  size_t M, N;
-  if (trans_x) {
-    M = dims_x[ndims_x - 1];
-  } else {
-    M = dims_x[ndims_x - 2];
-  }
-  if (trans_y) {
-    N = dims_y[ndims_y - 2];
-  } else {
-    N = dims_y[ndims_y - 1];
-  }
-
-  std::vector<int64_t> new_dims;
-  if (ndims_x > ndims_y) {
-    new_dims.assign(dims_x.begin(), dims_x.end() - 2);
-  } else if (ndims_x < ndims_y) {
-    new_dims.assign(dims_y.begin(), dims_y.end() - 2);
+  if (x.dtype() == phi::DataType::INT8) {
+    out->set_dtype(phi::DataType::INT32);
   } else {
-    new_dims.reserve(ndims_x);
-    for (size_t i = 0; i < ndims_x - 2; ++i) {
-      new_dims.push_back(std::max(dims_x[i], dims_y[i]));
-    }
-  }
-  if (!x_broadcasted) {
-    new_dims.push_back(M);  // NOLINT
-  }
-  if (!y_broadcasted) {
-    new_dims.push_back(N);  // NOLINT
+    out->set_dtype(x.dtype());
   }
-
-  auto ddim_out = phi::make_ddim(new_dims);
-
-  out->set_dims(ddim_out);
-  out->set_dtype(phi::DataType::INT32);
   out->set_layout(x.layout());
 }
 
@@ -2314,7 +2248,11 @@ void MatmulWithFlattenInferMeta(const MetaTensor& x,
   }
 
   out->set_dims(phi::make_ddim(output_dims));
-  out->set_dtype(x.dtype());
+  if (x.dtype() == phi::DataType::INT8) {
+    out->set_dtype(phi::DataType::INT32);
+  } else {
+    out->set_dtype(x.dtype());
+  }
   out->share_lod(x);
 }
 
 
@@ -347,12 +347,6 @@ void MatmulInferMeta(const MetaTensor& x,
                      bool trans_y,
                      MetaTensor* out);
 
-void MatmulInt8InferMeta(const MetaTensor& x,
-                         const MetaTensor& y,
-                         bool trans_x,
-                         bool trans_y,
-                         MetaTensor* out);
-
 void MatmulWithFlattenInferMeta(const MetaTensor& x,
                                 const MetaTensor& y,
                                 int x_num_col_dims,
 
@@ -106,7 +106,7 @@ class AutoTuneBase {
     float min_time = std::numeric_limits<float>::max();
 
     // Time cost test estabulished in default stream.
-    for (int i = 0; i < kernels_.size(); ++i) {
+    for (size_t i = 0; i < kernels_.size(); ++i) {
       auto time = RunAndMeasureKernel<Context>(ctx, i, args...);
       if (time < min_time) {
         min_time = time;