Zenn
🌊

Fastly Compute を便利に使う Tips (4) Wizer で Wasm モジュールの起動速度を最適化する

2024/12/30に公開
performance
WebAssembly
Fastly
Fastly Compute
tech

この記事は 2024 年の Fastly Compute 一人アドベントカレンダー 25 日目の記事です。前回の 23 日目の記事から引き続き大幅に遅れて（5日遅れ）の投稿になってしまいましたが、今年のアドベントカレンダー最終回となる今回は Wizer を使った WebAssembly モジュールのローディング速度最適化の tips について書いてみたいと思います。

Wizer とは

Wizer は Wasm モジュールの事前初期化ツールです。入力として Wasm モジュール(例: main.wasm など)を受け付け、そのモジュールの初期化処理完了後の線形メモリのスナップショットをデータセクションとして保持する新しい Wasm モジュールを出力してくれます。いかにも初期化処理が重いケースで利用価値が高そうなツールですが、Fastly の場合は JavaScript SDK の中でデフォルトで有効になっていて、去年のアドベントカレンダーのこちらの記事で以前も簡単に触れたことがあります。

インタプリタの初期化のようなた明らかに大きな処理以外でも、例えば大きめのテキストやバイナリデータを予め読み込んで parse してメモリに展開しておくだけの用途でも十分に便利なツールで、実際前回の記事で紹介したベクトル検索のデモにおいても、検索に必要な HNSW グラフの一部やベクトルの配列を予めメモリに読み込んでおく用途で活用されてたりします（ちなみにそのデモのコードで実際にどの程度 wizer によって起動かかるローディング時間を短縮しているか計測してみたところ、ローカルの環境では 10ms 前後は短縮されているようでした）。

使い方

wizer のリポジトリの README に基本的な使い方の解説があり、この内容を見ていきます。まず wizer で最適化する対象の Wasm モジュールは wizer.initialize という名前の初期化関数を export しておく必要があるようです（以下は Rust での例）

#[export_name = "wizer.initialize"]
pub extern "C" fn init() {
    // Your initialization code goes here...
}

README を読む限り他に難しい条件はなく、あとはいま最適化したい対象の Wasm モジュールのファイル名が input.wasm だとすると、下記コマンドを実行すれば OK なようです。

$ wizer input.wasm -o initialized.wasm

各 SDK での使い方

次に具体的に Fastly の各言語 SDK でどのように使えるかについてコード例を通して見ていきます。

JavaScript の場合

前述のように JavaScript SDK の場合は以下のようにデフォルトで wizer が組み込まれて有効になってるので、Fastly で JS でコードを書いてデプロイするとみなさん意識をするかしないかに関わらず、基本的に wizer の恩恵に預かっていることになります。

https://github.com/fastly/js-compute-runtime/blob/56aa96d5a223d45d79327f0a19bcfd93f1e90363/src/compileApplicationToWasm.js#L155-L165

なので、JS の場合グローバルのスコープで変数定義するだけで特に追加の設定など不要ですぐに使えます。以下は wizer により top レベルでの変数宣言の場合は乱数生成が効かなくなることを示したサンプルです。（fiddle で動作を確認したい方はこちらに作成してあります）

addEventListener("fetch", (event) => event.respondWith(handleRequest(event)));

console.log("stdout at top level scope");
let rootValue = Math.random();

async function handleRequest(event) {
  console.log("stdout inside handleRequest()... fyi:", Math.random());
  return new Response(rootValue, { status: 200 });
}

これを実行すると以下のような標準出力がログで得られます。以下は fastly compute serve で起動したローカル環境上の Compute サービスに対して http request を複数回送信した際の実行結果です。Global レベルのスコープで代入されている Math.random() の実行結果を格納しているはずの rootValue 変数に毎回同じ値が設定されている様子が確認できます。

$ curl localhost:7676
0.017223164904862642
$ curl localhost:7676
0.017223164904862642
$ curl localhost:7676
0.017223164904862642

一方で、サーバ側の標準出力で記録されたログを確認すると、以下のようなログが確認できます。このログから、top レベルで実行された console.log() メソッドの呼び出しによる標準出力のログが出力されていないことが確認できますし、handleRequest() 関数内で実行された Math.random() であれば毎回の http request の送信の度に乱数が正しく生成されていることが確認できます。

stdout inside handleRequest()... fyi: 0.507978729903698
stdout inside handleRequest()... fyi: 0.744825457688421
stdout inside handleRequest()... fyi: 0.5817289529368281

上記例のように乱数を固定したいというニーズは現実にはあまりないかもしれませんが、同じ原理を利用して大きめのサイズのテキストで構成されるマスタ系やインデックスなどの json などのデータソースを予め読み込んで parse まで済ませておく、といったことであれば現実的に活用できる場面もあるかもしれません。

Rust の場合

上記 JS の例を読んでいただくと、なんとなく「メインとなるエントリポイントに入る前にグローバル変数に値をセットできれば使えるイメージなのかな」という感覚を持って頂けるかなと思います。次に Rust でどのようにすればエントリポイントに入る前にグローバル変数に値をセットできるかを見ていきます。

具体的なコード例を 3 つ挙げます。一つ目は、wizer のリードコントリビュータの Nick Fitzgerald が 2021 年の WebAssembly Summit の講演の際に使用したと思われるリポジトリからの抜粋です。init() 関数が wizer.initialize という名前で export されていますので、この関数が wizer によって初期化関数として呼び出される想定がされているようです。この init() 関数の終わりで LEXICON.set() というメソッド呼び出しによってグローバル変数の値がセットされているのが分かります。

https://github.com/fitzgen/wasm-summit-2021/blob/main/spelling-corrector/src/main.rs#L8-L23

次の例は Shopify が Ruby の実行環境として開発している ruvy が wizer を利用している例です。先ほどと同じように wizer.initialize という名前で export されている関数がありますのでこの load_user_code() 関数が初期化関数に該当しそうです。こちらの例も先ほどと似ていて、USER_CODE.set() というメソッド呼び出しでグローバル変数の値をセットしていることが確認できます。

https://github.com/Shopify/ruvy/blob/fb1be2345c7044cdd83f8f16b5e9c35b2bf9fbc7/crates/core/src/main.rs#L4-L26

三つ目の例は前回紹介したベクトル検索のコード例です。少し長いので主要部分のみ抜粋すると以下のようになります。今度は VEC_DB というグローバル変数が使われていて、Lazy::force()という遅延評価処理の中でグローバル変数への値のセット（確定）が行われているようです。

use once_cell::sync::Lazy;
...(中略)...
pub static VEC_DB: Lazy<Vectors<DIMENSIONS>> = Lazy::new(|| {
  ...(中略)...
});

#[export_name = "wizer.initialize"]
pub extern "C" fn init() {
    Lazy::force(&VEC_DB);
}

以上いくつかの具体的なグローバル変数への値セットの方法を見てきましたが、どのパターンの場合であっても、wizer に対して明示的に初期化関数を指定することで値のセットが可能そうであることが見ていただけたかなと思います。

ここまで見てきたように、Rust においてグローバル変数への値セットの方法が複数ある理由については以下の記事が大変分かりやすかったです。この辺りの経緯も踏まえつつ、どのメソッドが適切かを（必要があれば）状況に応じて使い分けてグローバル変数の値をセットしていくことになりそうです。

https://zenn.dev/frozenlib/articles/lazy_static_to_once_cell

自分なりの最小限のコードもメモ書き代わりに残しておきます。(コード中の combined.bincode や crate::common は、それぞれ前回の記事のこのバイナリやこの common.rs をコピーして持ってくれば動きます)

mod common;

use crate::common::{EmbeddingCluster,find_embedding_point_by_id};
use fastly::http::StatusCode;
use fastly::{Error, Request, Response};
use std::time;
use once_cell::sync::Lazy;

static CLUSTERS: Lazy<Vec<EmbeddingCluster<5>>> = Lazy::new(|| bincode::deserialize(include_bytes!("./combined.bincode")).unwrap());

#[export_name = "wizer.initialize"]
pub extern "C" fn init() {
    Lazy::force(&CLUSTERS);
}

#[fastly::main]
fn main(req: Request) -> Result<Response, Error> {
    let start = time::Instant::now();
    let ids: &Vec<u32> = &vec![84958, 85002, 85047, 85073, 85177];
    let embeddings_for_ids: Vec<Vec<f32>> = ids
        .iter()
        .filter_map(|id| find_embedding_point_by_id(&CLUSTERS, *id))
        .map(|embedding_point| embedding_point.into())
        .collect();
    println!("combined.bincode loaded in {:?}", start.elapsed());
    Ok(Response::from_status(StatusCode::OK).with_body(format!("{:?}",embeddings_for_ids)))
}

Go の場合

TinyGo では以前から wasm ターゲットで関数の export が可能でしたし、今年リリースされた Go の 1.24 からも go:wasmexport ディレクティブがサポートされたとのことなので、これを使えば wizer を使った初期化も可能であるに違いありません。と思って wizer の利用を試してみましたが、結論としては本稿執筆時点では私の試した限りではうまく動作させることに成功していません。以下試した内容のログです。（試したソースは TinyGo のこちらの Issue で利用されていたものをサンプルとして引用しています）

$ go version # go1.24rc1 を使ってテスト
go version go1.24rc1 darwin/arm64
$ cat main.go # 試したコード
package main

var theText string

//go:wasmexport wizer.initialize
func WizerInitialize() {
  theText = "initialized!"
}

func main() {
  println("theText: ", theText)
}
$ GOARCH=wasm GOOS=wasip1 go build -o bin/main.wasm . # wasm ターゲット向けにビルド
$ wasm-tools print bin/main.wasm | grep wizer.initialize | head -n 2 # wizer.initialize 関数は export されているように見える
  (func $wizer.initialize (;1048;) (type 10)
  (export "wizer.initialize" (func $wizer.initialize))
$ wizer --allow-wasi --wasm-bulk-memory=true bin/main.wasm -o bin/main.wizer.wasm
Error: the `wizer.initialize` function trapped

Caused by:
    0: error while executing at wasm backtrace:
           0: 0xea2f6 - <unknown>!wizer.initialize
    1: memory fault at wasm address 0xfffffff8 in linear memory of size 0x1d0000
    2: wasm trap: out of bounds memory access

上記の go build コマンド部分を tinygo build -target=wasi -o bin/main.wasm ./ に変えても、エラーメッセージが若干変わりますが同様に wizer の実行時にエラーとなりました。

$ wizer --allow-wasi --wasm-bulk-memory=true bin/main.wasm -o bin/main.wizer.wasm
panic: runtime error: //go:wasmexport function called before runtime initialization
Error: the `wizer.initialize` function trapped

Caused by:
    0: error while executing at wasm backtrace:
           0: 0x1bc3 - main!runtime.runtimePanicAt
           1: 0x18fe - main!runtime.runtimePanic
           2: 0x4701 - main!main.WizerInitialize#wasmexport
       note: using the `WASMTIME_BACKTRACE_DETAILS=1` environment variable may show more debugging information
    1: wasm trap: wasm `unreachable` instruction executed

私の環境やテスト方法に誤りがあり実は何か足りていないステップがあるのかなど、調べきれていない部分も多いので Go での wizer の利用については来年以降どこかで時間を見つけて調べてみたいと思います。

まとめ

本稿では Wizer を使った WebAssembly モジュールのローディング速度最適化の方法について、各言語での具体的な利用方法を中心に紹介しました。また時間の都合で本稿では紹介できませんでしたが、このような AOT コンパイラの領域は wizer 以外にも例えば weval [1] などの新しい手法の開発も進んでいる興味深い領域です。

このような AOT での最適化が一つのツールで実質的に言語を跨いで実現できることは WebAssembly をビルドのターゲットにして開発する場合の醍醐味の一つだと思いますし、WebAssembly をターゲットにした開発を検討される場合の観点の一つとして覚えておくと役に立つ場面もあるかと思います。今後もこういったツール周りで開発に役立ちそうな内容があれば取り上げて記事にしてみたいと思います。

脚注

  1. インタプリタ型言語での利用を想定した AOT 最適化ツール。参考: 開発者の Chris による解説ブログ ↩︎

Discussion