| 77 | } // namespace |
| 78 | |
| 79 | Status LowerFunctionalOpsPass::Run( |
| 80 | const GraphOptimizationPassOptions& options) { |
| 81 | if (options.partition_graphs != nullptr) { |
| 82 | return errors::Internal( |
| 83 | "Lowering If/While ops should happen before partitioning."); |
| 84 | } |
| 85 | if (options.graph == nullptr) { |
| 86 | return Status::OK(); |
| 87 | } |
| 88 | |
| 89 | Graph* g = options.graph->get(); |
| 90 | if (g == nullptr) { |
| 91 | return errors::Internal( |
| 92 | "Lowering While op requires a graph to be available."); |
| 93 | } |
| 94 | |
| 95 | FunctionLibraryDefinition* flib_def = options.flib_def; |
| 96 | if (flib_def == nullptr) { |
| 97 | return errors::Internal( |
| 98 | "Lowering If op requires a FunctionLibraryDefinition to be available."); |
| 99 | } |
| 100 | |
| 101 | // Lower function calls only if it's explicitly enabled in session options. |
| 102 | const bool lower_function_calls = |
| 103 | options.session_options && options.session_options->config.graph_options() |
| 104 | .optimizer_options() |
| 105 | .do_function_inlining(); |
| 106 | |
| 107 | // If graph is a function instantiation, it will have `_Arg` and `_Retval` |
| 108 | // nodes for input and output tensors. Otherwise it's unsafe to remove any of |
| 109 | // the nodes, because they might be later used as fetches. |
| 110 | // |
| 111 | // When we do not keep lowered nodes fetchable, we still add a NoOp node to |
| 112 | // the graph with the same name as lowered node, because it might be used as a |
| 113 | // control output source, and it's currently not expressed in a graph. |
| 114 | bool keep_lowered_nodes_fetchable = keep_lowered_nodes_fetchable_.has_value() |
| 115 | ? *keep_lowered_nodes_fetchable_ |
| 116 | : !HasArgsOrRetvals(*g); |
| 117 | |
| 118 | // We disable lowering control flow to switch/merge variants for the |
| 119 | // single-threaded executor, which does not support it. |
| 120 | const bool functional_control_flow = |
| 121 | options.session_options && |
| 122 | (options.session_options->config.experimental().executor_type() == |
| 123 | "SINGLE_THREADED_EXECUTOR"); |
| 124 | |
| 125 | // Lower all If, Case, While ops that have the `kLowerUsingSwitchMergeAttr` |
| 126 | // attr set and inline all function calls into the graph. |
| 127 | // We start at `i` = 2 to skip the source and sink nodes. |
| 128 | // Note that `g->num_node_ids()` may change in the for body if a matching If, |
| 129 | // Case, While node is lowered. Since new graph nodes are always added to the |
| 130 | // end of the list of nodes it is ensured that nested If/Case/While nodes will |
| 131 | // be lowered as well. |
| 132 | for (int i = 2; i < g->num_node_ids(); ++i) { |
| 133 | Node* n = g->FindNodeId(i); |
| 134 | if (n == nullptr) continue; // deleted node |
| 135 | if (MarkedForTpuCompilation(n)) continue; |
| 136 | if (MarkedForXlaCompilation(n)) continue; |
nothing calls this directly
no test coverage detected