[PIR] pir onednn add conv2d test, and support fused_conv2d (#60974)

wanghuancoder · web-flow · commit 166aced98db2 · 2024-01-23T10:47:19.000+08:00
* pir onednn add conv2d test, and support fused_conv2d
diff --git a/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.cc b/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.cc
@@ -314,11 +314,33 @@ OneDNNPhiKernelInstruction::OneDNNPhiKernelInstruction(
             .at("extra_args")
             .dyn_cast<pir::ArrayAttribute>()
             .AsVector();
-    std::vector<std::string> extra_args;
+    auto& op_normalizer = paddle::translator::OpNameNormalizer::instance();
+    std::string fluid_op_name = yaml_info_parser.GetOriginOpName();
+
     for (auto& attr : extra_args_attr) {
       auto attr_name = attr.dyn_cast<pir::StrAttribute>().AsString();
       extra_attr_[attr_name] = ConvertPirAttribute2RuntimeAttribute(
           op_attributes.at(attr_name), attr_name, yaml_info_parser);
+      auto legacy_attr_name =
+          op_normalizer.GetLegacyAttrName(fluid_op_name, attr_name);
+      if (legacy_attr_name != attr_name) {
+        extra_attr_[legacy_attr_name] = extra_attr_[attr_name];
+      }
+    }
+    auto attr_name_list = yaml_info_parser.AttrParams(true);
+    for (auto& attr : attr_name_list) {
+      auto attr_name = attr;
+      if (!op_attributes.count(attr_name)) {
+        // In PIR, IntArray attr will be input, but not attr.
+        continue;
+      }
+      ctx_attr_[attr_name] = ConvertPirAttribute2RuntimeAttribute(
+          op_attributes.at(attr_name), attr_name, yaml_info_parser);
+      auto legacy_attr_name =
+          op_normalizer.GetLegacyAttrName(fluid_op_name, attr_name);
+      if (legacy_attr_name != attr_name) {
+        ctx_attr_[legacy_attr_name] = ctx_attr_[attr_name];
+      }
     }
   }
   TensorNameMap(op, *value_exec_info_, yaml_info_parser, inputs_, outputs_);
@@ -336,6 +358,9 @@ void OneDNNPhiKernelInstruction::Run() {
       size_t(0), kernel_context_.InputsSize());
   for (size_t i = 0; i < inputs.size(); ++i) {
     auto input = inputs[i];
+    if (input == nullptr) {
+      continue;
+    }
     if (input->layout() != phi::DataLayout::ONEDNN) {
       phi::DataLayout from_layout = input->layout();
 
@@ -375,6 +400,9 @@ void OneDNNPhiKernelInstruction::Run() {
   for (auto& attr : extra_attr_) {
     one_dnn_ctx->SetDnnAttr(attr.first, attr.second);
   }
+  for (auto& attr : ctx_attr_) {
+    one_dnn_ctx->SetDnnAttr(attr.first, attr.second);
+  }
   one_dnn_ctx->SetInputsName(inputs_);
   one_dnn_ctx->SetOutputsName(outputs_);
 
diff --git a/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.h b/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.h
@@ -71,6 +71,7 @@ class OneDNNPhiKernelInstruction : public InstructionBase {
   std::set<int> data_format_tensors_{};
   phi::DataLayout input_layout_{phi::DataLayout::kAnyLayout};
   std::map<std::string, phi::Attribute> extra_attr_{};
+  std::map<std::string, phi::Attribute> ctx_attr_{};
   std::map<std::string, std::vector<std::string>> inputs_{};
   std::map<std::string, std::vector<std::string>> outputs_{};
 };
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -1990,6 +1990,15 @@ def OpGenerator(
                 if first_file:
                     op["is_onednn_only"] = True
                     onednn_only_op_list.append("\"" + op['name'] + "\"")
+                    if op['name'] in ops_onednn_extra_map:
+                        onednn_item = ops_onednn_extra_map[op['name']]
+                        op["is_onednn_only"] = onednn_item["is_onednn_only"]
+                        op["extra_args"] = onednn_item["extra_args"]
+                        op["data_format_tensors"] = onednn_item[
+                            "data_format_tensors"
+                        ]
+                        op["dynamic_fallback"] = onednn_item["dynamic_fallback"]
+                        op["attrs"] = op["attrs"] + onednn_item["attrs"]
                 elif op['name'] in ops_onednn_extra_map:
                     onednn_item = ops_onednn_extra_map[op['name']]
                     op["is_onednn_only"] = onednn_item["is_onednn_only"]
diff --git a/paddle/fluid/pir/dialect/operator/ir/onednn.yaml b/paddle/fluid/pir/dialect/operator/ir/onednn.yaml
@@ -8,6 +8,16 @@
     func : dequantize
     data_type : input
 
+- op : fused_conv2d
+  args : (Tensor input, Tensor filter, Tensor bias, Tensor residual_param, int[] strides={1, 1}, int[] paddings={0, 0}, str padding_algorithm="EXPLICIT", int[] dilations={1, 1}, int groups=1, str data_format="NCHW", str mkldnn_data_type="float32", str fuse_activation="", bool fuse_residual_connection=false, bool force_fp32_output=false)
+  output : Tensor(output)
+  infer_meta :
+    func : FusedConvInferMeta
+  kernel :
+    func : fused_conv2d
+    data_type : input
+  optional : bias, residual_param
+
 - op : quantize
   args : (Tensor input, bool is_negative_input=false, float scale=1.0, float shift=0.0, str output_format="NHWC", bool bfloat16=false)
   output : Tensor(output)
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -19,6 +19,10 @@
   extra_args : bool is_test=false
   data_format_tensors : input, out_grad
 
+- op : fused_conv2d
+  extra_args : float fuse_alpha = 0.0, float fuse_beta = 0.0, float scale_in=1.0, float scale_out=1.0, float scale_in_eltwise=1.0, float[] scale_weights={1.0f}
+  data_format_tensors : input
+
 - op : lrn
   extra_args : bool is_test=false
   data_format_tensors : x
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -1307,6 +1307,12 @@
     reserve_space: ReserveSpace
 
 - op : fused_conv2d
+  inputs :
+    {input : Input, filter : Filter, bias : Bias, residual_param : ResidualData}
+  outputs :
+    {output : Output}
+  attrs :
+    {scale_in : Scale_in, scale_out : Scale_out, scale_in_eltwise : Scale_in_eltwise, scale_weights : Scale_weights}
   extra :
     attrs : [bool use_cudnn = false, float fuse_alpha = 0.0f, float fuse_beta = 0.0f, float Scale_in = 1.0f,
              float Scale_out = 1.0f, float Scale_in_eltwise = 1.0f, 'float[] Scale_weights = {1.0f}', bool use_mkldnn = true, str mkldnn_data_type = "float32"]
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -567,6 +567,7 @@
     func : ConvInferMeta
   kernel :
     func : conv2d
+    data_type : input
   backward : conv2d_grad
 
 - op : conv3d
@@ -576,6 +577,7 @@
     func : Conv3DInferMeta
   kernel :
     func : conv3d
+    data_type : input
   backward : conv3d_grad
 
 - op : conv3d_transpose
@@ -713,6 +715,7 @@
     func : DepthwiseConvInferMeta
   kernel :
     func : depthwise_conv2d
+    data_type : input
   backward : depthwise_conv2d_grad
 
 - op : det
diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -827,7 +827,7 @@ void FusedConvInferMeta(const MetaTensor& input,
                         bool fuse_residual_conn,
                         bool force_fp32_output,
                         MetaTensor* out,
-                        MetaConfig config);
+                        MetaConfig config = MetaConfig());
 
 void MoeInferMeta(const MetaTensor& x,
                   const MetaTensor& gate,
diff --git a/paddle/phi/kernels/onednn/conv_function.h b/paddle/phi/kernels/onednn/conv_function.h
@@ -60,6 +60,10 @@ static dnnl::memory::data_type GetDstType(
           NAME, ::paddle::DataType::FLOAT32, float, __VA_ARGS__)          \
       PD_PRIVATE_CASE_TYPE(                                               \
           NAME, ::paddle::DataType::INT8, int8_t, __VA_ARGS__)            \
+      PD_PRIVATE_CASE_TYPE(NAME,                                          \
+                           ::paddle::DataType::BFLOAT16,                  \
+                           ::phi::dtype::bfloat16,                        \
+                           __VA_ARGS__)                                   \
       default:                                                            \
         PD_THROW("function " #NAME " is not implemented for data type `", \
                  __dtype__,                                               \
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -1065,7 +1065,7 @@ set_tests_properties(test_sigmoid_cross_entropy_with_logits_op
                      PROPERTIES TIMEOUT 120)
 set_tests_properties(test_imperative_optimizer_v2 PROPERTIES TIMEOUT 150)
 set_tests_properties(test_partial_sum_op PROPERTIES TIMEOUT 120)
-set_tests_properties(test_cond PROPERTIES TIMEOUT 120)
+set_tests_properties(test_cond PROPERTIES TIMEOUT 240)
 set_tests_properties(test_sgd_op PROPERTIES TIMEOUT 250)
 set_tests_properties(test_parallel_executor_seresnext_base_gpu
                      PROPERTIES TIMEOUT 120)
diff --git a/test/legacy_test/op_test.py b/test/legacy_test/op_test.py
@@ -431,6 +431,7 @@ def setUpClass(cls):
         cls.check_prim = False
         cls.check_prim_pir = False
         cls._check_cinn = False
+        cls.check_pir_onednn = False
 
         np.random.seed(123)
         random.seed(124)
diff --git a/test/legacy_test/test_conv2d_op.py b/test/legacy_test/test_conv2d_op.py
@@ -499,7 +499,10 @@ def test_check_output(self):
         place = core.CUDAPlace(0) if self.has_cuda() else core.CPUPlace()
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         self.check_output_with_place(
-            place, atol=1e-5, check_dygraph=(not self.use_mkldnn)
+            place,
+            atol=1e-5,
+            check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad(self):
@@ -515,6 +518,7 @@ def test_check_grad(self):
             'Output',
             max_relative_error=0.02,
             check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_filter(self):
@@ -531,6 +535,7 @@ def test_check_grad_no_filter(self):
             max_relative_error=0.02,
             no_grad_set={'Filter'},
             check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_input(self):
@@ -546,6 +551,7 @@ def test_check_grad_no_input(self):
             'Output',
             no_grad_set={'Input'},
             check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def init_test_case(self):
@@ -824,7 +830,10 @@ def test_check_output(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         place = core.CUDAPlace(0) if self.has_cuda() else core.CPUPlace()
         self.check_output_with_place(
-            place, atol=1e-5, check_dygraph=(not self.use_mkldnn)
+            place,
+            atol=1e-5,
+            check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad(self):
@@ -838,6 +847,7 @@ def test_check_grad(self):
             'Output',
             max_relative_error=0.02,
             check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_filter(self):
@@ -852,6 +862,7 @@ def test_check_grad_no_filter(self):
             max_relative_error=0.02,
             no_grad_set={'Filter'},
             check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_input(self):
@@ -865,6 +876,7 @@ def test_check_grad_no_input(self):
             'Output',
             no_grad_set={'Input'},
             check_dygraph=(not self.use_mkldnn),
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def init_test_case(self):
diff --git a/test/legacy_test/test_elementwise_add_op.py b/test/legacy_test/test_elementwise_add_op.py
@@ -30,7 +30,6 @@
 class TestElementwiseAddOp(OpTest):
     def init_kernel_type(self):
         self.use_mkldnn = False
-        self.check_pir_onednn = False
 
     def setUp(self):
         self.op_type = "elementwise_add"
diff --git a/test/mkldnn/test_conv2d_bf16_mkldnn_op.py b/test/mkldnn/test_conv2d_bf16_mkldnn_op.py
@@ -50,6 +50,7 @@ def setUp(self):
         self.init_data_type()
         self.init_force_fp32_output()
         self.init_infer_or_train()
+        self.check_pir_onednn = True
 
         self.conv2d_param = {
             'stride': self.stride,
@@ -117,7 +118,9 @@ def setUp(self):
         self.init_additional_attrs()
 
     def test_check_output(self):
-        self.check_output_with_place(core.CPUPlace())
+        self.check_output_with_place(
+            core.CPUPlace(), check_pir_onednn=self.check_pir_onednn
+        )
 
     def test_check_grad(self):
         pass
@@ -186,6 +189,7 @@ def test_check_grad(self):
             "Output",
             user_defined_grads=[dx, dweights],
             user_defined_grad_outputs=[convert_float_to_uint16(dout)],
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_filter(self):
@@ -202,6 +206,7 @@ def test_check_grad_no_filter(self):
             {'Filter'},
             user_defined_grads=[dx],
             user_defined_grad_outputs=[convert_float_to_uint16(dout)],
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_input(self):
@@ -218,6 +223,7 @@ def test_check_grad_no_input(self):
             {'Input'},
             user_defined_grads=[dweights],
             user_defined_grad_outputs=[convert_float_to_uint16(dout)],
+            check_pir_onednn=self.check_pir_onednn,
         )
 
 
diff --git a/test/mkldnn/test_conv2d_int8_mkldnn_op.py b/test/mkldnn/test_conv2d_int8_mkldnn_op.py
@@ -48,6 +48,7 @@ def setUp(self):
         self.init_fuse_activation()
         self.init_fuse_residual()
         self.init_data_type()
+        self.check_pir_onednn = True
 
         conv2d_param = {
             'stride': self.stride,
@@ -184,7 +185,10 @@ def test_check_output(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         # the atol for integer tests should be 1
         self.check_output_with_place(
-            core.CPUPlace(), atol=1, check_dygraph=False
+            core.CPUPlace(),
+            atol=1,
+            check_dygraph=False,
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad(self):
diff --git a/test/mkldnn/test_conv2d_mkldnn_op.py b/test/mkldnn/test_conv2d_mkldnn_op.py
@@ -17,9 +17,6 @@
 import numpy as np
 from op_test import OpTest, skip_check_grad_ci
 from test_conv2d_op import TestConv2DOp, TestConv2DOp_v2
-from utils import compare_legacy_with_pt
-
-from paddle.base import core
 
 
 def conv2d_bias_naive(out, bias):
@@ -64,6 +61,7 @@ def setUp(self):
         self.input_residual_size = None
 
         TestConv2DOp.setUp(self)
+        self.check_pir_onednn = True
 
         output = self.outputs['Output']
 
@@ -144,6 +142,7 @@ def setUp(self):
         self.input_residual_size = None
 
         TestConv2DOp.setUp(self)
+        self.check_pir_onednn = True
 
         output = self.outputs['Output']
 
@@ -195,14 +194,6 @@ def setUp(self):
 
         self.outputs['Output'] = output
 
-    @compare_legacy_with_pt
-    def test_check_output(self):
-        place = core.CUDAPlace(0) if self.has_cuda() else core.CPUPlace()
-        # TODO(wangzhongpu): support mkldnn op in dygraph mode
-        self.check_output_with_place(
-            place, atol=1e-5, check_dygraph=(not self.use_mkldnn)
-        )
-
 
 @skip_check_grad_ci(
     reason="Fusion is for inference only, check_grad is not required."
