Rustでwasm用カスタムアロケータを書く

Wasm とホスト間での関数呼び出しでは数値しか受け渡すことができないので、文字列や配列をやり取りするときには wasm インスタンスのメモリを経由する必要がある。例えば文字列をホスト(JavaScript)から wasm に渡したいときは一度 wasm インスタンスのメモリ領域に文字列を書いてから、書いた場所のアドレスとサイズを関数に渡す。これは数値以外の情報をやり取りする場合は wasm 側にメモリアロケータが必要であることを意味している。

Rust が生成する wasm がメモリ割り当てをどうやっているのかを調べて、自分でナイーブなアロケータを作るところまでやったので、その過程で理解したことをエントリにしておこうと思う。

ここではホスト環境としてブラウザ(JavaScript)を想定している。また wasm のインスタンス化の際にはメモリを import せず wasm 内で管理することを前提とする。

WebAssembly のメモリ

WebAssembly は単純なメモリモデルを採用していて、メモリは連続したバイト列として表現される。初期化時にある一定のサイズが確保されて、必要に応じて領域を広げることができる。ただし一度確保した領域を縮小することはできない。メモリ領域の伸張は 64KB のページ単位で行う。

Wasm インスタンスのメモリは JavaScript からは ArrayBuffer として見える。例として[1,2,3]という配列を JS から wasm に&[u32]として渡したいとしよう。JS 側のコードはこんな感じになる。

// `wasm` is an WebAssembly.Instance.
const arg = [1, 2, 3];
const arr = new Uint32Array(wasm.memory.buffer);
const addr = wasm.allocBytes(arg.length * 4);
arr.set(arg, addr / 4);
wasm.someFunc(addr, arg.length);
wasm.freeBytes(addr);

wasm.memory.bufferから見える ArrayBuffer が wasm インスタンスのメモリ。wasm 側のアロケータの仕事は、wasm.memory.bufferをヒープ領域として扱い、そのヒープを適宜伸張しながら malloc()/free()相当の機能を提供すること。上記ではそれらをwasm.allocBytes()やwasm.freeBytes()として定義している。

wasm からメモリ領域を伸張するにはmemory.grow命令を使う。この命令は引数を二つとる。一つは新たに確保するページの数。もう一つは対象となるメモリ領域を指定するインデックス。このインデックスには常に 0 を指定する。仕様の上では wasm インスタンスは複数のメモリ領域を参照できるようになっているが、これは将来の拡張のために用意されているもの。現時点では一つの wasm インスタンスにつきメモリ領域は一つしか存在しない。

アロケータの実装

Rust 側での実装に話を進めていく。Rust から WebAssembly のmemory.grow命令を呼ぶにはwasm32::memory_grow()を使う。以下の Rust のコードをコンパイルすると、

#![cfg_attr(target_arch = "wasm32", feature(stdsimd))]
use core::arch::wasm32;
fn f() -> usize {
  const MEMORY_INDEX: usize = 0;
  unsafe { wasm32::memory_grow(MEMORY_INDEX, 1) }
}

以下の wasm が生成される。

(func (;1;) (type 1) (result i32)
  i32.const 1
  memory.grow)

memory_grow()は成功すると伸張前のメモリのサイズをページ単位で返す。言い換えれば、新規に確保されたメモリの開始アドレスは戻り値に 64KB を掛けることで求められる。また伸長に失敗したときはusize::max_value()を返す。以下のようなラッパーを書くと malloc()に近い使い勝手になる。

unsafe fn grow_pages(num_pages: usize) -> *mut u8 {
    let num_pages = wasm32::memory_grow(MEMORY_INDEX, num_pages);
    if num_pages == usize::max_value() {
        process::abort();
    }
    (num_pages * PAGE_SIZE) as *mut u8
}

下準備ができたのでアロケータを書いていく。今回はアロケータを作るうえで必要なパーツを説明するのが主眼なので、Memory Allocators 101という記事を参照にして単純な first-fit のアロケータを作成した。コードは Github にあげているのでここではインタフェースのみ載せておく。

struct Heap {
    // ...
}

impl Heap {
    unsafe fn alloc(&mut self, size: usize) -> *mut u8 { /* ... */ }
    unsafe fn dealloc(&mut self, ptr: *mut u8) { /* ... */ }
}

ポインタを安全でない方法でキャストしたりするので割り切って全体を unsafe にしている。もっと Rust 力があれば安全に書けるかもしれないが、後述する GlobalAlloc が unsafe なのであまり頑張っても意味がないかもしれない。

GlobalAlloc

さらに、さきほど作成したアロケータを Rust のVec<T>やStringなどの標準ライブラリのメモリ割り当てにも使うようにする。ホスト環境とのデータのやり取りのみに自作のアロケータを使う場合はこの作業は不要だが、デフォルトのアロケータはバイナリサイズが大きく(10KB ぐらいになる)、使用するアロケータは一つに統一しておいたほうが良い。

標準ライブラリのアロケータを自前のものに差し替えるには以下の二つを行う。

GlobalAllocトレイトを実装する型を作る
その型の static なグローバル変数を#[global_allocator]属性をつけて定義する

GlobalAllocを実装するには以下の二つのメソッドを実装する。Layoutはサイズとアライメントを保持する構造体。

unsafe fn alloc(&self, layout: Layout) -> *mut u8;
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);

これら二つのメソッド、&mut selfではなく&selfを取る。これは GlobalAlloc を実装する型が static なグローバル変数として使われることからくる制約だろう。

とはいえメモリの割り当てと解放にはアロケータの状態を変更しなければならない。こういう場合の Rust のイディオムは内部可変性を使うことだろう。ここではUnsafeCell<T>を使う。内部可変性を提供する型を書く際は排他制御をして各種の競合に注意しなければならないが、Rust は WebAssembly 向けのスレッドはまだサポートしていないので今回はシングルスレッドを前提にする。

struct CustomAlloc {
    heap: UnsafeCell<Heap>,
}

impl GlobalAlloc for CustomAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        (*self.heap.get()).alloc(layout.size())
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        (*self.heap.get()).dealloc(ptr)
    }
}

unsafe impl Sync for CustomAlloc {}

static ALLOC: CustomAlloc = CustomAlloc {
    heap: UnsafeCell::new(Heap { /* ... */ }),
};

UnsafeCell<T>はSyncを実装しないので、シングルスレッドで使う前提でSyncを明示的に実装している。これを忘れるとCustomAllocを static な変数として使おうとしたときにコンパイラに怒られる。

これで実装は一通り完了したのであとは実際に使ってみる。

wasm-bindgen

冒頭に説明したように wasm とホスト環境で文字列や配列をやり取りする際には面倒な手順を踏む必要がある。文字列を渡す場合を例に具体的なステップを見てみると、(1) ホスト側からメモリ領域を確保するよう wasm インスタンスに要求する、(2)ホスト側で wasm のメモリ領域に文字列の内容を書く、(3)データを書いたアドレスとサイズを wasm の関数に渡す、(4)wasm の関数では渡されたアドレスとサイズから String なり &str に変換する、という作業が必要になる。

この部分も自作しようかと考えたが、結構煩雑になるのでこの変換部分は wasm-bindgen に頼ることにした。wasm-bindgen を使うとホスト(JavaScript)と wasm 側両方のラッパーを作ってくれるのでとても楽になる。

extern crate wasm_bindgen;
use wasm_bindgen::prelude::*;

#[wasm_bindgen]
pub fn to_uppercase(s: &str) -> String {
    s.to_uppercase()
}

あとはビルドして wasm-bindgen の CLI で後処理をすれば最終的な wasm/JS ができる。wasm-bindgen はデフォルトだと Webpack で処理する前提の ES modules 形式で出力するが、オプションを指定することでブラウザで直接読み込める形式にすることもできる。

$ cargo build --release --target wasm32-unknown-unknown
$ wasm-bindgen --browser --no-modules target/wasm32-unknown-unknown/release/wasm_custom_allocator_example.wasm --out-dir public/dist

wee_alloc

アロケータを差し替える主な動機はバイナリサイズの削減だろう。今回は理解を深める目的で自分でアロケータを書いてみたが、実用的には独自アロケータを書くよりもwee_allocを使うのが良いと思う。アロケータは unsafe なコードがどうしても多くなるし、unsafe なコードをきちんと書くのは相当難しい。

kanejaku.org