[Inference] Add remove_redundant_transpose_pass (#63357)

Sunny-bot1 · web-flow · commit 3dd6a795ace6 · 2024-04-15T10:53:56.000+08:00
* extract RemoveRedundantTransposePattern from identity_op_clean_pass and use it as remove_redundant_transpose_pass

* fix
diff --git a/paddle/fluid/inference/api/paddle_pass_builder.cc b/paddle/fluid/inference/api/paddle_pass_builder.cc
@@ -614,7 +614,8 @@ const std::vector<std::string> kPirGpuPasses{
     "fc_elementwise_layernorm_fuse_pass",
     "matmul_scale_fuse_pass",
     "matmul_transpose_fuse_pass",
-    "transpose_flatten_concat_fuse_pass"};
+    "transpose_flatten_concat_fuse_pass",
+    "remove_redundant_transpose_pass"};
 
 const std::vector<std::string> kPirXpuPasses{// Functional pass
                                              "map_op_to_another_pass",
diff --git a/paddle/fluid/pir/transforms/general/identity_op_clean_pass.cc b/paddle/fluid/pir/transforms/general/identity_op_clean_pass.cc
@@ -245,39 +245,6 @@ class ReplaceDropoutWithScalePattern : public paddle::drr::DrrPatternBase {
   }
 };
 
-class RemoveRedundantTransposePattern : public paddle::drr::DrrPatternBase {
- public:
-  std::string name() const override {
-    return "RemoveRedundantTransposePattern";
-  }
-
-  void operator()(paddle::drr::DrrPatternContext *ctx) const override {
-    paddle::drr::SourcePattern pat = ctx->SourcePattern();
-    const auto &transpose1 =
-        pat.Op("pd_op.transpose", {{"perm", pat.Attr("perm_1")}});
-    const auto &transpose2 =
-        pat.Op("pd_op.transpose", {{"perm", pat.Attr("perm_2")}});
-
-    pat.Tensor("ret") = transpose2(transpose1(pat.Tensor("arg_transpose")));
-
-    paddle::drr::ResultPattern res = pat.ResultPattern();
-    const auto &new_perm_attr = res.ComputeAttr(
-        [](const paddle::drr::MatchContext &match_ctx) -> std::vector<int> {
-          const auto &perm1 = match_ctx.Attr<std::vector<int>>("perm_1");
-          const auto &perm2 = match_ctx.Attr<std::vector<int>>("perm_2");
-          std::vector<int> new_perm;
-          for (int v : perm2) {
-            new_perm.emplace_back(perm1[v]);
-          }
-          return new_perm;
-        });
-    const auto &transpose_continuous =
-        res.Op("pd_op.transpose", {{"perm", new_perm_attr}});
-
-    res.Tensor("ret") = transpose_continuous(res.Tensor("arg_transpose"));
-  }
-};
-
 class IdentityOpCleanPass : public pir::PatternRewritePass {
  public:
   IdentityOpCleanPass()
@@ -292,7 +259,6 @@ class IdentityOpCleanPass : public pir::PatternRewritePass {
     ps.Add(paddle::drr::Create<RemoveRedundantCastPattern>(context));
     ps.Add(paddle::drr::Create<DeleteDropoutOpPattern>(context));
     ps.Add(paddle::drr::Create<ReplaceDropoutWithScalePattern>(context));
-    ps.Add(paddle::drr::Create<RemoveRedundantTransposePattern>(context));
     return ps;
   }
 };
diff --git a/paddle/fluid/pir/transforms/general/remove_redundant_transpose_pass.cc b/paddle/fluid/pir/transforms/general/remove_redundant_transpose_pass.cc
@@ -0,0 +1,79 @@
+// 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.
+// You may obtain a copy of the License at
+//
+//     https://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/fluid/pir/transforms/general/remove_redundant_transpose_pass.h"
+
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/drr/include/drr_pattern_base.h"
+
+#include "paddle/pir/include/pass/pass.h"
+#include "paddle/pir/include/pass/pass_registry.h"
+
+namespace {
+
+class RemoveRedundantTransposePattern : public paddle::drr::DrrPatternBase {
+ public:
+  std::string name() const override {
+    return "RemoveRedundantTransposePattern";
+  }
+
+  void operator()(paddle::drr::DrrPatternContext *ctx) const override {
+    paddle::drr::SourcePattern pat = ctx->SourcePattern();
+    const auto &transpose1 =
+        pat.Op("pd_op.transpose", {{"perm", pat.Attr("perm_1")}});
+    const auto &transpose2 =
+        pat.Op("pd_op.transpose", {{"perm", pat.Attr("perm_2")}});
+
+    pat.Tensor("ret") = transpose2(transpose1(pat.Tensor("arg_transpose")));
+
+    paddle::drr::ResultPattern res = pat.ResultPattern();
+    const auto &new_perm_attr = res.ComputeAttr(
+        [](const paddle::drr::MatchContext &match_ctx) -> std::vector<int> {
+          const auto &perm1 = match_ctx.Attr<std::vector<int>>("perm_1");
+          const auto &perm2 = match_ctx.Attr<std::vector<int>>("perm_2");
+          std::vector<int> new_perm;
+          for (int v : perm2) {
+            new_perm.emplace_back(perm1[v]);
+          }
+          return new_perm;
+        });
+    const auto &transpose_continuous =
+        res.Op("pd_op.transpose", {{"perm", new_perm_attr}});
+
+    res.Tensor("ret") = transpose_continuous(res.Tensor("arg_transpose"));
+  }
+};
+
+class RemoveRedundantTransposePass : public pir::PatternRewritePass {
+ public:
+  RemoveRedundantTransposePass()
+      : pir::PatternRewritePass("remove_redundant_transpose_pass", 2) {}
+
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add(paddle::drr::Create<RemoveRedundantTransposePattern>(context));
+    return ps;
+  }
+};
+
+}  // namespace
+
+namespace pir {
+
+std::unique_ptr<Pass> CreateRemoveRedundantTransposePass() {
+  return std::make_unique<RemoveRedundantTransposePass>();
+}
+}  // namespace pir
+
+REGISTER_IR_PASS(remove_redundant_transpose_pass, RemoveRedundantTransposePass);
diff --git a/paddle/fluid/pir/transforms/general/remove_redundant_transpose_pass.h b/paddle/fluid/pir/transforms/general/remove_redundant_transpose_pass.h
@@ -0,0 +1,26 @@
+// 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.
+// You may obtain a copy of the License at
+//
+//     https://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.
+
+#pragma once
+
+#include <memory>
+#include "paddle/pir/include/core/dll_decl.h"
+
+namespace pir {
+
+class Pass;
+
+IR_API std::unique_ptr<Pass> CreateRemoveRedundantTransposePass();
+
+}  // namespace pir
diff --git a/paddle/fluid/pir/transforms/passes.h b/paddle/fluid/pir/transforms/passes.h
@@ -39,6 +39,7 @@ USE_PIR_PASS(embedding_eltwise_layernorm_fuse_pass);
 USE_PIR_PASS(add_norm_fuse_pass);
 USE_PIR_PASS(fused_dot_product_attention_pass);
 USE_PIR_PASS(fused_flash_attn_pass);
+USE_PIR_PASS(remove_redundant_transpose_pass);
 
 #ifdef PADDLE_WITH_DNNL
 USE_PIR_PASS(depthwise_conv_onednn_pass);
diff --git a/test/ir/pir/test_pir_identity_op_clean_pass.py b/test/ir/pir/test_pir_identity_op_clean_pass.py
@@ -209,41 +209,6 @@ def setUp(self):
             self.places.append(paddle.CUDAPlace(0))
 
 
-class TestRemoveRedundantTransposePattern(PassTest):
-    def is_program_valid(self, program=None):
-        return True
-
-    def sample_program(self):
-        for perm1_shape in [[1, 2, 0]]:
-            for perm2_shape in [[0, 2, 1]]:
-                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=[2, 3, 4], dtype="float32"
-                        )
-                        out = paddle.transpose(
-                            paddle.transpose(x, perm1_shape), perm2_shape
-                        )
-                        out = paddle.assign(out)
-                        self.pass_list = ['identity_op_clean_pass']
-                        self.feeds = {
-                            "x": np.random.random((2, 3, 4)).astype("float32")
-                        }
-                        self.fetch_list = [out]
-                        self.valid_op_map = {"pd_op.transpose": 1}
-                        yield [main_prog, start_prog], False
-
-    def test_check_output(self):
-        self.check_pass_correct()
-
-    def setUp(self):
-        self.places.append(paddle.CPUPlace())
-        if core.is_compiled_with_cuda():
-            self.places.append(paddle.CUDAPlace(0))
-
-
 class TestDeleteDropoutOpPatternPattern(PassTest):
     def is_program_valid(self, program=None):
         return True
diff --git a/test/ir/pir/test_remove_redundant_transpose_pass.py b/test/ir/pir/test_remove_redundant_transpose_pass.py
@@ -0,0 +1,62 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     https://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.
+
+import unittest
+
+import numpy as np
+from fused_pass.pass_test import PassTest
+
+import paddle
+from paddle.base import core
+
+paddle.enable_static()
+
+
+class TestRemoveRedundantTransposePattern(PassTest):
+    def is_program_valid(self, program=None):
+        return True
+
+    def sample_program(self):
+        for perm1_shape in [[1, 2, 0]]:
+            for perm2_shape in [[0, 2, 1]]:
+                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=[2, 3, 4], dtype="float32"
+                        )
+                        out = paddle.transpose(
+                            paddle.transpose(x, perm1_shape), perm2_shape
+                        )
+                        out = paddle.assign(out)
+                        self.pass_list = ['remove_redundant_transpose_pass']
+                        self.feeds = {
+                            "x": np.random.random((2, 3, 4)).astype("float32")
+                        }
+                        self.fetch_list = [out]
+                        self.valid_op_map = {"pd_op.transpose": 1}
+                        yield [main_prog, start_prog], False
+
+    def test_check_output(self):
+        self.check_pass_correct()
+
+    def setUp(self):
+        self.places.append(paddle.CPUPlace())
+        if core.is_compiled_with_cuda():
+            self.places.append(paddle.CUDAPlace(0))
+
+
+if __name__ == "__main__":
+    unittest.main()
