| 3061 | std::pair<int, const HloInstruction*>)>; |
| 3062 | template <typename Visitor> |
| 3063 | static Status PostOrderDFS(HloInstruction* root, Visitor* visitor, |
| 3064 | const InternalCompareFunction* operand_order, |
| 3065 | bool ignore_control_predecessors) { |
| 3066 | // Calculating the instruction count within a module can be expensive on large |
| 3067 | // models so only do it if the visit state is empty. This will help when the |
| 3068 | // same visitor is reused across many computations of a single module. |
| 3069 | if (visitor->VisitStateCapacity() == 0) { |
| 3070 | visitor->ReserveVisitStates(root->GetModule()->instruction_count()); |
| 3071 | } |
| 3072 | |
| 3073 | // dfs_stack holds pairs of <HloInstruction*->unique_id(), HloInstruction*>. |
| 3074 | // |
| 3075 | // We need to keep track of both the id and the instruction because |
| 3076 | // instructions can get deleted while they are on the stack, so we |
| 3077 | // can't always use the (potentially dead) instruction object to grab |
| 3078 | // its id. |
| 3079 | DFSStack dfs_stack; |
| 3080 | dfs_stack.emplace_back(root->unique_id(), root); |
| 3081 | |
| 3082 | do { |
| 3083 | DCHECK(!dfs_stack.empty()); |
| 3084 | |
| 3085 | int current_id = dfs_stack.back().first; |
| 3086 | HloInstruction* current_node = dfs_stack.back().second; |
| 3087 | CHECK_GE(current_id, 0) << current_id << ": " << current_node |
| 3088 | << ": instruction may not have parent computation"; |
| 3089 | typename Visitor::VisitState visit_state = |
| 3090 | visitor->GetVisitState(current_id); |
| 3091 | if (visit_state == Visitor::kVisited) { |
| 3092 | dfs_stack.pop_back(); |
| 3093 | VLOG(3) << "Not visiting HLO (id = " << current_id |
| 3094 | << ") as it was already visited."; |
| 3095 | continue; |
| 3096 | } |
| 3097 | |
| 3098 | if (visit_state == Visitor::kVisiting) { |
| 3099 | dfs_stack.pop_back(); |
| 3100 | |
| 3101 | TF_RETURN_IF_ERROR(visitor->Preprocess(current_node)); |
| 3102 | VLOG(2) << "Visiting HLO %" << current_node->name(); |
| 3103 | TF_RETURN_IF_ERROR(current_node->Visit(visitor)); |
| 3104 | visitor->SetVisitState(current_id, Visitor::kVisited); |
| 3105 | TF_RETURN_IF_ERROR(visitor->Postprocess(current_node)); |
| 3106 | continue; |
| 3107 | } |
| 3108 | |
| 3109 | visitor->SetVisitState(current_id, Visitor::kVisiting); |
| 3110 | |
| 3111 | const size_t old_dfs_stack_size = dfs_stack.size(); |
| 3112 | for (HloInstruction* child : current_node->operands()) { |
| 3113 | if (!TF_PREDICT_TRUE(PushDFSChild(visitor, &dfs_stack, child))) { |
| 3114 | PrintCycle(child, &dfs_stack); |
| 3115 | return FailedPrecondition( |
| 3116 | "A cycle is detected while visiting instruction %s", |
| 3117 | current_node->ToString()); |
| 3118 | } |
| 3119 | } |
| 3120 |
no test coverage detected