SIMD in Velox - BigintValuesUsingHashTable

介绍

SIMD使用CPU中的特殊寄存器同时操作多个primitive data。在一些基本情况下，编译器能够为我们将紧密循环转换为SIMD指令，但通常需要显式调用SIMD内在函数。Velox中有几个地方明确使用SIMD来获得更好的性能。

Velox使用SIMD的一个非常典型的例子就是BigintValuesUsingHashTable::testValues方法，BigintValuesUsingHashTable是Velox的common::Filter的一个子类，Filter用于TableScan的数据过滤（oerling出品，我也有所贡献:) ），在这篇SIMD Usage in Velox的官方文章中也简单介绍了该方法的实现，BigintValuesUsingHashTable::testValues使用SIMD来同时检查多个values是否在一个哈希表中。在哈希表中使用特殊的empty marker来标识值缺失：

1. 如果所有值都超出了范围，直接返回false。
2. 如果有empty marker插入哈希表中，则回退到逐个检查值的方式。
3. 使用SIMD乘法和取模计算所有有效值的哈希值，然后使用maskGather获取哈希表中对应的状态。
4. 如果状态为空标记，则表示值缺失；如果状态等于值，则表示找到了该值。否则，我们遇到了哈希冲突，需要查看哈希表中的下一个位置。如果没有发生冲突，我们可以立即返回结果。
5. 对于每个发生冲突的值，使用SIMD一次性推进多个位置，直到找到匹配的值或空标记。

我们首先过一下SIMD最基础的概念，然后整体看一下BigintValuesUsingHashTable的构造函数（即hash表build过程）和testValues方法（即probe过程）的实现，最后分析这些过程中底层的SIMD intrinsics。Velox使用xsimd作为内在函数的零成本抽象，以解决可移植性问题，这里我们假设是AVX/AVX2架构，寄存器宽度256，实际上oerling大神最初的实现版本就是基于avx2的，后面代码中会看到适配avx2设计上的trick。在文章末尾也给出用到的SIMD intrinsics的说明。

SIMD一些基本概念和指令可以看这篇SIMD Basic文章。

BigintValuesUsingHashTable分析

BigintValuesUsingHashTable构造函数

BigintValuesUsingHashTable::BigintValuesUsingHashTable(
    int64_t min,
    int64_t max,
    const std::vector<int64_t>& values,
    bool nullAllowed)
    : Filter(true, nullAllowed, FilterKind::kBigintValuesUsingHashTable),
      min_(min),
      max_(max),
      values_(values) {
  constexpr int32_t kPaddingElements = 4;
  VELOX_CHECK(min < max, "min must be less than max");
  VELOX_CHECK(values.size() > 1, "values must contain at least 2 entries");

  // Size the hash table to be 2+x the entry count, e.g. 10 entries
  // gets 1 << log2 of 50 == 32. The filter is expected to fail often so we
  // wish to increase the chance of hitting empty on first probe.
  auto size = 1u << (uint32_t)std::log2(values.size() * 5);
  hashTable_.resize(size + kPaddingElements);
  sizeMask_ = size - 1;
  std::fill(hashTable_.begin(), hashTable_.end(), kEmptyMarker);
  for (auto value : values) {
    if (value == kEmptyMarker) {
      containsEmptyMarker_ = true;
    } else {
      auto position = ((value * M) & sizeMask_);
      for (auto i = position; i < position + size; i++) {
        uint32_t index = i & sizeMask_;
        if (hashTable_[index] == kEmptyMarker) {
          hashTable_[index] = value;
          break;
        }
      }
    }
  }
  // Replicate the last element of hashTable kPaddingEntries times at 'size_' so
  // that one can load a full vector of elements past the last used index.
  for (auto i = 0; i < kPaddingElements; ++i) {
    hashTable_[sizeMask_ + 1 + i] = hashTable_[sizeMask_];
  }
  std::sort(values_.begin(), values_.end());
}

• 从BigintValuesUsingHashTable这个名字不难看出这个哈希表的数据类型是int64_t，其成员values_就是一个int64_t类型的vector，是构建这个哈希表的数据，min_，max_分别是values_的最小与最大值。
• Line 17就是一个经验公式用于计算哈希表的size（一般是2的n次方），比如10个entries时size为32，为什么不是16？注释中提到“The filter is expected to fail”，这是因为在计算引擎中的TableScan的场景下，Filter预期的行为是能够过滤掉很多数据，所以希望通过更大的哈希表size来增加第一次probe就立马失败（遇到empty marker）的几率，一定程度的空间换时间（oerling大神对细节的把控真是极致，详见PR-#587）。
• Line 18是resize哈希表的实际size，注意，这里实际size增加了kPaddingElements（kPaddingElements = 4），这是给了simd batch操作预留空间，因为在设计这个类的时候erling默认使用avx/avx2指令集，寄存器宽度256，而BigintValuesUsingHashTable的数据类型是int64_t，正好批处理step是4，详见后面probe部分。
• Line19-34，初始化哈希表默认值为empty marker，使用开放寻找+线性探测的方式构建哈希表，如果values中有empty marker则设置*containsEmptyMarker_*为true（与介绍那一节中的步骤2对应），填充部分用哈希表最后一个slot的value填充。

BigintValuesUsingHashTable::testValues

xsimd::batch_bool<int64_t> BigintValuesUsingHashTable::testValues(
    xsimd::batch<int64_t> x) const { // xsimd::batch<long long, xsimd::fma3<xsimd::avx2>>
  // outOfRange = {xsimd::batch_bool<long long, xsimd::fma3<xsimd::avx2>>} 
  auto outOfRange = (x < xsimd::broadcast<int64_t>(min_)) |
      (x > xsimd::broadcast<int64_t>(max_)); //return _mm256_set1_epi64x(val), _mm256_cmpgt_epi64(other, self), _mm256_or_si256

  // _mm256_movemask_pd(reinterpret_cast<__m256d>(mask.data))
  if (simd::toBitMask(outOfRange) == simd::allSetBitMask<int64_t>()) {
    return xsimd::batch_bool<int64_t>(false);
  }
  if (containsEmptyMarker_) {
    return Filter::testValues(x);
  }
  // broadcast _mm256_set1_epi64x
  auto allEmpty = xsimd::broadcast<int64_t>(kEmptyMarker);
  // Temporarily casted to unsigned to suppress overflow error.
  auto indices = simd::reinterpretBatch<int64_t>(
      simd::reinterpretBatch<uint64_t>(x) * M & sizeMask_);
  // ~outOfRangne: _mm256_xor_si256
  // maskGather : _mm256_mask_i64gather_epi64
  auto data =
      simd::maskGather(allEmpty, ~outOfRange, hashTable_.data(), indices);
  // The lanes with kEmptyMarker missed, the lanes matching x hit and the other
  // lanes must check next positions.
  auto result = x == data; // _mm256_cmpeq_epi64
  auto resultBits = simd::toBitMask(result);
  auto missed = simd::toBitMask(data == allEmpty);
  static_assert(decltype(result)::size <= 16);
  // allSetBitMask : bits::lowMask(xsimd::batch_bool<T, A>::size);
  uint16_t unresolved = simd::allSetBitMask<int64_t>() ^ (resultBits | missed);
  if (!unresolved) {
    return result;
  }
  constexpr int kAlign = xsimd::default_arch::alignment();
  constexpr int kArraySize = xsimd::batch<int64_t>::size;
  alignas(kAlign) int64_t indicesArray[kArraySize];
  alignas(kAlign) int64_t valuesArray[kArraySize];
  (indices + 1).store_aligned(indicesArray);
  // store_aligned -> broadcast -> _mm256_set1_epi64x
  x.store_aligned(valuesArray);
  while (unresolved) {
    auto lane = bits::getAndClearLastSetBit(unresolved);
    // Loop for each unresolved (not hit and
    // not empty) until finding hit or empty.
    int64_t index = indicesArray[lane];
    int64_t value = valuesArray[lane];
    auto allValue = xsimd::broadcast<int64_t>(value);
    for (;;) {
      // _mm256_loadu_si256
      auto line = xsimd::load_unaligned(hashTable_.data() + index);

      if (simd::toBitMask(line == allValue)) {
        resultBits |= 1 << lane;
        break;
      }
      if (simd::toBitMask(line == allEmpty)) {
        resultBits &= ~(1 << lane);
        break;
      }
      index += line.size;
      if (index > sizeMask_) {
        index = 0;
      }
    }
  }
  return simd::fromBitMask<int64_t>(resultBits);
}

• xsimd::batch<T, A>是xsimd对SIMD寄存器的封装，其底层data是一个SIMD数据类型，比如xsimd::batch<int64_t>底层数据类型是__m256i，本文默认是avx family指令集，256位寄存器。
• Line 4-5，计算outRange（xsimd::batch<long long, xsimd::fma3<xsimd::avx2>>），xsimd::broadcast底层是_mm256_set1_epi64x，将min_或max_的值广播到目标simd寄存器所有lane上（本文场景4个）。x与广播后到min，max比较的底层是_mm256_cmpgt_epi64，比较两个packed signed 64-bit整数，保持结果到目标寄存器，大于为0xFFFFFFFFFFFFFFFF，否则为0。|底层是_mm256_or_si256，是bitwise or两个256 bits。
• Line 8-13，simd::toBitMask底层是_mm256_movemask_pd将一个batch转成bit mask，然后比较，如果都被set了说明x中所有数据都超过了range，直接返回。如果values中有empty marker回退到逐个检查值的方式。
• Line 14-22，这里开始进入了SIMD实现hash probe的首次探测部分
  • 生成一个empty marker的batch，allEmpty
  • 计算x的哈希值在哈希表中的index的batch，indices
  • 将~outRange（_mm256_xor_si256操作）中非0的lane（本文场景4个lane）在indices中的index索引的hashTable_.data()load到data对应的lane中，其它lane从allEmpty中load。
  • 上述就是一个经典的smid gather操作，底层是__m256i _mm256_mask_i64gather_epi64 (__m256i src, __int64 const* base_addr, __m256i vindex, __m256i mask, const int scale)，通过vindex索引oad其lane对应的mask MSB为1的base_addr地址开始的数据，其它lane则load在src中对于lane的数据。
  • 最后的结果是变量data，它是一个batch，包含x中没有超出range的整数第一次探测到的哈希表数据（可能会有冲突），以及超出range而load的empty marker
• Line 23-33，这里开始计算需要进行后续探测的lane，其中
  • x == data底层是_mm256_cmpeq_epi64得到实际值对比结果的batch，即result，首次探测成功，也就是实际值相同的lane的所有bit都是1（本文场景为0xFFFFFFFFFFFFFFFF），然后转成bit mask也就是resultBits
  • missed = simd::toBitMask(data == allEmpty)是首次探测即失败（其index对应的hash表的slot中的value位emtpy marker)
  • unresolved = simd::allSetBitMask<int64_t>() ^ (resultBits | missed)就可以得到首次探测没有失败，但是实际值不同（也就是哈希冲突）的lane转成的bit mask
  • 如果没有unresolved的lane那么直接返回结果result
• Line 34-40，开始后续探测的准备，前面提到该类的哈希实现是开放寻找+线性探测，所以发生冲突之后需要index后移一位后继续探测，
  • (indices + 1).store_aligned(indicesArray)初始化indicesArray，x.store_aligned(valuesArray)初始化valuesArray
  • 这里都用到了store_aligned，它本质就是broadcast其底层是_mm256_set1_epi64x
  • 这里之所以转成数组是因为后续探测需要一个一个lane的操作
• Line 41-65，获取unresolved当前最后一位为1的bit位并清零，算出其对应的lane（本文场景中x有4个lane），然后用lane获取当前需要继续探测的index和value
  • allValue = xsimd::broadcast(value)把需要继续探测的value加载到一个batch中，即allValues
  • 把需要探测的index索引的哈希表的值load到一个batch中，即line，这里底层是_mm256_loadu_si256
  • 这里比较line与value，如果结果不是0就设置resultBits对应的bit位，如果line == allEmpty则设置resultBits对应的bit位位0
  • 这个循环里面有个优化是，_mm256_loadu_si256加载的是index索引的哈希表的数据和它之后的几个数据（本文是3），实际上是一次进行了多次探测，所以后面更新index是index += line.size，这里也对应了前面提到的那个kPaddingElements填充

一些intrinsics的解释

__m256i _mm256_set1_epi64x (long long a)

Synopsis

__m256i _mm256_set1_epi64x (long long a)
#include <immintrin.h>
Instruction: Sequence
CPUID Flags: AVX

Description
Broadcast 64-bit integer a to all elements of dst. This intrinsic may generate the vpbroadcastq.

Operation

FOR j := 0 to 3
	i := j*64
	dst[i+63:i] := a[63:0]
ENDFOR
dst[MAX:256] := 0

__m256i _mm256_cmpgt_epi64 (__m256i a, __m256i b)

Synopsis
__m256i _mm256_cmpgt_epi64 (__m256i a, __m256i b)
#include <immintrin.h>
Instruction: vpcmpgtq ymm, ymm, ymm
CPUID Flags: AVX2

Description
Compare packed signed 64-bit integers in a and b for greater-than, and store the results in dst.

Operation

FOR j := 0 to 3
	i := j*64
	dst[i+63:i] := ( a[i+63:i] > b[i+63:i] ) ? 0xFFFFFFFFFFFFFFFF : 0
ENDFOR
dst[MAX:256] := 0

__m256i _mm256_mask_i64gather_epi64 (__m256i src, __int64 const* base_addr, __m256i vindex, __m256i mask, const int scale)

Synopsis

__m256i _mm256_mask_i64gather_epi64 (__m256i src, __int64 const* base_addr, __m256i vindex, __m256i mask, const int scale)
#include <immintrin.h>
Instruction: vpgatherqq ymm, vm64x, ymm
CPUID Flags: AVX2

Description

Gather 64-bit integers from memory using 64-bit indices. 64-bit elements are loaded from addresses starting at base_addr and offset by each 64-bit element in vindex (each index is scaled by the factor in scale). Gathered elements are merged into dst using mask (elements are copied from src when the highest bit is not set in the corresponding element). scale should be 1, 2, 4 or 8.

Operation

FOR j := 0 to 3
i := j64
m := j64
IF mask[i+63]
addr := base_addr + vindex[m+63:m] * ZeroExtend64(scale) * 8
dst[i+63:i] := MEM[addr+63:addr]
ELSE
dst[i+63:i] := src[i+63:i]
FI
ENDFOR
mask[MAX:256] := 0
dst[MAX:256] := 0

__m256i _mm256_cmpeq_epi64 (__m256i a, __m256i b)

Synopsis

__m256i _mm256_cmpeq_epi64 (__m256i a, __m256i b)
#include <immintrin.h>
Instruction: vpcmpeqq ymm, ymm, ymm
CPUID Flags: AVX2

Description

Compare packed 64-bit integers in a and b for equality, and store the results in dst.
Operation
FOR j := 0 to 3
i := j*64
dst[i+63:i] := ( a[i+63:i] == b[i+63:i] ) ? 0xFFFFFFFFFFFFFFFF : 0
ENDFOR
dst[MAX:256] := 0

int _mm256_movemask_pd (__m256d a)

Synopsis

int _mm256_movemask_pd (__m256d a)
#include <immintrin.h>
Instruction: vmovmskpd r32, ymm
CPUID Flags: AVX

Description
Set each bit of mask dst based on the most significant bit of the corresponding packed double-precision (64-bit) floating-point element in a.
Operation
FOR j := 0 to 3
i := j*64
IF a[i+63]
dst[j] := 1
ELSE
dst[j] := 0
FI
ENDFOR
dst[MAX:4] := 0

__m256i _mm256_loadu_si256 (__m256i const * mem_addr)

Synopsis
__m256i _mm256_loadu_si256 (__m256i const * mem_addr)
#include <immintrin.h>
Instruction: vmovdqu ymm, m256
CPUID Flags: AVX
Description
Load 256-bits of integer data from memory into dst. mem_addr does not need to be aligned on any particular boundary.
Operation
dst[255:0] := MEM[mem_addr+255:mem_addr]
dst[MAX:256] := 0

__m256i _mm256_lddqu_si256 (__m256i const * mem_addr)

Synopsis

__m256i _mm256_lddqu_si256 (__m256i const * mem_addr)
#include <immintrin.h>
Instruction: vlddqu ymm, m256
CPUID Flags: AVX

Description

Load 256-bits of integer data from unaligned memory into dst. This intrinsic may perform better than _mm256_loadu_si256 when the data crosses a cache line boundary.
Operation
dst[255:0] := MEM[mem_addr+255:mem_addr]
dst[MAX:256] := 0

References

• https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html
• https://www.codeproject.com/Articles/874396/Crunching-Numbers-with-AVX-and-AVX
• https://blog.triplez.cn/posts/avx-avx2-learning-notes/
• https://github.com/Triple-Z/AVX-AVX2-Example-Code