[Paddle-TRT] support sub block: while op, condition_block op

paddle/fluid/framework/ir/CMakeLists.txt

@@ -167,6 +167,7 @@ if(WITH_TENSORRT)
   pass_library(preln_embedding_eltwise_layernorm_fuse_pass inference)
   pass_library(split_layernorm_to_math_ops_pass inference)
   pass_library(trt_remove_amp_strategy_op_pass inference)
+  pass_library(set_subgraph_edge_pass inference)
 endif()
 
 if(WITH_GPU OR WITH_ROCM)

paddle/fluid/framework/ir/conv_bn_fuse_pass.cc

@@ -312,6 +312,14 @@ void ConvBNFusePass::ApplyImpl(ir::Graph* graph) const {
       graph, platform::errors::InvalidArgument("Graph cannot be nullptr."));
   FusePassBase::Init(name_scope_, graph);
 
+  VLOG(3) << "Running conv_bn_fuse_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running conv_bn_fuse_pass is: 0(main_graph)";
+  } else {
+    VLOG(3) << "The ID of block running conv_bn_fuse_pass is: "
+            << graph->GetBlockId();
+  }
+
   auto* scope = param_scope();
   PADDLE_ENFORCE_NOT_NULL(
       scope, platform::errors::InvalidArgument("Scope cannot be nullptr."));
@@ -612,6 +620,15 @@ void ConvEltwiseAddBNFusePass::ApplyImpl(ir::Graph* graph) const {
       graph, platform::errors::InvalidArgument("Graph cannot be nullptr."));
   FusePassBase::Init(name_scope_, graph);
 
+  VLOG(3) << "Running conv_eltwiseadd_bn_fuse_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running conv_eltwiseadd_bn_fuse_pass is: "
+               "0(main_graph)";
+  } else {
+    VLOG(3) << "The ID of block running conv_eltwiseadd_bn_fuse_pass is: "
+            << graph->GetBlockId();
+  }
+
   auto* scope = param_scope();
   PADDLE_ENFORCE_NOT_NULL(
       scope, platform::errors::InvalidArgument("Scope cannot be nullptr."));

paddle/fluid/framework/ir/delete_quant_dequant_linear_op_pass.cc

@@ -95,6 +95,17 @@ void DeleteQuantDequantLinearOpPass::ApplyImpl(ir::Graph* graph) const {
       scope,
       platform::errors::InvalidArgument(
           "Scope in DeleteQuantDequantLinearOpPass should not be null."));
+
+  VLOG(3) << "Running delete_quant_dequant_linear_op_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running delete_quant_dequant_linear_op_pass "
+               "is: 0(main_graph)";
+  } else {
+    VLOG(3)
+        << "The ID of block running delete_quant_dequant_linear_op_pass is: "
+        << graph->GetBlockId();
+  }
+
   std::unordered_map<std::string, std::vector<float>> var_quant_scales{};
 
   // Create pattern

paddle/fluid/framework/ir/delete_weight_dequant_linear_op_pass.cc

@@ -36,7 +36,17 @@ void DeleteWeightDequantLinearOpPass::ApplyImpl(ir::Graph* graph) const {
                     true,
                     platform::errors::InvalidArgument(
                         "Graph must have kParamScopeAttr attribute."));
-  VLOG(3) << "Handle delete weight dequant linear op pass ...";
+
+  VLOG(3) << "Running delete_weight_dequant_linear_op_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running delete_weight_dequant_linear_op_pass "
+               "is: 0(main_graph)";
+  } else {
+    VLOG(3)
+        << "The ID of block running delete_weight_dequant_linear_op_pass is: "
+        << graph->GetBlockId();
+  }
+
   auto& scope = graph->Get<framework::Scope>(kParamScopeAttr);
   bool is_int8 = false;
 

paddle/fluid/framework/ir/fuse_multi_transformer_layer_pass.cc

@@ -291,6 +291,16 @@ void FuseMultiTransformerLayerPass::ApplyImpl(Graph* graph) const {
       scope,
       platform::errors::Fatal("During the fuse_multi_transformer_layer pass, "
                               "The scope should not be null."));
+
+  VLOG(3) << "Running fuse_multi_transformer_layer_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running fuse_multi_transformer_layer_pass is: "
+               "0(main_graph)";
+  } else {
+    VLOG(3) << "The ID of block running fuse_multi_transformer_layer_pass is: "
+            << graph->GetBlockId();
+  }
+
   int fusion_count = BuildFusion(graph, name_scope_, scope);
 
   AddStatis(fusion_count);

paddle/fluid/framework/ir/fused_multi_transformer_decoder_pass.cc

@@ -1671,6 +1671,16 @@ void FusedMultiTransformerDecoderPass::ApplyImpl(Graph* graph) const {
       platform::errors::Fatal("During the multi_transformer pass, "
                               "The scope should not be null."));
 
+  VLOG(3) << "Running fused_multi_transformer_decoder_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running fused_multi_transformer_decoder_pass "
+               "is: 0(main_graph)";
+  } else {
+    VLOG(3)
+        << "The ID of block running fused_multi_transformer_decoder_pass is: "
+        << graph->GetBlockId();
+  }
+
   int fusion_count = BuildFusion(graph, name_scope_, scope);
   if (fusion_count > 0) {
     graph->Set(kFusedMultiTransformerDecoderPass, new bool(true));
@@ -2376,6 +2386,17 @@ void FusedMultiTransformerDecoderFuseQKVPass::ApplyImpl(Graph* graph) const {
       platform::errors::Fatal("During the fused_multi_transformer_decoder "
                               "pass, The scope should not be null."));
 
+  VLOG(3) << "Running fused_multi_transformer_decoder_fuse_qkv_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3)
+        << "The ID of block running "
+           "fused_multi_transformer_decoder_fuse_qkv_pass is: 0(main_graph)";
+  } else {
+    VLOG(3) << "The ID of block running "
+               "fused_multi_transformer_decoder_fuse_qkv_pass is: "
+            << graph->GetBlockId();
+  }
+
   int fusion_count = BuildFusion(graph, name_scope_, scope);
   if (fusion_count > 0) {
     graph->Set(kFusedMultiTransformerDecoderFuseQKVPass, new bool(true));
@@ -3146,6 +3167,19 @@ void MultiDevicesFusedMultiTransformerDecoderFuseQKVPass::ApplyImpl(
       platform::errors::Fatal("During the fused_multi_transformer_decoder "
                               "pass, The scope should not be null."));
 
+  VLOG(3)
+      << "Running multi_devices_fused_multi_transformer_decoder_fuse_qkv_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running "
+               "multi_devices_fused_multi_transformer_decoder_fuse_qkv_pass "
+               "is: 0(main_graph)";
+  } else {
+    VLOG(3)
+        << "The ID of block running "
+           "multi_devices_fused_multi_transformer_decoder_fuse_qkv_pass is: "
+        << graph->GetBlockId();
+  }
+
   int fusion_count = BuildFusion(graph, name_scope_, scope);
   if (fusion_count > 0) {
     graph->Set(kFusedMultiTransformerDecoderFuseQKVPass, new bool(true));

paddle/fluid/framework/ir/fused_multi_transformer_encoder_pass.cc

@@ -2400,6 +2400,16 @@ void FusedMultiTransformerEncoderPass::ApplyImpl(Graph* graph) const {
       platform::errors::Fatal(
           "During the multi_transformer pass, The scope should not be null."));
 
+  VLOG(3) << "Running fused_multi_transformer_encoder_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running fused_multi_transformer_encoder_pass "
+               "is: 0(main_graph)";
+  } else {
+    VLOG(3)
+        << "The ID of block running fused_multi_transformer_encoder_pass is: "
+        << graph->GetBlockId();
+  }
+
   int fusion_count = BuildFusion(graph, name_scope_, scope);
   if (fusion_count > 0) {
     graph->Set(kFusedMultiTransformerEncoderPass, new bool(true));
@@ -3211,6 +3221,17 @@ void FusedMultiTransformerEncoderFuseQKVPass::ApplyImpl(Graph* graph) const {
           "During the fused_multi_transformer_encoder pass, "
           "The scope should not be null."));
 
+  VLOG(3) << "Running fused_multi_transformer_encoder_fuse_qkv_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3)
+        << "The ID of block running "
+           "fused_multi_transformer_encoder_fuse_qkv_pass is: 0(main_graph)";
+  } else {
+    VLOG(3) << "The ID of block running "
+               "fused_multi_transformer_encoder_fuse_qkv_pass is: "
+            << graph->GetBlockId();
+  }
+
   int fusion_count = BuildFusion(graph, name_scope_, scope);
   if (fusion_count > 0) {
     graph->Set(kFusedMultiTransformerEncoderFuseQKVPass, new bool(true));
@@ -4013,6 +4034,17 @@ void MultiDevicesFusedMultiTransformerEncoderPass::ApplyImpl(
       platform::errors::Fatal(
           "During the multi_transformer pass, The scope should not be null."));
 
+  VLOG(3) << "Running multi_devices_fused_multi_transformer_encoder_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running "
+               "multi_devices_fused_multi_transformer_encoder_pass is: "
+               "0(main_graph)";
+  } else {
+    VLOG(3) << "The ID of block running "
+               "multi_devices_fused_multi_transformer_encoder_pass is: "
+            << graph->GetBlockId();
+  }
+
   int fusion_count = BuildFusion(graph, name_scope_, scope);
   if (fusion_count > 0) {
     graph->Set(kFusedMultiTransformerEncoderPass, new bool(true));
@@ -4872,6 +4904,19 @@ void MultiDevicesFusedMultiTransformerEncoderFuseQKVPass::ApplyImpl(
           "During the fused_multi_transformer_encoder pass, "
           "The scope should not be null."));
 
+  VLOG(3)
+      << "Running multi_devices_fused_multi_transformer_encoder_fuse_qkv_pass.";
+  if (graph->IsMainGraph()) {
+    VLOG(3) << "The ID of block running "
+               "multi_devices_fused_multi_transformer_encoder_fuse_qkv_pass "
+               "is: 0(main_graph)";
+  } else {
+    VLOG(3)
+        << "The ID of block running "
+           "multi_devices_fused_multi_transformer_encoder_fuse_qkv_pass is: "
+        << graph->GetBlockId();
+  }
+
   int fusion_count = BuildFusion(graph, name_scope_, scope);
   if (fusion_count > 0) {
     graph->Set(kMultiDevicesFusedMultiTransformerEncoderFuseQKVPass,

paddle/fluid/framework/ir/graph.cc

@@ -22,6 +22,10 @@ PADDLE_DEFINE_EXPORTED_bool(convert_all_blocks,
                             true,
                             "Convert all blocks in program into SSAgraphs");
 
+PADDLE_DEFINE_EXPORTED_bool(all_blocks_convert_trt,
+                            false,
+                            "Convert all blocks'Ops into TensorRT Ops");
+
 namespace paddle {
 namespace framework {
 namespace ir {

paddle/fluid/framework/ir/graph.h

@@ -28,6 +28,7 @@ limitations under the License. */
 #include "paddle/utils/flags.h"
 
 PD_DECLARE_bool(convert_all_blocks);
+PD_DECLARE_bool(all_blocks_convert_trt);
 
 namespace paddle {
 namespace framework {
@@ -383,6 +384,8 @@ class Graph {
 
   bool IsMainGraph() const { return main_graph_ == nullptr; }
 
+  const Graph *GetMainGraph() const { return main_graph_; }
+
   Graph *GetSubGraph(const size_t idx) const {
     PADDLE_ENFORCE_EQ(
         this->IsMainGraph(),

paddle/fluid/framework/ir/graph_helper.cc

@@ -34,6 +34,7 @@ PHI_DECLARE_bool(dynamic_static_unified_comm);
 #endif
 #include "paddle/fluid/platform/flags.h"
 PD_DECLARE_bool(convert_all_blocks);
+PD_DECLARE_bool(all_blocks_convert_trt);
 PADDLE_DEFINE_EXPORTED_string(print_sub_graph_dir,
                               "",
                               "FLAGS_print_sub_graph_dir is used "

paddle/fluid/framework/ir/graph_pattern_detector.cc

@@ -4937,6 +4937,13 @@ void patterns::MulMatmulMatmulV2::operator()(
   ops->LinksTo({ops_out});
 }
 
+// subgraph_edge_pattern
+PDNode *patterns::SubgraphEdgePattern::operator()(
+    const std::unordered_set<std::string> &ops_type) {
+  auto ops = pattern->NewNode(ops_repr())->assert_is_ops(ops_type);
+  return ops;
+}
+
 PDNode *patterns::ConvBNAddAct::operator()(
     const std::unordered_set<std::string> &act_types,
     bool shortcut,

paddle/fluid/framework/ir/graph_pattern_detector.h

@@ -2296,6 +2296,14 @@ struct AddSupportInt8 : public PatternBase {
   PATTERN_DECL_NODE(quant_out);
 };
 
+// subgraph_edge_pattern
+struct SubgraphEdgePattern : public PatternBase {
+  SubgraphEdgePattern(PDPattern* pattern, const std::string& name_scope)
+      : PatternBase(pattern, name_scope, "subgraph_edge_pattern") {}
+  PDNode* operator()(const std::unordered_set<std::string>& ops_type);
+  PATTERN_DECL_NODE(ops);
+};
+
 // The following patterns are used to fuse feedforward in forward
 // 1. layer_norm -> linear1 -> activation -> dropout1 -> linear2 -> dropout2
 // -> residual_add (pre_layer_norm)

paddle/fluid/framework/ir/node.h

@@ -254,6 +254,14 @@ class Node {
   // so expose it is a good idea
   static constexpr int NO_DESC_ORDER = INT_MAX;
 
+  // Set whether the node is an edge of the subgraph.
+  void SetSubgraphOutput() { subgraph_output_ = true; }
+  void SetSubgraphInput() { subgraph_input_ = true; }
+
+  // Get whether the node is an edge of the subgraph.
+  bool IsSubgraphOutput() { return subgraph_output_; }
+  bool IsSubgraphInput() { return subgraph_input_; }
+
  protected:
   std::string name_;
   std::unique_ptr<VarDesc> var_desc_;
@@ -269,6 +277,10 @@ class Node {
   uint64_t original_desc_id_{0};
   int graph_id_{-1};
 
+  // Is it the edge of the subgraph.
+  bool subgraph_output_ = false;
+  bool subgraph_input_ = false;
+
  private:
   // ID can only set by a Graph.
   void SetId(int id) { id_ = id; }

paddle/fluid/framework/ir/pass.cc

@@ -39,6 +39,10 @@ namespace ir {
 static const char kParamScopeAttr[] = "__param_scope__";  // NOLINT
 
 static const std::vector<std::string> support_subgraph_passes = {
+    "feed_fetch_subgraph_pass",
+    "set_subgraph_edge_pass",
+    "trt_map_ops_to_matrix_multiply_pass",
+    "tensorrt_subgraph_pass",
     "simplify_with_basic_ops_pass",
     "fused_multi_transformer_encoder_pass",
     "fused_multi_transformer_decoder_pass",
@@ -128,7 +132,8 @@ Graph *Pass::Apply(Graph *graph) const {
   } else {
     subgraph_passes = support_subgraph_passes;
   }
-  if (graph->IsMainGraph() &&
+  if (FLAGS_all_blocks_convert_trt && FLAGS_convert_all_blocks &&
+      graph->IsMainGraph() &&
       (std::count(subgraph_passes.begin(), subgraph_passes.end(), Type()) ||
        std::count(support_subgraph_generate_passes.begin(),
                   support_subgraph_generate_passes.end(),