Go ARM64 Map优化小记

Go ARM64 Map优化小记

背景

Go内置了map类型，而其中重要的哈希算法是一个cityhash的变种。

同时，为了避免 哈希冲突攻击（collision attack） 和加速哈希计算速度，Keith Randall于Go1.0中就添加了x86_64支持的有硬件加速的 AESHASH算法 。 我搜遍了互联网，惊讶地发现，这个算法仅仅在Go里面有实现，这思路真是绝了。

这就被我这个四处搜索ARM64 Go runtime待优化点的人找到了：ARM64也支持AESHASH的硬件加速指令，但是Go并没有用上。

我嘴角又微微地一笑，满心欢喜准备加代码。可我并不知道，这看似平静的海面下不知道藏着什么妖怪……

开工

打蛇打七寸，看代码实现自然要先看头。 初始化代码在 runtime/alg.go

if (GOARCH == "386" || GOARCH == "amd64") &&
	GOOS != "nacl" &&
	support_aes && // AESENC
	support_ssse3 && // PSHUFB
	support_sse41 { // PINSR{D,Q}
	useAeshash = true
	algarray[alg_MEM32].hash = aeshash32
	algarray[alg_MEM64].hash = aeshash64
	algarray[alg_STRING].hash = aeshashstr
	// Initialize with random data so hash collisions will be hard to engineer.
	getRandomData(aeskeysched[:])
	return
}

可以看到，通过替换algarrary中的hash函数成aeshash，就完成了这个加速替换， 充分考虑了未来其他平台的跟进，赞叹同时感到这个简直就是盛情邀请我来完成接下来的工作。

先看看最简单的 aeshash64的具体实现

//func aeshash64(p unsafe.Pointer, h uintptr) uintptr
TEXT runtime·aeshash64(SB),NOSPLIT,$0-24
	MOVQ	p+0(FP), AX	// ptr to data
	MOVQ	h+8(FP), X0	// seed
	PINSRQ	$1, (AX), X0	// data
	AESENC	runtime·aeskeysched+0(SB), X0
	AESENC	runtime·aeskeysched+16(SB), X0
	AESENC	runtime·aeskeysched+32(SB), X0
	MOVQ	X0, ret+16(FP)
	RET

注释也很清晰，AX载入了数据的指针，然后，把map的种子和数据载入X0中等待计算。 这里要说明一下，每个hashmap在初始化的时候都会随机分配一个种子，防止黑客找到系统的种子而发起哈希冲突攻击。 在接下来的几步中，使用程序初始化时生成的随机种子 runtime·aeskeysched 对数据进行加密，最后把结果返回。

更复杂的aeshash也只是加载各种长度进行计算而已。

到这，我只能感叹，这太简单了，便花了两个周末就写完了大体的代码，还碰到了以下问题:

1. 平台差异
2. Smhasher
3. 冲突（Collision）
4. Go编译器bug

平台差异

首先发现的问题是，ARM64并没有X86的AESENC，而是分成两个指令，AESE和AESMC。

先看了看 AES的介绍 标准AES加密分成了4步：

1. AddRoundKey
2. SubBytes
3. ShiftRows
4. MixColumns

X86的AESENC指令等价于：

1. SubBytes
2. ShiftRows
3. MixColumns
4. data XOR Key

但是……ARM64的AESE指令等价于：

1. AddRoundKey
2. SubBytes
3. ShiftRows

所以，要是单纯模仿X86的 AESENC X0,X0 写法时……数据就会被清空掉……（摔。 无奈，只好另寻出路，用系统随机种子加密数据，代码思路如下：

// 把系统种子载入 V1
// 再将种子和数据载入 V2
AESE	V1.B16, V2.B16

SMhasher&Collision

解决了上个问题，开始进行测试。 Go使用的是Smhasher，一个Hash函数是否达标需要通过以下测试：

1. Sanity，必须处理整个key
2. AppendedZeros，填零，长度不同
3. SmallKeys，所有小key（< 3字节）组合
4. Cyclic，循环，例如：11211->11121
5. Sparse，稀疏，例如：0b00001和 0b00100
6. Permutation，组合，每个block排列组合顺序不同
7. Avalanche，翻转每个位
8. Windowed，例如产生的hash值是32位的，那么就20位相同，结果也需要不同
9. Seed，种子变化会影响结果

每一次Smhasher报错，我都开始看代码哪里出了问题，一般都是放错寄存器这些低级错误…… Go还会对map中bucket分配情况进行测试，如果某个bucket放了过多的数据，一样也会出错。 不过，这部分还是比较顺利的。

Go编译器bug

搞定了这个，我又发现另一个坑。 为了减少指令，我希望直接把寄存器中的数据直接载入到 ARM64 Vector Lane