[PIR] Fix layout transformation (#66977)

LLee233 · web-flow · commit 50c237ee0e78 · 2024-08-07T14:22:42.000+08:00
* fix layout transformation

* fix copyright

* fix copyright

* delete useless continue
diff --git a/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.cc b/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2024 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.
@@ -429,18 +429,19 @@ void OneDNNPhiKernelInstruction::Run() {
       tmp_holders.emplace_back(std::make_shared<phi::DenseTensor>(*input));
       auto transed_tensor = tmp_holders.back().get();
 
-      std::set<std::string> elementwise_kernels = {
-          "add", "subtract", "multiply", "divide"};
-      if (elementwise_kernels.count(kernel_name_)) {
+      std::set<std::string> elementwise_kernels = {"onednn_op.add",
+                                                   "onednn_op.subtract",
+                                                   "onednn_op.multiply",
+                                                   "onednn_op.divide"};
+
+      if (elementwise_kernels.count(phi_op_name_)) {
         if (phi::OneDNNContext::tls().get_cur_paddle_data_layout() ==
                 phi::DataLayout::kNHWC &&
             !(kernel_key_.dtype() == phi::DataType::COMPLEX64 ||
               kernel_key_.dtype() == phi::DataType::COMPLEX128)) {
+          from_layout = phi::DataLayout::kNHWC;
           phi::funcs::MatchShapeToLayout(
               transed_tensor, from_layout, phi::DataLayout::ONEDNN);
-          from_layout = phi::DataLayout::kNHWC;
-        } else {
-          continue;
         }
       } else {
         //  Handle 'layout_transform' in
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -1,4 +1,4 @@
-/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+/* Copyright (c) 2024 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.
@@ -1647,9 +1647,8 @@ void DotInferMeta(const MetaTensor& x, const MetaTensor& y, MetaTensor* out) {
 
 void ElementwiseInferMeta(const MetaTensor& x,
                           const MetaTensor& y,
-                          MetaTensor* out,
-                          MetaConfig config) {
-  return ElementwiseRawInferMeta(x, y, -1, out, config);
+                          MetaTensor* out) {
+  return ElementwiseRawInferMeta(x, y, -1, out);
 }
 
 void BitwiseShiftInferMeta(const MetaTensor& x,
diff --git a/paddle/phi/infermeta/binary.h b/paddle/phi/infermeta/binary.h
@@ -1,4 +1,4 @@
-/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+/* Copyright (c) 2024 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.
@@ -324,8 +324,7 @@ void DropoutNdInferMeta(const MetaTensor& x,
 
 TEST_API void ElementwiseInferMeta(const MetaTensor& x,
                                    const MetaTensor& y,
-                                   MetaTensor* out,
-                                   MetaConfig config = MetaConfig());
+                                   MetaTensor* out);
 
 void ElementwiseRawInferMeta(const MetaTensor& x_meta,
                              const MetaTensor& y_meta,
diff --git a/test/ir/pir/fused_pass/onednn/test_onednn_placement.py b/test/ir/pir/fused_pass/onednn/test_onednn_placement.py
@@ -183,5 +183,70 @@ def test_check_output(self):
         self.check_pass_correct()
 
 
+# This case is for testing layout transformation
+class TestConv2dAddPlacmentPass(PassTest):
+    def is_program_valid(self, program=None):
+        return True
+
+    def build_ir_program(self):
+        with paddle.pir_utils.IrGuard():
+            main_prog = paddle.static.Program()
+            start_prog = paddle.static.Program()
+            with paddle.pir.core.program_guard(main_prog, start_prog):
+                x = paddle.static.data(
+                    name='x', shape=[1, 4, 4, 5], dtype='float32'
+                )
+                bias_attr = paddle.ParamAttr(
+                    learning_rate=0.0,
+                    initializer=paddle.nn.initializer.Normal(mean=0.0, std=2.0),
+                )
+                bias = paddle.static.create_parameter(
+                    shape=[1, 1, 1, 3],
+                    dtype='float32',
+                    attr=bias_attr,
+                    is_bias=False,
+                )
+                w_attr = paddle.ParamAttr(
+                    learning_rate=0.0,
+                    initializer=paddle.nn.initializer.Normal(mean=0.0, std=2.0),
+                )
+                conv2d = paddle.nn.Conv2D(
+                    in_channels=5,
+                    out_channels=3,
+                    kernel_size=[1, 1],
+                    groups=1,
+                    stride=[1, 1],
+                    padding=[1, 1, 1, 1],
+                    dilation=[1, 1],
+                    data_format='NHWC',
+                    bias_attr=False,
+                    weight_attr=w_attr,
+                )
+
+                out = paddle.add(conv2d(x), bias)
+
+                out = paddle.assign(out)
+                self.pass_attr_list = [{'onednn_placement_pass': {}}]
+                self.feeds = {
+                    "x": np.random.random((1, 4, 4, 5)).astype("float32"),
+                    "bias": np.random.random((1, 1, 1, 3)).astype("float32"),
+                }
+                self.fetch_list = [out]
+                self.valid_op_map = {
+                    "pd_op.conv2d": 0,
+                    "pd_op.add": 0,
+                }
+                return [main_prog, start_prog]
+
+    def sample_program(self):
+        yield self.build_ir_program(), False
+
+    def setUp(self):
+        self.places.append(paddle.CPUPlace())
+
+    def test_check_output(self):
+        self.check_pass_correct()
+
+
 if __name__ == "__main__":
     unittest.main()
