| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from torch import fx |
| 4 | |
| 5 | from ..vllm_inductor_pass import VllmInductorPass |
| 6 | |
| 7 | |
| 8 | class PostCleanupPass(VllmInductorPass): |
| 9 | """This pass performs cleanup after custom passes. |
| 10 | It topologically sorts the graph and removes unused nodes. |
| 11 | This is needed because the pattern matcher does not guarantee producing |
| 12 | a topologically sorted graph, and there may be unused nodes left around. |
| 13 | """ |
| 14 | |
| 15 | @VllmInductorPass.time_and_log |
| 16 | def __call__(self, graph: fx.Graph) -> None: |
| 17 | from torch._inductor.pattern_matcher import stable_topological_sort |
| 18 | |
| 19 | stable_topological_sort(graph) |
| 20 | graph.eliminate_dead_code() |