PyTorch/FX Graph

FX:按逆拓扑顺序删除无用计算

借助计算依赖图与使用者计数,清理既不影响输出也无副作用的节点。

v2.5.1图处理与编译核验 2026-10-08
相关概念

工程问题#

模型图变换之后会残留无用算子;随意删除节点可能破坏依赖或丢失副作用。

实际实现#

FX Graph.lint 先校验节点已按拓扑次序排列,再逆序遍历;仅删除没有使用者且被判定无副作用的节点,删除会更新其输入的使用者集合。

处理流程#

  1. 验证依赖先定义后使用
  2. 从图尾向前遍历节点
  3. 检查使用者数量与副作用
  4. 删除符合条件的节点并更新依赖

源码片段#

固定版本 · 原文摘录
torch/fx/graph.pyv2.5.1 · L1641–L1659
查看原文件
1641        # Lint the graph first to make sure its topologically sorted, otherwise1642        # DCE below will not behave as expected.1643        self.lint()16441645        def has_side_effect(node):1646            if is_impure_node is not None:1647                return is_impure_node(node)1648            return node.is_impure()16491650        # Reverse iterate so that when we remove a node, any nodes used as an1651        # input to that node have an updated user count that no longer reflects1652        # the removed node.1653        changed = False1654        for node in reversed(self.nodes):1655            if not has_side_effect(node) and len(node.users) == 0:1656                self.erase_node(node)1657                changed = True16581659        return changed
Graph.eliminate_dead_code

lint 只验证图已满足拓扑顺序,并不会进行拓扑排序。逆序遍历让删除节点后的使用者计数及时传递给上游;仅移除无使用者且无副作用的节点。源码明确警告默认副作用判定不全面,需要纯函数图或自定义可靠判定才能安全使用。

收益与取舍#

工程分析

带来的收益

  • 删除后的依赖计数及时影响上游判断
  • 为后续代码生成清理冗余图节点

付出的代价

  • 默认副作用判断不覆盖所有操作
  • 图重写者必须维持合法依赖与拓扑顺序

可比较的选择

从必要输出反向标记活跃节点另一类死代码分析方式,同样需要正确识别必须保留的副作用根节点。

版本与适用范围#

这里利用并校验已有拓扑顺序,并非 lint 自动执行拓扑排序。只在操作纯函数化或提供可靠副作用判断时安全应用死代码删除。

参考来源#

源码、官方文档与原始论文
01
FX Graph 校验与死代码删除源码

torch/fx/graph.py · Graph.lint / Graph.eliminate_dead_code

源码v2.5.1

相关实现与比较#

4 个相关案例
5 个匹配案例