[ pass_enhance ]transpose_flatten_concat_fuse_pass (#33744)

PaddlePaddle · Jun 29, 2021 · 2e97faf · 2e97faf
1 parent 07eeb36
commit 2e97faf
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
diff --git a/paddle/fluid/framework/ir/transpose_flatten_concat_fuse_pass.cc b/paddle/fluid/framework/ir/transpose_flatten_concat_fuse_pass.cc
@@ -19,7 +19,50 @@ namespace paddle {
 namespace framework {
 namespace ir {
 
-void RunTransposeFlattenConcatFuse(ir::Graph *graph, int times) {
+TransposeFlattenConcatFusePass::TransposeFlattenConcatFusePass() {
+  AddOpCompat(OpCompat("transpose2"))
+      .AddInput("X")
+      .IsTensor()
+      .End()
+      .AddOutput("Out")
+      .IsTensor()
+      .End()
+      .AddOutput("XShape")
+      .IsTensor()
+      .End()
+      .AddAttr("axis")
+      .IsType<std::vector<int>>()
+      .End();
+  AddOpCompat(OpCompat("flatten2"))
+      .AddInput("X")
+      .IsTensor()
+      .End()
+      .AddOutput("Out")
+      .IsTensor()
+      .End()
+      .AddOutput("XShape")
+      .IsTensor()
+      .End()
+      .AddAttr("axis")
+      .IsNumGE(0)
+      .End();
+  AddOpCompat(OpCompat("concat"))
+      .AddInput("X")  // Input("X"): vector<tensors>
+      .End()
+      .AddInput("AxisTensor")
+      .IsTensor()
+      .IsOptional()
+      .End()
+      .AddOutput("Out")
+      .IsTensor()
+      .End()
+      .AddAttr("axis")
+      .IsIntIn({0, 1})
+      .End();
+}
+
+void TransposeFlattenConcatFusePass::RunTransposeFlattenConcatFuse(
+    ir::Graph *graph, int times) const {
   const std::string pattern_name =
       "transpose_flatten" + std::to_string(times) + "_concat_fuse";
 
@@ -37,6 +80,10 @@ void RunTransposeFlattenConcatFuse(ir::Graph *graph, int times) {
 
   auto handler = [&](const GraphPatternDetector::subgraph_t &subgraph,
                      Graph *g) {
+    if (!IsCompat(subgraph, g)) {
+      LOG(WARNING) << "Pass in op compat failed.";
+      return;
+    }
     const int kNumFields = 5;
     const int kTransOffset = 1;
     const int kTransOutOffset = 2;

diff --git a/paddle/fluid/framework/ir/transpose_flatten_concat_fuse_pass.h b/paddle/fluid/framework/ir/transpose_flatten_concat_fuse_pass.h
@@ -16,7 +16,6 @@
 #include <memory>
 
 #include "paddle/fluid/framework/ir/fuse_pass_base.h"
-#include "paddle/fluid/framework/ir/graph_pattern_detector.h"
 
 namespace paddle {
 namespace framework {
@@ -28,10 +27,14 @@ namespace ir {
 // structure.
 class TransposeFlattenConcatFusePass : public FusePassBase {
  public:
+  TransposeFlattenConcatFusePass();
   virtual ~TransposeFlattenConcatFusePass() {}
 
  protected:
   void ApplyImpl(ir::Graph* graph) const override;
+
+ private:
+  void RunTransposeFlattenConcatFuse(ir::Graph* graph, int times) const;
 };
 
 }  // namespace ir