HEC：基于动态规则生成的MLIR等价性验证工具

1. 项目概述

HEC（Hybrid Equivalence Checker）是一款基于动态规则生成的MLIR等价性验证工具，专门用于验证编译器优化前后的代码语义一致性。在编译器优化领域，确保变换前后代码的语义等价性至关重要，但传统方法往往难以处理复杂的控制流变换场景。

HEC的核心创新在于将静态数据路径重写规则与动态规则生成技术相结合，利用e-graph数据结构实现高效的语义等价验证。与现有工具相比，HEC能够处理更广泛的优化场景，特别是那些涉及循环变换和控制流修改的情况。

提示：MLIR（Multi-Level Intermediate Representation）是近年来兴起的一种编译器中间表示，它采用模块化设计，支持多级抽象，特别适合硬件加速和高性能计算场景下的代码优化。

1.1 核心需求解析

程序等价性验证面临两大核心挑战：

1. 静态规则的局限性：传统方法依赖预定义的静态重写规则，难以覆盖编译器优化中可能出现的所有变换模式，特别是那些与具体程序上下文相关的变换。
2. 验证效率问题：完全依赖形式化方法（如SMT求解器）进行验证会导致严重的性能问题，无法扩展到实际规模的代码。

HEC通过以下方式解决这些挑战：

• 对数据路径变换，使用基于数学恒等式（如德摩根定律、算术结合律）的静态规则
• 对控制流变换，采用动态规则生成技术，仅在必要时调用Z3 SMT求解器进行验证
• 利用e-graph的共享结构特性，避免重复计算，提高验证效率

2. 技术架构与原理

2.1 e-graph基础结构

e-graph（equivalence graph）是HEC的核心数据结构，它能够高效表示和操作程序表达式之间的等价关系。e-graph由以下部分组成：

1. e-node（等价节点）：表示程序中的基本操作（如加载、存储、算术运算）
2. e-class（等价类）：包含语义等价的一组e-node
3. 合并操作：将两个e-class合并，表示它们语义等价

e-graph的关键特性是同余闭包（congruence closure）：如果两个表达式的所有子表达式都等价，则这两个表达式也等价。这一特性使得e-graph能够自动发现和传播等价关系。

2.1.1 e-graph构造算法

HEC的e-graph构造过程如算法1所示：

1. 对输入程序的图表示进行拓扑排序
2. 按照从叶节点到根节点的顺序插入操作
3. 为每个操作生成唯一标识符并建立输入输出关系

这种构造方式确保了子节点总是先于父节点被处理，为后续的等价性验证奠定了基础。

2.2 动态规则生成机制

HEC的规则生成器能够自动识别潜在的变换模式并生成相应的验证规则。动态规则生成分为三个阶段：

1. 模式识别：分析程序图表示，识别可能的变换模式（如循环展开、循环融合）
2. 规则生成：为候选操作创建动态重写规则
3. 条件验证：使用Z3 SMT求解器验证变换条件的数学正确性

特别值得注意的是，HEC将Z3的使用限制在动态规则条件验证环节，而不是整个饱和过程，这显著提高了工具的可扩展性。测试表明，动态规则生成（包括Z3检查）通常在1秒内完成，仅占整个验证运行时的一小部分。

2.3 验证流程详解

HEC的验证流程通过图7的案例可以得到清晰展示。考虑一个嵌套循环展开的场景：

1. 初始验证：仅使用静态规则，如果差异在静态规则处理范围内，直接报告等价性
2. 动态规则应用：对于无法通过静态规则处理的差异，创建伪组合节点绑定候选循环操作
3. 迭代验证：通过多次迭代逐步合并等价类，直至确认等价性或耗尽所有可能性

在每次迭代中，HEC都会：

• 将e-graph转换回图表示以便进一步处理
• 应用新生成的动态规则
• 检查是否已达到等价状态

这种迭代方法确保了验证的完备性，同时通过动态规则生成避免了不必要的计算开销。

3. 关键实现细节

3.1 循环展开的等价性验证

循环展开是编译器常用的优化技术，但传统的验证工具往往难以处理。HEC通过专门的动态规则处理这类变换。以图6的案例为例：

原始循环：

affine.for %arg2 = 0 to %0 { %1 = affine.load %arg1[%arg2] : memref<?xf64> }

展开后代码（展开因子2和3）：

affine.for %arg2 = 0 to #map1()[%0] step 6 { // 展开后的循环体 } affine.for %arg2 = #map1()[%0] to #map2()[%0] step 2 { // 剩余迭代处理 }

HEC的验证过程会：

1. 识别展开模式
2. 生成合并规则将多个循环合并为等效的单一循环
3. 验证迭代空间保持不变

3.2 内存访问模式验证

HEC特别关注内存操作（加载/存储）的语义一致性验证。对于如下内存操作序列：

%1 = affine.load %A[%i] : memref<?xf32> %2 = arith.addf %1, %cst : f32 affine.store %2, %A[%i] : memref<?xf32>

HEC会：

1. 跟踪内存访问地址
2. 分析读写依赖关系
3. 确保优化不改变内存访问顺序和次数

这一机制使得HEC能够检测出Listing 12中展示的内存RAW（Read-After-Write）违规问题。

3.3 边界条件处理

循环边界检查是循环变换中常见的错误来源。HEC通过动态规则生成专门处理各种边界情况，如：

• 循环上界小于下界时的行为
• 非整数倍展开因子的剩余迭代处理
• 跨循环迭代的数据依赖关系

在PolyBench测试集中，HEC成功识别出了mlir-opt编译器在Jacobi_1d和Seidel_2d基准测试中引入的循环边界检查错误。

4. 性能评估与实验结果

4.1 实验设置

评估环境：

• CPU: Intel Xeon Gold 6418H (48物理核心)
• 内存: 1024GB
• MLIR编译器: mlir-opt (LLVM 18.0.0)
• 基准测试集: PolyBench和PolyBench-NN

4.2 控制流变换验证

表4展示了HEC在各种变换配置下的性能：

• 循环平铺：验证时间稳定在7-28秒之间
• 循环展开：验证时间随展开因子增大而增加，但大多数情况仍在1分钟内完成
• 混合变换：同时包含平铺和展开的复杂变换，验证时间可达数分钟

关键发现：

1. e-class数量是影响验证时间的主要因素
2. 嵌套展开的验证时间呈指数增长（图9）
3. 动态规则生成时间可以忽略不计（<1秒）

4.3 数据路径变换评估

图10展示了HEC处理大规模数据路径变换的能力：

• 问题规模：15,000到90,000行代码
• 验证时间：与代码规模呈线性关系
• 最大测试案例（108,012 LOC）：验证时间2,305秒

这表明HEC能够高效处理实际规模的代码验证任务。

4.4 错误检测能力

HEC在PolyBench测试集中发现了mlir-opt编译器的两类关键错误：

1. 循环边界检查错误（Listing 10）：

   • 当循环上界小于下界时，原始代码不执行
   • 但展开后的代码错误地执行了剩余迭代
   • 影响Jacobi_1d和Seidel_2d基准测试

2. 内存RAW违规（Listing 12）：

   • 循环融合改变了内存访问顺序
   • 导致计算结果与原始代码不一致
   • 可能引发严重的数据一致性问题

5. 使用经验与最佳实践

5.1 性能优化技巧

基于实际使用经验，我们总结了以下优化建议：

1. 控制展开因子：嵌套展开的验证时间随展开因子呈指数增长（图8）。实践中，很少需要超过16的展开因子。

2. 增量验证：对于大型优化序列，建议分阶段验证，而不是一次性验证所有变换。

3. 规则优先级：为常见模式设置高优先级静态规则，可以减少动态规则生成的开销。

5.2 常见问题排查

在实际使用中，我们遇到了以下典型问题及解决方案：

问题1：验证时间异常长

• 检查e-class数量是否过多
• 尝试简化优化序列或降低展开因子
• 检查是否有无限循环的规则应用

问题2：误报不等价

• 确认所有相关动态规则已正确生成
• 检查Z3求解器是否超时
• 验证输入程序是否包含未定义行为

问题3：内存消耗过大

• 限制e-graph大小
• 启用更激进的垃圾回收
• 考虑分模块验证

5.3 工具集成建议

将HEC集成到现有编译流程时：

1. 作为优化pass：在关键优化阶段后插入验证pass
2. 调试辅助：当优化结果异常时，调用HEC定位问题变换
3. 持续集成：在CI流程中加入关键优化的等价性检查

6. 技术对比与优势分析

6.1 与现有工具对比

表5比较了HEC与主流验证工具的特性：

HEC的主要优势在于：

1. 唯一支持MLIR级别的动态控制流变换验证
2. 通过e-graph实现高效的等价类管理
3. 可扩展的规则生成机制

6.2 局限性

HEC目前存在以下限制：

1. 不完备性：无法验证所有可能的等价变换，特别是那些需要特殊规则的情况
2. 资源消耗：极端情况下（如大展开因子）验证成本较高
3. 模式覆盖：需要人工定义新的变换模式模板

这些限制也是未来改进的方向。在实际使用中，我们发现HEC已经能够覆盖绝大多数常见的编译器优化场景，特别是那些在HLS（高级综合）流程中频繁使用的变换。