推理引擎图优化 Pass 详细解析：常量折叠、死代码消除与公共子表达式消除的注册流程

推理引擎图优化 Pass 详细解析：常量折叠、死代码消除与公共子表达式消除的注册流程

一、当计算图成为性能瓶颈：图优化Pass的工程价值

在TFLite、ONNX Runtime、NCNN等推理引擎中，模型文件中的计算图是"忠实还原"训练框架导出结构的。但训练时的图结构和推理时的最优图结构之间存在显著差异——训练图包含梯度计算节点、BatchNorm的独立算子、未被剪枝的分支路径等推理阶段无用的结构。一个未经优化的MobileNetV2计算图可能包含约450个操作节点，而经过标准图优化Pass后的节点数可降至约300个——约33%的节点纯粹是冗余的。

图优化Pass（Graph Optimization Pass）的价值在于对计算图进行语义等价变换，在不改变模型输出的前提下消除冗余计算、合并算子、简化常量表达式。本文以三种最基础的优化Pass——常量折叠、死代码消除、公共子表达式消除——为例，深入分析其遍历算法、注册机制和相互作用关系。

二、三种优化Pass的算法原理与执行流程

2.1 常量折叠（Constant Folding）

常量折叠在编译期替换输入全部为常量的操作，将其输出预计算为常量节点。在推理引擎中，这直接减少了运行时的计算量。

flowchart TD subgraph 优化前["优化前：含常量操作的子图"] CW["权重 (Const)<br>shape: [64,32,3,3]"] --> CONV["Conv2D"] Input["输入 (Placeholder)"] --> CONV CONV --> ADD["Add"] CBias["偏置 (Const)<br>shape: [64]"] --> ADD ADD --> MUL["Mul"] CSCALE["缩放因子 (Const)<br>shape: [64]"] --> MUL MUL --> Output["输出"] end subgraph 优化后["常量折叠后：预计算的等价图"] CW2["已折叠权重 (Const)<br>= weight × scale"] --> CONV2["Conv2D"] Input2["输入 (Placeholder)"] --> CONV2 CBias2["已折叠偏置 (Const)<br>= bias × scale"] --> ADD2["Add (可选合并)"] CONV2 --> ADD2 ADD2 --> Output2["输出"] end style CW fill:#ffd,stroke:#d80 style CBias fill:#ffd,stroke:#d80 style CSCALE fill:#ffd,stroke:#d80

核心算法：对图中每个节点，检查其所有输入是否均为Const节点。若是，则通过该节点的算子实现计算输出值，将结果替换为一个新的Const节点，同时更新所有下游节点的输入引用。这是一个典型的worklist算法——一个折叠可能使下游节点变为全常量输入，触发连锁折叠。

2.2 死代码消除（Dead Code Elimination）

死代码消除移除那些输出不被任何下游节点使用的子图。在推理图中，死代码的常见来源包括：

• 训练阶段的辅助操作（如Loss计算、梯度更新）在模型导出时未被正确剥离；
• 常量折叠后，原来的Const节点变为无引用；
• 分支结构中被剪枝的路径。

flowchart TD subgraph 优化前["死代码示例：未连接的输出路径"] Input2["输入"] --> Conv1["Conv2D_1"] Conv1 --> ReLU1["ReLU_1"] Conv1 --> DeadBranch["Conv2D_dead"] DeadBranch --> DeadOutput["DeadOutput<br>（无消费者）"] ReLU1 --> FC["FullyConnected"] FC --> RealOutput["模型输出"] end subgraph 优化后["DCE后：移除无引用路径"] Input3["输入"] --> Conv3["Conv2D_1"] Conv3 --> ReLU3["ReLU_1"] ReLU3 --> FC3["FullyConnected"] FC3 --> RealOutput3["模型输出"] end style DeadBranch fill:#fee,stroke:#e33 style DeadOutput fill:#fee,stroke:#e33

算法以输出节点为根进行反向BFS/DFS遍历，标记所有可达节点。未标记的节点即为死代码，从图中安全移除。

2.3 公共子表达式消除（Common Subexpression Elimination）

CSE识别计算图中语义相同的多个子表达式（输入相同、操作相同、属性相同），将它们合并为一个，冗余引用全部指向合并后的节点。

flowchart TD subgraph 优化前["CSE前：重复的子表达式"] CX["x (输入)"] --> SIG1["Sigmoid"] CX --> SIG2["Sigmoid<br>（重复）"] SIG1 --> ADD1["Add"] SIG2 --> ADD2["Add<br>（重复）"] CY["y (输入)"] --> ADD1 CY --> ADD2 ADD1 --> OUT1["输出1"] ADD2 --> OUT2["输出2"] end subgraph 优化后["CSE后：共享子表达式"] CX2["x (输入)"] --> SIG3["Sigmoid<br>（共享）"] SIG3 --> ADD3["Add<br>（共享）"] CY2["y (输入)"] --> ADD3 ADD3 --> OUT3["输出1"] ADD3 --> OUT4["输出2"] end style SIG2 fill:#fee,stroke:#e33 style ADD2 fill:#fee,stroke:#e33

CSE的核心是节点哈希——将（操作类型，输入引用的哈希，属性）三元组作为键，检测重复。对于大型图，需要设计高效的哈希函数以在O(N)时间内完成去重。

2.4 三种Pass的相互作用与执行顺序

Pass之间存在依赖关系：DCE应在常量折叠之后执行（折叠产生的孤立Const节点需要被消除），CSE应在所有图结构变换之后执行。推荐的标准执行顺序：

flowchart LR A[原始计算图] --> B["1. 常量折叠<br>(消除运行时常量计算)"] B --> C["2. 死代码消除<br>(移除折叠产生的孤立节点)"] C --> D["3. 公共子表达式消除<br>(去重共享子图)"] D --> E["4. 常量折叠 (二次)<br>(CSE合并可能产生新常量路径)"] E --> F["5. 死代码消除 (二次)"] F --> G[优化后计算图] style A fill:#fff0f0 style G fill:#f0fff0

三次折叠与两次DCE的交替是必要的：CSE合并可能使之前需要运行时计算的路径变为全常量路径，需要第二轮折叠来捕获这些新机会。

三、图优化Pass框架的代码实现

以下代码在Python中实现一个轻量级的图优化Pass框架，使用类似TFLite的图表示：

""" 推理引擎图优化Pass框架 实现常量折叠、死代码消除、公共子表达式消除 """ from typing import Dict, List, Set, Tuple, Callable, Optional, Any from enum import Enum from dataclasses import dataclass, field from collections import deque # ===================== 图数据结构定义 ===================== class NodeType(Enum): CONST = 0 # 常量节点 PLACEHOLDER = 1 # 输入占位符 OPERATION = 2 # 可执行操作节点 @dataclass class GraphNode: """计算图节点""" id: int name: str node_type: NodeType op_type: str = "" # 操作类型，如 "Conv2D", "Add", "Mul" inputs: List[int] = field(default_factory=list) # 输入节点ID列表 outputs: List[int] = field(default_factory=list) # 输出节点ID列表 attrs: Dict[str, Any] = field(default_factory=dict) # 节点属性 data: Optional[Any] = None # Const节点的数据或中间计算结果 @dataclass class ComputeGraph: """计算图容器""" nodes: Dict[int, GraphNode] = field(default_factory=dict) inputs: List[int] = field(default_factory=list) # 图的输入节点ID outputs: List[int] = field(default_factory=list) # 图的输出节点ID def add_node(self, node: GraphNode) -> None: self.nodes[node.id] = node def get_node(self, node_id: int) -> Optional[GraphNode]: return self.nodes.get(node_id) # ===================== Pass 注册与执行框架 ===================== @dataclass class GraphPass: """图优化Pass定义""" name: str execute: Callable[[ComputeGraph], int] """执行优化并返回修改的节点数量""" class PassRegistry: """ Pass注册中心 管理所有优化Pass的注册、排序和执行 """ def __init__(self): self._passes: List[GraphPass] = [] def register(self, pass_instance: GraphPass) -> None: """注册一个优化Pass""" if any(p.name == pass_instance.name for p in self._passes): raise ValueError(f"Pass '{pass_instance.name}' 已注册") self._passes.append(pass_instance) def run_all(self, graph: ComputeGraph, max_iterations: int = 5) -> Dict[str, int]: """ 按注册顺序执行所有Pass，迭代至收敛 参数: graph: 待优化的计算图 max_iterations: 最大迭代次数（防止无限循环） 返回: {pass_name: total_changes} 每个Pass的总修改节点数 """ stats: Dict[str, int] = {} for iteration in range(max_iterations): total_changes = 0 for graph_pass in self._passes: changes = graph_pass.execute(graph) pass_name = graph_pass.name stats[pass_name] = stats.get(pass_name, 0) + changes total_changes += changes if changes > 0: # 每个Pass后运行死代码消除（如果当前Pass不是DCE自身） if "dead_code" not in pass_name.lower(): _pass = GraphPass( name="dead_code_elimination_auto", execute=dead_code_elimination ) auto_changes = _pass.execute(graph) stats["dce_auto"] = stats.get("dce_auto", 0) + auto_changes if total_changes == 0: # 收敛：本轮无任何修改 break return stats # ===================== Pass 1: 常量折叠 ===================== def constant_folding(graph: ComputeGraph) -> int: """ 常量折叠：输入全为Const的操作 → 预计算为Const节点 返回: 折叠的节点数量 """ # 操作类型到计算函数的映射 FOLD_OPS: Dict[str, Callable] = { 'Add': lambda a, b: a + b, 'Sub': lambda a, b: a - b, 'Mul': lambda a, b: a * b, 'Div': lambda a, b: a / b if b != 0 else float('nan'), 'Reshape': lambda tensor, shape: tensor, # 简化：不实际reshape } changes = 0 # 使用worklist算法，因为折叠可能使下游节点变为可折叠 changed = True while changed: changed = False # 收集当前可折叠的节点（避免在迭代中修改字典） foldable: List[int] = [] for node_id, node in graph.nodes.items(): if node.node_type != NodeType.OPERATION: continue if node.op_type not in FOLD_OPS: continue if not node.inputs: continue # 检查所有输入是否都是Const节点 all_const = True for input_id in node.inputs: input_node = graph.get_node(input_id) if input_node is None or input_node.node_type != NodeType.CONST: all_const = False break if all_const: foldable.append(node_id) # 执行折叠 for node_id in foldable: node = graph.nodes[node_id] try: # 收集所有Const输入的数据 input_data = [] for input_id in node.inputs: input_node = graph.get_node(input_id) input_data.append(input_node.data) # 执行操作 fold_func = FOLD_OPS[node.op_type] result = fold_func(*input_data) # 将操作节点转换为Const节点 node.node_type = NodeType.CONST node.data = result node.op_type = "" node.inputs = [] # Const节点无输入 node.attrs = {} changes += 1 changed = True # 可能产生新的可折叠节点 except Exception as e: # 折叠失败（如除零），保持原样 print(f" 警告: 折叠节点 '{node.name}' 失败: {e}") continue return changes # ===================== Pass 2: 死代码消除 ===================== def dead_code_elimination(graph: ComputeGraph) -> int: """ 死代码消除：从输出节点反向标记可达节点，删除不可达节点 返回: 删除的节点数量 """ if not graph.nodes: return 0 # 步骤1: 从所有输出节点反向BFS，标记可达节点 reachable: Set[int] = set() queue = deque() for output_id in graph.outputs: if output_id in graph.nodes: queue.append(output_id) while queue: node_id = queue.popleft() if node_id in reachable: continue reachable.add(node_id) node = graph.get_node(node_id) if node is None: continue for input_id in node.inputs: if input_id in graph.nodes and input_id not in reachable: queue.append(input_id) # 步骤2: 删除不可达节点 removed = 0 to_remove = [nid for nid in graph.nodes if nid not in reachable] for node_id in to_remove: # 清理其他节点的引用 node = graph.nodes[node_id] for input_id in node.inputs: input_node = graph.get_node(input_id) if input_node and node_id in input_node.outputs: input_node.outputs.remove(node_id) for output_id in node.outputs: output_node = graph.get_node(output_id) if output_node and node_id in output_node.inputs: output_node.inputs.remove(node_id) del graph.nodes[node_id] removed += 1 return removed # ===================== Pass 3: 公共子表达式消除 ===================== def _node_hash(graph: ComputeGraph, node_id: int) -> Optional[int]: """ 计算节点的语义哈希 相同操作+相同输入+相同属性 → 相同哈希 返回: 哈希值，节点信息不完整时返回None """ node = graph.get_node(node_id) if node is None: return None if node.node_type == NodeType.CONST: # Const节点按数据值比较（简化：使用id或数据哈希） return hash(('__CONST__', id(node.data))) elif node.node_type == NodeType.OPERATION: # 操作节点：类型 + 输入ID集合 + 关键属性的哈希 input_tuple = tuple(sorted(node.inputs)) attr_tuple = tuple(sorted( (k, v) for k, v in node.attrs.items() if isinstance(v, (int, float, str, bool, tuple)) )) return hash((node.op_type, input_tuple, attr_tuple)) else: # Placeholder 节点不合并（每个输入是独立的） return None def common_subexpression_elimination(graph: ComputeGraph) -> int: """ 公共子表达式消除：合并语义重复的节点 返回: 消除的重复节点数量 """ # 构建语义哈希 → 节点ID的映射 hash_to_node: Dict[int, int] = {} replacements: Dict[int, int] = {} # 待替换映射: old_id → new_id changes = 0 for node_id, node in graph.nodes.items(): h = _node_hash(graph, node_id) if h is None or h not in hash_to_node: if h is not None: hash_to_node[h] = node_id continue # 找到重复节点：记录替换映射 canonical_id = hash_to_node[h] canonical_node = graph.get_node(canonical_id) if canonical_node is None: continue # 验证语义等价（深度校验，防止哈希碰撞） if node.node_type != canonical_node.node_type: continue if node.node_type == NodeType.OPERATION: if (node.op_type != canonical_node.op_type or set(node.inputs) != set(canonical_node.inputs)): continue # 标记替换：所有引用node_id的地方改为引用canonical_id replacements[node_id] = canonical_id # 执行替换 for old_id, new_id in replacements.items(): old_node = graph.nodes[old_id] # 更新所有节点的输入引用 for _, node in graph.nodes.items(): if old_id in node.inputs: # 替换引用 new_inputs = [new_id if i == old_id else i for i in node.inputs] node.inputs = new_inputs # 更新图输出引用 new_outputs = [new_id if o == old_id else o for o in graph.outputs] graph.outputs = new_outputs # 删除重复节点 del graph.nodes[old_id] changes += 1 return changes # ===================== 使用示例 ===================== def example_usage(): """演示三种Pass的组合使用""" # 构建示例图 graph = ComputeGraph() # 常量节点 c1 = GraphNode(1, "const_2", NodeType.CONST, data=2.0) c3 = GraphNode(3, "const_3", NodeType.CONST, data=3.0) c5 = GraphNode(5, "const_5", NodeType.CONST, data=5.0) graph.add_node(c1); graph.add_node(c3); graph.add_node(c5) # 输入节点 inp = GraphNode(0, "input", NodeType.PLACEHOLDER) graph.add_node(inp); graph.inputs = [0] # 操作节点 add1 = GraphNode(10, "add_1", NodeType.OPERATION, op_type="Add", inputs=[1, 3]) mul1 = GraphNode(11, "mul_1", NodeType.OPERATION, op_type="Mul", inputs=[10, 5]) add2 = GraphNode(12, "add_2", NodeType.OPERATION, op_type="Add", inputs=[0, 11]) # 死代码节点: 输出不被使用 dead = GraphNode(99, "dead_op", NodeType.OPERATION, op_type="Mul", inputs=[3, 5]) graph.add_node(add1); graph.add_node(mul1); graph.add_node(add2) graph.add_node(dead) graph.outputs = [12] print("=== 原始图 ===") print(f"节点数: {len(graph.nodes)}, 输出: {graph.outputs}") # 注册Pass registry = PassRegistry() registry.register(GraphPass("constant_folding", constant_folding)) registry.register(GraphPass("cse", common_subexpression_elimination)) registry.register(GraphPass("dce", dead_code_elimination)) # 执行优化 stats = registry.run_all(graph) print(f"\n=== 优化结果 ===") for name, count in stats.items(): print(f" {name}: {count} 处修改") print(f"优化后节点数: {len(graph.nodes)}, 输出: {graph.outputs}") if __name__ == "__main__": example_usage()

四、图优化Pass的设计权衡与工程陷阱

Const数据的生命周期管理：常量折叠将运行时计算替换为Const节点，但Const节点的数据量不可忽略。一个全连接层的常量折叠可能产生数MB的权重矩阵。在资源受限的边缘设备上，折叠后的Const节点总量可能超过可用内存。解决方案是在折叠时增加数据量阈值检查——超过阈值的常量操作保留为运行时计算，以空间换时间。

节点哈希的碰撞风险：CSE的核心是哈希匹配。在大型图中（10K+节点），哈希碰撞可能导致语义不同的节点被错误合并，产生难以察觉的精度偏差。应在哈希匹配后进行严格的语义等价验证（操作类型、输入拓扑、属性全部一致），仅在确认等价后执行替换。

与量化Pass的顺序冲突：图优化Pass通常应在量化之前执行（折叠减少计算 → 量化减少精度）。但部分引擎在量化后执行的算子融合Pass（如Conv+ReLU合并）可能破坏已消除的公共子表达式。Pass顺序的管理需在框架层面统一编排，而非Ad-Hoc叠加。

优化收敛的无限循环风险：折叠 → DCE → CSE → 折叠的迭代可能在某些不规则图中形成无限循环（如折叠产生新Const → DCE删除 → CSE合并产生新折叠机会 → ...）。虽然实践中极少出现，但框架必须在PassRunner中设置迭代上限（如最大5轮或超时退出）。

五、总结

常量折叠、死代码消除和公共子表达式消除是推理引擎图优化的基础Pass，按照"折叠→DCE→CSE→二次折叠→二次DCE"的执行顺序可最大化收益。常量折叠将全Const操作预计算，DCE从输出反向标记删除死代码，CSE通过语义哈希去重子表达式。

工程落地建议：在推理引擎的模型加载阶段插入7-10个基础优化Pass（除本文介绍的三种外，还包括算子融合、布局转换、广播消除等）；在PassRunner中实现迭代至收敛的控制逻辑（设置迭代上限）；对每个Pass的输出添加节点计数日志，便于性能回归分析；在CI中对比优化前后的模型拓扑差异，确保优化变换的语义正确性。