fc_lstm_fuse_pass.cc modify use_mkldnn [fluid_ops] (#74550)

Author: co63oc

paddle/fluid/framework/ir/fc_lstm_fuse_pass.cc

@@ -195,7 +195,7 @@ int FCLstmFusePass::BuildFusion(Graph* graph,
                           Node* cell,
                           Node* xx,
                           Node* fc_bias,
-                          const bool use_mkldnn) {
+                          const bool use_onednn) {
     OpDesc op_desc;
     op_desc.SetType("fusion_lstm");
 #define SET_IN(Key, node__) op_desc.SetInput(#Key, {node__->Name()});
@@ -235,7 +235,7 @@ int FCLstmFusePass::BuildFusion(Graph* graph,
     op_desc.SetOutput("XX", {xx->Name()});
     op_desc.SetAttr("is_reverse", lstm->Op()->GetAttr("is_reverse"));
     op_desc.SetAttr("use_peepholes", lstm->Op()->GetAttr("use_peepholes"));
-    op_desc.SetAttr("use_mkldnn", use_mkldnn);
+    op_desc.SetAttr("use_onednn", use_onednn);
     // TODO(TJ): get from attr
     op_desc.SetAttr("use_seq", true);
 
@@ -300,8 +300,9 @@ int FCLstmFusePass::BuildFusion(Graph* graph,
     GET_IR_NODE_FROM_SUBGRAPH(Cell, Cell, lstm_pattern);
     GET_IR_NODE_FROM_SUBGRAPH(w, w, fc_pattern);
     GET_IR_NODE_FROM_SUBGRAPH(mul, mul, fc_pattern);
-    const bool use_mkldnn =
-        (mul->Op()->GetAttrIfExists<bool>("use_mkldnn") &&
+    const bool use_onednn =
+        ((mul->Op()->GetAttrIfExists<bool>("use_mkldnn") ||
+          mul->Op()->GetAttrIfExists<bool>("use_onednn")) &&
          lstm->Op()->GetAttrIfExists<std::string>("gate_activation") ==
              "sigmoid" &&
          lstm->Op()->GetAttrIfExists<std::string>("cell_activation") ==
@@ -323,7 +324,7 @@ int FCLstmFusePass::BuildFusion(Graph* graph,
                    Cell,
                    fc_out,
                    fc_bias,
-                   use_mkldnn);
+                   use_onednn);
       // Remove unneeded nodes.
       std::unordered_set<const Node*> marked_nodes(
           {mul, lstm, elementwise_add, mul_out, BatchGate, BatchCellPreAct});
@@ -339,7 +340,7 @@ int FCLstmFusePass::BuildFusion(Graph* graph,
                    Cell,
                    fc_out,
                    nullptr,
-                   use_mkldnn);
+                   use_onednn);
       // Remove unneeded nodes.
       std::unordered_set<const Node*> marked_nodes(
           {mul, lstm, BatchGate, BatchCellPreAct});

paddle/fluid/framework/ir/fuse_pass_base.cc

@@ -58,10 +58,16 @@ void FusePassBase::AddStatis(int count_of_fused) const {
 FuseOptions FusePassBase::FindFuseOption(const Node& node1,
                                          const Node& node2) const {
 #ifdef PADDLE_WITH_DNNL
-  bool node1_onednn = node1.Op()->HasAttr("use_mkldnn") &&
-                      PADDLE_GET_CONST(bool, node1.Op()->GetAttr("use_mkldnn"));
-  bool node2_onednn = node2.Op()->HasAttr("use_mkldnn") &&
-                      PADDLE_GET_CONST(bool, node2.Op()->GetAttr("use_mkldnn"));
+  bool node1_onednn =
+      (node1.Op()->HasAttr("use_mkldnn") &&
+       PADDLE_GET_CONST(bool, node1.Op()->GetAttr("use_mkldnn"))) ||
+      (node1.Op()->HasAttr("use_onednn") &&
+       PADDLE_GET_CONST(bool, node1.Op()->GetAttr("use_onednn")));
+  bool node2_onednn =
+      (node2.Op()->HasAttr("use_mkldnn") &&
+       PADDLE_GET_CONST(bool, node2.Op()->GetAttr("use_mkldnn"))) ||
+      (node2.Op()->HasAttr("use_onednn") &&
+       PADDLE_GET_CONST(bool, node2.Op()->GetAttr("use_onednn")));
   if (node1_onednn && node2_onednn)
     return FUSE_ONEDNN;
   else if (!node1_onednn && !node2_onednn)

paddle/fluid/framework/ir/gpu_cpu_map_matmul_to_mul_pass.cc

@@ -444,6 +444,9 @@ void GpuCpuMapMatmulV2ToMatmulPass::ApplyImpl(ir::Graph* graph) const {
     if (matmul_v2_op->Op()->HasAttr("use_mkldnn")) {
       desc.SetAttr("use_mkldnn", matmul_v2_op->Op()->GetAttr("use_mkldnn"));
     }
+    if (matmul_v2_op->Op()->HasAttr("use_onednn")) {
+      desc.SetAttr("use_onednn", matmul_v2_op->Op()->GetAttr("use_onednn"));
+    }
     if (matmul_v2_op->Op()->HasAttr("enable_int8")) {
       desc.SetAttr("enable_int8", matmul_v2_op->Op()->GetAttr("enable_int8"));
       desc.SetAttr("Input_scale", matmul_v2_op->Op()->GetAttr("Input_scale"));

paddle/fluid/framework/op_desc.h

@@ -235,7 +235,7 @@ class TEST_API OpDesc {
   // attribute name => all original attrs
   AttributeMap attrs_;
   // runtime_attrs_ contains the attributes which used for dispatching kernel
-  // (use_mkldnn, use_cudnn, ...) or passing additional configuration for
+  // (use_onednn, use_cudnn, ...) or passing additional configuration for
   // special heterogeneous kernel (workspace_size_MB, ...).
   // The attributes in runtime_attrs_ are set by framework (such as PASS),
   // and not in the python api.

paddle/phi/kernels/fusion/onednn/fc_kernel.cc

@@ -411,7 +411,7 @@ void RunKernel(const phi::OneDNNContext& dev_ctx,
                const paddle::optional<DenseTensor>& bias,
                const int in_num_col_dims,
                const std::string& activation_type,
-               const bool use_mkldnn,
+               const bool use_onednn,
                const bool padding_weights,
                const bool use_quantizer,
                const std::string& mkldnn_data_type,