高性能Rust wasm最佳实践

在我司项目里写rust一小段时间了。记得在前年年底(2018)小组聚餐的时候我第一次给前manager提rust这门语言的时候,完全是想不到我会有今天这样的机会。目前这个语言用于wasm的加速上,也是主要由我推动的。预计到明年上半年之前这么语言会在性能优化的主题上起到至关重要的作用,几乎所有的cpu端的重型计算模块,都会被迁移进rust实现的wasm内,并由wasm-bindgen提供接口供中层建筑调用,其性能提升的关键指标原型都已经得到了认证。经过这一段时间,以及我长期业余时间内的rust开发经验。这里简单谈一下性能方面的最佳实践。

inline

强烈建议对所有小函数都标记#[inline],如果非常简短的函数调用可以标记#[inline(always)]。 比方说数学库,很多小函数,又比如说调用上的封装,都最好加上。虽然即便不加,在最高优化等级下,编译器也会自己做很多激进的inline,但是我遇到的几个case,只要我检查了调用上是否应该inline的地方,然后手工添加这些标记,大概都直接收获了10-20%的性能提升。

我是怎么发现关键的小函数没有inline呢,其实还是主要通过开发者工具profile工具的火焰图。正常情况下是不可能sample到小函数调用的,如果有,且整个逻辑不是非常复杂,那么就应该尝试inline。这里逻辑不是非常复杂指的是inline的结果如果造成函数体过大,因为代码cache命中降低的原因,反而会对性能产生影响。如果你的逻辑其实还是很复杂的,编译器完成inline以后,不一定能通过后续的优化,无论是死代码消除还是窥孔来减少代码量,还是会自行取消inline。

要特别注意高频率调用函数的inline情况,一般而言,如果有inline的可能,我们一定要确保inline。某些情况下这可能造成30-40%的性能差异。记得我的同事写了一段和排序相关的逻辑,结果比较原版的ts实现反而变慢了。我分析了wasm的变异结果,看到其比较函数没有inline,仔细分析比较函数的编译结果,发现逻辑过于复杂,其中包含了多个pattern match,所以为了触发inline,我简单的将比较函数内的pattern match移出,取而代之的编写多个比较函数来选择调用。这一做法的确奏效了,内联之后性能有了巨大提升。

inline是非常非常重要的事情,相当多的优化和分析都是函数内的,inline是触发后续优化的前提。除了上面函数体过大,还有一些原因使得inline无法发生也需要注意:

dynamic dispatch

我上周重构了一小段代码,但是写完发现性能稳定的衰退的20%。仔细看我的改动,发现造成性能衰退的原因是我改变了trait的抽象结构,导致其中一处高频的方法调用走了trait object的dynamic dispatch,而我改动之前是static dispatch。这真是教科书级别的例子(dynamic的性能慢于static dispatch)。

当然不是说dynamic dispatch一无是处,这个东西是非常好非常有用的,只不过我们要衡量好tradeoff,到底是不是我们愿意牺牲编译时间,代码体积翻倍膨胀,代码范型写的到处都是来换取20%的性能。如果是我这样的case,这个逻辑是如此的高频,我希望我为这个东西的每一种具体的实例化的类型都生成高度优化的代码,那么就要确保这个性能sensitive的地方要走static dispatch,编译器一定不能让我失望。而那些低频率的接口,dynamic其实是很有必要的,因为它完成了类型概括的功能,能让代码写的合理和简单。

我发现我对trait的摆弄,那些build nice abstractions的事情,其实无非是把原本存在于运行时的值中的信息编码到类型系统里,来大幅提升系统的类型上的安全保证和维护性。而很多时候你在代码中还是要抹去一些类型上的差异的,这就是dynamic dispatch存在的意义,它就是cpp的虚函数,它就是通过虚表来完成多个不同的类型间的统一。

dynamic的性能慢于static dispatch的原因很简单,就是不能内联了。因为代码调用走的是虚表,除非编译器有更全局的信息和推导,不到运行时是不能确定具体的函数调用。

heap allocation

每一个对性能上点心的人都不会忽视堆内存分配的成本。如果你的代码出发了堆内存的分配和释放,那么inline就不可能的,而分配这件事情成本也很高,更别说触发wasm memeory grow,触发系统调用。如果可以,禁止不必要的堆内存分配,比如随随便便Vec::new(),等随意变出个容器,用两下扔了?

避免堆分配的另一个重要话题是,如果可以,严格预分配正确大小的内存,检查所有collect的迭代器都实现size_hint。不然你在高频操作时push/insert容器会触发容器扩容,如果是vec则会memcopy,如果是hashmap则会是代价更加高昂的rehash。在我的实践中这大概能减少了1-2ms的时间。

尽可能避免高频使用智能指针操作

这里不是指不用,而是Rc, Arc, RefCell的某些调用开销比你想象的要高。如果追求性能,则不应在高频部分做引用计数增减,borrow等操作。这些操作事实上很多时候把内存读附加了一次内存写。inline废了,cache也受到影响。

如无必要,不要使用hashmap

rust标准库的hashmap有simd的加速,wasm simd还没做完。

即便非要用记得换更高性能但是没有密码学安全的hasher实现

避免不必要的大对象clone/move

wasm的memcopy没有simd支持,相关标准可能还在设计,要比native慢很多。

不过这个平时都很注意的,没啥可说的

严格禁止高频wasm IO

IO性能很差,千万别乱来。实践上还是走批处理的做法,获取指针位置/大小,从js侧建立数据的view,按布局直接批量直接读写,这个是最佳实践。其他要是小接口,或者没啥高频调用的,走wasm-bindgen也挺好。

Prefer struct of array

了解编译原理,了解编译结果,了解底层

建议对性能有追求的同学,学习编译原理后端优化方面的内容。这样你能很自信知道你代码里的漂亮写法是零开销抽象,你能知道哪些事情哪些做法不是free的。阅读编译结果,调整抽象,调整hint,和编译器一起完成工作。同样的,你对其他底层知识知道的越多,优化的空间就越大。

数据驱动

这里推荐一下rust一个benchmark的框架 https://github.com/bheisler/criterion.rs,看起来比较高端也好用。