[CustomPass] add support for outputting the intermediate variables (#55728)

ronny1996 · web-flow · commit da2589649e12 · 2023-07-27T14:16:57.000+08:00
* add support for outputting the intermediate variables

* fix fuse_rresnet_unit
diff --git a/paddle/fluid/framework/ir/generate_pass.cc b/paddle/fluid/framework/ir/generate_pass.cc
@@ -220,6 +220,16 @@ void InitGeneratePattern(const proto::PassDesc& pass_desc, PDPattern* pattern) {
       });
     }
   }
+
+  // The output of the pattern must be marked as AsOutput.
+  for (const proto::PassDesc::VarMap& var_map : pass_desc.var_maps()) {
+    PDNode* var_pdnode = pattern->RetrieveNode(var_map.pattern_var());
+    PADDLE_ENFORCE_NOT_NULL(
+        var_pdnode,
+        platform::errors::NotFound("Not found the var %s in the pattern.",
+                                   var_map.pattern_var()));
+    var_pdnode->AsOutput();
+  }
 }
 
 // There are some duplicate patterns.
@@ -314,10 +324,10 @@ GraphPatternDetector::handle_t GetGenerateRewrite(
         }
         // `var_node_maps` record the mapping of variable to the pattern
         // subgraph.
-        std::map<std::string, Node*> var_node_maps;
+        std::map<std::string, std::vector<Node*>> var_node_maps;
         for (const proto::PassDesc::VarMap& var_map : pass_desc.var_maps()) {
           Node* node = subgraph.at(pattern.RetrieveNode(var_map.pattern_var()));
-          var_node_maps.insert({var_map.replace_var(), node});
+          var_node_maps[var_map.replace_var()].emplace_back(node);
         }
         // Traverse all operators to create subgraph.
         for (int index = 0; index < pass_desc.replace_size(); ++index) {
@@ -330,17 +340,13 @@ GraphPatternDetector::handle_t GetGenerateRewrite(
             std::vector<std::string> arguments;
             for (const std::string& argument : var.arguments()) {
               // The input may be mapped on the operator of pattern subgraph.
-              Node* node = nullptr;
-              auto iter = var_node_maps.find(argument);
-              if (var_node_maps.end() == iter) {
+              if (var_node_maps[argument].size() == 0) {
                 VarDesc var_desc(patterns::UniqueKey(argument));
-                node = graph->CreateVarNode(&var_desc);
-                var_node_maps.insert({argument, node});
-              } else {
-                node = iter->second;
+                var_node_maps[argument].emplace_back(
+                    graph->CreateVarNode(&var_desc));
               }
-              in_nodes.push_back(node);
-              arguments.push_back(node->Name());
+              in_nodes.push_back(var_node_maps[argument][0]);
+              arguments.push_back(var_node_maps[argument][0]->Name());
             }
             op_desc.SetInput(var.parameter(), arguments);
           }
@@ -349,22 +355,20 @@ GraphPatternDetector::handle_t GetGenerateRewrite(
             std::vector<std::string> arguments;
             for (const std::string& argument : var.arguments()) {
               // The output may be mapped on the operator of pattern subgraph.
-              Node* node = nullptr;
-              auto iter = var_node_maps.find(argument);
-              if (var_node_maps.end() == iter) {
+              if (var_node_maps[argument].size() == 0) {
                 VarDesc var_desc(patterns::UniqueKey(argument));
-                node = graph->CreateVarNode(&var_desc);
-                var_node_maps.insert({argument, node});
+                var_node_maps[argument].emplace_back(
+                    graph->CreateVarNode(&var_desc));
+              }
+              if (in_nodes.end() == std::find(in_nodes.begin(),
+                                              in_nodes.end(),
+                                              var_node_maps[argument][0])) {
+                out_nodes.push_back(var_node_maps[argument][0]);
               } else {
-                if (in_nodes.end() ==
-                    std::find(in_nodes.begin(), in_nodes.end(), iter->second)) {
-                  node = iter->second;
-                } else {
-                  node = graph->CreateVarNode(iter->second->Var());
-                }
+                out_nodes.push_back(
+                    graph->CreateVarNode(var_node_maps[argument][0]->Var()));
               }
-              out_nodes.push_back(node);
-              arguments.push_back(node->Name());
+              arguments.push_back(var_node_maps[argument][0]->Name());
             }
             op_desc.SetOutput(var.parameter(), arguments);
           }
@@ -413,7 +417,40 @@ GraphPatternDetector::handle_t GetGenerateRewrite(
           remove_nodes.emplace(subgraph.at(pdnode.get()));
         }
         for (auto iter : var_node_maps) {
-          remove_nodes.erase(iter.second);
+          for (auto& node : iter.second) {
+            remove_nodes.erase(node);
+          }
+        }
+        GraphSafeRemoveNodes(graph, remove_nodes);
+
+        // Replace the redundant node by the first node in var_nodes_nmaps.
+        remove_nodes.clear();
+        for (auto iter : var_node_maps) {
+          auto var_node = iter.second[0];
+          for (size_t i = 1; i < iter.second.size(); ++i) {
+            auto replaced_var_node = iter.second[i];
+            for (auto op_node : replaced_var_node->outputs) {
+              auto index = std::find(op_node->inputs.begin(),
+                                     op_node->inputs.end(),
+                                     replaced_var_node) -
+                           op_node->inputs.begin();
+              op_node->inputs[index] = var_node;
+
+              auto& input_name_maps = *op_node->Op()->MutableInputs();
+              for (auto& item : input_name_maps) {
+                auto iter = std::find(item.second.begin(),
+                                      item.second.end(),
+                                      replaced_var_node->Name());
+                if (iter != item.second.end()) {
+                  item.second[iter - item.second.begin()] = var_node->Name();
+                  input_name_maps[item.first] = item.second;
+                  break;
+                }
+              }
+              op_node->Op()->Flush();
+            }
+            remove_nodes.emplace(replaced_var_node);
+          }
         }
         GraphSafeRemoveNodes(graph, remove_nodes);
       };
diff --git a/python/paddle/incubate/passes/fuse_resnet_unit_pass.py b/python/paddle/incubate/passes/fuse_resnet_unit_pass.py
@@ -90,7 +90,7 @@ def pattern_two_input(
         varZ,
     ):
         bnX = pattern_conv_bn(x, filterX, scaleX, biasX, meanX, varX)
-        bnZ = pattern_conv_bn(x, filterZ, scaleZ, biasZ, meanZ, varZ)
+        bnZ = pattern_conv_bn(z, filterZ, scaleZ, biasZ, meanZ, varZ)
         ewadd = ir.PassDesc.OP.elementwise_add(
             X=bnX.Output("Y"), Y=bnZ.Output("Y")
         )
diff --git a/python/paddle/incubate/passes/ir.py b/python/paddle/incubate/passes/ir.py
@@ -314,6 +314,36 @@ def Attr(self, name):
             return attr
 
     class OpHelper:
+        def _to_readable_code(self, skip_op_callstack=True):
+            assert isinstance(
+                skip_op_callstack, bool
+            ), "skip_op_callstack parameter's type is error, expect bool, received {}".format(
+                type(skip_op_callstack)
+            )
+            outputs_str = "{"
+            outputs_str += ", ".join(
+                [f"{k}={v}" for k, v in self._outputs.items()]
+            )
+            outputs_str += "}"
+
+            inputs_str = "{"
+            inputs_str += ", ".join(
+                [f"{k}={v}" for k, v in self._inputs.items()]
+            )
+            inputs_str += "}"
+
+            attrs_str = "{"
+            attrs_str += ", ".join([f"{k}={v}" for k, v in self._attrs.items()])
+            attrs_str += "}"
+
+            op_str = "{outputs} = {op_type}(inputs={inputs}, {attrs})".format(
+                outputs=outputs_str,
+                op_type=self._type,
+                inputs=inputs_str,
+                attrs=attrs_str,
+            )
+            return op_str
+
         def __init__(self, type=None):
             self._type = type
 

Original file line number	Diff line number	Diff line change
`@@ -90,7 +90,7 @@ def pattern_two_input(`
`90`	`90`	`varZ,`
`91`	`91`	`):`
`92`	`92`	`bnX = pattern_conv_bn(x, filterX, scaleX, biasX, meanX, varX)`
`93`		`- bnZ = pattern_conv_bn(x, filterZ, scaleZ, biasZ, meanZ, varZ)`
	`93`	`+ bnZ = pattern_conv_bn(z, filterZ, scaleZ, biasZ, meanZ, varZ)`
`94`	`94`	`ewadd = ir.PassDesc.OP.elementwise_add(`
`95`	`95`	`X=bnX.Output("Y"), Y=bnZ.Output("Y")`
`96`	`96`	`)`