求个平均值，为什么编译器先写起了循环？

系列目录 · 数值与量化 · Read in English

“这个算子就一行：求平均。”同事把纸推过来。纸上确实只有一行。过了半小时，白板已经出现轴、步长、临时缓冲区和两个循环标签。有人笑：“平均值还没算，先把工作量平均分配一下？”

谜底是：模型描述的是数学集合，向量引擎接收的是地址和固定形状。把“沿某一轴”翻译成“从哪里读、隔多少字节读、在哪停止”，才是编译器的工作。

一、同一个平均值，有三种完全不同的走路方式

设连续存储的教学张量为 [P,R,Q]=[3,5,7]，沿中间轴 R 求平均。目标为：

1

out[p,q] = sum(input[p,r,q], r=0..R-1) / R

固定 p、q 后，相邻参与归约的元素相隔 Q 个元素。它们在数学上是一组，在物理地址上却不是相邻的一排。

若沿末轴 Q 归约，读取的是连续片段；若沿 R、Q 同时归约，可以把两轴合成一个连续片段；若仅沿 R 归约，就必须保留 Q 的各个位置。这三种情况的输入元素数量相同，访问模式却不同。一个只擅长尾部归约的引擎，不会因为高层算子名字相同就自动支持第三种。

历史实现最早尝试在代码生成端处理这些差别：把相邻归约轴组成一段，为中间轴生成运行时循环，为不连续的多轴安排多趟计算。后续差异逐步收紧底层支持范围，把非原生轴转换交给图层。这个过程证明设计方向发生过变化，不证明第一种方案在所有硬件上都慢。

二、循环真正维护的是地址不变量

假设某个内核一次处理一个 p 平面，输入元素占 b 字节，输出元素占 d 字节。每轮结束后的步进应为：

1
2

next_input  = current_input  + R * Q * b
next_output = current_output + Q * d

注意，输入步长和输出步长不一样，元素字节数也可能不一样。把“输入轴被压缩了”忘在输出地址里，往往得到第一块正确、第二块错位的结果。只测 P=1 会让错误躲得很好。

循环不变量可以写得非常朴素：开始第 p 轮时，输入指向该平面的第一个元素，输出指向它对应的 Q 个结果，循环计数等于尚未处理的平面数。每条地址更新都应能从这句话推出来。分支的停止条件也要明确：最后一次计算之后先退出还是先推进？越过尾部的地址即使没被读取，也可能影响带校验的描述符或后续复用。

历史代码增加过分支和地址步进支持，也出现过后来删除归约循环配置的改动。这提醒我们：一个新算子看似只添一个数学操作，实际可能扩展了整套指令生成基础设施。基础设施一旦存在，还要承担标签唯一性、寄存器占用和可重复输出的维护成本。

三、多趟平均，在实数里成立，在定点里未必逐位相同

对规则矩形数据，先沿一轴求平均，再沿另一轴求平均，在精确实数运算下等于一次对所有对应元素求平均。但定点计算通常在每趟结束时舍入，甚至饱和。

例如两组教学整数为 [0,1] 和 [0,3]。精确总平均为 1。若每组平均先采用向下取整，得到 0 和 1，再向下取整求平均得到 0。这里的偏差来自两次舍入，而不是平均公式错误。真实实现应以自己的舍入规则分析；此例只是展示“分阶段舍入不具有结合性”。

量化平均的基本关系为：

1
2

real_input = s_in * (q_in - z_in)
q_out ≈ round[(s_in / s_out) * sum(q_in - z_in) / R] + z_out

倒数 1/R 和尺度比都需要表示。它们是先各自近似再相乘，还是合并成一个固定点系数，会影响误差和寄存器约束。历史差异中能看到平均倒数与输入输出尺度的参数设置被重新安排，因此参数“存在”不够，还得放在执行单元实际使用的位置。

为多趟算法设计中间类型时，至少要回答：是否保留更宽累加值？每趟是否应用零点？中间结果的尺度是谁决定的？某趟饱和后，后续缩小数值能否挽回信息？最后一个问题的答案通常是否定的。

四、为什么后来把复杂性搬到了图里

另一条路线是先重排轴，把非原生归约转成引擎支持的尾部归约，然后恢复输出形状。底层只需检查有限的合法模式，生成一套参数。

这种分层的好处不是“代码短，所以快”。它让形状转换出现在可检查的 IR 中，让已有的重排优化、缓冲区规划和调度有机会参与。相反，把多趟运算藏在一个代码生成函数里，可能让调度器只看到一个不透明大节点，难以识别内部临时张量的生存期。

代价同样清楚：显式重排可能搬动大量数据。能否把重排吸收到生产者、消费者或布局解释中，决定它是不是值得。只减少底层分支数，不能推导运行时间下降。

一个可用于选型的成本模型为：

1
2

T_loop    ≈ T_read + P * (T_reduce + T_address + T_branch)
T_reorder ≈ T_permute_in + T_native_reduce + T_permute_out

如果重排需要完整读写，而归约本身很小，重排可能占主导；如果循环每次工作极少，指令和参数加载成本也可能显著。应把这两个判断变成可测假设，而不是凭算子个数投票。

五、回归测试要故意让地址和数值“不舒服”

建议的测试矩阵如下：

历史中有编译脚本和 IR 变换测试；仅凭这些文件不能宣称数值或硬件性能全部验证通过。结构测试回答“生成了什么”，参考计算回答“算得对不对”，性能实验回答“为什么快或慢”。三种证据必须各自站稳。

最后再回头看那张只有一行公式的纸：公式没有撒谎，它只是没写出内存的生活习惯。编译器工程的乐趣，就在于把这种习惯翻译得既准确，又能被下一位维护者看懂。

六、做一次不用硬件的手工审查

可以把一轮地址计算写成表，专门挑输入输出位宽不同的情况。仍取外层三块、每块五行七列，输入每元素一字节，输出每元素两字节。第一轮输入从块首开始，下一块输入应前进三十五字节；输出每轮七个元素，下一块输出应前进十四字节。若两边都加三十五，第一轮结果仍可能正确，但后两轮会留下空洞甚至越界；若输出也按一字节算，则后续写入会覆盖上一轮部分结果。

这张表还要标清地址单位。有些描述符接受字节地址，有些字段接受元素数或固定粒度的块数。两个变量都叫 stride并不能证明单位相同。最稳的做法是在推导旁边写单位，并在转换边界集中换算。这样很容易发现“某处乘了元素字节数，调用者又乘一遍”的双重换算。

再审查多趟计划时，不要只记录当前 shape，还要记录当前每一轴对应哪些原始轴。先归约后面一组再处理前面一组，是为了让未处理轴的位置较容易保持；若选择另一种顺序，就必须同步更新轴映射。单纯把原轴编号列表复用于已经降秩的张量，会让算法在某个特殊形状上碰巧正确。

量化信息也应放进同一张表：每趟输入尺度、输出尺度、零点、临时位宽、是否舍入。若某趟沿用输入尺度，应问这是为了保持分辨率，还是只是方便复制类型；如果某趟直接使用最终尺度，则检查是否过早丢失小量。把形状计划和数值计划分两份记录，再在每个中间张量处对齐，比在代码中来回搜索 multiplier更容易发现缺口。

最后给所有计划加一个共同断言：每个输出元素的贡献集合与原始归约定义相同，且每个输入贡献恰好出现一次。地址正确但贡献集合重叠，仍然会算错；元素数相等但分组方式不同，也不是同一个平均值。这个集合视角能把循环方案和置换方案放到同一张审查表上。

返回系列目录 · 下一篇

• ← Previous Post
• Next Post →