rust
WebAssembly
OriginalWACULDay 5

Rust + WebAssembly でpngデコードを行うnode_moduleを作ってみる

この記事は「WACUL Advent Calendar 2017」の5日目です。
WACULでフロントエンドエンジニアをしている@bokuwebと申します。

表題の通りですがRustWebAssemblyを使用してpngデコードを行うnode_moduleを作ってみました。

モチベーション

  • あるmoduleで使用しているpngjsによるデコード処理が時間を食っておりwasmで高速化できないかの調査
  • wasm32-unknown-unknownを使ってnode_moduleを作るとこまで体験しときたい

リポジトリ

https://github.com/bokuweb/rust-wasm-png-decoder-example

RustとWebAssembly

これまではwasm32-unknown-emscriptenを指定して、emscriptenを介してwasmを出力する必要があったんですが、先日のリリースにおいて1.24.0-nightlywasm32-unknown-unknownというtargetが指定可能となりました。これによりemscriptenなしにRustからwasmを吐けるようになりました。めでたい :tada:

まずは簡単なモジュールを作ってみる

まずは引数に1を加算して返すだけの単純なモジュールを作ってみます。

Rust側

セットアップ。

rustup update
rustup target add wasm32-unknown-unknown --toolchain nightly
cargo init --bin add

main.rsを編集します。

add/src/main.rs
fn main() {}

#[no_mangle]
pub fn add_one(x: i32) -> i32 {
    x + 1
}

ビルド。

rustc +nightly --target wasm32-unknown-unknown src/main.rs -o main.wasm

main.wasmが出力されていたら成功です。
以下は蛇足で今回の調査で知ったんですが、wasm-gcというRust製のツールがあるようで、こいつを使用することで未使用関数の削除などが行われサイズダウンができるようです。

cargo install --git https://github.com/alexcrichton/wasm-gc
wasm-gc main.wasm main.min.wasm

これにより本ケースでは15kiBが108Bまで落ちました。(この時点では、めっちゃおちるじゃないですか。やだー。ってなってたんですが、現実は甘くなかったです。後述。)

JS側

add/index.js
const fs = require('fs');

const wasm = fs.readFileSync('./main.wasm');
const mod = new WebAssembly.Instance(new WebAssembly.Module(wasm));

console.log(mod.exports.add_one(2)); // -> 3
node index.js

良さそうです。

pngのデコードを行う

では本題のpngデコードを行ってみます。

Rust側

emscriptenがないためalloc,freeは自分で用意する必要があります。
JSとのやり取りも含め以下のリポジトリを参考にさせていただきました。

https://github.com/KOBA789/rust-wasm

今回初めて知ったんですが、mem::forget()を使用しているのがポイントらしく、これによりスコープを抜けた後も確保したヒープが解放されないようです。freeのほうは指定領域のヒープを再確保してやりDrop時に解放してもらう感じでしょうか。

pngのデコードには以下のcrateを使用していて、JS側からもらったバッファをこいつに食わせることでデコードを行っています。

https://github.com/PistonDevelopers/image-png

後はwidthやheightといった情報とデコード結果へのポインタをシリアライズしてJS側に返しています。この辺シリアライズ・デシリアライズのコストも勿体無いのでなんとかしたいのですが、いい方法がわかっていません。

src/main.rs
extern crate png;

use std::os::raw::{c_char, c_void};
use std::mem;
use std::ffi::CString;

#[macro_use]
extern crate serde_derive;
extern crate serde;
extern crate serde_json;

#[no_mangle]
pub fn alloc(size: usize) -> *mut c_void {
    let mut buf = Vec::with_capacity(size);
    let ptr = buf.as_mut_ptr();
    mem::forget(buf);
    return ptr as *mut c_void;
}

#[no_mangle]
pub fn free(ptr: *mut c_void, size: usize) {
    unsafe {
        let _buf = Vec::from_raw_parts(ptr, 0, size);
    }
}

fn main() {}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct DecodingResult {
    ptr: u32,
    len: u32,
    width: u32,
    height: u32,
}

#[no_mangle]
pub fn decode(ptr: *mut u8, len: usize) -> *mut c_char {
    let buf: &[u8] = unsafe { std::slice::from_raw_parts_mut(ptr, len) };
    let decoder = png::Decoder::new(buf);
    let (info, mut reader) = match decoder.read_info() {
        Ok(i) => i,
        Err(why) => panic!(why.to_string()),
    };
    let mut img_data = vec![0; info.buffer_size()];
    reader.next_frame(&mut img_data).unwrap();
    let result = DecodingResult {
        ptr: img_data.as_mut_ptr() as u32,
        len: info.buffer_size() as u32,
        width: info.width,
        height: info.height,
    };
    mem::forget(img_data);
    let res = serde_json::to_string(&result).unwrap();
    let c_str = CString::new(res).unwrap();
    c_str.into_raw()
}

JS側

JSONをパースしたりポインタからデータを取り出したりしてるだけですね。注意点はfreeし忘れないことですね。JS書いてる気分でいると抜けます。

index.js
'use strict';

const fs = require('fs');
const wasm = fs.readFileSync('./wasm/main.wasm');
const mod = new WebAssembly.Instance(new WebAssembly.Module(wasm));

module.exports = (buf) => {
    const heap = new Uint8Array(mod.exports.memory.buffer);
    const ptr = mod.exports.alloc(buf.length);
    heap.set(buf, ptr);

    const resultPtr = mod.exports.decode(ptr, buf.length);
    mod.exports.free(ptr, buf.length)
    const resultBuf = new Uint8Array(mod.exports.memory.buffer, resultPtr);
    const getSize = (buf) => {
        let i = 0;
        while (buf[i] !== 0) i++;
        return i;
    }
    const resultSize = getSize(resultBuf);
    const json = String.fromCharCode.apply(null, resultBuf.slice(0, resultSize));
    mod.exports.free(resultPtr, resultSize)
    const result = JSON.parse(json);
    const data = new Uint8Array(result.len);
    const target = new Uint8Array(mod.exports.memory.buffer, result.ptr, result.len);
    data.set(target);
    const ret = {
        width: result.width,
        height: result.height,
        buf: data
    }
    mod.exports.free(result.ptr, result.len);
    return ret;
}

サイズ

今回サンプルの最終的なサイズは以下になりました。なんだかんだ、大きくなりますね。。。かなしい。

項目 サイズ
rust -> wasm後 370KiB
wasm-gc後 354KiB
wasm-gc後にgzip 73KiB

ベンチマーク

それでは速度比較を行ってみます。
対象となるpngはpng crateに含まれていたテスト用pngから10枚選択しています。
結果は以下。

log
## 0.png
rust-png wasm x 10,912 ops/sec ±0.87% (90 runs sampled)
pngjs x 2,882 ops/sec ±1.30% (85 runs sampled)
Fastest is rust-png wasm

## 1.png
rust-png wasm x 6,913 ops/sec ±14.55% (62 runs sampled)
pngjs x 2,839 ops/sec ±1.11% (89 runs sampled)
Fastest is rust-png wasm

## 2.png
rust-png wasm x 12,349 ops/sec ±1.83% (85 runs sampled)
pngjs x 2,839 ops/sec ±1.42% (90 runs sampled)
Fastest is rust-png wasm

## 3.png
rust-png wasm x 8,360 ops/sec ±0.84% (90 runs sampled)
pngjs x 2,801 ops/sec ±1.61% (88 runs sampled)
Fastest is rust-png wasm

## 4.png
rust-png wasm x 8,407 ops/sec ±1.08% (92 runs sampled)
pngjs x 6,935 ops/sec ±2.71% (85 runs sampled)
Fastest is rust-png wasm

## 5.png
rust-png wasm x 5,732 ops/sec ±1.49% (91 runs sampled)
pngjs x 2,414 ops/sec ±1.14% (90 runs sampled)
Fastest is rust-png wasm

## 6.png
rust-png wasm x 6,643 ops/sec ±1.01% (91 runs sampled)
pngjs x 7,138 ops/sec ±2.16% (80 runs sampled)
Fastest is pngjs

## 7.png
rust-png wasm x 2,220 ops/sec ±0.88% (90 runs sampled)
pngjs x 1,957 ops/sec ±1.28% (89 runs sampled)
Fastest is rust-png wasm

## 8.png
rust-png wasm x 11,944 ops/sec ±1.01% (91 runs sampled)
pngjs x 2,929 ops/sec ±1.12% (90 runs sampled)
Fastest is rust-png wasm

## 9.png
rust-png wasm x 5,229 ops/sec ±0.92% (87 runs sampled)
pngjs x 10,119 ops/sec ±1.12% (92 runs sampled)
Fastest is pngjs

スクリーンショット 2017-12-04 20.39.42.png

*キャプチャだと単位等いろいろぬけているので下記リンクを参照してください
https://bokuweb.github.io/rust-wasm-png-decoder-example/index.html

wasmの方が高速でよかったですねー。と終わりたかったんですが、JS版がwasm版より倍近く速いケースが見つかりました。(9.png)。僅差でJSの方が速い(6.pngのような)ケースは発生しうるんじゃないかと思っていましたが、なぜこれほどの差が発生するのかは分かっておらず、要調査。

まとめ

wasm32-unknown-unknownemscriptenなしにnode_moduleを作製してみました。グルーコードを自身で用意する必要はありますが、emscripten分の容量を削減できたり、emscriptenの古いnodeに苦しめられることもなくなるので可能な限りこちらを使用したいですね。

速度に関しては、「wasmの方が速いよー」とは現時点では言うことができなそうで、要調査。

また、今回のサンプルモジュールは以下で使用できます。

npm i https://github.com/bokuweb/rust-wasm-png-decoder-example
const fs = require('fs');
const decode = require("wasm-png-decoder");
console.log(decode(fs.readFileSync('hoge.png')));