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 行),工作逻辑如下:

  1. 搜集归约内建
    • 遍历函数中的所有指令,识别出所有 IntrinsicInst,检查其 IntrinsicID 是否属于 vector_reduce_* 族。
    • 对每个找到的归约内建,调用 TTI->shouldExpandReduction(II) 判断目标是否需要展开(例如目标后端不支持或性能更优)。
  2. 按类型分支展开
    • 浮点加/乘 (vector_reduce_fadd、vector_reduce_fmul):
      • 若开启了允许重结合(FastMathFlags::allowReassoc),且向量长度是 2 的幂,则用递归的 Shuffle(二叉树或扇形方式)实现;
      • 否则,调用有序归约 getOrderedReduction(按元素顺序循环累加)以保证语义。
    • 整数按位运算 (and/or/xor):
      • 对于布尔型向量(i1),先 bitcast 到整型,再比较与全零或全一;
      • 其它情况同样用 Shuffle + 二元指令实现。
    • 普通整数加/乘/最值 & 浮点最值
      • 要求向量长度为 2 的幂;
      • 对最值还要求 noNaNs(浮点最值),否则跳过;
      • 一律使用 Shuffle 形式的二叉折半归约。
  3. 替换与删除
    • 用 IRBuilder 在原位置插入上述 Shuffle/循环代码,计算出标量结果 Rdx;
    • 将原有归约内建的所有使用者重定向到 Rdx,然后删除该 IntrinsicInst;
    • 标记函数被修改,最终由 PassManager 更新 CFG 和分析状态。

补充:shuffle:在 LLVM 向量化归约展开中,“shuffle” 指的就是对向量中各元素进行重新排列/混合(permute)的操作,核心就是 IR 里的 shufflevector 指令。

为什么要用 shuffle?
  • 并行度最大化:借助 SIMD 指令库里的向量乱序/跨 Lane 传输指令,一步就能把高半区搬到低半区并行处理;
  • 指令集映射:许多后端(如 x86 AVX、ARM SVE)都有对应的 vperm / swizzle 指令,用来高效地重排向量;
  • 减少依赖链长度:树形归约比顺序循环更短的依赖深度,能获得更好的吞吐。