| 2344 | } |
| 2345 | |
| 2346 | void PrintCycle(const HloInstruction* child, DFSStack* dfs_stack) { |
| 2347 | // This set contains HloInstructions from the top of `DFSStack` that might |
| 2348 | // belong to the cycle, i.e. if DFSStack :=[back,...,child,...,top], then |
| 2349 | // `subgraph` := {child,...,top}. |
| 2350 | absl::flat_hash_set<const HloInstruction*> subgraph; |
| 2351 | while (!dfs_stack->empty() && dfs_stack->back().second != child) { |
| 2352 | subgraph.insert(dfs_stack->back().second); |
| 2353 | dfs_stack->pop_back(); |
| 2354 | } |
| 2355 | // Start dfs at `child` and find a cycle with all nodes in `subgraph`. |
| 2356 | absl::flat_hash_set<const HloInstruction*> visited; |
| 2357 | absl::InlinedVector<const HloInstruction*, 16> dfs; |
| 2358 | dfs.push_back(child); |
| 2359 | while (!dfs.empty()) { |
| 2360 | bool found_next_instr = false; |
| 2361 | for (const auto& user : dfs.back()->users()) { |
| 2362 | if (user == child) { |
| 2363 | dfs.push_back(child); |
| 2364 | LOG(INFO) << "\n\nDirected cycle:\n " |
| 2365 | << absl::StrJoin( |
| 2366 | dfs, "\n ", |
| 2367 | [](std::string* out, const HloInstruction* instr) { |
| 2368 | out->append(instr->name()); |
| 2369 | }); |
| 2370 | return; |
| 2371 | } |
| 2372 | if (!subgraph.contains(user) || visited.contains(user)) { |
| 2373 | continue; |
| 2374 | } |
| 2375 | visited.insert(user); |
| 2376 | dfs.push_back(user); |
| 2377 | found_next_instr = true; |
| 2378 | } |
| 2379 | if (!found_next_instr) { |
| 2380 | dfs.pop_back(); |
| 2381 | } |
| 2382 | } |
| 2383 | } |
| 2384 | |
| 2385 | } // namespace |
| 2386 | |