分类 汇编

使用AVX2指令集加速推荐系统MMR层余弦相似度计算

发布于
分类 汇编
标签 Go

1. 背景

前一段时间公司上线了一套Go实现的推荐系统,上线后发现MMR层虽然只有纯计算但耗时十分离谱,通过pprof定位问题所在之后进行了优化,虽然降低了非常多但是我们认为其中还有优化空间。

image.png

可以看到日常平均耗时126ms,P95 360ms。

MMR层主要耗时集中在了余弦相似度的计算部分,这部分我们使用的gonum库进行计算,其底层在x86平台上利用了SSE指令集进行了加速。

SSE指令集已经非常古老了,xmm寄存器只能存储两个双精度浮点数,每次只能并行进行两个双精度浮点数的计算,而AVX2指令集可以并行计算四个,理论上可以获得两倍的性能提升,因此我们决定自己使用AVX2指令集手写汇编的方式替代掉gonum库。

1.1 余弦相似度算法

余弦相似度的计算公式为

image.png

对应的代码为

import "gonum.org/v1/gonum/floats"

func CosineSimilarity(a, b []float64) float64 {
    dotProduct := floats.Dot(a, b) // 计算a和b的点积
    normA := floats.Norm(a, 2) // 计算向量a的L2范数
    normB := floats.Norm(b, 2) // 计算向量b的L2范数
    return dotProduct / (normA * normB)
}

2. Dot点积计算加速

gonum点积计算Dot的部分汇编代码如下:

TEXT ·DotUnitary(SB), NOSPLIT, $0
    ...
loop_uni:
	// sum += x[i] * y[i] unrolled 4x.
	MOVUPD 0(R8)(SI*8), X0
	MOVUPD 0(R9)(SI*8), X1
	MOVUPD 16(R8)(SI*8), X2
	MOVUPD 16(R9)(SI*8), X3
	MULPD  X1, X0
	MULPD  X3, X2
	ADDPD  X0, X7
	ADDPD  X2, X8

	ADDQ $4, SI   // i += 4
	SUBQ $4, DI   // n -= 4
	JGE  loop_uni // if n >= 0 goto loop_uni

    ...

end_uni:
	ADDPD    X8, X7
	MOVSD    X7, X0
	UNPCKHPD X7, X7
	ADDSD    X0, X7
	MOVSD    X7, sum+48(FP) // Return final sum.
	RET

可以看到其中使用xmm寄存器并行计算两个双精度浮点数,并且还采用了循环展开的优化手段,一个循环中同时进行4个元素的计算。

我们利用AVX2指令集并行计算四个双精度浮点数进行加速

loop_uni:
	// sum += x[i] * y[i] unrolled 8x.
	VMOVUPD 0(R8)(SI*8), Y0 // Y0 = x[i:i+4]
	VMOVUPD 0(R9)(SI*8), Y1 // Y1 = y[i:i+4]
	VMOVUPD 32(R8)(SI*8), Y2 // Y2 = x[i+4:i+8]
	VMOVUPD 32(R9)(SI*8), Y3 // Y3 = x[i+4:i+8]
	VMOVUPD 64(R8)(SI*8), Y4 // Y4 = x[i+8:i+12]
	VMOVUPD 64(R9)(SI*8), Y5 // Y5 = y[i+8:i+12]
	VMOVUPD 96(R8)(SI*8), Y6 // Y6 = x[i+12:i+16]
	VMOVUPD 96(R9)(SI*8), Y7 // Y7 = x[i+12:i+16]
	VFMADD231PD Y0, Y1, Y8 // Y8 = Y0 * Y1 + Y8
	VFMADD231PD Y2, Y3, Y9
	VFMADD231PD Y4, Y5, Y10
	VFMADD231PD Y6, Y7, Y11
	ADDQ $16, SI   // i += 16
	CMPQ DI, SI
	JG  loop_uni // if len(x) > i goto loop_uni

可以看到我们每个循环中同时用到8个ymm寄存器即一次循环计算16个数,而且还用到了VFMADD231PD指令同时进行乘法累积的计算。

最终Benchmark结果:

BenchmarkDot 一个循环中计算8个数
BenchmarkDot-2          14994770                78.85 ns/op
BenchmarkDot16 一个循环中计算16个数
BenchmarkDot16-2        22867993                53.46 ns/op
BenchmarkGonumDot Gonum点积计算
BenchmarkGonumDot-2      8264486               144.4 ns/op

可以看到点积部分我们得到了大约2.7倍的性能提升

3. L2范数计算加速

gonum库中进行L2范数计算的算法并不是常规的a1^2 + a2^2 ... + aN^2这种计算,而是采用了Netlib算法,减少了溢出和下溢,其Go源码如下:

func L2NormUnitary(x []float64) (norm float64) {
	var scale float64
	sumSquares := 1.0
	for _, v := range x {
		if v == 0 {
			continue
		}
		absxi := math.Abs(v)
		if math.IsNaN(absxi) {
			return math.NaN()
		}
		if scale < absxi {
			s := scale / absxi
			sumSquares = 1 + sumSquares*s*s
			scale = absxi
		} else {
			s := absxi / scale
			sumSquares += s * s
		}
	}
	if math.IsInf(scale, 1) {
		return math.Inf(1)
	}
	return scale * math.Sqrt(sumSquares)
}

其汇编代码比较晦涩难懂,但管中窥豹再结合Go源码可以看出来没有用到并行能力,每次循环只计算一个数

TEXT ·L2NormUnitary(SB), NOSPLIT, $0
    ...
loop:
	MOVSD   (X_)(IDX*8), ABSX // absxi = x[i]
	...

我们优化之后的核心代码如下:

loop:
	VMOVUPD 0(R8)(SI*8), Y0 // Y0 = x[i:i+4]
	VMOVUPD 32(R8)(SI*8), Y1 // Y1 = y[i+4:i+8]
	VMOVUPD 64(R8)(SI*8), Y2 // Y2 = x[i+8:i+12]
	VMOVUPD 96(R8)(SI*8), Y3 // Y3 = x[i+12:i+16]
	VMOVUPD 128(R8)(SI*8), Y4 // Y4 = x[i+16:i+20]
	VMOVUPD 160(R8)(SI*8), Y5 // Y5 = y[i+20:i+24]
	VMOVUPD 192(R8)(SI*8), Y6 // Y6 = x[i+24:i+28]
	VMOVUPD 224(R8)(SI*8), Y7 // Y7 = x[i+28:i+32]
	VFMADD231PD Y0, Y0, Y8 // Y8 = Y0 * Y0 + Y8
	VFMADD231PD Y1, Y1, Y9
	VFMADD231PD Y2, Y2, Y10
	VFMADD231PD Y3, Y3, Y11
	VFMADD231PD Y4, Y4, Y12
	VFMADD231PD Y5, Y5, Y13
	VFMADD231PD Y6, Y6, Y14
	VFMADD231PD Y7, Y7, Y15

	ADDQ $32, SI // i += 32
	CMPQ DI, SI
	JG  loop // if len(x) > i goto loop

我们采用原始的算法计算以利用到并行计算的能力,并且循环展开,一次循环中同时计算32个数,最终Benchmark结果:

BenchmarkAVX2L2Norm
BenchmarkAVX2L2Norm-2          29381442                40.99 ns/op
BenchmarkGonumL2Norm
BenchmarkGonumL2Norm-2           1822386               659.4 ns/op

可以看到得到了大约16倍的性能提升

4. 总结

通过这次优化我们在余弦相似度计算部分最终得到了(144.4 + 659.4 * 2) / (53.46 + 40.99 * 2) = 10.8倍的性能提升,效果还是非常显著的。相较于《记一次SIMD指令优化计算的失败经历》这次失败的初次尝试,本次还是非常成功的,切实感受到了SIMD的威力。

另外在本次优化过程中也涨了不少姿势

AVX-512指令降频问题

AVX-512指令因为并行度更高理论上性能也更高,但AVX-512指令会造成CPU降频,因此业界使用非常慎重,这一点可以参考字节的json解析库sonic的这个issue: https://github.com/bytedance/sonic/issues/319

循环展开优化

在一次循环中做更多的工作,优点有很多:

  • 减少循环控制的开销,循环变量的更新和条件判断次数更少,降低了分支预测失败的可能性
  • 增加指令并行性,更多的指令可以在流水线中并行执行

但一次循环使用过多的寄存器从实际Benchmark看性能确实更好,但是否存在隐患我没有看到相关的资料,希望这方面的专家可以指教一下。

...

阅读全文 »

记一次SIMD指令优化计算的失败经历

发布于
分类 汇编
标签 Go

1. 前言

书接上回 《统计一个数字二进制位1的个数》,现在我们已经知道如何快速计算出一个int64数字的二进制位1的个数,那么回到我们最初的需求,我们的目的是快速统计一个bitmap中二进制位1的个数,假设我们使用[]uint64来实现bitmap,那么如果要统计这个bitmap中二进制位1的个数,我们可以遍历每个元素,计算出每个uint64元素二进制位1的个数,最后加起来,代码大概如下:

type Bitmap []uint64

func (bitmap Bitmap) OnesCount() (count int) {
	for _, v := range bitmap {
		count += OnesCount64(v)
	}

	return
}

const m0 = 0x5555555555555555 // 01010101 ...
const m1 = 0x3333333333333333 // 00110011 ...
const m2 = 0x0f0f0f0f0f0f0f0f // 00001111 ...

// 计算出x中二进制位1的个数,该函数上篇文章有详细解释,看不懂可以再回去看下
func OnesCount64(x uint64) int {
	const m = 1<<64 - 1
	x = x>>1&(m0&m) + x&(m0&m)
	x = x>>2&(m1&m) + x&(m1&m)
	x = (x>>4 + x) & (m2 & m)
	x += x >> 8
	x += x >> 16
	x += x >> 32
	return int(x) & (1<<7 - 1)
}

这种实现方式在bitmap元素过多,切片长度过长的情况下,计算十分耗时。那么如何优化这段代码呢?

2. 优化

现代CPU一般都支持SIMD指令,通过SIMD指令可以并行执行多个计算,以加法运算为例,如果我们要计算{A0,A1,A2,A3}四个数与{B0,B1,B2,B3}的和,不使用SIMD指令的话,需要挨个计算A0+B0A1+B1A2+B2A3+B3的和。使用SIMD指令的话,可以将{A0,A1,A2,A3}{A0,A1,A2,A3}四个数加载到xmm(128bit)/ymm(256bit)/zmm(512bit)寄存器中,然后使用一条指令就可以同时计算对应的和。这样理论上可以获得N倍的性能提升。

image.png

我们可以采用SIMD指令将OnesCount64函数并行化,并行计算4个uint64数字的结果,代码实现如下:

在popcnt.go文件中定义SimdPopcntQuad函数

package popcnt

func SimdPopcntQuad(nums [4]uint64) [4]uint64

在popcnt.s文件中我们使用汇编实现SimdPopcntQuad函数

#include "textflag.h"

TEXT ·SimdPopcntQuad(SB),NOSPLIT,$0-64
    VMOVDQU nums+0(FP), Y0 // Y0 = x,将四个uint64数字加载到Y0寄存器
    MOVQ $0x5555555555555555, AX
    MOVQ AX, X9
    VPBROADCASTQ X9, Y5 // Y5 = m0 // 上面三行代码将4个m0加载到Y5寄存器
    MOVQ $0x3333333333333333, AX
    MOVQ AX, X9
    VPBROADCASTQ X9, Y6 // Y6 = m1 // 上面三行代码将4个m1加载到Y6寄存器
    MOVQ $0x0f0f0f0f0f0f0f0f, AX
    MOVQ AX, X9
    VPBROADCASTQ X9, Y7 // Y7 = m2 // 上面三行代码将4个m2加载到Y7寄存器
    MOVQ $0x7f, AX
    MOVQ AX, X9
    VPBROADCASTQ X9, Y8 // Y8 = m;上面三行代码将4个m3加载到Y8寄存器
    VPSRLQ $1, Y0, Y1 // Y1 = x>>1;Y0寄存器上四个uint64数字并行右移1位
    VPAND Y1, Y5, Y1 // Y1 = x>>1&m0;Y1寄存器上四个uint64数字并行与Y5寄存器上的四个m0并行与,结果存到Y1寄存器
    VPAND Y0, Y5, Y2 // Y2 = x&m0
    VPADDQ Y1, Y2, Y0 // x = x>>1&m0 + x&m0
    VPSRLQ $2, Y0, Y1 // Y1 = x>>2
    VPAND Y1, Y6, Y1 // Y1 = x>>2&m1
    VPAND Y0, Y6, Y2 // Y2 = x&m1
    VPADDQ Y1, Y2, Y0 // x = x>>2&m1 + x&m1
    VPSRLQ $4, Y0, Y1 // Y1 = x>>4
    VPAND Y1, Y7, Y1 // Y1 = x>>4&m2
    VPAND Y0, Y7, Y2 // Y2 = x&m2
    VPADDQ Y1, Y2, Y0 // x = x>>2&m2 + x&m2
    VPSRLQ $8, Y0, Y1 // Y1 = x >> 8
    VPADDQ Y1, Y0, Y0 // x += x >> 8
    VPSRLQ $16, Y0, Y1 // Y1 = x >> 16
    VPADDQ Y1, Y0, Y0 // x += x >> 16
    VPSRLQ $32, Y0, Y1 // Y1 = x >> 32
    VPADDQ Y1, Y0, Y0 // x += x >> 32
    VPAND Y0, Y8, Y0 // x & (1<<7-1)
    VMOVDQU Y0, ret+32(FP) // 将结果加载到内存中返回值的位置
    RET

Benchmark

理论上讲如此优化之后我们应该可以获得四倍的性能提升,所以我们写个基准测试验证下:

// 优化之后的并行计算测试
func BenchmarkSimdPopcntQuad(b *testing.B) {
        // 使用随机数防止编译阶段被编译器预先计算出来
	rand.Seed(time.Now().UnixNano())
	nums := [4]uint64{rand.Uint64(), rand.Uint64(), rand.Uint64(), rand.Uint64()}
	for i := 0; i < b.N; i++ {
		SimdPopcntQuad(nums)
	}
}

// 优化之前的顺序计算测试
func BenchmarkSerial(b *testing.B) {
        // 使用随机数防止编译阶段被编译器预先计算出来
	rand.Seed(time.Now().UnixNano())
	nums := [4]uint64{rand.Uint64(), rand.Uint64(), rand.Uint64(), rand.Uint64()}
	for i := 0; i < b.N; i++ {
		serialPopcntQuad(nums)
	}
}

func serialPopcntQuad(nums [4]uint64) [4]uint64 {
	return [4]uint64{uint64(bits.OnesCount64(nums[0])), uint64(bits.OnesCount64(nums[1])), uint64(bits.OnesCount64(nums[2])), uint64(bits.OnesCount64(nums[3]))}
}

运行后结果如下

# go test -bench=. -v
=== RUN   TestSimdPopcntQuad
--- PASS: TestSimdPopcntQuad (0.00s)
goos: linux
goarch: amd64
pkg: github.com/Orlion/popcnt
cpu: Intel Core Processor (Broadwell, no TSX)
BenchmarkSimdPopcntQuad
BenchmarkSimdPopcntQuad-8        3693530               330.8 ns/op
BenchmarkSerial
BenchmarkSerial-8               539924296                2.232 ns/op
PASS
ok      github.com/Orlion/popcnt        2.993s

可以看到优化后的并行计算比原始的顺序计算慢了150倍😭,失败~

image.png

3. 分析

虽然优化失败了,但是我们还是要分析复盘下其中的原因,从中汲取一些经验,下面我们从两方面来分析下。

3.1 未优化函数为什么快?

首先我们可以看到未优化的函数serialPopcntQuad计算四个数字竟然只花了2ns,根据Numbers Everyone Should Know一文,访存的时间大概是100ns,这就有点离谱了,计算竟然不从内存加载我们的参数?

下面我们写段main函数,使用随机数来调用下serialPopcntQuad函数,然后反汇编看下汇编代码分析下。

func main() {
	rand.Seed(time.Now().UnixNano())
	nums := [4]uint64{rand.Uint64(), rand.Uint64(), rand.Uint64(), rand.Uint64()}
	results := serialPopcntQuad(nums)
	fmt.Println(results)
}

func serialPopcntQuad(nums [4]uint64) [4]uint64 {
	return [4]uint64{uint64(bits.OnesCount64(nums[0])), uint64(bits.OnesCount64(nums[1])), uint64(bits.OnesCount64(nums[2])), uint64(bits.OnesCount64(nums[3]))}
}

编译后反汇编:

image.png

从汇编代码中可以看到在调用bits.OnesCount64之前会判断cpu是否支持popcnt指令,如果支持则使用popcnt指令来计算而不是调用bits.OnesCount64来计算,恰好我机器支持popcnt指令,省略了bits.OnesCount64中的一堆计算,因此计算速度非常快。

3.2 优化后为什么慢?

正如3.1中所提到的,相较于cpu计算,访存的代价是非常高的,大概是100ns,而我们汇编代码中为了使用SIMD指令实现统计算法有大量的访存操作。

受限于本人对汇编掌握程度,上面的汇编代码质量应该是很差的,并不能证明SIMD性能差,可能有性能更高的实现,请各位大佬指点。

而且当前Go汇编在不指定编译参数的情况下只能采用旧函数调用约定,必须采用内存传参,所以导致最终基准测试的结果很差。

4. 收获

这一通瞎折腾虽然最终结果失败,但还是有很多收获的。首先真实的体会到了访存有多慢,所以日后在进行性能优化时就会注意这一点,尽量使代码能命中CPU缓存。

再一个就是之前并没有使用过SIMD指令,也没有接触过这种级别的优化,这次算是入门了。

后端选手,水平有限,各位计算机科学家见笑了。

image.png

5. 参考资料

  1. 玩转SIMD指令编程

...

阅读全文 »

GNU 汇编器的语法

发布于
分类 汇编
标签 AT&T汇编
标签 编译器

学习汇编语法的目的

为什么要学习汇编语法呢?原因是我最近在做一个面向对象语言的编译器(地址:https://github.com/Orlion/Mizar),目前已经完成了parser部分,即已经生成了AST,下一步要做的就是语义分析了,而语义分析之后要做的就是生成AT&T汇编代码了,所以有必要提前了解下汇编语法看在语义分析的实现阶段能否有所指导。

先看一段代码

首先我们有这样一段c语言代码:

#include <stdio.h>

char msg[14] = "Hello,world!\n";
 
int main(void)
{
    puts(msg);
    return 0;
}

运行 gcc -S -Os hello.c

    .file   "hello.c"
    .section    .text.startup,"ax",@progbits
    .globl  main
    .type   main, @function
main:
.LFB0:
    .cfi_startproc
    pushq   %rax
    .cfi_def_cfa_offset 16
    movl    $msg, %edi
    call    puts
    xorl    %eax, %eax
    popq    %rdx
    .cfi_def_cfa_offset 8
    ret
    .cfi_endproc
.LFE0:
    .size   main, .-main
    .globl  msg
    .data
    .align 8
    .type   msg, @object
    .size   msg, 14
msg:
    .string "Hello,world!\n"
    .ident  "GCC: (GNU) 4.8.5 20150623 (Red Hat 4.8.5-39)"
    .section    .note.GNU-stack,"",@progbits

接下来解释下AT&T汇编的语法

指令

指令是直接由CPU负责处理的命令,不以.开头的行首缩进的行都是指令行。

    movl    $msg, %edi
    call    puts
    xorl    %eax, %eax

指令由操作符和作为参数的操作数组成,以 movl $msg, %edi 为例,movl 为操作符, $msg%edi 为操作数,操作数以逗号来间隔。

汇编伪操作

. 开头末尾没有:的行都是汇编伪操作。例如,.file "hello.c", .globl main。汇编伪操作是由汇编器而非CPU处理的指令。一般用于在目标文件中记录元数据(meta data)或者设定指定的属性等。例如 .string 是用来定义字符串常量的汇编伪操作。

标签(labal)

以冒号: 结尾的行都是标签行,例如:.LFB0:,main:。 标签具有为汇编伪操作生成的数据或者指令命名(标上符号)的功能,这样就可以在其他地方调用通过标签定义的符号。 标签可以以.开头

注释

支持两种注释:

# xxx

/* xxx
xxx */

助记符后缀

刚才提到的movlsubl为助记符,更准确的说movsub为助记符,末尾的l是后缀,llong的缩写,表示操作对象的数据大小。类似这样的后缀还有b,w,l

|后缀|操作对象的大小|缩写| |-|-|-| |b|8位|byte| |w|16位|word| |l|32位|long|

操作数

操作数有四种:

  1. 立即数
  2. 寄存器
  3. 直接内存引用
  4. 间接内存引用

1. 立即数

立即数就是C语言中的字面量,机器语言中立即数以整数的形式出现,能高速访问。像$27这样,立即数用$来标识,如果漏掉了$就成了直接内存引用了。立即数有8位,16位,32位。

2. 寄存器

GUN汇编器规定寄存器以%开头,例如eax寄存器写作%eax

3. 直接内存引用

直接访问固定内存地址的方式。GNC汇编器会将任何整数解释为内存地址并访问。比起使用数字,更常用符号(symbol)直接访问内存。例如.LFE0就是访问符号.LFE0所指向的地址。符号在汇编和链接的过程中会被置换为实际内存地址。

4. 间接内存引用

是将寄存器的值作为内存地址访问的方式。间接内存引用中最通用的就是下方的形式:

disp(base, index, scale)

其中任何一者都可以省略。

上述指令访问disp + (base + index * scale)的地址。 下面详细讲解,首先最简单的间接引用的形式如下:

(%eax)

即只指定基地址(base)的形式。上述表达式将eax寄存器中的值作为内存地址访问。 接着带有disp的形式如下。disp是displacement(偏移)的简称。

4(%eax)

上述就是访问 4 + (eax寄存器中值) 这个内存地址。在C语言中用来访问如下结构体中成员y的情况:

struct point {
    int x; // x占4个字节,即4个内存地址
    int y;
}

最后使用index和scale的情况如下所示:

(%ebx, %eax, 4)

上面访问的就是(ebx寄存器中的值 + eax寄存器中的值 * 4)内存地址。在C语言中用来访问数组,例如访问元素大小为4字节(例如int)的数组中元素的第%ebx个元素时就可以用这种方式。当并非所有的数组访问都可以只靠间接内存引用来表示,因为scale只能是1、2、4、8之一。

2020-09-22更

突然意识到如果要将一个复杂工程的AST编译为汇编代码必须具备能够用汇编实现这个复杂工程的能力才行,这太难了... 所以暂时放弃吧,先编译到了LLVM IR再说,嗯。

2020-10-07更

不如趁此机会学习下汇编以及后续的链接装载,有助于建立宏观的了解。所以还是继续学习吧。 😸

未完待续...

...

阅读全文 »