| 63 | } |
| 64 | |
| 65 | StatusOr<bool> ConditionalToSelect::Run(HloModule* module) { |
| 66 | std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module); |
| 67 | bool did_mutate = false; |
| 68 | VLOG(1) << "Running conditional-to-select pass"; |
| 69 | TF_RETURN_IF_ERROR( |
| 70 | call_graph->VisitNodes([&](const CallGraphNode& node) -> Status { |
| 71 | std::vector<HloInstruction*> ToInline; |
| 72 | if (node.context() != CallContext::kParallel) { |
| 73 | return Status::OK(); |
| 74 | } |
| 75 | for (const CallSite& callsite : node.callsites()) { |
| 76 | if (callsite.instruction()->opcode() == HloOpcode::kConditional) { |
| 77 | VLOG(1) << "Visiting conditional: " << callsite.ToString(); |
| 78 | HloInstruction* conditional = callsite.instruction(); |
| 79 | TF_ASSIGN_OR_RETURN(bool result, |
| 80 | DoConditionalToSelect(conditional)); |
| 81 | did_mutate |= result; |
| 82 | } |
| 83 | } |
| 84 | return Status::OK(); |
| 85 | })); |
| 86 | return did_mutate; |
| 87 | } |
| 88 | |
| 89 | } // namespace xla |
nothing calls this directly
no test coverage detected