所有文章 > AI驱动 > 向量检索研究系列:本地向量检索
向量检索研究系列:本地向量检索

向量检索研究系列:本地向量检索

向量检索研究系列:本地向量检索(上)

1 背景

广告推荐业务场景中,如果召回服务的向量检索时延过高,会导致召回服务整体时延达到上限,很多请求超时以至于没有广告返回给上游服务。同时粗排服务对召回服务返回的广告列表进行自定义向量相似度计算过滤,传统的数学公式计算非常耗时和耗资源,导致粗排服务压力很大,上游召回服务又想召回更多广告给到粗排服务进行再次过滤以提高召回精度。因此关于向量相关的检索和计算需要进行优化以缓解线上服务压力,助力业务发展。

在数据量不大但检索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汇编代码如下:

#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
RET

 Golang如何调用Plan9汇编呢?

  1. 先将汇编函数保存到后缀为.s的汇编文件中。
  2. 然后在同级目录下新建一个.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,首先需要包含那些定义了数据类型和函数的头文件。

#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))

查看机器支持的指令集两种方式

lscpu // 查看flags标志中支持的所有指令集
gcc -mavx2 -dM -E - < /dev/null|egrep AVX // 查看是否支持AVX
gcc -msse4 -dM -E - < /dev/null|egrep SSE // 查看是否支持SSE

2.2.1 内积距离计算

使用SIMD AVX2指令进行256维向量的内积距离计算,计算公式如下:

SIMD代码如下,256位寄存器一次可加载8个32位浮点数。

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代码如下:

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代码如下:

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函数调用。

/*
#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汇编,编译命令如下:

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内积运算的汇编代码如下:

  .globl        VecInnerProductAVX2     # -- Begin function VecInnerProductAVX2
.p2align 4, 0x90
.type VecInnerProductAVX2,@function
VecInnerProductAVX2: # @VecInnerProductAVX2
# %bb.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
# %bb.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路径。

git clone git@github.com:minio/c2goasm.git
go mod init c2goasm
go build

 2. 下载asm2plan9s可执行文件,并添加到PATH路径。

go install github.com/minio/asm2plan9s

3. 使用c2goasm工具将SIMD的汇编文件simd.s转成plan9汇编simd_avx2.s 。

c2goasm -a simd.s simd_avx2.s

4. 将SIMD内积运算的C汇编代码通过c2goasm转成Plan9汇编如下,默认会在函数名前加一个下划线。

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
RET

 asm2plan9s生成以上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位的,若有返回值,返回值的名字不能省略。

//go:noescape
func _VecInnerProductAVX2(x, y, z *float32)

 其它业务函数直接调用该函数即可,尝试过这些距离计算函数直接返回结果,最终拿不到结果,而有些函数可以直接返回结果,暂时还没有发现c2goasm转换后的调用规律,有兴趣的小伙伴可以研究讨论。

2.4 对比测试

实验环境

运行环境CPU类型操作系统CPU内存
支持AVX2的CVMIntel XeonCentOS 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汇编方案。

但实际上向量检索的流程还有前置的向量过滤(可选流程)和后置的检索结果排序,这两个方面也有进一步优化的空间,以及整体优化后的效果将在下一篇文章《向量检索研究系列:本地向量检索(下)》中进行详细介绍。

向量检索研究系列:本地向量检索(下)

1 背景

上一篇文章介绍了如何加快向量相似度计算,但是一般的向量检索流程还包括对计算结果进行排序,以及有必要的话,在计算相似度之前可以对向量库中的向量进行过滤筛选(可选流程)。

本地向量检索在过滤和排序这两个方面也有进一步优化的空间,本文将介绍一下过滤方案和排序方案。

2 过滤优化

本地向量计算是先把向量加载内存再计算,经过SIMD优化,计算速度快了,但是否还能进一步减少待计算相似度的向量集呢?

举个例子,一个用户向量本来要和向量集所有1000个向量进行相似度计算,是否可以在内存中通过对向量进行属性过滤,让用户向量只需要和向量集中500个向量进行相似度计算,这样可以加快总体的向量检索速度。

2.1 向量过滤

把广告通过模型转成向量后,向量应该关联广告的一些基本信息,广告检索条件是基于这些广告属性的,检索的时候可以根据检索条件在向量关联的广告信息中进行向量的筛选过滤。

广告信息和检索条件:

  • 模型版本
  • 冷启动或非冷启动创意
  • 平台
  • 模板
  • 媒体

基于内存进行向量过滤暂时有想到如下三种方案:

方案一:内存对象

  • 将广告信息存储为对象属性,向量也是其中一个属性,遍历广告对象,根据对象属性进行过滤。

方案二:内存Bitmap

  • 每个广告属性的取值都生成一个Bitmap,广告ID为下标,如平台属性中为iOS平台和安卓平台各生成一个bitmap,检索条件对应着多个bitmap,对这些bitmap进行集合运算即可得到满足条件的广告。bitmap举例如下:

方案三:内存倒排索引

  • 使用两个两层Map结构存储广告信息,第一个Map存储索引信息,一级Key :“模型版本_冷启动或非冷启动创意”,二级Key :“平台_模板_媒体”,值为广告ID列表。第二个Map存储向量信息,一级Key :“模型版本”,二级Key :广告ID,值为广告向量。
  • 检索时把检索条件在第一个Map中查询到满足检索条件的广告ID列表,再根据ID列表从第二个Map中取出对应向量列表。
  • 大致结构可以参考2.2中向量存储方案图。

对这三种方案的QPS和资源占用情况进行了测试,测试结果如下图:

  • QPS:倒排索引 > Bitmap> 对象
  • CPU资源:Bitmap > 对象 > 倒排索引
  • 时延(微秒级别,此处没有展示):对象 > Bitmap > 倒排索引

100万数据量以下方案三倒排索引的综合性能较优。

2.2 向量存储

线上倒排索引需要考虑向量存储,实现方案分为离线刷入数据到Redis在线从Redis读取数据到内存两个阶段。

在离线刷入数据到Redis阶段,有两种刷入方案:

方案一:如下图左侧所示,使用单个Hash存储,Hash的Key和Field存储条件,Value存储向量列表,同时对这些向量列表进行zip和base64压缩,浮点数的压缩率不高,仅2倍的压缩率。因为有些广告会在多个条件中出现,因此向量也会在多个Filed中出现,所以会存在向量冗余。

方案二:如下图右侧所示,使用一个Hash存储索引条件和广告ID列表,用多个单独的Key/value存储广告ID和对应的向量。若在Redis把这些单独的向量Key用一个Hash进行存储,则会出现大Key,请求这些大Key会导致某些节点压力过高,响应速度变慢,而使用单独的Key存储可以分散请求压力,提高后台服务请求Redis速度。

后台服务从Redis读取向量数据到内存,实验10万个广告,使用方案二,存储向量需要内存270M,存储倒排索引3M。如果线上4个版本的向量进行AB实验,则内存总占用约1G

Redis中多个单独的Key和Value读到内存后被存储在一个两层的Map中。

综合刷入数据和读取数据两个阶段,两种方案的优缺点如下:

方案一

  • 优点:查询Redis次数少
  • 缺点:解析复杂、冗余向量内存占用大、不便增量更新

方案二

  • 优点:解析简单、内存占用小、方便增量更新
  • 缺点:查询Redis次数比方案一要多

因此方案二的Redis存储方式更有利于线上服务存储和更新广告向量。

3 排序优化

向量过滤和相似度计算完之后,对结果取TopK的排序是否耗时?能否优化?

3.1 全量排序

把Golang官方的排序算法(快排+堆排序+插入排序)时间和SIMD相似度计算时间进行对比,如下图:

可见,排序运算时间 > 内积运算时间,Golang默认的排序算法不满足需求。

向量是浮点数数组,内积计算的结果是浮点数,浮点数结果排序方案对比:

  • Go官方排序(快排+堆排序+插入排序)
  • 堆排序(TopK问题常用算法)
  • 浮点数基数排序(非比较型排序)
  • 并行浮点数基数排序(分而治之)

基数排序常用于整数排序,基于浮点数的基数排序也是本小节的重点,其改造核心思想如下:

  • 浮点数转二进制分段
  • 多次分桶排序
  • 处理负数

浮点数基数排序的大致流程如下,可参考下图数字表标识顺序:

  1. 将待排序的浮点数转成二进制,并分成多段。
  2. 将所有浮点数的第1段映射到桶里面,段的二进制位数决定了桶的大小,如8位二进制段对应的桶大小为256。
  3. 在桶里面确定浮点数的相对位置。
  4. 根据这个相对位置再进行浮点数第2段排序,重复步骤2~3。
  5. 直至所有分段都分桶完成并确定元素相对位置后已经得到浮点数的大致顺序,因为负数带符号位,最高位为1,负数会在数组的后面,需要将负数反转至数组头部即可得到最终排序好的浮点数数组。

上面提到需要对浮点数的二进制进行分段,到底分多少段比较合适呢?

根据算法流程,得出时间复杂度公式:O(d*(n+2^(32/d))+n),其中d为浮点数分段个数,n为待排序数据量,括号中三个时间的相加,分别代表着分桶、确定元素相对位置、将原数组元素按顺序放到新数组中。32表示是32位的浮点数。

浮点数分段数时间复杂度待排序数量的合适取值范围
12n+2^32n > 2,147,418,112
24n+2^1732512 < n <= 2,147,418,112
48n+2^10124 < n <= 32512
816n+2^74 < n <= 124
1632n+2^61 =< n <= 4
3264n+2^60

注意:这仅是理论上的估算值,对分段趋势的一个大概判断。同时也在代码层面对分2段、4段、8段进行了测试,其排序时间对比如下图:

可以看出,数据量越大,分段数越少排序越快,这和表格中的分段趋势估算一致。在5万数据量以下,分4段的效果最好,大于5万时,分2段的效果较好。

数据量非常大的时候是否能并行排序?

并行浮点数基数排序思想:

  • 多协程分批排序后各取TopK
  • 子结果汇总后再取TopK

对于上面提到的几种排序算法进行了测试,同样也和SIMD内积运算时间进行对比,其测试结果如下:

上图中各排序方案性能:并行浮点数基数排序 > 浮点数基数排序 > Golang官方排序 > 堆排序

浮点数基数排序是Golang官方排序的4~5倍,并行浮点数基数排序是Golang官方排序的1~11倍。并行浮点数基数排序在数据量比较小的时候反而性能没有浮点数基数排序好,因为并行也存在性能损耗。

此时可以看出浮点数基数排序时间已经比SIMD相似度计算时间要短,已经满足我们的业务需求。

3.2 局部排序

前面提到的排序都是对全量的数据进行排序,然后对结果取TopK,如果只对部分数据进行排序拿到TopK结果,不关心其它数据顺序,因此可以考虑对现有排序算法进行局部排序改造。

局部排序改造思想

方案一:冒泡排序

  • 冒泡排序每次循环都会找到一个最大或最小的数值,循环TopK次就可以找到最终的TopK结果,退出算法即可。
  • 时间复杂度:O(n*k)

方案二:快速排序.

  • 利用二分法的思想,对于每次找到支点pivot时,判断pivot位置若大于TopK,则在pivot的左半部分继续递归,若pivot位置小于TopK,则在pivot的右半部分继续递归,直至pivot等于TopK退出递归,数组的前TopK个数就是最终的TopK结果。
  • 时间复杂度:O(n*logn)

方案三:堆排序

  • 取出数组的前TopK个数构建的一个小顶堆,然后遍历原数组第TopK之后所有的数,依次和堆顶进行比较,若比堆顶大,则插入堆中,进行堆调整。遍历完之后的堆即可TopK结果。
  • 时间复杂度:O(n*logk)

对这几种局部排序在不同的数据量下取TopK(k=1000)进行了测试,结果如下:

算法\数据量20001万2万5万10万100万
冒泡排序TopK2.3μs12μs114μs205.087ms321.765ms2530ms
快速排序TopK1403μs8505μs17135μs43.705ms88.822ms883ms
堆排序TopK59μs246μs335μs0.436ms0.551ms1.364ms

从表格中可以得出以下结论:

  • 冒泡排序在2万数据量以下,其排序时间比其它局部排序要短。
  • 快速排序的整体性能不佳,从其时间复杂度可以得知原因。
  • 堆排序的性能比较稳定,在5万及以上的数据量时,其排序性能比较好
  • 堆排序对比之前的浮点数基数排序和并行浮点数基数排序,在10万以下数据量时,性能相差不大,在10万数据量时还是堆排序的性能较优。

根据当前的业务数据量和数据增长趋势,选择堆排序的局部排序算法比较合适

4. 方法论

实践过程中的经验不仅能优化当前业务,也能沉淀成可复制的方法论,应用到更广的业务场景,为更多的业务赋能。

  1. 本地向量检索方案可以为100万以下数据量的业务提供快速、高性能且稳定的向量检索方案。
  2. SIMD自定义编程可以在应用到其它偏数学计算的业务,加速计算。
  3. 倒排索引和Bitmap的内存过滤方案可以为其它数据过滤场景提供思路。
  4. 浮点数基数排序局部排序算法可应用到业务的其它排序场景,加速排序。

5 总结

经本地向量检索和计算优化后,召回和粗排服务的时延都有大幅度下降,随着QPS和广告数的增长,线上服务仍能轻松处理请求,可支撑更大规模的业务发展。

文章转自微信公众号@IEG增长中台技术团队

#你可能也喜欢这些API文章!