创建掩蔽 kreg 值的有效方法

Efficient way to create masking kreg values

Intel 的 AVX-512 扩展的一个好处是几乎所有的操作都可以通过提供除向量寄存器之外的 kreg 来屏蔽,它指定了要应用的屏蔽操作:掩码排除的元素可以设置为零或保留其先前的值。

kreg 的一个特别常见的用途是创建一个掩码,该掩码排除向量开头或结尾的 N 个连续元素,例如,作为向量化循环中的第一次或最后一次迭代,其中小于完整向量会被处理。例如,对于超过 121 int32_t 值的循环,前 112 个元素可以由 7 个完整的 512 位向量处理,但剩下的 9 个元素可以由仅对前 9 个进行操作的掩码操作处理元素。

所以问题是,给定一个(运行时赋值的)整数 r,它是 0 - 16 范围内的某个值,表示剩余元素,加载 16 位 kreg 的最有效方法是什么低 r 位已设置,其余位未设置? KSHIFTLW 似乎不适合这个目的,因为它只需要立即数。

BMI2 bzhi 完全符合您的要求:从指定位位置开始的零高位。到目前为止,每个 CPU 的 AVX512 都有 BMI2。

__mmask16 k = _bzhi_u32(-1UL, r);

这需要 2 条指令,都是单 uop:mov-immediate 和 bzhi。它甚至是单周期延迟。 (或 KNL 上的 3 个周期)

  • 对于 r=0,它将所有给出 0 的位清零。
  • 对于r=1,它只留下低位(位#0)给出1
  • 对于 r=12,它将第 12 位及更高位清零,留下 0x0FFF(设置 12 位)
  • 对于 r>=32 BZHI 保留所有 32 位设置(并设置 CF)

INDEX 由第二个源操作数7:0 的位

指定

如果你有一个在一个展开的向量循环之后运行的一次一个向量的清理循环,你可以甚至在每个循环迭代中使用它,计算剩余的长度下降到零,而不是单独的最后一个向量清理。它将所有位设置为高长度。但这在循环内花费了 2 微指令,包括端口 5 kmovw,这意味着您的主循环将不得不使用屏蔽指令。这只适用于 r<=255 因为它只查看低字节,而不是完整的整数索引。但是 mov reg, -1 可以被吊起,因为 bzhi 不会破坏它。


PS。通常我认为你会想要安排你的清理来处理 1..16 个元素,(或者 0..15 如果你分支可能会跳过它)。但是,如果此清理还处理根本不进入主循环的小长度,则完整的 17 种可能性 0..16 是有意义的,并且 len=0 是可能的。 (并且您的主循环以剩余长度 = 1..16 退出,因此最终迭代可以是无条件的)