| 135 | } |
| 136 | |
| 137 | StatusOr<bool> CallInliner::Run(HloModule* module) { |
| 138 | std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module); |
| 139 | // Because call graph nodes are visited in post-order (callees before callers) |
| 140 | // we'll always inline kCalls into their callers in the appropriate order. |
| 141 | bool did_mutate = false; |
| 142 | TF_RETURN_IF_ERROR( |
| 143 | call_graph->VisitNodes([&](const CallGraphNode& node) -> Status { |
| 144 | VLOG(1) << "Visiting node: " << node.ToString(); |
| 145 | for (HloInstruction* instruction : |
| 146 | node.computation()->MakeInstructionPostOrder()) { |
| 147 | if (instruction->opcode() == HloOpcode::kCall && |
| 148 | (!single_call_site_ || |
| 149 | call_graph->GetNode(instruction->to_apply()) |
| 150 | .caller_callsites() |
| 151 | .size() == 1)) { |
| 152 | TF_RETURN_IF_ERROR(Inline(instruction).status()); |
| 153 | did_mutate = true; |
| 154 | } |
| 155 | } |
| 156 | return Status::OK(); |
| 157 | })); |
| 158 | if (did_mutate) { |
| 159 | // Run DCE to remove called computations which are now becoming unused. |
| 160 | // This can result then in problems if within the called computation, there |
| 161 | // were send/recv instructions, which the module group verifier will flag as |
| 162 | // error findingthe same channel ID used for multiple send/recv |
| 163 | // instructions. |
| 164 | TF_RETURN_IF_ERROR(HloDCE().Run(module).status()); |
| 165 | } |
| 166 | return did_mutate; |
| 167 | } |
| 168 | |
| 169 | } // namespace xla |
nothing calls this directly
no test coverage detected