Go ARM64 Map优化小记

发表评论

背景

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

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

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

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

Fishman on Danang Beach

开工

打蛇打七寸，看代码实现自然要先看头。
初始化代码在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

但是Go的编译器没办法正确编译不同的lane index。
例如下面两条指令，最终产生的指令码是一样……

VMOV R1, V2.D[0]
VMOV R1, V2.D[1]

只好先报个bug。

cmd/asm: wrong implement vmov/vld on arm64

然后用原生字节码先顶着了，想知道如何做到可以直接拉到Tips

妖怪

终于，所有runtime的Smhasher和Hash测试都通过了，我开始试着运行src/all.bash来构建Go。

这时我拉到了海底的那只妖怪……

SpongeBob

构建日志

$ ./all.bash
Building Go cmd/dist using /usr/lib/go-1.6.
Building Go toolchain1 using /usr/lib/go-1.6.
Building Go bootstrap cmd/go (go_bootstrap) using Go toolchain1.
Building Go toolchain2 using go_bootstrap and Go toolchain1.
Building Go toolchain3 using go_bootstrap and Go toolchain2.
# runtime
duplicate type..hashfunc.struct { "".full "".lfstack; "".empty "".lfstack; "".pad0 [64]uint8; "".wbufSpans struct { "".lock "".mutex; "".free "".mSpanList; "".busy "".mSpanList }; _ uint32; "".bytesMarked uint64; "".markrootNext uint32; "".markrootJobs uint32; "".nproc uint32; "".tstart int64; "".nwait uint32; "".ndone uint32; "".alldone "".note; "".helperDrainBlock bool; "".nFlushCacheRoots int; "".nDataRoots int; "".nBSSRoots int; "".nSpanRoots int; "".nStackRoots int; "".markrootDone bool; "".startSema uint32; "".markDoneSema uint32; "".bgMarkReady "".note; "".bgMarkDone uint32; "".mode "".gcMode; "".userForced bool; "".totaltime int64; "".initialHeapLive uint64; "".assistQueue struct { "".lock "".mutex; "".head "".guintptr; "".tail "".guintptr }; "".sweepWaiters struct { "".lock "".mutex; "".head "".guintptr }; "".cycles uint32; "".stwprocs int32; "".maxprocs int32; "".tSweepTerm int64; "".tMark int64; "".tMarkTerm int64; "".tEnd int64; "".pauseNS int64; "".pauseStart int64; "".heap0 uint64; "".heap1 uint64; "".heap2 uint64; "".heapGoal uint64 }
build failed

咦……编译器构建出错？？？ 测试都通过了啊！?

再运行一次all.bash，发现出错的地方还不一样？？？

用gdb断点在我写的asm_arm64.s:aeshash，跟踪执行流程，在对长字符串（>129)哈希也没有问题。
难道是编译器的bug？

所以我开始跟踪编译器的动作，发现只有符号表（symbol）使用了map的功能。
恶补了编译器的基本原理和Go实现后，我才意识到，这个符号表也仅仅用了aeshashstr（对字符做hash）的功能。
我把Smhasher中对aeshash的全部改装成了aeshashstr，发现还是能奇迹般地通过测试！
手动校验了一次，发现就算把aeshash32和aeshash64都搞得和x86实现一样,包括结果，还是报错！

于是我把这怪异的问题发邮件，发帖子，发群里问遍了所有人。还是无解。

就这样折腾了1个月业余的时间，基本看完了编译器的相关代码，
发现明明是两个不同的符号（symbol）还是会被认为是同一个。
最后还蛋疼地想用钱看看有没有人愿意帮debug一下。
态度惹怒了不少人。我想我是被这bug整得脑子进水了吧……

出坑

直到最近，我才突然意识到没准Smhasher测试并没有覆盖完所有情况？
果然，仔细检查后在aeshash129plus这一段有

SUBS $1, R1, R1
BHS aesloop

这个SUBS是SUB再对比，在R1=1的时候，就退出了。
但Smhasher仅仅对128个字节做了完整的测试，所以测试能通过，但是编译不了。
而这个bug仅仅在256个字节以上才会触发（摔。

改进后是

SUB $1, R1, R1
CBNZ R1, aesloop

提速效果

最终可以提交CL

runtime: implement aeshash for arm64 platform

注意，如果使用PRFM指令，速度能加快30-40MB左右(Hash1024)。可能是下次优化的重点（对齐和Cache)

name old speed new speed delta
Hash5 97.0MB/s ± 0% 97.0MB/s ± 0% -0.03% (p=0.008 n=5+5)
Hash16 329MB/s ± 0% 329MB/s ± 0% ~ (p=0.302 n=4+5)
Hash64 858MB/s ±20% 890MB/s ±11% ~ (p=0.841 n=5+5)
Hash1024 3.50GB/s ±16% 3.57GB/s ± 7% ~ (p=0.690 n=5+5)
Hash65536 4.54GB/s ± 1% 4.57GB/s ± 0% ~ (p=0.310 n=5+5)

如何用GNU汇编语言生成原生ARM64指令字节码？

$cat vld.s
ld1 {v2.b}[14], [x0]
$as -o vld.o -al vld.lis vld.s
AARCH64 GAS vld.s page 1
1 0000 0218404D ld1 {v2.b}[14], [x0]

其中第三列就是生成的字节码，复制到go中就OK了

WORD $0x4D401802

其实还有工具asm2plan9s, 只是目前这个工具没办法编译ARM64

感谢

最后非常感谢

  • Wei Xiao
  • Fangming
  • Keith Randall

对于我细致的帮助

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

此站点使用Akismet来减少垃圾评论。了解我们如何处理您的评论数据。