1 背景
当广告推荐业务峰值QPS已经达到10万以上，向量检索QPS峰值就会就会达到30万以上，召回服务的向量检索P99时延和平均时延已经超出了能接受的正常范围，导致召回服务整体时延达到上限，很多请求超时以至于没有广告返回给上游服务。同时粗排服务对召回服务返回的广告列表进行自定义向量相似度计算过滤，传统的数学公式计算非常耗时和耗资源，导致粗排服务压力很大，上游召回服务又想召回更多广告给到粗排服务进行再次过滤以提高召回精度。因此关于向量相关的检索和计算需要进行优化以缓解线上服务压力，助力业务发展。
在数据量不大但检索QPS非常高的场景下，使用第三方的向量检索产品可能并不一定是最佳选择，像开源的Faiss、HNSWliib和ScaNN这些优秀的向量检索库比较适用于上亿数量级，而且第三方服务毕竟存在网络请求开销和不稳定性因素，高并发场景下容易导致检索平均时延上升和出现很多毛刺现象。
而百万以内的数据是可以接受在业务服务本身内存中存储，这样可以省去很多网络请求时延，而且在服务本身做向量检索，不依赖第三方服务，检索性能相对稳定。但是在业务服务本身做向量检索会消耗比较多的CPU资源和内存资源，CPU资源是比较稀缺的，而且普通的向量检索效率比较低，时延比较长，如何减少资源消耗和加快向量检索效率成为了优化目标。
2 解决方案
在探索向量检索优化方案的过程中，想到向量检索是一个数学运算的过程，业务服务是Golang写的，Golang是否有开源的做过数学计算优化的库，然后在Github上发现了开源项目Gonum，作为Golang的一个科学计算包，实现了很多常见的数学运算，因此它成为了优化的切入点。
2.1 Gonum计算
向量检索的过程是两个向量按照一定的相似度计算公式进行运算，比如做内积、余弦或欧式距离计算。Gonum包里面有float64和float32的内积运算函数，float64的可以直接调用，float32的没有提供接口，需要在其源码中拷贝到项目中。
Github地址：https://github.com/gonum/gonum
float32的内积运行函数如下，采用的是Plan9汇编实现的。关于Plan9汇编的介绍在这暂时不展开，网上有很多资源可参考，如：https://xargin.com/plan9-assembly/
内积运算公式：
Gonum中内积运算Plan9汇编代码如下：
代码语言：javascript复制#include "textflag.h"
#define HADDPS_SUM_SUM LONG $0xC07C0FF2 // @ HADDPS X0, X0
#define X_PTR SI
#define Y_PTR DI
#define LEN CX
#define TAIL BX
#define IDX AX
#define SUM X0
#define P_SUM X1
// func DotUnitary32(x, y []float32) (sum float32)
TEXT ·DotUnitary32(SB), NOSPLIT, $0
MOVQ x_base 0(FP), X_PTR // X_PTR = &x
MOVQ y_base 24(FP), Y_PTR // Y_PTR = &y
PXOR SUM, SUM // SUM = 0
MOVQ x_len 8(FP), LEN // LEN = min( len(x), len(y) )
CMPQ y_len 32(FP), LEN
CMOVQLE y_len 32(FP), LEN
CMPQ LEN, $0
JE dot_end
XORQ IDX, IDX
MOVQ Y_PTR, DX
ANDQ $0xF, DX // Align on 16-byte boundary for MULPS
JZ dot_no_trim // if DX == 0 { goto dot_no_trim }
SUBQ $16, DX
dot_align: // Trim first value(s) in unaligned buffer do {
MOVSS (X_PTR)(IDX*4), X2 // X2 = x[i]
MULSS (Y_PTR)(IDX*4), X2 // X2 *= y[i]
ADDSS X2, SUM // SUM = X2
INCQ IDX // IDX
DECQ LEN
JZ dot_end // if --TAIL == 0 { return }
ADDQ $4, DX
JNZ dot_align // } while --DX > 0
dot_no_trim:
PXOR P_SUM, P_SUM // P_SUM = 0 for pipelining
MOVQ LEN, TAIL
ANDQ $0xF, TAIL // TAIL = LEN % 16
SHRQ $4, LEN // LEN = floor( LEN / 16 )
JZ dot_tail4_start // if LEN == 0 { goto dot_tail4_start }
dot_loop: // Loop unrolled 16x do {
MOVUPS (X_PTR)(IDX*4), X2 // X_i = x[i:i 1]
MOVUPS 16(X_PTR)(IDX*4), X3
MOVUPS 32(X_PTR)(IDX*4), X4
MOVUPS 48(X_PTR)(IDX*4), X5
MULPS (Y_PTR)(IDX*4), X2 // X_i *= y[i:i 1]
MULPS 16(Y_PTR)(IDX*4), X3
MULPS 32(Y_PTR)(IDX*4), X4
MULPS 48(Y_PTR)(IDX*4), X5
ADDPS X2, SUM // SUM = X_i
ADDPS X3, P_SUM
ADDPS X4, SUM
ADDPS X5, P_SUM
ADDQ $16, IDX // IDX = 16
DECQ LEN
JNZ dot_loop // } while --LEN > 0
ADDPS P_SUM, SUM // SUM = P_SUM
CMPQ TAIL, $0 // if TAIL == 0 { return }
JE dot_end
dot_tail4_start: // Reset loop counter for 4-wide tail loop
MOVQ TAIL, LEN // LEN = floor( TAIL / 4 )
SHRQ $2, LEN
JZ dot_tail_start // if LEN == 0 { goto dot_tail_start }
dot_tail4_loop: // Loop unrolled 4x do {
MOVUPS (X_PTR)(IDX*4), X2 // X_i = x[i:i 1]
MULPS (Y_PTR)(IDX*4), X2 // X_i *= y[i:i 1]
ADDPS X2, SUM // SUM = X_i
ADDQ $4, IDX // i = 4
DECQ LEN
JNZ dot_tail4_loop // } while --LEN > 0
dot_tail_start: // Reset loop counter for 1-wide tail loop
ANDQ $3, TAIL // TAIL = TAIL % 4
JZ dot_end // if TAIL == 0 { return }
dot_tail: // do {
MOVSS (X_PTR)(IDX*4), X2 // X2 = x[i]
MULSS (Y_PTR)(IDX*4), X2 // X2 *= y[i]
ADDSS X2, SUM // psum = X2
INCQ IDX // IDX
DECQ TAIL
JNZ dot_tail // } while --TAIL > 0
dot_end:
HADDPS_SUM_SUM // SUM = sum{ SUM[i] }
HADDPS_SUM_SUM
MOVSS SUM, sum 48(FP) // return SUM
RETGolang如何调用Plan9汇编呢?
- 先将汇编函数保存到后缀为.s的汇编文件中。
- 然后在同级目录下新建一个.go文件，在文件中声明函数，如以上汇编函数声明如下，业务代码直接调用该函数即可。
func DotUnitary32(x, y []float32) (sum float32)Gonum的计算性能怎么样呢？采用了并行计算，内积运算性能是原来的8倍，是满足要求的，具体测试结果在2.4章节中会统一进行对比测试。那既然性能满足要求，是不是就可以了？
因为Gonum的计算函数有限，并不能完全覆盖到我们需要的一些函数，如余弦和欧式距离计算，或者在标准的计算过程中加一些自定义的计算公式，Gonum是做不到的。
受到Gonum并行计算的启发，想到是否可以使用SIMD（单指令多数据流）指令集来加速计算。
2.2 SIMD计算
SIMD单指令流多数据流(SingleInstruction Multiple Data,SIMD)是一种采用一个控制器来控制多个处理器，同时对一组数据（又称“数据向量”）中的每一个分别执行相同的操作从而实现空间上的并行性的技术。在微处理器中，单指令流多数据流技术则是一个控制器控制多个平行的处理微元，例如Intel的MMX或SSE以及AMD的3D Now!技术。
目前Intel处理器支持的SIMD技术包括MMX,SSE,AVX.，MMX提供了8个64bit的寄存器进行SIMD操作，SSE系列提供了128bit的8个寄存器进行SIMD指令操作，AVX指令则支持256bit的SIMD操作。目前SIMD指令可以有四种方法进行使用分别是汇编语言，C 类，编译器Intrisincs和自动矢量化。SIMD intrinsics有些类似于C语言中的函数，可以被其它的代码直接调用，相比汇编语言来说更容易使用。
Intel官网的SIMD intrinsics指令集查询
https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html
为了使用与SIMD技术相关的intrinsics，首先需要包含那些定义了数据类型和函数的头文件。
代码语言：javascript复制#include <mmintrin.h> //MMX __m64 定义
#include <xmmintrin.h> //SSE(include mmintrin.h) __m128 定义
#include <emmintrin.h> //SSE2(include xmmintrin.h) __m128i ,__m128d 类型 定义
#include <pmmintrin.h> //SSE3(include emmintrin.h)
#include <tmmintrin.h>//SSSE3(include pmmintrin.h)
#include <smmintrin.h>//SSE4.1(include tmmintrin.h)
#include <nmmintrin.h>//SSE4.2(include smmintrin.h)
#include <wmmintrin.h>//AES(include nmmintrin.h)
#include <immintrin.h>//AVX(include wmmintrin.h) __m256 ,__m256i ,__m256d 类型定义
#include <intrin.h>//(include immintrin.h) 所有架构关键的SIMD AVX2指令集（256位寄存器），可以利用这些常见的指令进行自定义计算。
- __m256（声明寄存器变量）
- _mm256_loadu_ps（加载数据到寄存器）
- _mm256_mul_ps（寄存器相乘）
- _mm256_add_ps（寄存器相加）
- _mm256_extractf128_ps（获取高128位或低128位）
- _mm_hadd_ps（水平相加，如x的第1位和第2位相加结果放在新数组第1位，y的第1位和第2位相加结果放在新数组第2位，然后x和y下标移动两位依次重复以上操作将结果追加到新数组后面）
- _mm_storeu_ps（取出寄存器值赋值）
- _mm_sqrt_ps（开平方）
- _mm_log1p_ps（log(1 p)）
查看机器支持的指令集两种方式
代码语言：javascript复制lscpu // 查看flags标志中支持的所有指令集
gcc -mavx2 -dM -E - < /dev/null|egrep AVX // 查看是否支持AVX
gcc -msse4 -dM -E - < /dev/null|egrep SSE // 查看是否支持SSE2.2.1 内积距离计算
使用SIMD AVX2指令进行256维向量的内积距离计算，计算公式：
void VecInnerProductAVX2(const float* x, const float* y, float* z) {
int d = 256;
__m256 msum1 = _mm256_setzero_ps();
// 加载数据并计算乘积和累计和
while (d >= 8) {
__m256 mx = _mm256_loadu_ps(x);
x = 8;
__m256 my = _mm256_loadu_ps(y);
y = 8;
msum1 = _mm256_add_ps(msum1, _mm256_mul_ps(mx, my));
d -= 8;
}
// 将寄存器中的结果求和并赋值给新数组返回
__m128 msum2 = _mm256_extractf128_ps(msum1, 1);
msum2 = _mm_add_ps(msum2, _mm256_extractf128_ps(msum1, 0));
msum2 = _mm_hadd_ps(msum2, msum2);
msum2 = _mm_hadd_ps(msum2, msum2);
_mm_storeu_ps(z, msum2);
}2.2.2 欧式距离计算
使用SIMD AVX2指令进行256维向量的欧式距离计算，计算公式如下：
SIMD代码如下：
代码语言：javascript复制void VecEuclideanDistanceAVX2(const float* x, const float* y, float* z) {
int d = 256;
__m256 msum1 = _mm256_setzero_ps();
// 加载数据并计算差值平方累计和
while (d >= 8) {
__m256 mx = _mm256_loadu_ps(x);
x = 8;
__m256 my = _mm256_loadu_ps(y);
y = 8;
__m256 sub = _mm256_sub_ps(mx, my);
msum1 = _mm256_add_ps(msum1, _mm256_mul_ps(sub, sub));
d -= 8;
}
// 将寄存器中的结果求和并开平方根，最终结果赋值给新数组返回
__m128 msum2 = _mm256_extractf128_ps(msum1, 1);
msum2 = _mm_add_ps(msum2, _mm256_extractf128_ps(msum1, 0));
msum2 = _mm_hadd_ps(msum2, msum2);
msum2 = _mm_hadd_ps(msum2, msum2);
_mm_storeu_ps(z, _mm_sqrt_ps(msum2));
}2.2.3 余弦距离计算
使用SIMD AVX2指令进行256维向量的余弦距离计算，计算公式如下：
SIMD代码如下：
代码语言：javascript复制void VecCosineAVX2(const float* x, const float* y, float* z) {
int d = 256;
__m256 msum1 = _mm256_setzero_ps();
__m256 msumx = _mm256_setzero_ps();
__m256 msumy = _mm256_setzero_ps();
// 加载数据并计算x和y的内积、模长
while (d >= 8) {
__m256 mx = _mm256_loadu_ps(x);
x = 8;
__m256 my = _mm256_loadu_ps(y);
y = 8;
msum1 = _mm256_add_ps(msum1, _mm256_mul_ps(mx, my));
msumx = _mm256_add_ps(msumx, _mm256_mul_ps(mx, mx));
msumy = _mm256_add_ps(msumy, _mm256_mul_ps(my, my));
d -= 8;
}
// 将寄存器msum1中的内积结果求和
__m128 msum2 = _mm256_extractf128_ps(msum1, 1);
msum2 = _mm_add_ps(msum2, _mm256_extractf128_ps(msum1, 0));
msum2 = _mm_hadd_ps(msum2, msum2);
msum2 = _mm_hadd_ps(msum2, msum2);
// 将寄存器msumx中的x向量模长结果求和
__m128 msumx2 = _mm256_extractf128_ps(msumx, 1);
msumx2 = _mm_add_ps(msumx2, _mm256_extractf128_ps(msumx, 0));
msumx2 = _mm_hadd_ps(msumx2, msumx2);
msumx2 = _mm_hadd_ps(msumx2, msumx2);
// 将寄存器msumy中的y向量模长结果求和
__m128 msumy2 = _mm256_extractf128_ps(msumy, 1);
msumy2 = _mm_add_ps(msumy2, _mm256_extractf128_ps(msumy, 0));
msumy2 = _mm_hadd_ps(msumy2, msumy2);
msumy2 = _mm_hadd_ps(msumy2, msumy2);
// 根据内积、x模长和y模长计算余弦距离
__m128 result = _mm_div_ps(msum2, _mm_mul_ps(_mm_sqrt_ps(msumx2), _mm_sqrt_ps(msumy2)));
_mm_storeu_ps(z, result);
}一个寄存器一次加载8个32位的浮点数，理论上性能应该是原来的8倍，实际上经过测试这个猜想也得到了验证，详细数据在2.4节中给出。
为什么这些函数不直接返回结果，而把结果存在一个数组中呢？若C或C 调用这些函数可以直接返回结果，但是若使用Golang进行调用，需要进行一些转换，为什么要这么做？接下来将介绍Golang如何调用SIMD函数。
2.3 Golang调用SIMD
2.3.1 CGO调用
SIMD函数是使用C编写的，Golang调用C函数，最容易想到的就是采用Golang提供的CGO方式进行C函数调用。
代码语言：javascript复制/*
#cgo CFLAGS: -mavx2
#include <immintrin.h>
float vecInnerProduct(float* x, float* y) {
// 此次省略内积运算过程，返回值可以直接返回，不用放到额外的数组中
// _mm_storeu_ps(z, msum2);
return _mm_cvtss_f32(msum2);
}
*/
import "C"
func VecInnerProductByCGo(userVector []float32, itemVector []float32) float32 {
// Golang的floa和C的float互转
return float32(C.vecInnerProduct((*C.float)(&userVector[0]), (*C.float)(&itemVector[0])))
}CGO这种方式确实是可以的，但是存在Golang和C之间不同语言的上下文切换，存在性能问题，肯定不能完全发挥出SIMD所能提升的8倍，经过测试，CGO这种方式只能是未优化的普通计算的2倍左右，这种方式不能满足我们的业务需求，详细数据在2.4节中给出。
2.3.1 Plan9汇编调用
Golang是可以直接调用Plan9汇编的，但是C写的SIMD函数怎么转Plan9汇编呢？
在Github上发现一个开源的项目c2goasm，它可以将C函数汇编转成Plan9汇编，c2goasm的本质也是调用asm2plan9s工具将C的汇编转成Plan9汇编。
c2goasm项目地址：https://github.com/minio/c2goasm
asm2plan9s项目地址：https://github.com/minio/asm2plan9s
（1）C转C汇编
可以将SIMD函数先使用Clang编译成C的汇编，如将simd.c编译成simd.s汇编，编译命令如下：
代码语言：javascript复制clang -S -O1 -mavx2 -mfma -masm=intel -mno-red-zone -mstackrealign -mllvm -inline-threshold=1000 -fno-asynchronous-unwind-tables -fno-exceptions -fno-rtti -c simd.c -o simd.s注意：
- GCC优化标识必须为O1或者Os，否则最终的结果会不正确，关于编译的代码优化标识解释，可以参考：https://blog.csdn.net/liang_baikai/article/details/110137374
- 安装clang 7.0.0版本，可执行文件在根目录下的bin目录，其它版本（高于10.0版本）可能不支持-masm=intel参数。下载链接：https://releases.llvm.org/download.html
SIMD内积运算的汇编代码如下：
代码语言：javascript复制 .globl VecInnerProductAVX2 # -- Begin function VecInnerProductAVX2
.p2align 4, 0x90
.type VecInnerProductAVX2,@function
VecInnerProductAVX2: # @VecInnerProductAVX2
# �.0:
push rbp
mov rbp, rsp
and rsp, -8
vxorps xmm0, xmm0, xmm0
mov eax, 264
xor ecx, ecx
.p2align 4, 0x90
.LBB1_1: # =>This Inner Loop Header: Depth=1
vmovups ymm1, ymmword ptr [rdi rcx]
vmulps ymm1, ymm1, ymmword ptr [rsi rcx]
vaddps ymm0, ymm0, ymm1
add rcx, 32
add eax, -8
cmp eax, 15
ja .LBB1_1
# �.2:
vextractf128 xmm1, ymm0, 1
vaddps xmm0, xmm1, xmm0
vhaddps xmm0, xmm0, xmm0
vhaddps xmm0, xmm0, xmm0
vmovups xmmword ptr [rdx], xmm0
mov rsp, rbp
pop rbp
vzeroupper
ret
.Lfunc_end1:
.size VecInnerProductAVX2, .Lfunc_end1-VecInnerProductAVX2
# -- End function（2）C汇编转Plan9汇编
1. 编译c2goasm，生成可执行文件c2goasm，并添加到PATH路径。
代码语言：javascript复制git clone git@github.com:minio/c2goasm.git
go mod init c2goasm
go build2. 下载asm2plan9s可执行文件，并添加到PATH路径。
代码语言：javascript复制go install github.com/minio/asm2plan9s
yum install -y yasm3. 使用c2goasm工具将SIMD的汇编文件simd.s转成plan9汇编simd_avx2.s 。
代码语言：javascript复制c2goasm -a simd.s simd_avx2.s4. 将SIMD内积运算的C汇编代码通过c2goasm转成Plan9汇编如下，默认会在函数名前加一个下划线。
代码语言：javascript复制TEXT ·_VecInnerProductAVX2(SB), $0-24
MOVQ x 0(FP), DI
MOVQ y 8(FP), SI
MOVQ z 16(FP), DX
LONG $0xc057f8c5 // vxorps xmm0, xmm0, xmm0
LONG $0x000108b8; BYTE $0x00 // mov eax, 264
WORD $0xc931 // xor ecx, ecx
LBB2_1:
QUAD $0x0000000f8c10fcc5; BYTE $0x00 // vmovups ymm1, yword [rdi rcx]
QUAD $0x0000000e8c59f4c5; BYTE $0x00 // vmulps ymm1, ymm1, yword [rsi rcx]
LONG $0xc158fcc5 // vaddps ymm0, ymm0, ymm1
LONG $0x20c18348 // add rcx, 32
WORD $0xc083; BYTE $0xf8 // add eax, -8
WORD $0xf883; BYTE $0x0f // cmp eax, 15
JA LBB2_1
LONG $0x197de3c4; WORD $0x01c1 // vextractf128 xmm1, ymm0, 1
LONG $0xc058f0c5 // vaddps xmm0, xmm1, xmm0
LONG $0xc07cfbc5 // vhaddps xmm0, xmm0, xmm0
LONG $0xc07cfbc5 // vhaddps xmm0, xmm0, xmm0
LONG $0x4211f8c5; BYTE $0x10 // vmovups oword [rdx], xmm0
VZEROUPPER
RETasm2plan9s生成以上Plan9汇编指令的源代码参考：https://github.com/minio/asm2plan9s/blob/master/yasm.go中的函数toPlan9sYasm。
关键Plan9汇编指令
- MOVQ（搬运8个字节）
- BYTE、WORD、LONG、QUAD（将1、2、4、8字节数据放入指令流）
- JA (EFLAGS寄存器的标志位大于则跳转)
- VZEROUPPER（YMM寄存器高位置零）
（3）Golang调用Plan9汇编
需要提前在创建一个与目标汇编文件（simd_avx2.s）同名的go文件（如simd_avx2.go），声明C语言中的函数（带下划线），函数入参个数与原来C源码中的入参个数相等，参数需要是64位的，若有返回值，返回值的名字不能省略。
代码语言：javascript复制//go:noescape
func _VecInnerProductAVX2(x, y, z *float32)其它业务函数直接调用该函数即可，尝试过这些距离计算函数直接返回结果，最终拿不到结果，而有些函数可以直接返回结果，暂时还没有发现c2goasm转换后的调用规律，有兴趣的小伙伴可以研究讨论。
2.4 对比测试
实验环境
运行环境 | CPU类型 | 操作系统 | CPU | 内存 |
|---|---|---|---|---|
支持AVX2的CVM | Intel Xeon | CentOS Linux release 8.2.2.2004 (Core) | 4核 | 8G |
根据Gonum32、CGO方式调用SIMD，Golang调用Plan9汇编SIMD三种计算方式，对比未优化的普通内积计算，计算能力对比如下图：
内积计算能力和时延相比未优化的普通内积计算均有提升，结果如下：
- Gonum： 8倍
- SIMD-Cgo： 2倍
- SIMD-Plan9汇编：8.7倍
注意：
- 此处是单协程测试数据，多协程测试的时延更低，处理能力更高。
- 在AMD架构的机器上进行相同的测试，和Intel架构的机器测试结果没有明显差异。
对未优化的比普通内积计算，CPU资源使用对比如下图：
从图中看出，SIMD-Plan9汇编的内积运算CPU资源使用最低。
另外，对于余弦距离和欧式距离也进行了相同测试，测试结果与内积距离的性能提升结果基本一致。
综合计算能力、时延和CPU资源等方面，SIMD-Plan9汇编方案综合性能较优，因此可以采用此方案对线上服务进行优化。
3 小结
本文主要介绍了在当前的向量检索业务挑战的背景下，研究了如何在内存中进行本地向量检索的探索流程，对探索的多种方案也进行了压测，最终得出了综合性能较优的SIMD-Plan9汇编方案。
但实际上向量检索的流程还有前置的向量过滤（可选流程）和后置的检索结果排序，这两个方面也有进一步优化的空间，以及整体优化后的效果将在下一篇文章《向量检索研究系列：本地向量检索（下）》中进行详细介绍。
