深度干货|PolarDB分布式版基于向量化SIMD指令的探索

bashendan
发布于 2023-10-25 11:34
浏览
0收藏

1. 背景

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

深度干货|PolarDB分布式版基于向量化SIMD指令的探索-鸿蒙开发者社区

PolarDB分布式版正在全面自研列存节点Columnar,负责提供列式存储数据,结合行列混存 + 分布式计算节点构建HTAP新架构。近期即将正式上线公有云,未来也会同步发布开源。另外,在面向列存场景最典型的就是向量化,SIMD指令作为向量化中的关键一环,已经被诸多主流AP引擎用来提升计算速度。然而由于Java语言本身的限制,PolarDB分布式版CN计算引擎无法在JDK 17前主动调用SIMD指令,而在JDK 17版本中Java官方提供了对SIMD指令的封装,即Vector API。

本文将介绍PolarDB分布式版对于向量化SIMD指令的探索和实践,包括基本用法及实现原理,以及在具体算子实现中的思考和沉淀。

2. SIMD简介

SIMD(Single Instruction Multiple Data)是一种处理器指令类型,即单个指令可以同时处理多个数据。以下为加法的标量(Scalar)与SIMD(Vector)两种执行方式:

深度干货|PolarDB分布式版基于向量化SIMD指令的探索-鸿蒙开发者社区

为了支持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的值。

3. 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倍

深度干货|PolarDB分布式版基于向量化SIMD指令的探索-鸿蒙开发者社区

结果分析:

● 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操作的封装,我们只需要传入对应的参数即可。

深度干货|PolarDB分布式版基于向量化SIMD指令的探索-鸿蒙开发者社区

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
    for (; indexPos < indexVectorLimit; indexPos += laneSize) {
        IntVector av = IntVector.fromArray(INTEGER_VECTOR_SPECIES, source, 0, indexes, indexPos);
        av.intoArray(target, indexPos);
    }
    if (indexPos < indexLimit) {
        scalarPrimitives.gather(source, indexes, indexPos, indexLimit - indexPos, target, targetPos);
    }
}

2. Scatter实现:与Gather同理,Vector的intoArray运算提供了对Scatter运算的封装。

深度干货|PolarDB分布式版基于向量化SIMD指令的探索-鸿蒙开发者社区

a.标量实现

public void scatter(long[] source, long[] target, int[] scatterMap, int copySize) {
    for (int i = 0; i < copySize; i++) {
        target[scatterMap[i]] = source[i];
    }
}

b.SIMD实现

public void scatter(int[] source, int[] target, int[] scatterMap, int copySize) {
    int laneSize = SIMDHandles.INT_VECTOR_LANE_SIZE; //每次SIMD能处理的位数
    final int indexVectorLimit = copySize / laneSize * laneSize;
    int index = 0;
    for (; index < indexVectorLimit; index += laneSize) {
        IntVector dataInVector = IntVector.fromArray(INTEGER_VECTOR_SPECIES, source, index); //从source[index]位置取出K个数字
        dataInVector.intoArray(target, 0, scatterMap, index); 
    }
    if (index < copySize) {
        scalarPrimitives.scatter(source, target, scatterMap, index, copySize - index);
    }
}

Vector API实现原理

以下面的向量化相加为例,我们探索一下add函数的实现:

LongVector va = LongVector.fromArray(SPECIES, longArr, i);
LongVector vb = LongVector.fromArray(SPECIES, longArr2, i);
LongVector vc = va.add(vb);

JDK层面

在JDK层面Java并没有做任何的优化,其底层实现就是对Vector中的每个元素调用了apply函数,而apply函数指向了一个绑定的函数,该函数的实现为标量加法。显然,这么做甚至会增加执行的开销,那Vector API的高性能从何谈起呢?

1. add函数的实现

@Override
@ForceInline
public final LongVector add(Vector<Long> v) {
    return lanewise(ADD, v);
}

2. 最终调用b0pTemplate函数进行计算

@ForceInline
final
LongVector bOpTemplate(Vector<Long> o,
                                 FBinOp f) {
    long[] res = new long[length()];
    long[] vec1 = this.vec();
    long[] vec2 = ((LongVector)o).vec();
    for (int i = 0; i < res.length; i++) {
        res[i] = f.apply(i, vec1[i], vec2[i]);
    }
    return vectorFactory(res);
}

3. apply绑定了标量执行的函数实现

LongVector lanewiseTemplate(VectorOperators.Binary op,
                                      Vector<Long> v) {
    LongVector that = (LongVector) v;
    that.check(this);
    ....
    int opc = opCode(op);
    return VectorSupport.binaryOp(
        opc, getClass(), long.class, length(),
        this, that,
        BIN_IMPL.find(op, opc, (opc_) -> {
          switch (opc_) {
            case VECTOR_OP_ADD: return (v0, v1) ->
                    v0.bOp(v1, (i, a, b) -> (long)(a + b));
            case VECTOR_OP_SUB: return (v0, v1) ->
                    v0.bOp(v1, (i, a, b) -> (long)(a - b));
            case VECTOR_OP_MUL: return (v0, v1) ->
                    v0.bOp(v1, (i, a, b) -> (long)(a * b));
            case VECTOR_OP_DIV: return (v0, v1) ->
                    v0.bOp(v1, (i, a, b) -> (long)(a / b));
            ....
            }}));
}

JVM层面


  1. 前置知识:JIT与C2编译器


这里需要简单讲一下Java的即时编译(JIT)。Java的执行过程整体可以分为解释执行和编译执行(JIT),第一步由javac将源码编译成字节码并进行解释执行,在解释执行的过程中,JVM会对程序的运行信息进行收集,对于其中的热点代码通过编译器(默认为C2)进行编译,将字节码直接转化为机器码,然后进行编译执行。

怎么样才会被认为是热点代码呢?JVM中会设置一个阈值,当方法或者代码块的在一定时间内的调用次数超过这个阈值时就会被认为是热点代码。

JIT的执行流程如下:

深度干货|PolarDB分布式版基于向量化SIMD指令的探索-鸿蒙开发者社区

  1. Vector API在JVM层面的实现


通过查找,我们找到了Vector API的first commit,从该commit中我们发现Vector API的实现中有大量对JVM内核的修改。 

首先,其在c2compiler中增加了Vector API相关的intrinsic

(src/hotspot/share/opto/c2compiler.cpp)

深度干货|PolarDB分布式版基于向量化SIMD指令的探索-鸿蒙开发者社区

当触发JIT时,Vector API相关的代码会替换为JVM层面的intrinsic,并使用SIMD指令优化。具体来说,当Vector API的代码被JIT后,其会将语法树中原本的IR节点替换为intrinsic的实现,即将method替换为intrinsic的方法。

//---------------------------make_vm_intrinsic----------------------------
CallGenerator* Compile::make_vm_intrinsic(ciMethod* m, bool is_virtual) {
  vmIntrinsicID id = m->intrinsic_id();

  C2Compiler* compiler = (C2Compiler*)CompileBroker::compiler(CompLevel_full_optimization);
  bool is_available = false;

  methodHandle mh(THREAD, m->get_Method());
  is_available = compiler != NULL && compiler->is_intrinsic_supported(mh, is_virtual) &&
                   !C->directive()->is_intrinsic_disabled(mh) &&
                   !vmIntrinsics::is_disabled_by_flags(mh);

  if (is_available) {
    return new LibraryIntrinsic(m, is_virtual,
                                vmIntrinsics::predicates_needed(id),
                                vmIntrinsics::does_virtual_dispatch(id),
                                id);
  } else {
    return NULL;
  }
}

最终的SIMD指令在C2的汇编器中生成实现

void Assembler::addpd(XMMRegister dst, XMMRegister src) {
  NOT_LP64(assert(VM_Version::supports_sse2(), ""));
  InstructionAttr attributes(AVX_128bit, /* rex_w */ VM_Version::supports_evex(), /* legacy_mode */ false, /* no_mask_reg */ true, /* uses_vl */ true);
  attributes.set_rex_vex_w_reverted();
  int encode = simd_prefix_and_encode(dst, dst, src, VEX_SIMD_66, VEX_OPCODE_0F, &attributes);
  emit_int16(0x58, (0xC0 | encode));
}



文章转载自公众号:阿里云瑶池数据库

分类
标签
已于2023-10-25 11:34:48修改
收藏
回复
举报
回复
    相关推荐