|
| 1 | +// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +#include "paddle/fluid/framework/ir/conv_elementwise_add_mkldnn_fuse_pass.h" |
| 16 | +#include <functional> |
| 17 | +#include <utility> |
| 18 | + |
| 19 | +#include "paddle/fluid/framework/ir/graph_traits.h" |
| 20 | + |
| 21 | +namespace paddle { |
| 22 | +namespace framework { |
| 23 | +namespace ir { |
| 24 | +namespace { |
| 25 | + |
| 26 | +// The function keeps the graph consistent by replacing |
| 27 | +// a node 'from' in the set of inputs nodes |
| 28 | +// of the visited node by a node 'to'. |
| 29 | +void CorrectGraphEdges(Graph* graph, Node* from, Node* to) { |
| 30 | + for (auto& node : GraphTraits::DFS(*graph)) { |
| 31 | + auto from_in_inputs = |
| 32 | + std::find(std::begin(node.inputs), std::end(node.inputs), from); |
| 33 | + |
| 34 | + if (from_in_inputs != std::end(node.inputs)) { |
| 35 | + IR_NODE_LINK_TO(to, (&node)); |
| 36 | + |
| 37 | + auto inputs = node.Op()->Inputs(); |
| 38 | + |
| 39 | + using input_type = VariableNameMap::value_type; |
| 40 | + |
| 41 | + std::for_each(std::begin(inputs), std::end(inputs), |
| 42 | + [from, to, &node](const input_type& i) -> void { |
| 43 | + auto param_names = i.second; |
| 44 | + auto pi = std::find(std::begin(param_names), |
| 45 | + std::end(param_names), from->Name()); |
| 46 | + |
| 47 | + if (pi != std::end(param_names)) { |
| 48 | + node.Op()->SetInput(i.first, {to->Name()}); |
| 49 | + } |
| 50 | + }); |
| 51 | + } |
| 52 | + } |
| 53 | +} |
| 54 | +} // namespace |
| 55 | +using graph_ptr = std::unique_ptr<ir::Graph>; |
| 56 | + |
| 57 | +graph_ptr ConvElementwiseAddMKLDNNFusePass::ApplyImpl(graph_ptr graph) const { |
| 58 | + FusePassBase::Init(name_scope_, graph.get()); |
| 59 | + |
| 60 | + GraphPatternDetector gpd; |
| 61 | + auto pattern = gpd.mutable_pattern(); |
| 62 | + |
| 63 | + patterns::Conv conv_pattern{pattern, name_scope_}; |
| 64 | + auto conv_output = conv_pattern(); |
| 65 | + |
| 66 | + patterns::ElementwiseAdd elementwise_add_pattern{pattern, name_scope_}; |
| 67 | + elementwise_add_pattern(conv_output); |
| 68 | + |
| 69 | + conv_output->AsIntermediate(); |
| 70 | + |
| 71 | + auto conv_op_has_bias = [](const Node& conv_op) -> std::pair<bool, Node*> { |
| 72 | + auto bias_input_names = conv_op.Op()->Inputs(); |
| 73 | + auto bias_it = bias_input_names.find("Bias"); |
| 74 | + |
| 75 | + if (bias_it != std::end(bias_input_names)) { |
| 76 | + bool has_bias = !bias_it->second.empty(); |
| 77 | + |
| 78 | + if (has_bias) { |
| 79 | + auto conv_bias_names = bias_it->second; |
| 80 | + auto conv_bias_names_it = |
| 81 | + std::find_if(std::begin(conv_op.inputs), std::end(conv_op.inputs), |
| 82 | + [&conv_bias_names](Node* n) -> bool { |
| 83 | + return n->Name() == conv_bias_names[0]; |
| 84 | + }); |
| 85 | + return std::make_pair(has_bias, *conv_bias_names_it); |
| 86 | + } |
| 87 | + } |
| 88 | + |
| 89 | + return std::make_pair(false, nullptr); |
| 90 | + }; |
| 91 | + |
| 92 | + auto handler = [&](const GraphPatternDetector::subgraph_t& subgraph, |
| 93 | + Graph* g) { |
| 94 | + GET_IR_NODE_FROM_SUBGRAPH(conv_op, conv_op, conv_pattern); |
| 95 | + GET_IR_NODE_FROM_SUBGRAPH(conv_input, conv_input, conv_pattern); |
| 96 | + GET_IR_NODE_FROM_SUBGRAPH(conv_filter, conv_filter, conv_pattern); |
| 97 | + GET_IR_NODE_FROM_SUBGRAPH(conv_output, conv_output, conv_pattern); |
| 98 | + GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_op, elementwise_add_op, |
| 99 | + elementwise_add_pattern); |
| 100 | + GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_x, elementwise_add_x, |
| 101 | + elementwise_add_pattern); |
| 102 | + GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_out, elementwise_add_out, |
| 103 | + elementwise_add_pattern); |
| 104 | + |
| 105 | + if (FindFuseOption(*conv_op, *elementwise_add_op) != FUSE_MKLDNN) return; |
| 106 | + |
| 107 | + OpDesc op_desc; |
| 108 | + op_desc.SetType("conv2d"); |
| 109 | + |
| 110 | + op_desc.SetInput("Input", {conv_input->Name()}); |
| 111 | + op_desc.SetInput("Filter", {conv_filter->Name()}); |
| 112 | + op_desc.SetInput("ResidualData", {elementwise_add_x->Name()}); |
| 113 | + op_desc.SetOutput("Output", {conv_output->Name()}); |
| 114 | + |
| 115 | + bool has_bias; |
| 116 | + Node* conv_bias; |
| 117 | + |
| 118 | + std::tie(has_bias, conv_bias) = conv_op_has_bias(*conv_op); |
| 119 | + |
| 120 | + if (has_bias) { |
| 121 | + op_desc.SetInput("Bias", {conv_bias->Name()}); |
| 122 | + } |
| 123 | + |
| 124 | + for (const auto& attr : conv_op->Op()->GetAttrMap()) { |
| 125 | + op_desc.SetAttr(attr.first, attr.second); |
| 126 | + } |
| 127 | + |
| 128 | + op_desc.SetAttr("fuse_residual_connection", true); |
| 129 | + |
| 130 | + auto fused_conv_op = g->CreateOpNode(&op_desc); |
| 131 | + |
| 132 | + IR_NODE_LINK_TO(conv_input, fused_conv_op); |
| 133 | + IR_NODE_LINK_TO(conv_filter, fused_conv_op); |
| 134 | + IR_NODE_LINK_TO(elementwise_add_x, fused_conv_op); |
| 135 | + IR_NODE_LINK_TO(fused_conv_op, conv_output); |
| 136 | + |
| 137 | + if (has_bias) { |
| 138 | + IR_NODE_LINK_TO(conv_bias, fused_conv_op); |
| 139 | + } |
| 140 | + |
| 141 | + CorrectGraphEdges(g, elementwise_add_out, conv_output); |
| 142 | + GraphSafeRemoveNodes(g, {elementwise_add_out, conv_op, elementwise_add_op}); |
| 143 | + }; |
| 144 | + |
| 145 | + gpd(graph.get(), handler); |
| 146 | + |
| 147 | + return graph; |
| 148 | +} |
| 149 | +} // namespace ir |
| 150 | +} // namespace framework |
| 151 | +} // namespace paddle |
| 152 | + |
| 153 | +REGISTER_PASS(conv_elementwise_add_mkldnn_fuse_pass, |
| 154 | + paddle::framework::ir::ConvElementwiseAddMKLDNNFusePass); |
0 commit comments