深度优化 | PolarDB-X 基于向量化 SIMD 指令的探索

背景

PolarDB-X作为一款云原生分布式数据库，具有在线事务及分析的处理能力(HTAP)、计算存储分离、全局二级索引等重要特性。在HTAP方面，PolarDB-X对于AP引擎的向量化已经有了诸多探索和实践，例如实现了列式内存布局，MPP，面向列存的执行器等高级特性(参考PolarDB-X 向量化执行引擎(1)及PolarDB-X 向量化引擎(2) )。

PolarDB-X正在全面自研列存节点Columnar，负责提供列式存储数据，结合行列混存 + 分布式计算节点构建HTAP新架构。近期即将正式上线公有云，未来也会同步发布开源。

另外，在面向列存场景最典型的就是向量化，SIMD指令作为向量化中的关键一环，已经被诸多主流AP引擎用来提升计算速度。然而由于Java语言本身的限制，PolarDB-X CN计算引擎无法在JDK 17前主动调用SIMD指令，而在JDK 17版本中Java官方提供了对SIMD指令的封装，即Vector API。本文将介绍PolarDB-X对于向量化SIMD指令的探索和实践，包括基本用法及实现原理，以及在具体算子实现中的思考和沉淀。

SIMD简介

SIMD(Single Instruction Multiple Data) 是一种处理器指令类型，即单个指令可以同时处理多个数据。

以下为加法的标量(Scalar)与SIMD(Vector)两种执行方式:

为了支持SIMD编程，CPU提供了一系列的特殊寄存器与指令

• 寄存器：
  • SSE指令集中的128位寄存器XMM0-XMM15
  • AVX指令集中的256位寄存器YMM0-YMM15
• 算术运算：PADDB，计算两组 8 bits 整型的和, 每组包含16 个 8 bits (SSE2)，可同时用于计算 unsigned 和 signed 类型
• 比较运算：PCMPEQB，比较两组 8 bits 是否相等, 每组包含16 个 8 bits(SSE2), 32 个 (AVX2), 64 个 (AVX512BW)
• 位运算
  • PAND：对两个寄存器的值作按位与(AND)
  • POR：同 PAND，OR 操作
  • PXOR：同 PAND，XOR 操作
• Load/Store指令
  • MOVAPS：每次移动128bits的值

Vector API简介

在JDK 17以前，Java并不能主动的调用SIMD指令，但是在JDK 17版本中，Java官方提供了对SIMD指令的封装- Vector API。Vector API提供了IntVector, LongVector等寄存器，其会根据底层CPU自动选择合适的指令集，这使得开发人员无需考虑具体的CPU架构来快速进行SIMD编程。同时它也提供了add, sub, xor, or等操作来进行SIMD运算，以及fromArray, intoArray等来一次性读取多位数据。

Vector API的基本用法

我们以一个数组相加的例子来快速入门Vector API

在Vector API中，每一个Vector代表一个寄存器，其可以存放若干个元素，取决于寄存器的大小和元素类型，例如当寄存器大小为128位时，可以存放4个int类型(每个int占32位)

public class LongSumBenchmark {
    //定义SPECIES，表示Vector的类型
    private static final VectorSpecies<Long> SPECIES = LongVector.SPECIES_PREFERRED;
    
    private int count;
    private long[] longArr;
    private long[] longArr2;
    private long[] longArr3;
    public void normalSum() {
        for (int i = 0; i < longArr.length; i++) {
            longArr3[i] = longArr[i] + longArr2[i];
        }
    }
    public void vectorSum() {
        int i;
        int batchSize = longArr.length;
        int end = SPECIES.loopBound(batchSize); //通过loopBound获取到对齐后的上限
        for (i = 0; i < end; i += SPECIES.length()) {
            //fromArray(SPECIES, longArr, i)表示从longArr的第i个位置元素开始取出SPECIES.length()个元素
            LongVector va = LongVector.fromArray(SPECIES, longArr, i);
            LongVector vb = LongVector.fromArray(SPECIES, longArr2, i);
            LongVector vc = va.add(vb); //调用add函数，使用SIMD指令求和
            //intoArray(longArr3, i)表示将vc寄存器中的内容存入longArr3中i偏移量开始的元素
            vc.intoArray(longArr3, i);
        }
        for(; i < batchSize; ++i) { //剩余的部分需要手动处理
            longArr3[i] = (longArr[i] + longArr2[i]);
        }
    }
}

FMA计算

为了展现Vector API对计算性能的提升，我们复现了FMA计算的例子

FMA加法是指：c = c + a[i] * b[i]。其中a和b都是float/double类型的数组

@Benchmark
public double normalSum() {
    double sum = 0;
    for (int i = 0; i < doubleArr.length; i++) {
        sum += doubleArr[i] * doubleArr2[i];
    }
    return sum;
}
@Benchmark
public double vectorSum() {
    var sum = DoubleVector.zero(SPECIES);
    int i;
    int batchSize = doubleArr.length;
    var upperBound = SPECIES.loopBound(doubleArr.length);
    for (i = 0; i < upperBound; i += SPECIES.length()) {
        DoubleVector va = DoubleVector.fromArray(SPECIES, doubleArr, i);
        DoubleVector vb = DoubleVector.fromArray(SPECIES, doubleArr2, i);
        sum = va.fma(vb, sum);
    }
    var c = sum.reduceLanes(VectorOperators.ADD);
    for(; i < batchSize; ++i) {
        doubleArr3[i] = (doubleArr[i] + doubleArr2[i]);
    }
    return c;
}

测试环境：随机生成1000/10w个双精度浮点数。

测试结果：向量化执行比标量执行快了2倍

结果分析：

• Vector API的执行结果：只有一条指令vfmadd231pd %ymm0,%ymm2,%ymm3: 将ymm2和ymm3中的双精度浮点数相乘，和ymm1中的数据相加，并把结果放到ymm1中

0x00007f764d363902:   vfmadd231pd %ymm0,%ymm2,%ymm3

• 标量执行的结果：将乘法和加法拆成了两条指令vmulsd和vaddsd

0x00007f000133050d:   vmulsd 0x10(%rax,%r13,8),%xmm0,%xmm0
0x00007f0001330514:   vaddsd %xmm0,%xmm1,%xmm1

由测试结果可以看出，对于FMA计算场景，Vector API将原本需要两条指令的vmulsd和vaddsd合并为了一条指令vfmadd。

但需要注意：FMA计算的优化无法用在数据库中，因为PolarDB-X是将乘法和加法拆为两个算子来执行的

使用Vector API实现基础SIMD操作

在这里我们将演示如何使用Vector API实现《Rethinking SIMD Vectorization for In-Memory Databases》论文中的Gather和Scatter运算

1.Gather：Vector API的fromArray操作提供了对Gather操作的封装，我们只需要传入对应的参数即可

a.标量实现

public void gather(int[] source, int[] indexes, int count, int[] target) {
    for (int i = 0; i < count; i++) {
        target[i] = source[indexes[i]];
    }
}

b.SIMD实现

public void gather(int[] source, int[] indexes, int count, int[] target) {
    final int laneSize = INTEGER_VECTOR_SPECIES.length();
    final int indexVectorLimit = count / laneSize * laneSize;
    int indexPos = 0