时间戳

首页 > 解决方案 > Can I use SIMD to bucket sort / categorize?

c - Can I use SIMD to bucket sort / categorize?

问题描述

I'm curious about SIMD and wondering if it can handle this use case.

Let's say I have an array of 2048 integers, like [0x018A, 0x004B, 0x01C0, 0x0234, 0x0098, 0x0343, 0x0222, 0x0301, 0x0398, 0x0087, 0x0167, 0x0389, 0x03F2, 0x0034, 0x0345, ...]

Note how they all start with either 0x00, 0x01, 0x02, or 0x03. I want to split them into 4 arrays:

I imagine I would have code like this:

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!
}

I imagine that my first inner loop doesn't need SIMD, it can be just a (x & 0xFF00FF00FF00FF00) instruction, but I wonder if we can make that second inner loop into a SIMD instruction.

Is there any sort of SIMD instruction for this "categorizing" action that I'm doing?

The "insert" instructions seem somewhat promising, but I'm a bit too green to understand the descriptions at https://software.intel.com/en-us/node/695331.

If not, does anything come close?

Thanks!

标签： carraysx86simdbucket-sort

解决方案

您可以使用 SIMD 来做到这一点，但它的速度取决于您有哪些可用的指令集，以及您在实现中的聪明程度。

一种方法是获取数组并“筛选”它以分离出属于不同存储桶的元素。例如，从包含 16 个 16 位元素的数组中抓取 32 个字节。使用一些cmpgt指令来获取一个掩码，该掩码确定每个元素是落入00 + 01桶中还是落入02 + 03桶中。然后使用某种“压缩”或“过滤”操作将所有被屏蔽的元素连续移动到一个寄存器的一端，然后对未屏蔽的元素进行相同的移动。

然后再重复一次以00从01和中整理02出来03。

使用 AVX2，您可以从这个问题开始，以获得有关“压缩”操作的灵感。使用 AVX512，您可以使用该vcompress指令来提供帮助：它完全执行此操作，但仅以 32 位或 64 位粒度执行，因此您至少需要对每个向量执行几次。

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

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

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

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

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

推荐阅读