Method FunctionalizeInternal

tensorflow/compiler/tf2xla/functionalize_cond.cc:1421–1521

}

Status FunctionalizeCond::FunctionalizeInternal() {
  // The general approach for converting a tf.cond (as lowered via switch/merge
  // nodes) to a functional if is as follows:
  // 1. Determine the topological order and collect all the switch and merge
  // nodes in the graph;
  // 2. Compute the predicates and dominance structure for all the nodes in the
  // graph - this includes which predicate must be true for a op to execute
  // (predicate values are considered directly rather than attempting to
  // determine deeper equivalence). We shall refer to this structure as the
  // CondState;
  // 3. Sort the merge nodes by nesting depth;
  // 4. Extract merge nodes together that have the same CondState and
  // AncestorState from the innermost to the outermost into IfOps;
  // Note: In the above only nodes that feed into a merge node will be
  // considered for functionalization.

  // Perform a DFS over the graph and
  // * Determine the reverse topological order of the nodes (there should be no
  //   cycles at this point so the post-order numbering corresponds to the
  //   reverse topological sorting);
  // * Record reverse topological for merge and switch nodes;
  std::vector<Node*> rev_topo_order;
  std::vector<Node*> merge_order;
  DFS(*graph_, nullptr, [&](Node* n) {
    if (IsSwitch(n)) {
      AddSwitchId(n->id());
    }
    if (IsMerge(n)) {
      merge_order.push_back(n);
    }
    if (n->IsOp()) {
      rev_topo_order.push_back(n);
    }
  });

  // No merges to functionalize.
  if (merge_order.empty()) {
    // No merges mean no switch values consumed (as only considering values
    // fetchable as output of merge);
    DeleteReachableAndDeadNodes(merge_order);
    return Status::OK();
  }

  TF_RETURN_IF_ERROR(DetermineStates(std::move(rev_topo_order)));
  if (VLOG_IS_ON(4)) DumpGraphWithCondState("id");

  // Sort the merge nodes from innermost outwards.
  SortMergeNodes(&merge_order);

  // Cluster merge nodes by (CondId, AncestorId, predicate) in order of
  // nesting. (CondId, AncestorId) is not enough, e.g.
  //   pred1 = array_ops.placeholder(dtypes.bool, name='pred1')
  //   pred2 = array_ops.placeholder(dtypes.bool, name='pred2')
  //   cond1 = control_flow_ops.cond(pred1, ...)
  //   cond2 = control_flow_ops.cond(pred2, ...)
  //   cond3 = control_flow_ops.cond(pred1, use cond1 and cond2)
  //   cond4 = control_flow_ops.cond(pred2, use cond1 and cond2)
  // cond3 and cond4 have the same (CondId, AncestorId), but they should not