如何在一个具有唯一元素的微小数组中尽快找到匹配元素？

Question

在 Erlang 运行时系统内部，如果持久哈希图较大，则表示为 hash-array-mapped-tries；如果较小，则表示为“平面映射”。

我最近被书呆子们吸引去寻找优化这个的方法。^_^'

平面地图具有以下特征：

• 最多有 32 个键（和 32 个值）；
• 它们无序存储在 C 数组中；
• 没有重复的键；
• 密钥已拆箱：我们可以直接将两者进行比较uint64_t以检查是否匹配。

当前的实现是：

uint64_t *original_flatmap_get(uint64_t *keys, uint64_t *vals, uint64_t key, uint64_t max_size) {
  uint64_t n = max_size;
  uint64_t i;

  for (i = 0; i < n; ++i) {
    if (keys[i] == key) {
      return &vals[i];
    }
  }
  return NULL;
}

（对原文进行了简化）

但这根本不使用上述信息。我尝试了如果编译器意识到会发生什么

• 最多有 32 个元素
• 返回“a”匹配而不是“第一个”就可以了；由于密钥是唯一的，因此最多只会出现一次匹配。

这导致了以下实施：

uint64_t *latereturn_flatmap_get(uint64_t *keys, uint64_t *vals, uint64_t key, uint64_t max_size) {
  uint64_t n = min(max_size, 32);
  uint64_t i;

  uint64_t *res = NULL;
  for (i = 0; i < n; ++i) {
    if (keys[i] == key) {
      res = &vals[i];
    }
  }
  return res;
}

查看 Compiler Explorer，我们可以看到 Clang 和 GCC 现在能够向量化和展开循环。 基准测试显示速度提高了 5-15%。

---

然而，现在的问题是：是否可以更进一步？

例如，是否可以以某种方式向编译器指示数组中的所有元素都是唯一的，这可能会实现更多优化？

或者是否有可能直接手动编写一些更快的 SIMD 指令？

Answer 1

我\xe2\x80\x99m 不确定它会变得多快，但这里\xe2\x80\x99s 手动矢量化了函数的AVX2 版本。

\n

uint64_t* flatmap_avx2( const uint64_t* keys, uint64_t* vals, uint64_t key, uint64_t max_size )\n{\n    const __m256i needle = _mm256_set1_epi64x( (int64_t)key );\n\n    const uint64_t* const keysEnd = keys + max_size;\n    const uint64_t* const keysEndAligned = keys + ( max_size / 4 ) * 4;\n\n    for( ; keys < keysEndAligned; keys += 4, vals += 4 )\n    {\n        __m256i src = _mm256_loadu_si256( ( const __m256i* )keys );\n        __m256i eq = _mm256_cmpeq_epi64( needle, src );\n        uint32_t mask = (uint32_t)_mm256_movemask_epi8( eq );\n        if( 0 == mask )\n            continue;\n        uint32_t byteIndex = _tzcnt_u32( mask );\n        // The index is multiple of 8, in assembly all addresses expressed in bytes,\n        // yet adding pointers in C adds elements not bytes, that's why casting\n        return (uint64_t*)( ( (uint8_t*)vals ) + byteIndex );\n    }\n\n    for( ; keys < keysEnd; keys++, vals++ )\n        if( *keys == key )\n            return vals;\n\n    return nullptr;\n}\n

\n

如果您使用 VC++ 构建此函数，最好在函数的#pragma loop( no_vector )第二个循环之前添加。for

\n

同样，如果您使用 gcc 或 clang 构建\xe2\x80\x99，最好__attribute__((optimize("no-tree-vectorize")))在整个函数之前添加。

\n

如果没有这些特定于编译器的恶作剧，编译器可能会决定自动将第二个for循环与剩余部分进行向量化，从而无缘无故地膨胀代码。

\n

另一个与性能相关的事情。如果可以的话，将你的按键指针对齐 32 个字节，会变得稍微快一些。

\n