SIMT 标量向量转换的魔法:谓词

谓词:从“做什么”到“做不做”的退化

硬件流水线最粗暴简单的划分方法则是 “指令处理”和“数据执行”单元,而并行处理器 SIMD 聚合了多条流水线的指令部分,从而大大降低了 scaling 的边界成本;但同时指令单元的减少也大大限制了控制流的灵活度,我们必须将一组向量(vector)执行单元视作一个整体调度,而无法细粒度控制控制每一个标量(scalar)。

但实际的应用并行度不总是和硬件宽度匹配,比如 Nvidia SM 底层基础调度单元 warp 为 32 宽度,而要执行 1000 个数加法相加,1000 和 32 并不能整除。每次执行 32 宽度加法,则最后一次并不能充满 32 宽度只有 8 个数据。我们想要在硬件层面引入一种自动化机制,消除 32 这个 “magic number” 给程序员的心智负担,隐藏这一底层细节。

因此需要在粗粒度的 SIMD 中引入一种细粒度的控制机制,Single Instruction 在牺牲了指令表达 Multiply Data 中的每个数据“做什么”的能力,但是我们仍可以表达“做还是不做”,True or False,我们通过微量的代价对每个单元分配一个 bit 表示这种布尔状态。一个返回值是布尔型的函数称作谓词(predicate),这个状态常常需要经过运算得到,因此这种特殊的存储单元也被叫做谓词寄存器 (predicate register),又因为其作用就像在向量单元上施加了掩码,所以也可以称作 predicate mask。谓词广泛应用在各类并行处理器中,比如 NVIDIA 中的特殊谓词寄存器、守卫指令以及 branch divergence 概念;RISCV Vector 扩展中的 v0 寄存器;ARM SVE 扩展的 p0-p15 寄存器;x 86 AVX-512 的 Opmask Registers 。

虽然前文从并行程序的边界条件引出了谓词,放更大来看,任何一个程序必然包含并行部分和串行部分,但并非意味着如果一个处理器只有向量单元便无法处理标量程序,SIMD 以“不做”(stall) 的代价获得了处理标量的能力,能做,只是 utilization 低一些。而另一种使向量获取标量能力的方法是同时拥有向量单元和标量单元,比如昇腾和基于 RISCV 改进的并行处理器。是牺牲 utilization 减少硬件设计复杂度,还是引入异构组件换取性能提升,则要细细 profile workload 串行部分算力占比以及是否会造成 bottleneck。

谓词的主体:指令还是数据?

具体而言,对于一个 32 宽度的向量单元,我们需要多个 bit 来表达这种“做不做”的状态呢?答案似乎显然是 32bit,而许多处理器也正是这个宽度,这是以“指令”作为主体设计谓词,也符合直觉“做不做”的对象。实际上,我们面对的情况可能更加复杂,也许 32 宽度加法 A+B=C,A 数据前 16 个有效,而 B 数据前 8 个有效,这时候似乎以指令为主体并不恰当。

另一种思路是以数据为主体设计谓词,具体来说一个 W 宽度的指令有 N 个输入操作数,那么就需要 \(W\times N\) bit 空间表示操作的谓词状态,并根据具体的指令操作生成输出数的谓词。将指令作为谓词主体可以看作以数据为主体的一种特殊情况:假设每个单元的输入操作数谓词状态相同。

对比这两种设计范式,额外的 bit 开销在向量计算里面基本忽略不计(如果是极低 bit 的并行单元另说),更重要的是设计范式的区别,以指令为主体设计谓词得到是守卫指令(Guard Instruction),则每个指令除了要指定输入输出操作数的寄存器,还要指定谓词寄存器;而以数据为主体设计谓词则是要修改访存结构,将谓词和 operation 融为一体。

常见处理器都是指令为主体,这种处理器在面对上述的特殊情况时会将谓词表达能力的缺失转移到分配额外寄存器空间弥补控制能力,比如 vector reduce operation。目前个人设计加速器时都是 DSA,不太考虑统一的指令抽象,一般以数据为主体设计谓词机制换取控制自由度。

谓词的生成:索引在软硬域的翻译

前文也提到谓词是指以返回布尔型的函数,这个函数实际到底对应什么硬件机制?前文举例子中有提到谓词一种用法是处理边界条件,将“1000-32” 这两个表示软件和硬件并行度的 int 数转换成表示 True or False 的布尔类型,输入 int 输出 bool 是比较函数。

这里用 lane 这个术语表示 scalar 控制粒度对应的硬件单元。以 NVIDIA 架构举例,表示软件并行度的 “1000” 从程序中来;而表示硬件并行度的 “32” 则对应一组特殊的寄存器,PTX 代码中是 %nctaid, %ctaid, %ntid, %tid 等等,其编码了 Grid/Block/Thread 不同粒度的硬化编码,比如 %tid 是一个 32 宽度的 32 位无符号寄存器,每个位置值是 0,1,2,..., 31,正好编码了每个处理单元的硬化索引。Grid/Block/Thread 可以看作一个多维索引,通过这一组多维索引可以找到硬件中唯一的 lane。在定义硬件 Mapping 方式时,我们定义了软件维度和硬件维度之间的映射关系,比如 id=blockId.x * blockDim.x + threadIdx.x,这里通过乘法加法操作,将硬件域的索引 block、thread 翻译到了软件域的索引 id,然后将 id 和软件并行度数字比较,便得到了一组布尔型谓词。

当然谓词还有不止 runtime 比较这一种生成方法,比如还可以通过提前编译好谓词塞入指令之中。不论谓词怎么生成,其生成的这个函数肩负着翻译软件索引和硬件索引的功能,是硬件向开发者隐藏底层粒度的核心机制。

粗粒度 AI 加速器中的硬化谓词设计

“粗粒度 AI 加速器” 是现捏的一个词,用以借代一类常见的单位指令宽度超长语义丰富的类 VLIW AI DSA 设计。这类加速器特征是通过限制支持算子类型从而限制了控制流的丰富度,将部分控制流的实现抽象从编译器转移到硬件开发者,从而提高在专用任务的收益。常见控制流硬化包括一些放射循环,当然也包括本文中提到的谓词处理。以下形式化分析以数据为主体推导 predicate mask 的变换。以下主要搬运自个人笔记,可能相对晦涩,好麻烦也懒得改了=.=

张量 Predicate 的形式化表征

AI 处理器的对象是张量,个人最受用张量的逻辑思维方法是以维度为中心,即爱因斯坦算子方式。对于一个张量其大小由多个常量定义,比如常见 4 维张量划分 \([B,C,H,W]\) ,这 4 个数字即代表了张量的大小,也是索引的边界。假设索引 \(b\) 从 0 开始,那么有 \(b\in [0, B-1]\)

我们用小写字母表示从 0 开始的变量;而用大写字母表示 Tensor 维度从 1 开始的变量。最常见的隐藏硬件细节是处理 pad 的边界条件,假设一个 Tensor 的多个维度边界的有序序列表征是 \(X=\{A, B, C,...\}\) ,其对应的有序 predicate 函数序列是 \(P=\{P_A(a), P_B(b), P_C(c),...\}\),其中每个函数满足:

\[P_N(n)=n \leq n_{p},\quad n\in[0,N-1],\quad n_p\in[0,N-1] \]

这里的维度代表已经 pad 处理后的维度,比如假设硬件每次至少处理 C 维度数量为 32,原始 C 维度为 48,pad 之后维度则是 64,那么有 \(P_C (c)=c\le c_p = c\le 47\)

从软件映射到硬件

我们现在定义了 Tensor 的 predicate 函数,可以得到每个 scalar 对应的 True or False 1 bit mask。但维度 \(N\) ,边界 \(n_p\) ,索引 \(n\) 都是软件域的数值,接下来我们要将形式化软件映射到硬件的过程,以及该过程中 predicate 函数的变化。这可以拆分成三种基础操作:

  • 交换维度次序,对应的 predicate 函数同样交换次序
  • Split 维度,比如 \(C=(A\quad B)\rightarrow A\quad B\)
  • Merge 维度,比如 \(A\quad B\rightarrow C=(A\quad B)\)

Split 前后维度满足 \(C=A\times B\),索引向量满足 \(b=c\mod B\)\(a=\lfloor\frac{c}{B}\rfloor\) , \(c=a\times B + b\)。带入原始函数,得到 \(P_C(c)=((a\times B +b) \leq c_{p})\)。定义 \(b_p\), \(c_p\) 满足 \(b_p=c_p\mod B\)\(a_p = \lfloor\frac{c_p}{B}\rfloor\)\(C_p=A_p\times B + b_p\)。则:

\[P_C (c)=((a\times B +b) \leq c_{p})= (a < a_p) \lor [(a = a_p)\land(b \leq b_p)] \]

对于 Mapping 过程中的 Temporal-Spatial 切分是一种特殊的 Split,满足 \(a_p=A-1,c_p\in (C-B,C-1]\),此时 \(\lnot(a<a_p) \equiv (a=a_p)\) 公式可退化成 \((a < a_p) \lor [\lnot (a < a_p)\land(b \leq b_p)]=(a<a_p)\lor (b\leq b_p)\)\((a\neq a_p)\lor(b\le b_p)\);但是对于 Spatial-Spatial 切分并不满足 \(a_p=A-1\) 关系,比如精度转换导致数据在 Spatial 切分维度的变化。

对于 Merge 前后维度同样满足 \(C=A\times B\),故索引向量满足 \(c=a\times B + b\)\(b=c\mod B\)\(a=\lfloor\frac{c}{B}\rfloor\)\(P_C(c) = P_A(a) \land P_B(b)\) ,同样定义 \(a_p\)\(b_p\),则

\[P_C(c) = \left( \left\lfloor \frac{c}{B} \right\rfloor < a_p \right) \land \left( (c \mod B) < b_p \right) \]

posted @ 2026-01-08 21:13  DevilXXL  阅读(132)  评论(0)    收藏  举报