👅

Component Model な Wasm ランタイムを作った

2024/12/22に公開

Unity / .NET で動く WebAssembly ランタイム WaaS を作りました。

https://github.com/ruccho/WaaS

本記事では WebAssembly ランタイムとして見た WaaS の設計について書いていこうと思います。WaaS のコンセプトや特長については別記事で紹介しているので、こちらもぜひご覧ください。

https://zenn.dev/ruccho/articles/11ad66aeeab0c2

前提

  • C# 9
  • .NET Standard 2.1
  • Unity 2022.3.14f1

まずは普通の WebAssembly ランタイムを作る

Component Model の話に入る前に、まずは普通の WebAssembly ランタイムを作っていきます。

ランタイムは全て C# 実装ですが、実装のとっかかりでは「RustでWasm Runtimeを実装する」がとても参考になりました。

https://zenn.dev/skanehira/books/writing-wasm-runtime-in-rust

デシリアライズ

効率の良いバイナリの読み方については MessagePack for C# などのシリアライザの作者として知られる neuecc さんの資料が参考になりました。

https://speakerdeck.com/neuecc/cedec-2023-modanhaipahuomansuc-number-2023-edition

デシリアライザの大部分は手書きしていますが、命令のような種類の多いデータについてはボイラープレートが発生しがちなので Source Generator によるデシリアライザの自動生成を行っています。各命令のデシリアライズ方法を属性で定義すると自動的にデシリアライザが生成される感じです。

[OpCode(0x23)] // オペコード指定
public partial class GlobalGet : Instruction
{
    [Operand(0)] public uint GlobalIndex { get; } // オペランド指定
}

VM

WaaS は iOS など JIT が禁止されているプラットフォームでも動的にロードした WebAssembly の実行をサポートしたかったので、スタックマシンベースのインタプリタを実装しています。可能なプラットフォームでは JIT や AOT もサポートしたいところですが、ひとまずはこのインタプリタをベースラインとしています。現状は少々ナイーブな実装ですが、いくつか最適化のアイデアはあるので今後取り組んでいきたいです。

一点特殊なのは、C# から Wasm に非同期メソッドをインポートできるようにしたことです。今回はゲーム向けの会話シーンなどを記述する組み込みスクリプト用途を主に想定していたので、Wasm から非同期関数を簡単に呼び出せるようにしたいと考えました。そこで、Wasm に C# の非同期メソッドをインポートすると、Wasm 内部からは同期的に見え、しかし実際に呼び出されると VM が自動的に中断と再開を行う仕組みを入れています。Wasm 視点ではグリーンスレッドともいえるかも。

バインディング

C# から Wasm の関数を呼び出す場合は、最も低レベルなインターフェースではこのようなコードになります。

// set args
Span<StackValueItem> args = stackalloc StackValueItem[3];
args[0] = new StackValueItem(1); // i32
args[1] = new StackValueItem(1.0); // f64
args[2] = new StackValueItem(1.0f); // f32

// invoke
context.Invoke(function, args);

// take results
Span<StackValueItem> results = stackalloc StackValueItem[1];
context.TakeResults(results);
Debug.Log(results[0]);

これではちょっと長ったらしすぎるので、次のように Wasm の引数をそのまま C# の引数として記述できるようにしています。

var result = CoreBinder.Instance.Invoke<int>(context, function, 1, 1.0, 1.0f);
Debug.Log(result);

さて、C#では「可変長で任意の型の組み合わせの引数をもてるメソッド」を定義しようとすると、ほぼ必然的に params object[] になってしまいます。これは boxing と配列のアロケーションコストの観点で望ましくありません。
そこで再び Source Generator の出番です。上記のように関数を呼び出すコードを書くと、Source Generator がシグネチャに応じてオーバーロードを自動生成します。この生成コードは boxing や配列のアロケーションを避けるようになっており、オーバーヘッドの少ない関数呼び出しを可能にします。

public static TResult Invoke<TResult>(this Binder binder, ExecutionContext context, IInvocableFunction function, int _0, double _1, float _2)
{
    /* ... */
}

同様のアプローチは Wasm に C# のメソッドをインポートする際にも使われています。

var module = await moduleAsset.LoadModuleAsync();
var instance = new Instance(module, new Imports()
{
    {
        "module name", new ModuleExports()
        {
            { "some function", CoreBinder.Instance.ToExternalFunction((Func<int, long, float, double, int>)SomeFunction) }
        }
    }
});

static int SomeFunction(int a, long b, float c, double d)
{
    return default;
}

WaaS では、このように C# のデリゲートを Wasm の関数に変換してインポートする形式をとります。この変換の際、デリゲートの型に合わせてコードが自動生成されます。これによってデリゲートに対する DynamicInvoke() を回避することができ、こちらもオーバーヘッドが抑えられます。

テスト

コア Wasm 仕様のテストではこちらの記事を参考にさせていただきました。

https://zenn.dev/ri5255/articles/bac96cf74f82f0

難しかったのはブロックが絡んだスタック状態のバリデーションで、結構細かい考慮事項が多く、なかなかテストが通りませんでした。特に unreachable が絡むパターンが難しい……!バリデーションについては詳しく解説されている記事があります:

https://labs.gree.jp/blog/2024/12/24431/

とはいえ今回は全部のケースを厳密に通したわけではなく、例えば一部の浮動小数点数の処理は C# でのストレートな実装だといくつか仕様と微妙な差異があったのをそのままにしていたりと、割と適当です。この辺りはパフォーマンスにも影響しそうなので、問題が起きたら考えることにします……

Component Model 対応

ここから本題です。
Component Model については詳しい記事がありますので、この記事では飛ばしていきます。

https://zenn.dev/newgyu/articles/82a181212c7bb2

経緯

Core Wasm なランタイムの実装が進むにつれ、これだと実用は厳しいだろうということがわかってきました。Core Wasm は引数・戻り値として直接やり取りできる型が数値しかなく、それ以外の複雑な型や大きなデータをやり取りする際は線型メモリを介する必要があります。

ただ、線型メモリを介した複雑なデータのやり取りの方法は十分に定義されておらず、例えば現状の wasm-bindgen のようにホスト・ゲストの両面でグルーコードを生成するようなアプローチが事実上必須となってきます。しかし様々なゲスト言語に向けたツールチェインを独自に用意するのは大変すぎます。

これらの問題を解決できそうな方法として注目したのが Component Model でした。Component Model で定義されている Canonical ABI では数値以外のリッチな型表現のデータをやり取りする方法を定めています。また、Component Model では「リソースハンドル」の仕組みがあり、Unity側のオブジェクトへの参照を WebAssembly に持ち込む方法としてぴったりでした。
Component Model で各ゲスト言語向けのバインディングを生成する wit-bindgen をはじめとする周辺ツールも精力的に開発されており、そのエコシステムに乗っかることができるというのも利点でした。

デシリアライザ

Component Model におけるコンポーネントのバイナリフォーマットは通常の Wasm Module バイナリをラップし、そこに様々なメタデータを付与するような形式をとっています。

Wasm Module では(命令のデコードを除いて)手でデシリアライザを書いていましたが、Component では登場するデータの種類や複雑性が格段に増しており、全体にわたって Source Generator によるデシリアライザの自動生成を行うことにしました。たとえば次の IValueTypeDefinition は Component Model における値の型を表現するデータですが、これは先頭1バイトを見て型の種類を見分けるフォーマットをとっており、このルールを C# の属性で表現することでデシリアライザが自動生成されます:

[GenerateFormatter]
[Variant(0x68, typeof(BorrowedType))]
[Variant(0x69, typeof(OwnedType))]
[Variant(0x6A, typeof(ResultType))]
[Variant(0x6B, typeof(OptionType))]
[Variant(0x6D, typeof(EnumType))]
[Variant(0x6E, typeof(FlagsType))]
[Variant(0x6F, typeof(TupleType))]
[Variant(0x70, typeof(ListType))]
[Variant(0x71, typeof(VariantType))]
[Variant(0x72, typeof(RecordType))]
[VariantFallback(typeof(PrimitiveValueType))]
public partial interface IValueTypeDefinition : ITypeDefinition, IUnresolvedValueType
{
}

Component Model には各データの種類に対してインデックス空間が作成されます。バイナリを先頭から読み取っていって、新たにデータが登場すると、そのインデックス空間の末尾にデータが追加されていきます。データが他のデータを参照する際はこのインデックスが使われ、あるデータは、バイナリの中でそれよりも前に登場したデータしか参照することができません。これによってインデックスの遅延バインド等は必要なく、循環参照も発生しません。デシリアライザとしては先頭から順番にデータを読みながら参照関係を構築していくことができるので、ありがたい構造です。

ただ、このデータの参照にはかなり複雑なケースがいろいろあります。Component Model には Core Wasm と同様に「インスタンス化」の概念があり、コンポーネントに対して関数やメモリ等をインポートすることで「インスタンス」が作成されます。ただし Component Model が Core Wasm と異なるのは、コンポーネントの定義をネストしたり、コンポーネントそのものをエクスポートしたり、インポートしたコンポーネントをインスタンス化してさらにエクスポートする……といった操作が可能になっている点です。コンポーネントのような静的なデータとインスタンスのような動的なデータが共通のインデックス空間の中に存在しており、どのデータがいつ解決されるべきなのかということが非常にわかりにくいパターンがあります。これに関しては Explainer を熟読したり、wasmtime の実装を読んだりしながらなんとなく理解していきました。

バインディング生成

Component Model では WIT という IDL を使ったバインディングの標準的な手順が定められています。WIT はあるコンポーネントがインポート・エクスポートする関数や型などを定義するもので、この WIT をソースとして各言語向けのバインディングを生成します。

ゲスト言語向けには wit-bindgen という Rust で実装された CLI ツールがあり、こちらを利用できます。対してホスト環境 (WaaS) むけには、この WIT から適切なバインディングを生成できるようにする必要がありました。

WIT のパースに関しては、wit-bindgen の内部でも使われている wit パッケージを利用することとし、必然的に WaaS 向けのバインディング生成ツールも Rust で実装することになりました。これはwit2waasとして公開しています。

https://crates.io/crates/wit2waas

また、細かいバインディングが大量にファイルとして生成されるのは好ましくなかったため、wit2waas では最低限の型定義 + 属性によるアノテーションのみを生成するに留め、そのほかの詳細なバインディングはまたもや Source Generator で生成しています。以下は実際に wit2waas が生成する型定義です。

// <auto-generated />
#nullable enable

namespace MyGame.MySequencer
{
    [global::WaaS.ComponentModel.Binding.ComponentInterface(@"env")]
    public partial interface IEnv
    {
        [global::WaaS.ComponentModel.Binding.ComponentApi(@"show-message")]
        global::System.Threading.Tasks.ValueTask ShowMessage(string @speaker, string @message);

        [global::WaaS.ComponentModel.Binding.ComponentApi(@"show-options")]
        global::System.Threading.Tasks.ValueTask<uint> ShowOptions(global::System.ReadOnlyMemory<string> @options);

    }
}

これにより、WIT を介さずに直接型定義を手書きするような運用も一応可能になっています。

この型定義をもとに Source Generator が生成するバインディングはかなり手厚く作っていて、ほとんど C# のメソッドを呼び出すのと変わらない感覚で Wasm と相互運用することができます。

// コンポーネントのロード
var component = await componentAsset.LoadComponentAsync();

// インスタンス化
var instance = component.Instantiate(null, new Dictionary<string, ISortedExportable>());

using var context = new ExecutionContext();

// Source Generator が生成したラッパーでインスタンスを包む
var wrapper = new IEnv.Wrapper(instance, context);

// 呼び出す
await wraapper.ShowMessage("ぼく", "こんにちは");

値の受け渡し

関数呼び出しにおける引数と戻り値のやり取りについてです。

Wasm Module の時は、型が i32 i64 f32 f64 の4つしかなったので、これらを union っぽくした構造体 StackValueItem を作り、引数や戻り値は Span<StackValueItem> としてやり取りできていました。また、引数や戻り値は一旦スタック上に作ったバッファに溜め込んでから丸ごと受け渡していました。

これに対し、Component Model の型表現はかなりリッチです。まずこれらのデータを boxing なしに受け渡すには、静的なコード生成によって型を決定することが必要でしょう。また、特にリストのようなデータは巨大になりうるので、一時バッファをスタックに作成できません。コピーの回数を減らしたいので、できればヒープにも一時バッファを作りたくありません。以上から、プリミティブ単位で値のプッシュ・プルを繰り返すような構造が理想的です。

var binder = function.GetBinder();
binder.Push(1.0);
binder.Push("hoge");
var recordBinder = binder.PushRecord();
recordBinder.Push(2.0);
recordBidner.Push("fuga");
binder.Push(3.0);
/* ... */

バインディングの Source Generator でこのような形式のコードを自動生成すれば静的にデータの構造を決定できますし、実際の値を直接的に関数に渡すこともできます。

この方法で一つ問題があるのは、プッシュとプルのどちらが主導権を握るのか、ということです。
例えば、C# からコンポーネントに引数を渡すときは上記のようなプッシュ型がいいですが、逆に戻り値を受け取るときは同様のコードをプル型で書きたいところです。

var a = binder.PullS64();
var b = binder.PullString();
var recordBinder = binder.PullRecord();
c.f0 = recordBinder.PullS64();
c.f1 = recordBidner.PullString();
var d = binder.PullS64();
/* ... */

プッシュ側が一つ値をプッシュしたら、値はバッファに溜めず、すぐプル側に渡される。両方を同期的に書くことはできないので、どちらか片方はステートマシンにして、『現在位置』を進めながら値を処理する必要があります。そして、C# ではそんなコードを自動生成する機能が標準で備わっています……そう、 async/await です。async/await は本来は非同期処理のためのコード生成機能ですが、今回のように単なるステートマシン展開にも使えます。ということで、プル側は async/await を使って書くことにしました。

var a = await binder.PullS64();
var b = await binder.PullString();
var recordBinder = await binder.PullRecord();
c.f0 = await recordBinder.PullS64();
c.f1 = await recordBidner.PullString();
var d = await binder.PullS64();
/* ... */

これによって実際には以下のようにプッシュ側とプル側で交代しながら処理が進みます。

  • プル側がawait binder.Pull()を呼んで中断し、プッシュ側に処理が移る
  • プッシュ側がbinder.Push(/* */)を呼ぶと、内部でプル側の処理が再開する
  • 以後繰り返し

つまり、一連の値の受け渡しは async/await でありながら(外から見ると)同期的に完了するということです。
これを実現するため、WaaS では STask という独自の Task-like 型を定義して使っています。STaskは処理をスレッドプールに投げたり SynchronizationContext.Post() したりすることが一切なく、完全に意図したタイミングで中断・再開される、ミニマルな Task-like 型です。

マーシャリング

さて、上記のプッシュ・プルの仕組みを使って、実際のマーシャリング処理を書いていきます。

Component Model の Canonical ABIでは、Core Wasm の関数を Component の関数に変換する「Lifting」と、逆に Component の関数を Core Wasm から呼び出せる関数に変換する「Lowering」という操作が定義されています。したがって、このような関数の変換の境界でプッシュ・プル時に値の変換が行えるようにすればいいということです。

https://zenn.dev/tanishiking/articles/2024-1209-canonical-abi

Lifting / Lowering の基本的な考え方は、受け渡す値の数が規定値より少なければ Wasm の引数・戻り値を直接使って受け渡し(Flattening)、規定値より多ければ線形メモリに格納してそのポインタを渡す(Load / Store)というシンプルなものです。あとは Component Model の各型に対して Flat, Load, Store の方法が定義されているのでそれに従って実装していきます。なかなか量が多くて骨が折れますが、頑張りました。

ふりかえり

感想

いまコミットログを遡ったら、コア Wasm 仕様のテストは着手から1週間で通ってますが、そこから半年くらい Component Model 対応に費やしてました(休日開発 & サボり含む)。Component Model の仕様は複雑で巨大です。正味 Canonical ABI さえあれば満足な自分としては、それ以外のものをかなりたくさん実装する必要があって大変でした。

フル C# での実装について

Unity で WebAssembly を動かしたい、となれば既成の Wasm エンジンを組み込むのが近道だったと思います。初めは勉強のつもりで C# で書きはじめたんですが、実装を進める中で Component Model を知り、ゲームへの組み込みのイメージが湧いてきました。現状 Component Model に対応するメジャーなランタイムは wasmtime くらいしかありませんが、wasmtime ではインタプリタが未実装なので、iOS でのスクリプトの動的ロードをカバーできません。また、非 C# な他のランタイムを Unity に組み込むにはランタイムをネイティブバイナリにビルドする必要がありますが、フル C# 実装ならばそれも必要ありません。なんなら NDA で保護されるコンソールでも動くはずです。

また、WaaS はフル C# だからこそインターフェースに余計な制約がありません。バインディング周りや非同期関数のインポートなどの機能はフル C# 実装のおかげで最大限の機能性が得られています。Source Generator を活用するタイミングがたくさんあったのも楽しかったです。

WaaS の存在意義について

非ブラウザ Wasm のユースケースとして、今回のようなクライアントアプリケーションでの組み込みスクリプト用途は時々言及されてはいつつも、まだそんなに流行ってない感があります。これは記事中でも言及した通りマーシャリングの規格が不足していたことが要因として大きかったのではないかと思いますが、今回 Component Model を採用したことでその点もカバーでき、一段と実用度の上がったソリューションを示せたのではないでしょうか。あとは Component Model が普及してくれれば最高ですね(そこが一番ハードル高そう)。

Discussion