| 104 | } |
| 105 | |
| 106 | bool CanRemoveNode(const utils::MutableNodeView& node_view, |
| 107 | const absl::flat_hash_set<absl::string_view>& function_names, |
| 108 | const OpRegistryInterface& op_registry) { |
| 109 | const auto* node = node_view.node(); |
| 110 | const bool no_fanins = node_view.NumRegularFanins() == 0 && |
| 111 | node_view.NumControllingFanins() == 0; |
| 112 | if (IsNoOp(*node) && no_fanins) { |
| 113 | return true; |
| 114 | } |
| 115 | if (IsConstant(*node) && no_fanins && node_view.NumRegularFanouts() == 0) { |
| 116 | return true; |
| 117 | } |
| 118 | if (RemovalIncreasesEdgeCount(node_view)) { |
| 119 | return false; |
| 120 | } |
| 121 | const string& device = node->device(); |
| 122 | for (const auto& regular_fanin : node_view.GetRegularFanins()) { |
| 123 | auto* fanin_node = regular_fanin.node_view()->node(); |
| 124 | if (fanin_node->device() != device) { |
| 125 | // Node is driven by a different device. |
| 126 | return false; |
| 127 | } else if (function_names.contains(fanin_node->op())) { |
| 128 | // Node input is a function call. |
| 129 | return false; |
| 130 | } else if (IsOutputPortRefValue(*fanin_node, regular_fanin.index(), |
| 131 | op_registry)) { |
| 132 | return false; |
| 133 | } |
| 134 | } |
| 135 | for (const auto& controlling_fanin : node_view.GetControllingFanins()) { |
| 136 | if (controlling_fanin.node_view()->GetDevice() != device) { |
| 137 | // Node is driven by a different device. |
| 138 | return false; |
| 139 | } |
| 140 | } |
| 141 | for (const auto& regular_fanouts : node_view.GetRegularFanouts()) { |
| 142 | for (const auto& regular_fanout : regular_fanouts) { |
| 143 | if (function_names.contains(regular_fanout.node_view()->GetOp())) { |
| 144 | // Node output is a function call. |
| 145 | return false; |
| 146 | } |
| 147 | } |
| 148 | } |
| 149 | return true; |
| 150 | } |
| 151 | |
| 152 | // ForwardFanins forwards fanins of a node to be removed to its fanouts. This |
| 153 | // currently is specific to nodes defined in `IsTrivialOp` under the assumption |
no test coverage detected