推理引擎图优化 Pass 详细解析常量折叠、死代码消除与公共子表达式消除的注册流程一、当计算图成为性能瓶颈图优化Pass的工程价值在TFLite、ONNX Runtime、NCNN等推理引擎中模型文件中的计算图是忠实还原训练框架导出结构的。但训练时的图结构和推理时的最优图结构之间存在显著差异——训练图包含梯度计算节点、BatchNorm的独立算子、未被剪枝的分支路径等推理阶段无用的结构。一个未经优化的MobileNetV2计算图可能包含约450个操作节点而经过标准图优化Pass后的节点数可降至约300个——约33%的节点纯粹是冗余的。图优化PassGraph Optimization Pass的价值在于对计算图进行语义等价变换在不改变模型输出的前提下消除冗余计算、合并算子、简化常量表达式。本文以三种最基础的优化Pass——常量折叠、死代码消除、公共子表达式消除——为例深入分析其遍历算法、注册机制和相互作用关系。二、三种优化Pass的算法原理与执行流程2.1 常量折叠Constant Folding常量折叠在编译期替换输入全部为常量的操作将其输出预计算为常量节点。在推理引擎中这直接减少了运行时的计算量。flowchart TD subgraph 优化前[优化前含常量操作的子图] CW[权重 (Const)brshape: [64,32,3,3]] -- CONV[Conv2D] Input[输入 (Placeholder)] -- CONV CONV -- ADD[Add] CBias[偏置 (Const)brshape: [64]] -- ADD ADD -- MUL[Mul] CSCALE[缩放因子 (Const)brshape: [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[DeadOutputbr无消费者] 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 EliminationCSE识别计算图中语义相同的多个子表达式输入相同、操作相同、属性相同将它们合并为一个冗余引用全部指向合并后的节点。flowchart TD subgraph 优化前[CSE前重复的子表达式] CX[x (输入)] -- SIG1[Sigmoid] CX -- SIG2[Sigmoidbr重复] SIG1 -- ADD1[Add] SIG2 -- ADD2[Addbr重复] CY[y (输入)] -- ADD1 CY -- ADD2 ADD1 -- OUT1[输出1] ADD2 -- OUT2[输出2] end subgraph 优化后[CSE后共享子表达式] CX2[x (输入)] -- SIG3[Sigmoidbr共享] SIG3 -- ADD3[Addbr共享] CY2[y (输入)] -- ADD3 ADD3 -- OUT3[输出1] ADD3 -- OUT4[输出2] end style SIG2 fill:#fee,stroke:#e33 style ADD2 fill:#fee,stroke:#e33CSE的核心是节点哈希——将操作类型输入引用的哈希属性三元组作为键检测重复。对于大型图需要设计高效的哈希函数以在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_factorylist) # 输入节点ID列表 outputs: List[int] field(default_factorylist) # 输出节点ID列表 attrs: Dict[str, Any] field(default_factorydict) # 节点属性 data: Optional[Any] None # Const节点的数据或中间计算结果 dataclass class ComputeGraph: 计算图容器 nodes: Dict[int, GraphNode] field(default_factorydict) inputs: List[int] field(default_factorylist) # 图的输入节点ID outputs: List[int] field(default_factorylist) # 图的输出节点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(fPass {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( namedead_code_elimination_auto, executedead_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, data2.0) c3 GraphNode(3, const_3, NodeType.CONST, data3.0) c5 GraphNode(5, const_5, NodeType.CONST, data5.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_typeAdd, inputs[1, 3]) mul1 GraphNode(11, mul_1, NodeType.OPERATION, op_typeMul, inputs[10, 5]) add2 GraphNode(12, add_2, NodeType.OPERATION, op_typeAdd, inputs[0, 11]) # 死代码节点: 输出不被使用 dead GraphNode(99, dead_op, NodeType.OPERATION, op_typeMul, 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如ConvReLU合并可能破坏已消除的公共子表达式。Pass顺序的管理需在框架层面统一编排而非Ad-Hoc叠加。优化收敛的无限循环风险折叠 → DCE → CSE → 折叠的迭代可能在某些不规则图中形成无限循环如折叠产生新Const → DCE删除 → CSE合并产生新折叠机会 → ...。虽然实践中极少出现但框架必须在PassRunner中设置迭代上限如最大5轮或超时退出。五、总结常量折叠、死代码消除和公共子表达式消除是推理引擎图优化的基础Pass按照折叠→DCE→CSE→二次折叠→二次DCE的执行顺序可最大化收益。常量折叠将全Const操作预计算DCE从输出反向标记删除死代码CSE通过语义哈希去重子表达式。工程落地建议在推理引擎的模型加载阶段插入7-10个基础优化Pass除本文介绍的三种外还包括算子融合、布局转换、广播消除等在PassRunner中实现迭代至收敛的控制逻辑设置迭代上限对每个Pass的输出添加节点计数日志便于性能回归分析在CI中对比优化前后的模型拓扑差异确保优化变换的语义正确性。