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

Question

在 Erlang 运行时系统中，持久性 HashMap 如果很大则表示为哈希数组映射尝试，如果它们很小则表示为“平面图”。

我最近被 Nerd 盯上了，开始寻找优化它的方法。 ^_^'

平面图具有以下特点:

• 最多有 32 个键(和 32 个值)；
• 它们无序地存储在 C 数组中；
• 没有重复的键；
• key 未装箱:我们可以直接比较两个 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;
}

(从 the original 简化)

但这根本不使用以上信息。我尝试了如果让编译器知道会发生什么

• 最多32个元素
• 可以返回“一个”匹配项而不是“第一个”匹配项；由于 key 是唯一的，因此最多只能匹配一次。

这导致以下实现:

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

Looking at Compiler Explorer我们可以看到 Clang 和 GCC 现在能够向量化和展开循环。 Benchmarking this显示 5-15% 的加速。

---

但是，现在问题来了:是否有可能走得更远？

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

或者有没有办法直接手动编写一些 SIMD 指令，速度更快？

Answer 1

我不确定它会变得多快，但这是您函数的手动矢量化 AVX2 版本。

uint64_t* flatmap_avx2( const uint64_t* keys, uint64_t* vals, uint64_t key, uint64_t max_size )
{
    const __m256i needle = _mm256_set1_epi64x( (int64_t)key );

    const uint64_t* const keysEnd = keys + max_size;
    const uint64_t* const keysEndAligned = keys + ( max_size / 4 ) * 4;

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

    for( ; keys < keysEnd; keys++, vals++ )
        if( *keys == key )
            return vals;

    return nullptr;
}

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

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

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

另一个与性能相关的东西。如果可以，将您的键指针对齐 32 个字节，会变得稍微快一些。