[PHI] Migrate matmul kernel

paddle/fluid/framework/operator.cc

@@ -3230,6 +3230,30 @@ void OperatorWithKernel::BuildPhiKernelContext(
   }
   VLOG(4) << "Done attributes";
 
+// Clear All old attrs before add new attrs,
+// because sometimes old attrs may be misused.
+#if defined(PADDLE_WITH_MKLDNN)
+  if (phi::OneDNNContext::classof(dev_ctx)) {
+    phi::OneDNNContext* one_dnn_ctx = static_cast<phi::OneDNNContext*>(dev_ctx);
+    one_dnn_ctx->ClearDnnAttr();
+  }
+#endif
+
+  // Note(YuanRisheng): Now, we can't open code below.
+  // Because some unittest run OLD dygraph and ExtraAttr is not supported in OLD
+  // dygraph. So, here we use trick that dev_ctx is a global object. We can
+  // store ExtraAttr in static graph and when unittest run OLD dygraph, it can
+  // obtain these ExtraAttr. We can open this code when OLD dygraph is no longer
+  // used.
+  /*
+  #if defined(PADDLE_WITH_CUDA)
+    if(phi::GPUContext::classof(dev_ctx)) {
+      phi::GPUContext* gpu_dnn_ctx = static_cast<phi::GPUContext*>(dev_ctx);
+      gpu_dnn_ctx->ClearDnnAttr();
+    }
+  #endif
+  */
+
   // For compatible with Op with extra attrs for specific backend
 #if defined(PADDLE_WITH_MKLDNN) || defined(PADDLE_WITH_CUDA)
   auto& runtime_attrs = RuntimeAttrs();

paddle/fluid/operators/mkldnn/matmul_v2_mkldnn_op.cc

@@ -383,7 +383,7 @@ void ExecuteMatMulV2(const ExecutionContext &ctx,
 }
 
 template <typename T>
-class MatMulV2MKLDNNKernel : public paddle::framework::OpKernel<T> {
+class MatMulMKLDNNKernel : public paddle::framework::OpKernel<T> {
  public:
   void Compute(const ExecutionContext &ctx) const override {
     if (ctx.HasAttr("head_number")) {
@@ -699,21 +699,13 @@ class MatMulGradMKLDNNKernel : public paddle::framework::OpKernel<T> {
 REGISTER_OP_KERNEL(matmul,
                    MKLDNN,
                    ::paddle::platform::CPUPlace,
-                   MatMulV2MKLDNNKernel<float>,
-                   MatMulV2MKLDNNKernel<paddle::platform::bfloat16>,
-                   MatMulV2MKLDNNKernel<int8_t>,
-                   MatMulV2MKLDNNKernel<uint8_t>);
+                   MatMulMKLDNNKernel<float>,
+                   MatMulMKLDNNKernel<paddle::platform::bfloat16>,
+                   MatMulMKLDNNKernel<int8_t>,
+                   MatMulMKLDNNKernel<uint8_t>);
 
 REGISTER_OP_KERNEL(matmul_grad,
                    MKLDNN,
                    ::paddle::platform::CPUPlace,
                    MatMulGradMKLDNNKernel<float>,
                    MatMulGradMKLDNNKernel<paddle::platform::bfloat16>);
-
-REGISTER_OP_KERNEL(matmul_v2,
-                   MKLDNN,
-                   ::paddle::platform::CPUPlace,
-                   MatMulV2MKLDNNKernel<float>,
-                   MatMulV2MKLDNNKernel<paddle::platform::bfloat16>,
-                   MatMulV2MKLDNNKernel<int8_t>,
-                   MatMulV2MKLDNNKernel<uint8_t>);

paddle/fluid/operators/ops_extra_info.h

@@ -98,6 +98,7 @@ const std::unordered_map<std::string, ExtraAttrPropertySet>
         {"fuse_alpha", ExtraAttrProperty::ONEDNN},
         {"fuse_beta", ExtraAttrProperty::ONEDNN},
         {"fuse_relu", ExtraAttrProperty::ONEDNN},
+        {"fused_output_scale", ExtraAttrProperty::ONEDNN},
         {"fuse_residual_connection", ExtraAttrProperty::ONEDNN},
         {"fuse_with_relu", ExtraAttrProperty::ONEDNN},
         {"fused_reshape_Out", ExtraAttrProperty::ONEDNN},
@@ -221,7 +222,8 @@ class ExtraInfoUtils {
   std::unordered_map<std::string, std::vector<std::string>>
       g_extra_input_names_map_ = {{"conv2d", {"Bias", "ResidualData"}},
                                   {"conv2d_transpose", {"Bias"}},
-                                  {"conv2d_grad", {"Bias"}}};
+                                  {"conv2d_grad", {"Bias"}},
+                                  {"matmul_v2", {"ResidualData"}}};
   std::vector<std::string> empty_extra_input_names_;
 };
 

paddle/phi/backends/gpu/gpu_context.cc

@@ -740,6 +740,8 @@ struct GPUContext::Impl {
     dnn_attrs_[attr_name] = attr;
   }
 
+  void ClearDnnAttr() { dnn_attrs_.clear(); }
+
   // use one flag for all handles?
   // they should be accessed consistently
   bool owned_{false};
@@ -1042,4 +1044,6 @@ void GPUContext::SetDnnAttr(const std::string& attr_name, Attribute attr) {
   return impl_->SetDnnAttr(attr_name, std::move(attr));
 }
 
+void GPUContext::ClearDnnAttr() { return impl_->ClearDnnAttr(); }
+
 }  // namespace phi

paddle/phi/backends/gpu/gpu_context.h

@@ -172,6 +172,7 @@ class PADDLE_API GPUContext : public DeviceContext,
   bool HasDnnAttr(const std::string& attr_name) const;
   const Attribute& GetDnnAttr(const std::string& attr_name) const;
   void SetDnnAttr(const std::string& attr_name, Attribute attr);
+  void ClearDnnAttr();
 
   static const char* name() { return "GPUContext"; }
 

paddle/phi/backends/onednn/onednn_context.cc

@@ -301,6 +301,8 @@ struct OneDNNContext::Impl {
     dnn_attrs_[attr_name] = attr;
   }
 
+  void ClearDnnAttr() { dnn_attrs_.clear(); }
+
   bool HasDnnInput(const std::string& input_name) const {
     return dnn_inputs_.count(input_name) != 0UL;
   }
@@ -429,6 +431,8 @@ bool OneDNNContext::HasDnnInput(const std::string& input_name) const {
   return impl_->HasDnnInput(input_name);
 }
 
+void OneDNNContext::ClearDnnAttr() { return impl_->ClearDnnAttr(); }
+
 const DenseTensor* OneDNNContext::GetDnnInput(
     const std::string& input_name) const {
   return impl_->GetDnnInput(input_name);

paddle/phi/backends/onednn/onednn_context.h

@@ -146,6 +146,8 @@ class OneDNNContext : public CPUContext {
   const DenseTensor* GetDnnInput(const std::string& input_name) const;
   void SetDnnInput(const std::string& input_name, const DenseTensor* input);
 
+  void ClearDnnAttr();
+
   void SetInputsName(const TensorNameMap& inputs_name);
 
   void SetOutputsName(const TensorNameMap& outputs_name);

paddle/phi/backends/onednn/onednn_reuse.h

@@ -1844,9 +1844,11 @@ class MatmulOneDNNHandler
     if (scale_out != 1.0f) {
       matmul_attrs.set_output_scales(0, {scale_out});
     }
+    const auto* residual_data = dev_ctx.HasDnnInput("ResidualData")
+                                    ? dev_ctx.GetDnnInput("ResidualData")
+                                    : nullptr;
 
-    if (dev_ctx.HasDnnInput("ResidualData")) {
-      auto* residual_data = dev_ctx.GetDnnInput("ResidualData");
+    if (residual_data) {
       auto residual_data_tz = vectorize(residual_data->dims());
       auto residual_data_md = memory::desc(residual_data_tz,
                                            OneDNNGetDataType<OT>(),
@@ -1863,9 +1865,11 @@ class MatmulOneDNNHandler
 
     AppendActivation(dev_ctx, post_operations);
 
-    if (dev_ctx.HasDnnAttr("fused_output_scale")) {
-      float scale_alpha =
-          PADDLE_GET_CONST(float, dev_ctx.GetDnnAttr("fused_output_scale"));
+    const float scale_alpha =
+        dev_ctx.HasDnnAttr("fused_output_scale")
+            ? PADDLE_GET_CONST(float, dev_ctx.GetDnnAttr("fused_output_scale"))
+            : 1.0f;
+    if (scale_alpha != 1.0f) {
       post_operations.append_eltwise(
           1.0, dnnl::algorithm::eltwise_linear, scale_alpha, 0.0f);
     }
@@ -1984,8 +1988,11 @@ void ExecuteMatmul(const OneDNNContext& dev_ctx,
       {DNNL_ARG_WEIGHTS, *weights_memory_p},
       {DNNL_ARG_DST, *dst_memory_p}};
 
-  if (dev_ctx.HasDnnInput("ResidualData")) {
-    auto* residual_data = dev_ctx.GetDnnInput("ResidualData");
+  const auto* residual_data = dev_ctx.HasDnnInput("ResidualData")
+                                  ? dev_ctx.GetDnnInput("ResidualData")
+                                  : nullptr;
+
+  if (residual_data) {
     const auto residual_data_memory_p = handler.AcquireSrcMemory(residual_data);
     matmul_args.insert({DNNL_ARG_ATTR_MULTIPLE_POST_OP(0) | DNNL_ARG_SRC_1,
                         *residual_data_memory_p});

paddle/phi/kernels/onednn/matmul_kernel.cc

@@ -0,0 +1,164 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/matmul_kernel.h"
+
+#include "paddle/phi/backends/onednn/onednn_reuse.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+DDim GetDimsForInput(const OneDNNContext &dev_ctx,
+                     DDim input_dims,
+                     std::string input_name) {
+  auto shape =
+      dev_ctx.HasDnnAttr("fused_reshape_" + input_name)
+          ? PADDLE_GET_CONST(std::vector<int>,
+                             dev_ctx.GetDnnAttr("fused_reshape_" + input_name))
+          : std::vector<int>();
+  auto axis = dev_ctx.HasDnnAttr("fused_transpose_" + input_name)
+                  ? PADDLE_GET_CONST(
+                        std::vector<int>,
+                        dev_ctx.GetDnnAttr("fused_transpose_" + input_name))
+                  : std::vector<int>();
+  if (!shape.empty() && !axis.empty()) {
+    return input_dims.reshape(shape).transpose(axis);
+  }
+  return input_dims;
+}
+
+void CalculateMatrixDims(const std::vector<int64_t> &x_dims,
+                         const std::vector<int64_t> &y_dims,
+                         std::vector<int64_t> *x_bd_dims,
+                         std::vector<int64_t> *y_bd_dims,
+                         DenseTensor *out,
+                         const bool is_output_fused) {
+  if (x_dims.size() == 1) {
+    (*x_bd_dims)[(*x_bd_dims).size() - 1] = x_dims[0];
+  } else if (x_dims.size() == 2) {
+    (*x_bd_dims)[(*x_bd_dims).size() - 1] = x_dims[1];
+    (*x_bd_dims)[(*x_bd_dims).size() - 2] = x_dims[0];
+  } else {
+    for (size_t i = 0; i < x_dims.size(); ++i) {
+      (*x_bd_dims)[(*x_bd_dims).size() - x_dims.size() + i] = x_dims[i];
+    }
+  }
+  if (y_dims.size() == 1) {
+    (*y_bd_dims)[(*x_bd_dims).size() - 2] = y_dims[0];
+  } else if (y_dims.size() == 2) {
+    (*y_bd_dims)[(*y_bd_dims).size() - 1] = y_dims[1];
+    (*y_bd_dims)[(*y_bd_dims).size() - 2] = y_dims[0];
+  } else {
+    for (size_t i = 0; i < y_dims.size(); ++i) {
+      (*y_bd_dims)[(*y_bd_dims).size() - y_dims.size() + i] = y_dims[i];
+    }
+  }
+
+  if (!is_output_fused && x_dims.size() > 2 && y_dims.size() > 2) {
+    auto out_dims = vectorize(out->dims());
+    for (size_t i = 0; i < (*x_bd_dims).size() - 2; ++i) {
+      PADDLE_ENFORCE_EQ(
+          (*x_bd_dims)[i] == (*y_bd_dims)[i] || (*x_bd_dims)[i] == 1 ||
+              (*y_bd_dims)[i] == 1,
+          true,
+          errors::InvalidArgument(
+              "Tensor dimensions are incorrect for broadcasting."
+              "Dimensions in X and Y must be same or equal to 1, but "
+              "received x_dim[%d]=%d and y_dims[%d]= %d",
+              i,
+              (*x_bd_dims)[i],
+              i,
+              (*y_bd_dims)[i]));
+      (out_dims)[i] = std::max((*x_bd_dims)[i], (*y_bd_dims)[i]);
+    }
+    out->Resize(make_ddim((out_dims)));
+  }
+}
+
+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 (dev_ctx.HasDnnAttr("head_number")) {
+    const auto head_number =
+        PADDLE_GET_CONST(int, dev_ctx.GetDnnAttr("head_number"));
+    PADDLE_ENFORCE_EQ(
+        head_number,
+        1,
+        errors::Unimplemented(
+            "oneDNN matmul doesn't support multiple heads. Expected "
+            "head_number=1. But received `head_number` is %d",
+            head_number));
+  }
+
+  constexpr bool is_int8 = funcs::is_int8<T>();
+  constexpr bool is_bfloat16 = funcs::is_bfloat16<T>();
+  const bool force_fp32_output =
+      dev_ctx.HasDnnAttr("force_fp32_output")
+          ? PADDLE_GET_CONST(bool, dev_ctx.GetDnnAttr("force_fp32_output"))
+          : false;
+
+  bool fuse_relu = false;
+  if (dev_ctx.HasDnnAttr("fuse_activation")) {
+    auto act_type =
+        PADDLE_GET_CONST(std::string, dev_ctx.GetDnnAttr("fuse_activation"));
+    if (act_type == "relu" || act_type == "relu6") {
+      fuse_relu = true;
+    }
+  }
+
+  auto x_dims = vectorize(GetDimsForInput(dev_ctx, x.dims(), "X"));
+  auto y_dims = vectorize(GetDimsForInput(dev_ctx, y.dims(), "Y"));
+
+  int ndims = std::max(x_dims.size(), y_dims.size());
+  ndims = std::max(ndims, 3);
+
+  std::vector<int64_t> x_bd_dims(ndims, 1);
+  std::vector<int64_t> y_bd_dims(ndims, 1);
+
+  CalculateMatrixDims(x_dims,
+                      y_dims,
+                      &x_bd_dims,
+                      &y_bd_dims,
+                      out,
+                      funcs::IsOutputFused(dev_ctx));
+
+  if (force_fp32_output || ((!is_int8) && (!is_bfloat16))) {
+    funcs::ExecuteMatmul<T, float>(
+        dev_ctx, x, y, x_bd_dims, y_bd_dims, transpose_x, transpose_y, out);
+  } else if (is_bfloat16) {
+    funcs::ExecuteMatmul<T, paddle::platform::bfloat16>(
+        dev_ctx, x, y, x_bd_dims, y_bd_dims, transpose_x, transpose_y, out);
+  } else if (fuse_relu) {
+    funcs::ExecuteMatmul<T, uint8_t>(
+        dev_ctx, x, y, x_bd_dims, y_bd_dims, transpose_x, transpose_y, out);
+  } else {
+    funcs::ExecuteMatmul<T, int8_t>(
+        dev_ctx, x, y, x_bd_dims, y_bd_dims, transpose_x, transpose_y, out);
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(matmul,
+                   OneDNN,
+                   ONEDNN,
+                   phi::MatmulKernel,
+                   float,
+                   phi::dtype::bfloat16,
+                   int8_t,
+                   uint8_t) {}