Merge pull request #14458 from tpatejko/tpatejko/mkldnn-skip-connections

[WIP] Correcting and extending MKLDNN residual connection fuse pass

paddle/fluid/framework/ir/conv_elementwise_add_mkldnn_fuse_pass.h

@@ -15,24 +15,119 @@
 #pragma once
 
 #include <string>
+#include <tuple>
+#include <utility>
 #include "paddle/fluid/framework/ir/fuse_pass_base.h"
 #include "paddle/fluid/framework/ir/graph.h"
 #include "paddle/fluid/framework/ir/graph_pattern_detector.h"
 
+#include <boost/optional.hpp>
+
 namespace paddle {
 namespace framework {
 namespace ir {
 
-class ConvElementwiseAddMKLDNNFusePass : public FusePassBase {
+using graph_ptr = std::unique_ptr<ir::Graph>;
+using GraphWithStats = std::pair<ir::Graph*, int>;
+
+void CorrectGraphEdges(Graph* graph, Node* from, Node* to);
+bool IsReachable(ir::Graph* graph, Node* from, Node* to);
+boost::optional<Node*> HasBias(const Node& op, const std::string& bias_name);
+
+class ResidualConnectionMKLDNNFusePass : public FusePassBase {
+ private:
+  GraphWithStats FuseConvAsX(const std::string& name_scope,
+                             const GraphWithStats& graph_with_stats) const;
+  GraphWithStats FuseConvAsY(const std::string& name_scope,
+                             const GraphWithStats& graph_with_stats) const;
+  GraphWithStats FuseProjectionConv(
+      const std::string& name_scope,
+      const GraphWithStats& graph_with_stats) const;
+
+  template <typename RetType>
+  using GetNodeFunc =
+      std::function<RetType(const GraphPatternDetector::subgraph_t& subgraph)>;
+  using IdentityConvFunc = GetNodeFunc<std::tuple<Node*, Node*, Node*, Node*>>;
+  using IdentityElementwiseAddFunc =
+      GetNodeFunc<std::tuple<Node*, Node*, Node*>>;
+
+  using ProjectionConvFunc = IdentityConvFunc;
+  using ProjectionElementwiseAddFunc = GetNodeFunc<std::tuple<Node*, Node*>>;
+
+  using CanFuseFunc = std::function<bool(Node*, Node*)>;
+
+  std::tuple<Node*, Node*, Node*, Node*> GetNodesFromConv(
+      const patterns::Conv& conv_pattern,
+      const GraphPatternDetector::subgraph_t& subgraph) const;
+
+  std::tuple<Node*, Node*, Node*, Node*> GetNodesFromProjectionConv(
+      const patterns::Conv& conv_pattern,
+      const GraphPatternDetector::subgraph_t& subgraph) const;
+
+  template <typename HandleType, typename... OpFuncs>
+  GraphWithStats ExecuteHandleOnGraph(GraphPatternDetector* gpd,
+                                      const GraphWithStats& graph_with_stats,
+                                      OpFuncs&&... op_funcs) const {
+    ir::Graph* graph;
+    int stats;
+
+    std::tie(graph, stats) = graph_with_stats;
+
+    auto can_fuse = [this](Node* op1, Node* op2) -> bool {
+      return this->FindFuseOption(*op1, *op2) == FUSE_MKLDNN;
+    };
+
+    auto fuse_handle = HandleType{can_fuse, std::forward<OpFuncs>(op_funcs)...};
+
+    (*gpd)(graph, fuse_handle);
+
+    return std::make_pair(graph, stats + fuse_handle.get_stats());
+  }
+
+  struct IdentityFuseHandle {
+    IdentityFuseHandle(
+        const CanFuseFunc& can_fuse_func,
+        const IdentityConvFunc& get_node_from_conv_op,
+        const IdentityElementwiseAddFunc& get_node_from_elementwise_add_op);
+
+    void operator()(const GraphPatternDetector::subgraph_t& subgraph,
+                    Graph* graph);
+    int get_stats() const { return *fusion_stats; }
+
+   private:
+    std::shared_ptr<int> fusion_stats;
+    CanFuseFunc can_fuse_func;
+    IdentityConvFunc get_node_from_conv_op;
+    IdentityElementwiseAddFunc get_node_from_elementwise_add_op;
+  };
+
+  struct ProjectionFuseHandle {
+    ProjectionFuseHandle(
+        const CanFuseFunc& can_fuse_func,
+        const ProjectionConvFunc& get_node_from_conv_x_op,
+        const ProjectionConvFunc& get_node_from_conv_y_op,
+        const ProjectionElementwiseAddFunc& get_node_from_elementwise_add_op);
+
+    void operator()(const GraphPatternDetector::subgraph_t& subgraph,
+                    Graph* graph);
+    int get_stats() const { return *fusion_stats; }
+
+   private:
+    std::shared_ptr<int> fusion_stats;
+    CanFuseFunc can_fuse_func;
+    ProjectionConvFunc get_node_from_conv_x_op;
+    ProjectionConvFunc get_node_from_conv_y_op;
+    ProjectionElementwiseAddFunc get_node_from_elementwise_add_op;
+  };
+
  public:
-  virtual ~ConvElementwiseAddMKLDNNFusePass() {}
+  virtual ~ResidualConnectionMKLDNNFusePass() {}
 
  protected:
-  std::unique_ptr<ir::Graph> ApplyImpl(std::unique_ptr<ir::Graph> graph) const;
+  std::unique_ptr<ir::Graph> ApplyImpl(graph_ptr graph) const;
 
-  const std::string name_scope_{"residual_connections_fuse_pass"};
+  const std::string name_scope_{"residual_connection_fuse_pass"};
 };
-
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/ir/conv_elementwise_add_mkldnn_fuse_pass_tester.cc

@@ -40,7 +40,7 @@ void SetOp(ProgramDesc* prog, const std::string& type,
   op->SetOutput(output.first, {output.second});
 }
 
-struct IsReachable {
+struct TestIsReachable {
   using func = std::function<bool(const std::string&, const std::string&)>;
 
   auto operator()(const std::unique_ptr<ir::Graph>& graph) -> func {
@@ -89,7 +89,9 @@ struct IsReachable {
   }
 };
 
-void AssertOpsCount(const std::unique_ptr<ir::Graph>& graph) {
+void AssertOpsCount(const std::unique_ptr<ir::Graph>& graph,
+                    int expected_conv_count,
+                    int expected_elementwise_add_count = 0) {
   int conv_count = 0;
   int elementwise_add_count = 0;
 
@@ -101,8 +103,8 @@ void AssertOpsCount(const std::unique_ptr<ir::Graph>& graph) {
       ++elementwise_add_count;
     }
   }
-  EXPECT_EQ(conv_count, 1);
-  EXPECT_EQ(elementwise_add_count, 0);
+  EXPECT_EQ(conv_count, expected_conv_count);
+  EXPECT_EQ(elementwise_add_count, expected_elementwise_add_count);
 }
 
 ProgramDesc BuildProgramDesc(const std::vector<std::string>& transient_vars,
@@ -127,22 +129,13 @@ ProgramDesc BuildProgramDesc(const std::vector<std::string>& transient_vars,
 
   return prog;
 }
-}  // namespace
-
-TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionWithElementwiseAddRelu) {
-  auto prog =
-      BuildProgramDesc({"a", "b", "c", "d", "e", "f"}, {"bias", "weights"});
-
-  SetOp(&prog, "conv2d",
-        {{"Input", "a"}, {"Bias", "bias"}, {"Filter", "weights"}},
-        {"Output", "b"});
-  SetOp(&prog, "elementwise_add", {{"X", "b"}, {"Y", "c"}}, {"Out", "d"});
-  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+void RunPassAndAssert(ProgramDesc* prog, const std::string& from,
+                      const std::string& to, int expected_conv_num) {
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(*prog));
 
-  IsReachable is_reachable;
-  EXPECT_TRUE(is_reachable(graph)("a", "relu"));
+  TestIsReachable is_reachable;
+  EXPECT_TRUE(is_reachable(graph)(from, to));
 
   auto pass =
       PassRegistry::Instance().Get("conv_elementwise_add_mkldnn_fuse_pass");
@@ -150,82 +143,87 @@ TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionWithElementwiseAddRelu) {
   graph = pass->Apply(std::move(graph));
   int current_nodes_num = graph->Nodes().size();
 
-  EXPECT_TRUE(is_reachable(graph)("a", "relu"));
+  EXPECT_TRUE(is_reachable(graph)(from, to));
 
   EXPECT_EQ(original_nodes_num - nodes_removed + nodes_added,
             current_nodes_num);
 
-  AssertOpsCount(graph);
+  AssertOpsCount(graph, expected_conv_num);
 }
+}  // namespace
 
-TEST(ConvElementwiseAddMKLDNNFusePass,
-     ConvolutionWithElementwiseAddReluNoBias) {
-  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"weights"});
-  SetOp(&prog, "conv2d", {{"Input", "a"}, {"Filter", "weights"}},
-        {"Output", "b"});
-  SetOp(&prog, "elementwise_add", {{"X", "b"}, {"Y", "c"}}, {"Out", "d"});
-  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
-
-  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionAsYWithElementwiseAddRelu) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"bias", "weights"});
 
-  IsReachable is_reachable;
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
+  SetOp(&prog, "conv2d",
+        {{"Input", "b"}, {"Bias", "bias"}, {"Filter", "weights"}},
+        {"Output", "c"});
 
-  EXPECT_TRUE(is_reachable(graph)("a", "relu"));
+  SetOp(&prog, "elementwise_add", {{"X", "a"}, {"Y", "c"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  auto pass =
-      PassRegistry::Instance().Get("conv_elementwise_add_mkldnn_fuse_pass");
-  int original_nodes_num = graph->Nodes().size();
-  graph = pass->Apply(std::move(graph));
-  int current_nodes_num = graph->Nodes().size();
+  RunPassAndAssert(&prog, "a", "relu", 1);
+}
 
-  EXPECT_TRUE(is_reachable(graph)("a", "relu"));
+TEST(ConvElementwiseAddMKLDNNFusePass,
+     ConvolutionAsYWithElementwiseAddReluNoBias) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"weights"});
 
-  EXPECT_EQ(original_nodes_num - nodes_removed + nodes_added,
-            current_nodes_num);
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
+  SetOp(&prog, "conv2d", {{"Input", "b"}, {"Filter", "weights"}},
+        {"Output", "c"});
+  SetOp(&prog, "elementwise_add", {{"X", "a"}, {"Y", "c"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  AssertOpsCount(graph);
+  RunPassAndAssert(&prog, "a", "relu", 1);
 }
 
-TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionElementwiseAdd) {
-  auto prog = BuildProgramDesc({"a", "b", "c", "d"}, {"bias", "weights"});
+TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionAsXWithElementwiseAddRelu) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"bias", "weights"});
+
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
   SetOp(&prog, "conv2d",
-        {{"Input", "a"}, {"Bias", "bias"}, {"Filter", "weights"}},
-        {"Output", "b"});
-  SetOp(&prog, "elementwise_add", {{"X", "b"}, {"Y", "c"}}, {"Out", "d"});
+        {{"Input", "b"}, {"Bias", "bias"}, {"Filter", "weights"}},
+        {"Output", "c"});
 
-  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+  SetOp(&prog, "elementwise_add", {{"X", "c"}, {"Y", "a"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  IsReachable is_reachable;
-  EXPECT_TRUE(is_reachable(graph)("a", "d"));
+  RunPassAndAssert(&prog, "a", "relu", 1);
+}
 
-  auto pass =
-      PassRegistry::Instance().Get("conv_elementwise_add_mkldnn_fuse_pass");
-  int original_nodes_num = graph->Nodes().size();
-  graph = pass->Apply(std::move(graph));
-  int current_nodes_num = graph->Nodes().size();
+TEST(ConvElementwiseAddMKLDNNFusePass,
+     ConvolutionAsXWithElementwiseAddReluNoBias) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"weights"});
 
-  EXPECT_FALSE(is_reachable(graph)("a", "d"));
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
+  SetOp(&prog, "conv2d", {{"Input", "b"}, {"Filter", "weights"}},
+        {"Output", "c"});
+  SetOp(&prog, "elementwise_add", {{"X", "c"}, {"Y", "a"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  EXPECT_EQ(original_nodes_num - nodes_removed + nodes_added,
-            current_nodes_num);
-  AssertOpsCount(graph);
+  RunPassAndAssert(&prog, "a", "relu", 1);
 }
 
-TEST(ConvElementwiseAddMKLDNNFusePass, SigmoidConvolutionAddElementwiseRelu) {
+TEST(ConvElementwiseAddMKLDNNFusePass, NoFusion) {
   auto prog =
-      BuildProgramDesc({"a", "b", "c", "d", "e", "f"}, {"bias", "weights"});
+      BuildProgramDesc({"a", "b", "c", "d", "e", "f", "g"}, {"weights"});
+
   SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
-  SetOp(&prog, "conv2d",
-        {{"Input", "b"}, {"Bias", "bias"}, {"Filter", "weights"}},
+  SetOp(&prog, "conv2d", {{"Input", "b"}, {"Filter", "weights"}},
         {"Output", "c"});
-  SetOp(&prog, "elementwise_add", {{"X", "c"}, {"Y", "d"}}, {"Out", "e"});
-  SetOp(&prog, "relu", {{"X", "e"}}, {"Out", "f"});
 
-  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+  SetOp(&prog, "conv2d", {{"Input", "d"}, {"Filter", "weights"}},
+        {"Output", "e"});
 
-  IsReachable is_reachable;
+  SetOp(&prog, "elementwise_add", {{"X", "c"}, {"Y", "e"}}, {"Out", "f"});
+  SetOp(&prog, "relu", {{"X", "f"}}, {"Out", "g"});
 
-  EXPECT_TRUE(is_reachable(graph)("a", "f"));
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+
+  TestIsReachable is_reachable;
+  EXPECT_TRUE(is_reachable(graph)("a", "g"));
 
   auto pass =
       PassRegistry::Instance().Get("conv_elementwise_add_mkldnn_fuse_pass");
@@ -233,11 +231,10 @@ TEST(ConvElementwiseAddMKLDNNFusePass, SigmoidConvolutionAddElementwiseRelu) {
   graph = pass->Apply(std::move(graph));
   int current_nodes_num = graph->Nodes().size();
 
-  EXPECT_TRUE(is_reachable(graph)("a", "f"));
+  EXPECT_TRUE(is_reachable(graph)("a", "g"));
+  EXPECT_EQ(original_nodes_num, current_nodes_num);
 
-  EXPECT_EQ(original_nodes_num - nodes_removed + nodes_added,
-            current_nodes_num);
-  AssertOpsCount(graph);
+  AssertOpsCount(graph, 2, 1);
 }
 
 }  // namespace ir

paddle/fluid/framework/ir/graph_pattern_detector.cc

@@ -1084,16 +1084,12 @@ PDNode *patterns::Conv::operator()() {
   return output_var;
 }
 
-PDNode *patterns::ElementwiseAdd::operator()(PDNode *x_var) {
+PDNode *patterns::ElementwiseAdd::operator()(PDNode *x_var, PDNode *y_var) {
   auto elementwise_add_op = pattern->NewNode(elementwise_add_op_repr())
                                 ->assert_is_op("elementwise_add");
 
-  x_var->assert_is_op_input("elementwise_add", "X");
-
-  auto y_var = pattern->NewNode(elementwise_add_x_repr())
-                   ->AsInput()
-                   ->assert_is_op_input("elementwise_add", "Y");
-
+  x_var->AsInput()->assert_is_op_input("elementwise_add", "X");
+  y_var->AsInput()->assert_is_op_input("elementwise_add", "Y");
   auto out_var = pattern->NewNode(elementwise_add_out_repr())
                      ->AsOutput()
                      ->assert_is_op_output("elementwise_add", "Out");

paddle/fluid/framework/ir/graph_pattern_detector.h

@@ -664,7 +664,7 @@ struct ElementwiseAdd : public PatternBase {
   ElementwiseAdd(PDPattern* pattern, const std::string& name_scope)
       : PatternBase(pattern, name_scope, "elementwise_add") {}
 
-  PDNode* operator()(PDNode* x_var);
+  PDNode* operator()(PDNode* x_var, PDNode* y_var);
 
   PATTERN_DECL_NODE(elementwise_add_op);
   PATTERN_DECL_NODE(elementwise_add_x);