[PIR] pir onednn support bn (#61170)

wanghuancoder · web-flow · commit a3dda38a5871 · 2024-01-29T12:20:08.000+08:00
* pir onednn support bn
diff --git a/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.cc b/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.cc
@@ -282,6 +282,7 @@ OneDNNPhiKernelInstruction::OneDNNPhiKernelInstruction(
 
   // Step2: build layout_transform information
   if (op_attributes.count("data_format_tensors")) {
+    VLOG(6) << "data_format_tensors is not empty";
     if (op_attributes.count("data_format")) {
       auto data_layout = op_attributes.at("data_format")
                              .dyn_cast<pir::StrAttribute>()
@@ -290,6 +291,7 @@ OneDNNPhiKernelInstruction::OneDNNPhiKernelInstruction(
     } else {
       input_layout_ = phi::OneDNNContext::tls().get_cur_paddle_data_layout();
     }
+    VLOG(6) << "input_layout = " << input_layout_;
 
     std::vector<pir::Attribute> data_format_tensors_attr =
         op->attributes()
@@ -298,11 +300,13 @@ OneDNNPhiKernelInstruction::OneDNNPhiKernelInstruction(
             .AsVector();
 
     for (auto& attr : data_format_tensors_attr) {
-      auto pair =
-          kernel_context_.InputRangeAt(yaml_info_parser.InputName2Id().at(
-              attr.dyn_cast<pir::StrAttribute>().AsString()));
+      auto data_format_tensor = attr.dyn_cast<pir::StrAttribute>().AsString();
+      auto pair = kernel_context_.InputRangeAt(
+          yaml_info_parser.InputName2Id().at(data_format_tensor));
+      VLOG(6) << "data_format_tensor = " << data_format_tensor;
       for (int i = pair.first; i < pair.second; ++i) {
         data_format_tensors_.insert(i);
+        VLOG(6) << data_format_tensor << " index = " << i;
       }
     }
   }
@@ -366,6 +370,10 @@ void OneDNNPhiKernelInstruction::Run() {
     if (input == nullptr) {
       continue;
     }
+    if (!input->initialized()) {
+      continue;
+    }
+    VLOG(6) << "input[" << i << "].layout() = " << input->layout();
     if (input->layout() != phi::DataLayout::ONEDNN) {
       phi::DataLayout from_layout = input->layout();
 
@@ -375,6 +383,7 @@ void OneDNNPhiKernelInstruction::Run() {
           input_layout_ != phi::DataLayout::kAnyLayout) {
         from_layout = input_layout_;
       }
+      VLOG(6) << "from_layout = " << from_layout;
 
       auto transed_tensor = const_cast<phi::DenseTensor*>(input);
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -11,6 +11,14 @@
 - op : add_triple_grad
   extra_args : str mkldnn_data_type="float32", bool use_quantizer=false, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0
 
+- op : batch_norm
+  extra_args : bool fuse_with_relu=false
+  data_format_tensors : x
+
+- op : batch_norm_grad
+  extra_args : bool fuse_with_relu=false
+  data_format_tensors : x, out_grad
+
 - op : conv2d
   extra_args : bool is_test=false
   data_format_tensors : input
@@ -54,11 +62,5 @@
 #     arg_name: cur_paddle_data_layout
 #     tensors: x, y
 
-# - op : batch_norm
-#   extra_args : bool fuse_with_relu=false
-#   layout_transform :
-#     arg_name: data_layout
-#     tensors: x
-
 # - op : prelu
 #   extra_args : bool is_test=false, str mkldnn_data_type="float32"
diff --git a/paddle/phi/kernels/cpu/onednn_to_paddle_layout_kernel.cc b/paddle/phi/kernels/cpu/onednn_to_paddle_layout_kernel.cc
@@ -55,13 +55,6 @@ void OneDNN2PaddleLayout(const Context& dev_ctx,
   VLOG(10) << " x: " << print_tensor_meta(x);
   VLOG(10) << " out: " << print_tensor_meta(*out) << " " << out;
 
-  if (src_layout != DataLayout::ONEDNN) {
-    out->ShareDataWith(x);
-    out->ShareInplaceVersionCounterWith(x);
-    out->set_layout(static_cast<DataLayout>(dst_layout));
-    return;
-  }
-
   DataLayout tmp_layout = static_cast<DataLayout>(dst_layout);
 
   if (tmp_layout == DataLayout::ANY) {
@@ -70,6 +63,13 @@ void OneDNN2PaddleLayout(const Context& dev_ctx,
 
   VLOG(4) << "src_layout: " << src_layout << ", tmp_layout: " << tmp_layout;
 
+  if (src_layout != DataLayout::ONEDNN || !x.storage_properties_initialized()) {
+    out->ShareDataWith(x);
+    out->ShareInplaceVersionCounterWith(x);
+    out->set_layout(static_cast<DataLayout>(tmp_layout));
+    return;
+  }
+
   // NOTE(zhiqiu): to handle the special case in ApplyDataTransform() in
   // data_transfer.cc
   if (!x.IsInitialized() && src_layout == DataLayout::ONEDNN &&
diff --git a/test/legacy_test/test_batch_norm_op.py b/test/legacy_test/test_batch_norm_op.py
@@ -462,9 +462,17 @@ def test_check_output(self):
         for place in places:
             for data_format in ["NCHW", "NHWC"]:
                 self.check_with_place(
-                    place, data_format, self.dtype, [2, 3, 4, 5]
+                    place,
+                    data_format,
+                    self.dtype,
+                    [2, 3, 4, 5],
+                )
+                self.check_with_place(
+                    place,
+                    data_format,
+                    self.dtype,
+                    [2, 3],
                 )
-                self.check_with_place(place, data_format, self.dtype, [2, 3])
                 self.check_with_place_without_scale_and_bias(
                     place, data_format, self.dtype, [2, 3, 4, 5]
                 )
@@ -493,9 +501,17 @@ def test_check_output(self):
             # for data_format in ["NCHW", "NHWC"]:
             for data_format in ["NCHW"]:
                 self.check_with_place(
-                    place, data_format, self.dtype, [2, 3, 4, 5]
+                    place,
+                    data_format,
+                    self.dtype,
+                    [2, 3, 4, 5],
+                )
+                self.check_with_place(
+                    place,
+                    data_format,
+                    self.dtype,
+                    [2, 3],
                 )
-                self.check_with_place(place, data_format, self.dtype, [2, 3])
 
 
 @unittest.skipIf(
@@ -516,9 +532,17 @@ def test_check_output(self):
             # for data_format in ["NCHW", "NHWC"]:
             for data_format in ["NCHW"]:
                 self.check_with_place(
-                    place, data_format, self.dtype, [2, 3, 4, 5]
+                    place,
+                    data_format,
+                    self.dtype,
+                    [2, 3, 4, 5],
+                )
+                self.check_with_place(
+                    place,
+                    data_format,
+                    self.dtype,
+                    [2, 3],
                 )
-                self.check_with_place(place, data_format, self.dtype, [2, 3])
 
 
 class TestBatchNormOpTraining(unittest.TestCase):
diff --git a/test/mkldnn/test_batch_norm_mkldnn_op.py b/test/mkldnn/test_batch_norm_mkldnn_op.py
@@ -16,7 +16,7 @@
 
 import numpy as np
 from mkldnn_op_test import check_if_mkldnn_batchnorm_primitives_exist_in_bwd
-from op_test import _set_use_system_allocator
+from op_test import _set_use_system_allocator, pir_executor_guard
 from test_batch_norm_op import (
     TestBatchNormOpInference,
     TestBatchNormOpTraining,
@@ -72,6 +72,11 @@ def ref_forward_backward(
             bias_grad,
         )
 
+    def test_forward_backward(self):
+        super().test_forward_backward()
+        with pir_executor_guard():
+            super().test_forward_backward()
+
 
 class TestMKLDNNBatchNormOpTraining_NHWC(TestMKLDNNBatchNormOpTraining):
     def init_kernel_type(self):
@@ -139,6 +144,11 @@ def test_check_output(self):
         self.check_with_place_without_scale_and_bias(
             place, data_format, self.dtype, [2, 3, 4, 5]
         )
+        with pir_executor_guard():
+            self.check_with_place(place, data_format, self.dtype, [2, 3, 4, 5])
+            self.check_with_place_without_scale_and_bias(
+                place, data_format, self.dtype, [2, 3, 4, 5]
+            )
 
 
 class TestMKLDNNBatchNormOpInference_NHWC(TestMKLDNNBatchNormOpInference):
@@ -160,6 +170,8 @@ def test_check_output(self):
         place = core.CPUPlace()
         data_format = "NCHW"
         self.check_with_place(place, data_format, self.dtype, [2, 3, 4, 5])
+        with pir_executor_guard():
+            self.check_with_place(place, data_format, self.dtype, [2, 3, 4, 5])
 
 
 if __name__ == '__main__':
