wasm-bindgenでマンデルブロ集合
はじめに
Rustで生成したWebAssemblyとJavaScriptをつなぎ合わせるインタフェースを生成してくれるwasm-bindgenというツールがあります。
wasm-bindgenを使うことで、若干面倒なWebAssembly周りのメモリ操作やデータの受け渡しを楽に行うことができます。
Rustの構造体やメソッドをJavaScriptのクラスとしてラップすることも可能です。
wasm-bindgenのデモとしてマンデルブロ集合を描いてみましょう。
誰もが描いたことがある(?)以下のような図ですね。
本稿の完成版+ズーム機能を付け加えたデモを以下で公開しています。
https://ionic-wasm-mandelbrot.likr-lab.com/
また、関連のソースコードはGitHubで公開しています。
https://github.com/likr/ionic-wasm-mandelbrot
なお本稿ではマンデルブロ集合自体の解説は省略しています。
方針
画像のバイト列をRust(WebAssembly)側で生成して、JavaScript側でCanvasを使って描画します。
Canvasに画像としてバイト列を読み込む方法として、CanvasRenderingContext2DのputImageDataメソッドがあります。
putImageDataはImageDataオブジェクトを引数として取り、ImageDataオブジェクトの生成には画像のバイト列を表すUint8ClampedArrayが必要です。
ImageDataの生成に必要なUint8ClampedArrayには、RGBAの画素が順番に格納されます。
つまり、Rust側ではそのバイト列を生成してやる必要があるので、Vec<u8>
に画素の情報を書き込んでいきます。
そのVectorの要素数は、画像の幅 * 画像の高さ * 4(RBGA各1バイトずつ)となります。
準備
Rust(バージョン1.30.0以降)とNode.jsはインストールされているものとします。
まずcrateを作成します。
$ cargo new --lib mandelbrot
ディレクトリ構成は以下のようになっています。
mandelbrot
├── Cargo.toml
└── src
└── lib.rs
wasm-bindgenを使用するために、Cargo.tomlに追記します。
[package]
name = "mandelbrot"
version = "0.1.0"
authors = ["Yosuke Onoue <onoue@likr-lab.com>"]
edition = "2018"
[lib]
crate-type = ["cdylib"]
[dependencies]
wasm-bindgen = "0.2"
Rust側
Rust側では、画像のバイト列をScreen構造体で表し、計算をmandelbrot関数で行うことにします。
次にそれぞれの詳細を説明します。
Screen構造体
Screen構造体の役割は、JavaScript側から指定された幅と高さの画像を表すバイト列を確保することです。
以下のようにScreen構造体を定義します。
#[wasm_bindgen]
pub struct Screen {
bytes: Vec<u8>,
#[wasm_bindgen(readonly)]
pub width: usize,
#[wasm_bindgen(readonly)]
pub height: usize,
}
#[wasm_bindgen]
によってJavaScript側へ公開するインタフェースを指定します。
publicなフィールドは、JavaScriptのプロパティとして読み書きすることができます。
#[wasm_bindgen(readonly)]
をつけることで読み取り専用のプロパティを作ることもできます。
JavaScript側では以下のようにScreen構造体を扱うことができます。
(正確には、上記のコードだけではまだnewでインスタンスを取得することができません)
const screen = new Screen(600, 400)
console.log(screen.width, screen.height)
次に、メソッドの実装を行います。
impl
に対して#[wasm_bindgen]
を適用することで、publicなメソッドをJavaScriptから読み出すことができるようになります。
ただし、メソッドの引数と戻り値はJavaScriptへ変換可能な型でなければなりません。
また、#[wasm_bindgen(constructor)]
をつけたメソッドは、JavaScript側でインスタンス生成を行う際のコンストラクタとして実行されます。
fn create_buffer(width: usize, height: usize) -> Vec<u8> {
let size = 4 * width * height;
let mut bytes = Vec::with_capacity(size);
bytes.resize(size, 0);
bytes
}
#[wasm_bindgen]
impl Screen {
#[wasm_bindgen(constructor)]
pub fn new(width: usize, height: usize) -> Screen {
Screen {
bytes: create_buffer(width, height),
width,
height,
}
}
pub fn pointer(&self) -> *const u8 {
self.bytes.as_ptr()
}
pub fn size(&self) -> usize {
self.bytes.len()
}
pub fn resize(&mut self, width: usize, height: usize) {
self.bytes = create_buffer(width, height);
self.width = width;
self.height = height;
}
}
JavaScriptでScreenを用いてCanvasに描画する例は以下のようになります。
(importによるモジュール読み込みについての説明は後回しにします)
const mod = import('mandelbrot')
mod.then(({ Screen, mandelbrot }) => {
const canvas = document.querySelector('canvas')
const screen = new Screen(canvas.width, canvas.height)
const bytes = new Uint8ClampedArray(memory.buffer, screen.pointer(), screen.size())
const image = new ImageData(bytes, screen.width, screen.height)
const ctx = canvas.getContext('2d')
ctx.putImageData(image, 0, 0)
})
mandelbrot関数
次に、マンデルブロ集合の計算を行う関数を作成します。
関数も、#[wasm_bindgen]
をつけることでJavaScript側に公開することができます。
メソッドと同様に、引数と戻り値はJavaScriptへ変換可能な型である必要があります。
fn repeat(a: f64, b: f64, limit: usize) -> usize {
let mut x0 = 0.;
let mut y0 = 0.;
for k in 0..limit {
let x = x0 * x0 - y0 * y0 + a;
let y = 2. * x0 * y0 + b;
if x * x + y * y >= 4. {
return k;
}
x0 = x;
y0 = y;
}
return limit;
}
#[wasm_bindgen]
pub fn mandelbrot(screen: &mut Screen, x0: f64, y0: f64, d: f64, limit: usize) {
let mut y = y0;
for i in 0..screen.height {
let mut x = x0;
for j in 0..screen.width {
let k = repeat(x, y, limit);
let v = (k * 255 / limit) as u8;
let offset = 4 * (screen.width * i + j);
screen.bytes[offset] = v;
screen.bytes[offset + 1] = v;
screen.bytes[offset + 2] = v;
screen.bytes[offset + 3] = 255;
x += d;
}
y += d;
}
}
先ほどのJavaScript側の例にmandelbrot関数の呼び出しを追加すると以下のようになります。
const mod = import('mandelbrot')
const bg = import('mandelbrot/mandelbrot_bg')
Promise.all([mod, bg]).then(([{ Screen, mandelbrot }, { memory }]) => {
const canvas = document.querySelector('canvas')
const screen = new Screen(canvas.width, canvas.height)
const bytes = new Uint8ClampedArray(memory.buffer, screen.pointer(), screen.size())
const image = new ImageData(bytes, screen.width, screen.height)
mandelbrot(screen, -3, -2, 0.01, 100)
const ctx = canvas.getContext('2d')
ctx.putImageData(image, 0, 0)
})
ビルドとパッケージ化
作成したcrateをNode.jsのパッケージとして公開する手順を簡単にするために、wasm-packを使います。
wasm-packのインストールを行います。
$ cargo install wasm-pack
ビルドとパッケージ化は以下のコマンドで行います。
$ wasm-pack build
pkgディレクトリが作成され、WebAssemblyの.wasm
ファイルとインタフェースの.js
ファイル、そしてNode.jsのパッケージ情報を記載したpackage.json
が生成されます。
wasm-pack publish
で、作成したパッケージをそのままnpmに公開することもできます。
JavaScript側
作成したパッケージを利用するために、新たにJavaScriptのプロジェクトを作成しましょう。
$ mkdir wasm-mandelbrot
$ cd wasm-mandelbrot
$ npm init -y
$ npm i path/to/mandelbrot/pkg
$ npm i -D webpack webpack-cli webpack-dev-server
index.jsを作成します。
上述のJavaScriptコードと同じものです。
const mod = import('mandelbrot')
const bg = import('mandelbrot/mandelbrot_bg')
Promise.all([mod, bg]).then(([{ Screen, mandelbrot }, { memory }]) => {
const canvas = document.querySelector('canvas')
const screen = new Screen(canvas.width, canvas.height)
const bytes = new Uint8ClampedArray(memory.buffer, screen.pointer(), screen.size())
const image = new ImageData(bytes, screen.width, screen.height)
mandelbrot(screen, -3, -2, 0.01, 100)
const ctx = canvas.getContext('2d')
ctx.putImageData(image, 0, 0)
})
Webpackを使うことでwasm-bindgenで作成したパッケージを読み込むことができます。
ただし、WebAssemblyプログラムをJavaScriptから呼び出すためには実行時の読み込みが必要なので、パッケージの読み込みにはDynamic importを使用します。
Dynamic importで読み込んだモジュールはPromiseになります。
mandelbrot
モジュールには、wasm-bindgenを使って定義されたJavaScript用のインタフェースが含まれます。
mandelbrot_bg
は、その内部で使用される低レベルのモジュールですが、今回のようにJavaScript側からRust側のメモリにアクセスする場合にはそれを読み込む必要があります。
残りの部分として、webpack.config.jsとindex.htmlを以下のように用意します。
const path = require('path')
module.exports = {
entry: './index.js',
output: {
path: path.resolve(__dirname, 'public'),
filename: 'bundle.js'
}
}
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8"/>
<title>Drawing Mandelbrot Set with Rust and WebAssembly</title>
</head>
<body>
<canvas width="600" height="400"></canvas>
<script src="bundle.js"></script>
</body>
</html>
ここまでのディレクトリ構成は以下の通りです。
wasm-mandelbrot
├── index.js
├── package.json
├── public
│ └── index.html
└── webpack.config.js
最後に、以下のコマンドでwebpack-dev-serverを起動し、http://localhost:8080/ にアクセスすると、画面中のCanvas要素にマンデルブロ集合の描画が行われるでしょう。
$ npx webpack-dev-server --content-base public --mode development
性能測定
最後に、JavaScriptで同じプログラムを実装した場合との性能比較を行いましょう。
JavaScriptによる実装は以下の通りです。
const repeat = (a, b, limit) => {
let x0 = 0
let y0 = 0
for (let k = 0; k < limit; ++k) {
const x = x0 * x0 - y0 * y0 + a
const y = 2 * x0 * y0 + b
if (x * x + y * y >= 4) {
return k
}
x0 = x
y0 = y
}
return limit
}
const jsMandelbrot = (screen, buffer, x0, y0, d, limit) => {
for (let i = 0; i < screen.height; ++i) {
const y = y0 + d * i
for (let j = 0; j < screen.width; ++j) {
const x = x0 + d * j
const k = repeat(x, y, limit)
const v = k * 255 / limit
buffer[4 * (screen.width * i + j)] = v
buffer[4 * (screen.width * i + j) + 1] = v
buffer[4 * (screen.width * i + j) + 2] = v
buffer[4 * (screen.width * i + j) + 3] = 255
}
}
}
const canvas = document.querySelector('canvas')
const screen = new Screen(canvas.width, canvas.height)
const bytes = new Uint8ClampedArray(memory.buffer, screen.pointer(), screen.size())
const image = new ImageData(bytes, screen.width, screen.height)
jsMandelbrot(screen, bytes, -3, -2, 0.01, 100)
const ctx = canvas.getContext('2d')
ctx.putImageData(image, 0, 0)
mandelbrot
およびjsMandelbrot
関数の呼び出しのみを1000回繰り返して経過時間を測ります。
また、参考のためcargo bench
によってネイティブコンパイル時の実行時間も測定します。
実行環境はMac mini (Late 2012), 2.6 GHz Intel Core i7 processor, 16 GB 1600MHz DDR memory上のChrome バージョン70です。
1回あたりの平均呼び出し時間は以下の通りになりました。
環境 | 時間(ミリ秒) |
---|---|
Rust+WebAssembly | 10.80 |
JavaScript | 25.43 |
Rust(ネイティブ) | 10.60 |
Rust+WebAssemblyがJavaScriptの約2.5倍速い結果となりました。
WebAssemblyのSIMDやマルチスレッドが活用できればもう少し差がつきそうですね。
また、ネイティブと比較してもほぼ同等の時間でした。
おわりに
性能面でも利便性でも、Rust+WebAssemblyがもうほぼ実用的に使えるレベルまで来ているのではないかと思います。
wasm-bindgenはかなり遊べるのでぜひお試しあれ。
最後にドキュメントへのリンクも載せておきます。