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 B0、A1 B1、A2 B2、A3 B3的和。使用SIMD指令的话，可以将{A0,A1,A2,A3}和{A0,A1,A2,A3}四个数加载到xmm(128bit)/ymm(256bit)/zmm(512bit)寄存器中，然后使用一条指令就可以同时计算对应的和。这样理论上可以获得N倍的性能提升。
我们可以采用SIMD指令将OnesCount64函数并行化，并行计算4个uint64数字的结果，代码实现如下：
代码语言：javascript复制在popcnt.go文件中定义SimdPopcntQuad函数
package popcnt
func SimdPopcntQuad(nums [4]uint64) [4]uint64代码语言：javascript复制在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) // 将结果加载到内存中返回值的位置
RETBenchmark
理论上讲如此优化之后我们应该可以获得四倍的性能提升，所以我们写个基准测试验证下：
代码语言：javascript复制// 优化之后的并行计算测试
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]))}
}运行后结果如下
代码语言：javascript复制# 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倍
