demo: https://github.com/aircloud/rust-wasm-demo

其他资料:入门 Rust 开发 WebAssembly

一般来说,使用 WebAssembly 能够在一定程度上提高性能,不过有的时候我们也许会发现,使用 WebAssembly 之后,有的时候我们不仅发现性能没有提升,反而下降了许多甚至数倍,实际上这是因为,使用 WebAssembly 需要非常谨慎,有很多细节都会大幅度影响性能,比如:

  • 我们编译采用的是 debug 还是 release 方式。
  • 最后编译的结果是否采用了不同级别的优化,如果使用了 opt-level = 's' 那么通常速度也会下降很多。
  • 是否在 JS 和 rust 之间存在大量的数据拷贝,因为很多代码是工具链生成的,也许有的时候我们会忽视这一点。

本文针对以上等一些问题特别是第三点,给出一个 wasm 优化的参考方案,并给出示例代码。

编译优化

我们在优化数据拷贝之前,对于编译我们可以做一些前置的简单的工作。

  • 检查 Cargo.toml 脚本中 [profile.release] 中的 opt-level 选项,确认我们所使用的值:
1
2
3
4
5
6
7
8
9
10
11
This flag controls the optimization level.

0: no optimizations, also turns on cfg(debug_assertions) (the default).
1: basic optimizations.
2: some optimizations.
3: all optimizations.
s: optimize for binary size.
z: optimize for binary size, but also turn off loop vectorization.
Note: The -O flag is an alias for -C opt-level=2.

The default is 0.

如果我们使用了 ‘s’ 或者 ‘z’,那么通常会牺牲一部分性能(对于 demo 而言,使用 ‘z’, wasm 的性能也只有 js 的 20%),因为其主要是对体积进行一定的优化,所以如果优化前的体积我们可以接受的话,通常不需要这样的优化。

在以上的前提下,我们使用 --release 的方式编译,通常就可以了。

减少拷贝

在这之前,我们需要有一个认知:

通过 rust 工具链编译的 wasm 代码,所有参数传入都是需要拷贝一次的,包括我们传入 ArrayBuffer 等 Buffer 类型的参数。这是由于 wasm 只能访问自己的线性内存,而这个拷贝,通常是我们在处理大规模计算的一个坎,有的时候虽然 wasm 计算快一点,但是拷贝的消耗还是比较大的,加之 js 有若干 v8 优化的加持,可能和 wasm 也相差不多。

所以我们要把计算移植到 wasm 中的话,首先要解决的就是大规模数据拷贝的问题。

这里的一般思路为:

  1. wasm 分配内存:调用 wasm 的方法,在 wasm 内存中分配空间,返回指针位置
  2. js 写入数据:js 端在 wasm 的 memory arraybuffer 上,按指针位置和数据量建立 view,把数据写入
  3. wasm 计算:调用 wasm 方法完成计算, 返回计算好的批量结果的指针位置和大小
  4. js 读取数据:js 端在 wasm 的 memory arraybuffer上,按指针位置和数据量建立 view,把数据读出

接下来,我们通过一个 demo 来完成以上几点,demo 的主要功能为:

  • 初始化一个 ImageData,内容随机。
  • 分别使用 js 和 WebAssembly 进行高斯模糊计算,并计算二者的时间,进行对比。

这里的 demo 只是辅助进行验证改方案的可行性并且给出一个示例,并不作为一个标准的 benchmark 去对比 js 和 WebAssembly 的性能,同时,也并没有 UI 展示,计算结果输出在控制台中。

最终笔者运行的结果为,js 比 WebAssembly 慢 30% 左右。

1. wasm 分配内存

这部分的通用做法,即我们在 wasm 的 rust 中分配一个数组(Vec),然后把其指针传递给 js:

1
2
3
4
5
6
7
8
9
10
// rust:
#[wasm_bindgen]
pub fn new_buffer(key: String, len: usize) -> *const u8 {
// GlobalBufferStorage 是一个 lazy_static
let mut global_buffer_storage = GlobalBufferStorage.lock().unwrap();
let mut buffer = vec![255; len];
let ptr = buffer.as_ptr();
global_buffer_storage.buffer_map.insert(key, buffer);
ptr
}

为了后续方便寻找到这段数据,我们可以使用一个 key 将这个 Vec 联系起来,并且在 Rust 中放入全局(可以使用 lazy_static!,因为这种类型的数据没有办法直接定义在全局),之后通过 key 来查找数据。

在 js 中,我们就可以建立各种 TypedArray 对其进行操作:

1
2
const ptr = this.wasm!.new_buffer(key, len);
const u8Arr = new Uint8ClampedArray(this.wasm!.get_wasm_buffer(), ptr, len);

这个时候,我们在 js 或 rust 任何一侧改了这个数据之后,都可以在另外一侧访问到。

实际上,在 js 侧的比如 ImageData 等一些对象中,也支持我们传递一个 TypedArray 进行初始化,这让我们在比如 canvas 等应用场景下,使用 wasm 分配的内存更为方便。

1
const imageData = new ImageData(u8Arr, width, height);

2. js 写入数据

如果我们需要在 js 侧写入数据,实际上这个时候我们得到的 TypedArray 已经和直接使用 js new 的 TypedArray 在使用上没有差别,可以正常按照数组的方式进行数据写入。

不过,这里需要注意的是,js 写入通过 wasm 分配内存建立的 TypedArray,有些场景下在一定程度上速度要慢于直接使用 js new 的 TypedArray(不过在笔者的测试数据中,wasm 分配的方式反而是更快的),所以如果我们是一个高频的数据写入的场景,比如帧数据等,这个时候最好进行一次对比测试。

3. wasm 计算

当我们真正需要进行计算的时候,我们可以调用 wasm 的计算函数,并且传入上文中定义的 key,这样 wasm 的 rust 函数可以直接找到这段数据,这里我们的 demo 为一段计算卷积的函数:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
#[wasm_bindgen]
pub fn convolution(key: String, width: usize, height: usize, kernel: Vec<i32>) {
let mut global_buffer_storage = GlobalBufferStorage.lock().unwrap();
let kernel_length = kernel.iter().sum::<i32>() as i32;
if let Some(buffer) = global_buffer_storage.buffer_map.get_mut(&key) {
for i in 1..width-1 {
for j in 1..height-1 {
let mut newR: i32 = 0;
let mut newG: i32 = 0;
let mut newB: i32 = 0;
for x in 0..3 { // 取前后左右共9个格子
for y in 0..3 {
newR += buffer[width * (j + y - 1) * 4 + (i + x - 1) * 4 + 0] as i32 * kernel[y * 3 + x] / kernel_length;
newG += buffer[width * (j + y - 1) * 4 + (i + x - 1) * 4 + 1] as i32 * kernel[y * 3 + x] / kernel_length;
newB += buffer[width * (j + y - 1) * 4 + (i + x - 1) * 4 + 2] as i32 * kernel[y * 3 + x] / kernel_length;
}
}
buffer[width * j * 4 + i * 4 + 0] = newR as u8;
buffer[width * j * 4 + i * 4 + 1] = newG as u8;
buffer[width * j * 4 + i * 4 + 2] = newB as u8;
}
}
} else {
return ();
}
}

因为这段函数对应操作的内存数据实际上已经在 wasm 和 js 之间共享了,所以也是不需要返回值的,等计算完成后 js 直接去读之前建立的 TypedArray,甚至直接使用通过 TypedArray 创建的 ImageData,进行绘制上屏等后续操作。

4. js 读取数据

在 demo 中,我们可以直接通过 CanvasRenderingContext2D.putImageData() 传入之前获取的 imageData,绘制上屏。

其他方案

实际上,我们如果目的是加速 js 计算,不仅仅有 WebAssembly 这一个方案可以选择,如果我们的环境中拥有可以访问 Node 的能力或者可以访问原生模块的能力(比如,我们的应用运行在 electron 环境,或者是一些移动客户端),也可以采用比如 addon 的方式来运行我们的计算部分,相比于 wasm,这部分的优缺点在于:

优点:

  • 通常可以更好的控制优化,甚至做到汇编级别的优化定制,性能提升空间更高(同样也可能会面临数据拷贝的问题,也需要一定方式减少拷贝)。
  • 在重 addon 的环境下(例如,其他大量功能也依赖 addon),可以更好的处理函数调用关系、依赖库使用等,一定程度上减少体积和增加开发的便捷性,而 wasm 会被编译成一个独立的二进制文件,处于沙盒环境中,无法直接调用其他的动态库。

缺点:

  • 无法做到像 wasm 一样跨平台,并且可以同时运行在网页、桌面环境、移动端等任何 Webview 存在的环境中。

不过总之,如果使用得当,二者的性能都是可以优于原生的 js,都可以作为优化方案考虑。