使用 C# Vector<T> SIMD 查找匹配元素的索引

Question

使用 C# 的Vector<T>，我们如何才能最有效地矢量化查找集合中特定元素的索引的操作？

作为约束，集合将始终是整数基元的Span<T>，并且最多包含 1 个匹配元素。

我想出了一个看起来不错的解决方案，但我很好奇我们是否可以做得更好。方法如下：

1. 在每个槽中创建一个仅包含目标元素的Vector<T>。
2. 在输入集向量和上一步的向量之间使用Vector.Equals()，以获得在单个匹配槽中包含 1 的掩码（如果没有匹配，则仅包含 0）。
3. 使用包含从 1 开始的索引 (1, 2, 3, 4, ...) 的预初始化向量，在该向量和上一步的掩码之间调用 Vector.Dot()。每个索引都将乘以 0，除了潜在的匹配索引，它将乘以 1。我们得到的是这些乘法的总和，即 0 或匹配元素的从 1 开始的索引。李>

4. 如果结果为 0，则返回 -1 表示不匹配。否则，从结果中减去 1 使其从 0 开始，然后返回。

       // One-time initialized vector containing { 1, 2, 3, 4, ... }
       Vector<ushort> indexes = MemoryMarshal.Cast<ushort, Vector<ushort>>(Enumerable.Range(1, Vector<ushort>.Count).Select(index => (ushort)index).ToArray())[0];
   
       // The input set and the element to search for
       Span<ushort> set = stackalloc ushort[]{ 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25 };
       ushort element = 22;
   
       // Interpret input set as a sequence of vectors (set is assumed to have length power of two for brevity)
       var setVectors = MemoryMarshal.Cast<ushort, Vector<ushort>>(set);
   
       // Create a vector that contains the target element in each slot
       var elementVector = new Vector<ushort>(element);
   
       // Loop per vector rather than per element
       foreach (var vector in setVectors)
       {
           // Get a mask that has a 1 in the single matching slot, or only 0s
           var mask = Vector.Equals(vector, elementVector);
   
           // Get the dot product of the mask and the indexes
           // This will multiple each index by 0, or by 1 if it is the matching one, and return their sum, i.e. the matching index or 0
           // Note that the indexes are deliberately 1-based, to distinguished from 0 (no match)
           var index = Vector.Dot(indexes, mask);
   
           // Either return 0 for no match, or reduce the index by 1 to get the 0-based index
           return index == 0 ? -1 : index - 1;
       }
   

Answer 1

正如我所见，简单的Span<char>.IndexOf 已经在使用 Intrinsics 来搜索简单的值。你甚至不需要转换为 char 来使用它，因为 MemoryExtensions.IndexOf 只关心大小和 Unsafe.SizeOf<ushort>() == sizeof(char) ！

同样在JsonReaderHelper.IndexOfOrLessThan 中，您会发现一个更复杂的用于搜索的矢量化示例。它使用字节搜索，但如果Span<ushort>.IndexOf 不合适，我相信您可以根据自己的需要进行调整。

Answer 2

您希望编译器生成的 x86 asm 是相等比较 (pcmpeqb)、pmovmskb 或 movmskps（向量到具有 1 字节或 4 字节元素的位掩码），然后如果掩码非零，对第一个设置位进行位扫描（bsf 或 tzcnt）。

这将比整数点积更有效！

您已经有了比较相等，我想我已经看到了其他具有矢量->位图内在特性的 C# Q&A。如果有人想编辑这个答案或使用 C# 将自己的答案编译/JITs 到这个 asm，请这样做。我不懂 C#，我只是为了 x86 SIMD。