我可以使用 SIMD 进行桶排序/分类吗?

Can I use SIMD to bucket sort / categorize?

我对 SIMD 很好奇,想知道它是否可以处理这个用例。

假设我有一个包含 2048 个整数的数组,比如 [0x018A, 0x004B, 0x01C0, 0x0234, 0x0098, 0x0343, 0x0222, 0x0301, 0x0398, 0x0087, 0x0167, 0x0389, 0x03F2, 0x0034, 0x0345, ...][=13]

请注意它们都是以 0x00、0x01、0x02 或 0x03 开头的。我想将它们分成 4 个数组:

我想我会有这样的代码:

int main() {
   uint16_t in[2048] = ...;

   // 4 arrays, one for each category
   uint16_t out[4][2048];

   // Pointers to the next available slot in each of the arrays
   uint16_t *nextOut[4] = { out[0], out[1], out[2], out[3] };

   for (uint16_t *nextIn = in; nextIn < 2048; nextIn += 4) {

       (*** magic simd instructions here ***)

       // Equivalent non-simd code:
       uint16_t categories[4];
       for (int i = 0; i < 4; i++) {
           categories[i] = nextIn[i] & 0xFF00;
       }
       for (int i = 0; i < 4; i++) {
           uint16_t category = categories[i];
           *nextOut[category] = nextIn[i];
           nextOut[category]++;
       }
   }
   // Now I have my categoried arrays!
}

我想我的第一个内循环不需要 SIMD,它可以只是一个 (x & 0xFF00FF00FF00FF00) 指令,但我想知道我们是否可以将第二个内循环变成 SIMD 指令。

我正在执行的这个 "categorizing" 操作是否有任何类型的 SIMD 指令?

"insert" 的说明似乎有点前途,但我有点太菜鸟了,无法理解 https://software.intel.com/en-us/node/695331 的说明。

如果没有,有什么接近的吗?

谢谢!

您可以使用 SIMD 执行此操作,但它的速度取决于您可用的指令集以及您在实施中的聪明程度。

一种方法是使用数组 "sift" 它来分离出属于不同桶的元素。例如,从包含 16 个 16 位元素的数组中获取 32 个字节。使用一些 cmpgt 指令获取掩码,其中确定每个元素是落入 00 + 01 桶还是 02 + 03 桶。然后使用某种 "compress" 或 "filter" 操作将所有屏蔽元素连续移动到寄存器的一端,然后对未屏蔽元素进行相同操作。

然后再重复一次,从 01 中找出 00,从 03 中找出 02

对于 AVX2,您可以从 开始,以获取有关 "compress" 操作的灵感。使用 AVX512,您可以使用 vcompress 指令来提供帮助:它完全执行此操作,但仅限于 32 位或 64 位粒度,因此您至少需要为每个向量执行几个操作。

您也可以尝试垂直方法,加载 N 个向量,然后在它们之间交换,使第 0 个向量具有最小元素,等等。此时,您可以在压缩阶段使用更优化的算法(例如,如果您对足够多的向量进行垂直排序,则末尾的向量可能完全以 0x00 等开头)。

最后,您还可以考虑以不同方式组织数据,无论是在源头还是作为预处理步骤:从有效负载字节中分离出始终为 0-3 的 "category" 字节。许多处理步骤只需要在一个或另一个上发生,因此您可以通过将它们分开来潜在地提高效率。例如,您可以对所有类别的 32 个字节进行比较操作,然后对 32 个有效负载字节进行压缩操作(至少在最后一步每个类别都是唯一的)。

这将导致字节元素数组,而不是 16 位元素数组,其中 "category" 字节是隐含的。您已将数据大小减半,这可能会加快您将来要对数据执行的所有其他操作。

如果您无法以这种格式生成源数据,则可以在将有效负载放入正确的存储桶时使用分桶作为删除标记字节的机会,因此输出为 uint8_t out[4][2048]; .如果您正在使用评论中讨论的 pshufb 字节洗牌进行 SIMD 左打包,您可以选择一个洗牌控制向量,它只将有效负载字节打包到低半部分。

(在 AVX512BW 之前,x86 SIMD 没有任何可变控制字洗牌,只有字节或双字,因此您已经需要一个字节洗牌,它可以在打包有效载荷的同时轻松地将有效载荷与标签分开字节到底部。)