| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Iterable, Iterator |
| 5 | |
| 6 | from torch import fx |
| 7 | from torch._higher_order_ops.auto_functionalize import auto_functionalized |
| 8 | from torch._ops import OpOverload, OpOverloadPacket |
| 9 | from torch.fx.node import Target |
| 10 | |
| 11 | |
| 12 | def is_func(node: fx.Node, target: Target) -> bool: |
| 13 | return bool(node.op == "call_function" and node.target == target) |
| 14 | |
| 15 | |
| 16 | # Returns the first auto_functionalized node with the given op (if it exists) |
| 17 | def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None: |
| 18 | for node in nodes: |
| 19 | if is_func(node, auto_functionalized) and node.args[0] == op: # noqa |
| 20 | return node |
| 21 | return None |
| 22 | |
| 23 | |
| 24 | # Returns the first auto_functionalized node with the given op |
| 25 | def find_auto_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node: |
| 26 | node = find_auto_fn_maybe(nodes, op) |
| 27 | assert node is not None, f"Could not find {op} in nodes {nodes}" |
| 28 | return node |
| 29 | |
| 30 | |
| 31 | # An auto-functionalization-aware utility for finding nodes with a specific op |
| 32 | # Also handles op overload packets and finds all overloads |
| 33 | def find_op_nodes( |
| 34 | op: OpOverload | OpOverloadPacket, graph: fx.Graph |
| 35 | ) -> Iterator[fx.Node]: |
| 36 | if isinstance(op, OpOverloadPacket): |
| 37 | for overload in op.overloads(): |
| 38 | overload_op = getattr(op, overload) |
| 39 | yield from find_op_nodes(overload_op, graph) |
| 40 | return |
| 41 | |
| 42 | assert isinstance(op, OpOverload) |
| 43 | |
| 44 | yield from graph.find_nodes(op="call_function", target=op) |
| 45 | |
| 46 | for n in graph.find_nodes(op="call_function", target=auto_functionalized): |
| 47 | if n.args[0] == op: |
| 48 | yield n |