| 241 | } |
| 242 | |
| 243 | void CallGraph::SetNodeDepths() { |
| 244 | std::queue<CallGraphNode*> worklist; |
| 245 | |
| 246 | // Initialize node depths to -1. |
| 247 | for (CallGraphNode& node : nodes_) { |
| 248 | node.set_depth(-1); |
| 249 | } |
| 250 | |
| 251 | // Initialize worklist with all roots of the call graph (computations without |
| 252 | // callers). |
| 253 | for (const HloComputation* computation : module_->computations()) { |
| 254 | CallGraphNode& node = GetNode(computation); |
| 255 | if (node.callers().empty()) { |
| 256 | node.set_depth(0); |
| 257 | worklist.push(&node); |
| 258 | } |
| 259 | } |
| 260 | |
| 261 | while (!worklist.empty()) { |
| 262 | CallGraphNode* node = worklist.front(); |
| 263 | worklist.pop(); |
| 264 | for (const HloComputation* callee : node->callees()) { |
| 265 | CallGraphNode& callee_node = GetNode(callee); |
| 266 | if (callee_node.depth() < node->depth() + 1) { |
| 267 | callee_node.set_depth(node->depth() + 1); |
| 268 | worklist.push(&callee_node); |
| 269 | } |
| 270 | } |
| 271 | } |
| 272 | |
| 273 | for (CallGraphNode& node : nodes_) { |
| 274 | CHECK_NE(node.depth(), -1); |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | /* static */ |
| 279 | std::unique_ptr<CallGraph> CallGraph::Build(const HloModule* module) { |