使用avx指令提升计算性能

前言

现在很多数据库内核引擎都使用向量化的技术来优化性能。例如StarRocks基于avx2指令做向量优化后,很多算子性能都得到了巨大的提升。
image.png

基础知识

AVX指令集与SIMD

SIMD是Single Instruction Stream, Multiple Data Stream的缩写,即单指令多数据流。AVX是advanced vector extensions的缩写。

AVX指令集是为了支持SIMD运算,CPU层面提供的指令集。像最新的是avx512,可以支持Skylake-server和Ice lake的intel芯片。

使用新指令集支持SIMD运算当然性能会更好,缺点就是会有一些额外的功耗开销,不过对于软件来说有机会使用SIMD运算总体肯定是利大于弊的。

确认硬件环境

我用的机器是mac(x86 cpu),使用如下命令确认我的机器支持avx指令集

1
2
3
4
5
# mac
sysctl -a | grep machdep.cpu.features

# linux
cat /proc/cpuinfo |grep avx

我得到的结果如下,可以看到支持的是avx1.0

1
machdep.cpu.features: FPU VME DE PSE TSC MSR PAE MCE CX8 APIC SEP MTRR PGE MCA CMOV PAT PSE36 CLFSH DS ACPI MMX FXSR SSE SSE2 SS HTT TM PBE SSE3 PCLMULQDQ DTES64 MON DSCPL VMX EST TM2 SSSE3 FMA CX16 TPR PDCM SSE4.1 SSE4.2 x2APIC MOVBE POPCNT AES PCID XSAVE OSXSAVE SEGLIM64 TSCTMR AVX1.0 RDRAND F16C

Java

java因为其设计本身的缘故,对操作系统一些native方法的调用就不是那么直接。Java要使用avx指令要么通过JNI调用low level api 要么就等JIT。JDK16开始官方已经开始在孵化vector api。可以看下这个油管视频Vector API: SIMD Programming in Java

示例代码

下面是一个vector api的应用。vector api具体的信息可以参考下JEP417

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
public static void scalarComputation(float[] a, float[] b, float[] c) {
for (int i = 0; i < a.length; i++) {
c[i] = (a[i] * a[i] + b[i] * b[i]) * -1.0f;
}
}
static final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_PREFERRED;

public static void vectorComputation(float[] a, float[] b, float[] c) {
int i = 0;
int upperBound = SPECIES.loopBound(a.length);
for (; i < upperBound; i += SPECIES.length()) {
// FloatVector va, vb, vc;
var va = FloatVector.fromArray(SPECIES, a, i);
var vb = FloatVector.fromArray(SPECIES, b, i);
var vc = va.mul(va)
.add(vb.mul(vb))
.neg();
vc.intoArray(c, i);
}

}

其他语言

本质上还是对intel提供的C++库的调用。指令集关联的C++方法可以参考:Intel® Intrinsics Guide

参考资料