HARDBOILED IR：面向张量计算的编译器优化设计

1. HARDBOILED IR设计精要解析

在编译器优化领域，中间表示（IR）的设计质量直接决定了后续优化的上限。HARDBOILED IR针对现代硬件加速器（如Intel AMX和NVIDIA WMMA）的特性，构建了一套专门面向张量计算的IR系统。其核心设计哲学可以概括为：通过最小化的语法结构捕获最大化的硬件特性，同时保持足够的抽象层级以支持跨平台优化。

1.1 语法结构设计

HARDBOILED的语法设计采用了典型的函数式风格，将程序表示为语句（Stmt）和表达式（Expr）的组合。这种设计在保持简洁性的同时，能够精确描述张量计算的关键特征：

// 存储操作：将向量表达式结果写入指定缓冲区的索引位置 Store(buffer_name, index_pattern, vector_expr) // 计算操作：对表达式求值（常用于有副作用的硬件指令） Evaluate(expr)

表达式系统特别强化了对向量化操作的支持：

• load/store显式标记内存操作
• cast处理不同精度张量间的类型转换
• call封装硬件加速指令（如tile_matmul）
• broadcast/ramp构建索引向量
• vector_reduce_add实现向量归约

提示：ramp(base, stride, length)实际上定义了一个算术序列向量，这在描述内存访问模式时特别有用。例如ramp(0, 1, 256)对应连续256个元素的线性访问。

1.2 硬件加速支持

IR通过Location类型显式区分不同存储层级：

Location ::= Mem | AMX | WMMA // 主存/AMX寄存器/WMMA寄存器

这种设计使得编译器可以：

1. 显式跟踪数据在存储层次结构中的位置
2. 自动插入必要的数据搬移指令
3. 针对不同存储类型选择最优的指令序列

例如，loc_to_loc(l1, l2, e)操作符明确表示了数据从位置l1到l2的移动，这对生成高效的DMA指令至关重要。

2. 四维优化规则体系

HARDBOILED的核心创新在于其基于egglog规则引擎的优化系统。规则被精心设计为四个相互协作的类别，形成完整的优化链条。

2.1 Lowering规则：硬件指令映射

Lowering规则负责将高级张量操作映射到具体的硬件指令。以矩阵乘法为例：

(rule ( ;; 矩阵乘法Lowering规则 (= e (Add C (VectorReduceAdd 512 (Mul (Cast (Float32 8192) A) (Cast (Float32 8192) B))))) (amx-A-tile A amx-A) (amx-B-tile B amx-B) ) ( (let new-e (Call "tile_matmul" (vec-of (Mem2AMX C) amx-A amx-B))) (union e (AMX2Mem new-e)) ))

该规则实现了以下转换：

1. 识别C + reduce_add(A * B)的数学模式
2. 通过amx-*-tile关系查询已注册的硬件加速方案
3. 生成tile_matmul硬件指令，并保留原始表达式等价性

注意事项：规则中的(union e (AMX2Mem new-e))是关键，它告诉优化器这两个表达式在语义上等价，但硬件实现不同。这种显式的等价关系声明是egglog优于传统编译器pass架构的地方。

2.2 应用特定规则：模式扩展

这类规则针对特定领域模式进行扩展优化。例如处理不同矩阵布局的B矩阵：

(rule ( ;; VNNI布局处理规则 (= orig-B (Load (BFloat16 8192) B B-index)) (= B-index (Broadcast (Ramp (Ramp (Ramp B-base 1 2) (Broadcast B-stride 2) 16) (Broadcast 2 32) 16) 16)) ) ( (let amx-B (Call "tile_load" (vec-of (Var B) B-base B-stride))) (amx-B-tile orig-B amx-B) ))

该规则专门处理Intel AMX的VNNI（Vector Neural Network Instructions）布局特征：

• 识别三维Ramp嵌套的特定访问模式
• 直接生成tile_load指令
• 通过amx-B-tile关系注册该优化方案

2.3 公理化规则：向量化优化

这类规则实现向量操作的代数变换，使模式匹配更鲁棒：

(rewrite (Broadcast (Broadcast x l1) l2) (Broadcast x (* l1 l2))) (rewrite (Add (Ramp base stride r-lanes) (Broadcast x b-lanes)) (Ramp (Add base (Broadcast x (/ b-lanes r-lanes))) stride r-lanes) :when ((= (% b-lanes r-lanes) 0)))

第一个规则将嵌套的Broadcast扁平化，第二个规则实现了Broadcast与Ramp的加法融合。这些规则虽然看似简单，但能有效处理编译器前端可能生成的各种语法变体。

2.4 支持规则：静态分析

支持规则提供类型检查和形状分析等基础功能：

(rule ((= e (Ramp inner s l)) (has-type inner ty)) ((has-type e (MultiplyLanes ty l)))) (function MultiplyLanes (Type i64) Type) (rewrite (MultiplyLanes (Float32 l) x) (Float32 (* l x)))

这些规则共同构成了类型推导系统，确保向量操作的维度一致性。例如MultiplyLanes函数会跟踪向量操作后的维度变化，这在生成硬件指令时至关重要——错误的向量长度会导致指令选择失败。

3. 卷积优化的实现艺术

HARDBOILED在图像卷积优化方面展现了强大的能力，其核心是将卷积运算转化为矩阵乘法。这种转换通过Toeplitz矩阵构造实现，在保持数学等价性的同时充分利用硬件矩阵加速单元。

3.1 卷积到矩阵乘法的转换

考虑以下1D卷积的Halide表达式：

conv[ramp(0, 1, 256)] = (float32x256)vector_reduce_add( float32x2048(I[ramp(ramp(rx, 1, 8), x8(1), 256)]) * x256(float32x8(K[ramp(rx, 1, 8)])) ) + conv[ramp(0, 1, 256)]

HARDBOILED通过模式匹配识别出该表达式符合卷积模式后，会执行以下转换：

1. 构造Toeplitz矩阵临时缓冲区
2. 生成ConvolutionShuffle内在函数调用
3. 将原卷积重写为矩阵乘法

(let new-K (ExprVar (Call "ConvolutionShuffle" (vec-of (Var K) base-K 16 8))))

3.2 WMMA指令生成

最终生成的NVIDIA WMMA指令序列如下：

wmma.mma.sync.aligned.row.row.m32n8k16.f32.f32( wmma.load.a.sync.aligned.row.m32n8k16.f16(I, rx, 8), wmma.load.b.sync.aligned.row.m32n8k16.f16(tempbuf, 0, 8), conv[ramp(0, 1, 8)])

这一转换过程的关键优势在于：

• 完全自动化，无需手动编写特定于硬件的优化
• 保持数学语义的严格等价
• 可适应不同硬件后端的指令集特性

4. 工程实现深度剖析

HARDBOILED的实际工程实现包含多个精妙的设计决策，这些决策直接影响最终生成的代码质量。

4.1 瓦片提取器（Tile Extractor）

瓦片提取器是连接Halide与egglog的关键组件，其主要工作流程：

1. 预处理：

   • 标记内存与加速器寄存器间的数据移动
   • 识别用户指定的加速器内存分配

2. 规则应用：

   egglog_process.run(` (rule ( (= e ...) ) ( ... )) (rewrite ...) `);

3. 结果提取：

   • 基于成本模型选择最优表达式
   • 将egglog输出转换回Halide IR
   • 处理ExprVar节点生成实际内存分配

实操技巧：成本模型当前采用简单的表达式大小评估，但在实际项目中可以扩展为考虑指令延迟、寄存器压力等更复杂的模型。

4.2 表达式变量（ExprVar）处理

ExprVar是HARDBOILED中处理临时缓冲区的关键抽象：

ExprVar ::= PointerTo(evaluation_result)

其处理流程包含以下关键步骤：

1. 生命周期分析：确定临时缓冲区的有效作用域
2. 内存分配：根据类型和维度计算缓冲区大小
3. 初始化代码生成：填充缓冲区内容
4. 提升优化：将分配尽可能外提以减少循环内开销

这种设计使得诸如矩阵转置、数据重排等操作可以延迟到优化阶段处理，为规则系统提供更大的优化空间。

5. 实战经验与性能考量

在实际部署HARDBOILED系统时，以下几个经验教训值得注意：

5.1 规则设计平衡术

编写egglog规则时需要特别小心以下平衡：

• 表达能力vs规则爆炸：例如实现加法交换律而不实现结合律
• 通用性vs特殊性：如VNNI布局的特殊处理规则
• 提前应用vs延迟应用：某些规则需要其他规则先创造匹配条件

一个典型的取舍案例是广播（Broadcast）规则：

(rewrite x (Broadcast x 1) :when ((IsExpr x)))

这条规则看似多余（因为x ≡ broadcast(x,1)），但它能帮助后续规则匹配到更统一的模式。

5.2 硬件特性适配

不同加速器需要不同的优化策略：

HARDBOILED通过Location类型和硬件特定规则优雅地处理这些差异。例如AMX的tile_load和WMMA的wmma.load虽然功能相似，但通过不同的规则实现。

5.3 调试与验证

验证优化正确性的有效方法包括：

1. 数学恒等验证：确保优化前后表达式数学等价
2. 边界条件测试：特别测试向量长度非倍数的情况
3. 逐级检查：在Lowering前后插入验证断言
4. 成本模型监控：观察优化过程中表达式成本的变化

例如，可以在关键规则后插入验证：

(rule ( (= e1 ...) ) ( (let e2 ...) (assert (math-equiv e1 e2)) (union e1 e2) ))

HARDBOILED IR系统通过精心设计的语法结构和规则体系，为张量计算提供了强大的优化能力。其核心价值在于将硬件特定的优化知识编码为可重用的规则，同时保持足够的灵活性以适应不同的计算模式和硬件架构。这种基于规则的中间表示设计范式，为下一代面向异构计算的编译器提供了有价值的参考。