后端Pass简介——ExpandReduction
ExpandReduction
在生成目标代码(codegen)之前,将 LLVM IR 中的向量化归约内建(vector_reduce_*)展开为具体的循环或 Shuffle 序列,以便让不直接支持这些内建的后端也能正确生成归约代码。
涉及的概念如下:
- 归约(Reduction):将一个向量或数组上的所有元素用某种二元运算(如加、乘、最值、按位运算等)累积到一个标量结果的操作。(算子计算里面比较常见)
- LLVM 归约内建:LLVM IR 提供了一组以
llvm.vector_reduce_*
命名的内建(intrinsic),如llvm.vector_reduce_add
、llvm.vector_reduce_fmax 等,用于在中间表示层面描述归约意图。 - 展开(Expansion):并非所有后端都直接支持这些归约内建,ExpandReductions Pass 就是在 IR 级别把它们替换成等价的、基于循环或shufflevector/extractelement 的实现。
下面来看一个例子,假设有如下 IR 片段,计算向量 %v 的所有元素之和,并加上初始累加值 %acc:
%sum = call <4 x float> @llvm.vector_reduce_fadd(<4 x float> %v, float %acc) [fast-math flags]
经过 ExpandReductions 展开后,可能变为(伪 IR):
; 先用 shuffle 两两相加
%v0 = shufflevector <4 x float> %v, <4 x float> undef, <2 x i32> <0,1>
%v1 = shufflevector <4 x float> %v, <4 x float> undef, <2 x i32> <2,3>
%r0 = fadd float extractelement(<2 x float> %v0, i32 0),
extractelement(<2 x float> %v0, i32 1)
%r1 = fadd float extractelement(<2 x float> %v1, i32 0),
extractelement(<2 x float> %v1, i32 1)
; 汇总
%r2 = fadd float %r0, %r1
; 加上初始累加值
%sum = fadd float %acc, %r2
其源码非常简单(190 行),工作逻辑如下:
- 搜集归约内建
- 遍历函数中的所有指令,识别出所有 IntrinsicInst,检查其 IntrinsicID 是否属于 vector_reduce_* 族。
- 对每个找到的归约内建,调用 TTI->shouldExpandReduction(II) 判断目标是否需要展开(例如目标后端不支持或性能更优)。
- 按类型分支展开
- 浮点加/乘 (vector_reduce_fadd、vector_reduce_fmul):
- 若开启了允许重结合(FastMathFlags::allowReassoc),且向量长度是 2 的幂,则用递归的 Shuffle(二叉树或扇形方式)实现;
- 否则,调用有序归约 getOrderedReduction(按元素顺序循环累加)以保证语义。
- 整数按位运算 (and/or/xor):
- 对于布尔型向量(i1),先 bitcast 到整型,再比较与全零或全一;
- 其它情况同样用 Shuffle + 二元指令实现。
- 普通整数加/乘/最值 & 浮点最值:
- 要求向量长度为 2 的幂;
- 对最值还要求 noNaNs(浮点最值),否则跳过;
- 一律使用 Shuffle 形式的二叉折半归约。
- 浮点加/乘 (vector_reduce_fadd、vector_reduce_fmul):
- 替换与删除
- 用 IRBuilder 在原位置插入上述 Shuffle/循环代码,计算出标量结果 Rdx;
- 将原有归约内建的所有使用者重定向到 Rdx,然后删除该 IntrinsicInst;
- 标记函数被修改,最终由 PassManager 更新 CFG 和分析状态。
补充:shuffle:在 LLVM 向量化归约展开中,“shuffle” 指的就是对向量中各元素进行重新排列/混合(permute)的操作,核心就是 IR 里的 shufflevector 指令。
为什么要用 shuffle?
- 并行度最大化:借助 SIMD 指令库里的向量乱序/跨 Lane 传输指令,一步就能把高半区搬到低半区并行处理;
- 指令集映射:许多后端(如 x86 AVX、ARM SVE)都有对应的 vperm / swizzle 指令,用来高效地重排向量;
- 减少依赖链长度:树形归约比顺序循环更短的依赖深度,能获得更好的吞吐。
评论