Julia で WebAssembly できるかな?


本日は

  • Julia で WebAssembly やりましょうという話です。要するにWebブラウザでJuliaのコードを動かしてJSと連携しようぜ。という話です。WebAssembly へのとりくみは Julia 以外ではC/C++そしてRustでもおこなわれています。Rustの入門書でも WebAssembly を取り扱う和書がでていますのでそっちをとりくんだほうが巨人の肩に乗れて幸せになれる人がおおいかもです。

Julia の青春の夢

  • さて、じつは Keno さんという Julia の開発に積極的に取り組まれている方がいましてね。

このリポジトリの README.md にある

にアクセスするとですね。

ほらご覧の通り動くんです。すごいでしょ?Julia1.3ですからね。去年にはできてたっぽいんです(震え声)。リポジトリおとして emcc とか導入してよなよなLLVMをビルドしてですねまぁ再現しようと思ったんですが。できませんでした。

ここでは

ここでは MikeInnes/WebAssembly.jl をつかって Julia の関数コンパイルし wasm で呼べるようにしてみましょう。MikeInnes さんは Flux.jl や MacroTools.jl など開発している方で、マクロの取扱い周りのツール、パッケージを多く開発しています。(数ヶ月前はバカンスに行ってたらしい。いいなぁ)

クイックスタートその1

  • wasm_with_julia を使ってみましょう。これは私が例として作成したリポジトリです. リポジトリのトップには すでに生成済みの twice.wasm というファイルと index.html があるはずです。まずはリポジトリをクローンしてHTTPサーバーをたててみましょう。
$ git clone https://github.com/terasakisatoshi/wasm_with_julia
$ cd wasm_with_julia
$ julia -e 'using Pkg; Pkg.add("LiveServer"); using LiveServer; serve()'
✓ LiveServer listening on http://localhost:8000/ ...
  (use CTRL+C to shut down)

localhost:8000 にアクセスするとでかでかと twice(3)=6 という文字が出ているはずです。(注意:index.html をクリックしてブラウザで閲覧するだと期待した動作は得られません。何かしらの方法でサーバーを立ててください)

const num = 3.0; の部分を変更すると対応する数値の2倍の値がえられるようになります。(いまはLiveServer.jlでサーバーを立てているため)

index.html
<!-- 一部抜粋 -->

.then(instance => {   
            console.log(instance); 
            const num = 3.0;
            let result = instance.exports.twice(num); 
            let msg = `twice(${num}) = ${result}`;
            document.getElementById("textcontent").innerHTML = msg;
         });

クイックスタートその2

あたえられた wasm を使って雰囲気をつかむことができました。では手元のPCで作ってみましょう(動作はLinuxとMacで確認済み)。
下記のようにして依存パッケージをついかします。

じゅんび

$ julia
julia> ENV["PYTHON"]=Sys.which("python3")
julia> using Pkg
julia> pkg"add WebAssembly#master"
julia> Pkg.add(["IRTools", "PyCall"])

pkg> juladd WebAssembly#master はつぎのようにGitのコミットハッシュ値を明示したほうが親切かもしれませんね。まぁ、アップデート最近されてないんでそこまで厳密にする必要はないと思いますけど念の為。

julia> using Pkg
julia> Pkg.add(url="https://github.com/MikeInnes/WebAssembly.jl", rev="d9260be1f5ef6c3f6fa5aa37152adde1c15adec9")

あとは wasmer で 生成した wasm ファイルを Python から呼び出せるようにします(動作テストを兼ねるという意味で)

pip3 install wasmer

うごかす

webassembly.jl を走らせればOKです。

$ rm twice.wasm
$ julia webassembly.jl
(module
 (type $0 (func (param f64) (result f64)))
 (export "twice" (func $0))
 (func $0 (; 0 ;) (type $0) (param $0 f64) (result f64)
  (f64.add
   (local.get $0)
   (local.get $0)
  )
 )
)

$ file twice.wasm # できたもの
twice.wasm: WebAssembly (wasm) binary module version 0x1 (MVP)

あとはクイックスタート1の方法に従ってWebサーバーを起動して見ましょう。

なにをやってるんだ?

とりあえず動いたことは確認できたでしょう。ここでは実際に何をしているかってのをみていきます。Julia のコードを IRTools.jl というものを使って一種の中間表現(IR)に直します。その中間表現をさらに WebAssembly.jl の関数に投げて wasm に変換するという流れになっています。

じつはこの webassembly.jl というファイルは jupytext によって ipynb のノートブックを .jl にしたものなので Jupyter Notebook で閲覧することができます。

$ pip3 install jupytext jupyter
$ jupyter notebook

さて、ノートブックの中身を見ていきましょう。下記のように依存する Julia パッケージをロードしていることがわかります。

using WebAssembly
using WebAssembly: f64, irfunc # i32 and so on
using IRTools.All
using PyCall

twice 関数を作っていきましょう。実際には x という名前の仮引数を入力として x + x を返すものです。この例は IRTools.jl のドキュメント から取ってきました。

f を Julia のREPLで定義し code_ir マクロを使って IR を生成します。

julia> using IRTools
julia> f(x) = x + x # define function
julia> IRTools.@code_ir f(3.0)
1: (%1, %2)
  %3 = %2 + %2
  return %3

初見だとわかりにくいですが %1 は関数 f そのものを指していて %2, %3 は関数の中身に出てくる変数を機械の世界に近い形に変換したものです。

WebAssembly.jl に投げるためには型の情報も込めて上記のようなIRを作る必要があります。code_ir で吐かれた情報をもとに式を作っていきます。手作りのぬくもりって大事ですよね。

twice = let 
    ir = IR()
    x = argument!(ir, f64)
    r = push!(ir, stmt(xcall(f64.add, x, x), type=f64))
    return!(ir, r)
end

今回はすごくシンプルな例だったので人間のぬくもりをつかってかけましたがループや分岐が入ると若干ふくざつになります。たとえば次のように定義した $x$ の $n$ 乗を計算する関数はどのようにIRやWASMのためにIRを書くべきでしょうか?

function pow(x, n) # A simple Julia function
         r = 1
         while n > 0
           n -= 1
           r *= x
         end
         return r
end

こたえは下記のリポジトリのREADMEやテストコードを参考にしてください。

https://github.com/FluxML/IRTools.jl
https://github.com/MikeInnes/WebAssembly.jl

さて、作ったあとは下記のように twice という名前で呼び出せるようにしていきます。

funcname = :twice
func = irfunc(funcname, twice)

mod = WebAssembly.Module(
    funcs=[func],
    exports=[WebAssembly.Export(funcname, funcname, :func)]
)
WebAssembly.binary(mod, "$(funcname).wasm")

上記のコードで twice.wasm が得られました。

disassembleすることで中身を眺めることもできます。

run(`$(WebAssembly.Binaryen.wasm_dis) $(funcname).wasm`);
 (type $0 (func (param f64) (result f64)))
 (export "twice" (func $0))
 (func $0 (; 0 ;) (type $0) (param $0 f64) (result f64)
  (f64.add
   (local.get $0)
   (local.get $0)
  )
 )
)

twice.wasm の動作確認のために Python から呼び出すこともできます。

pycode = """
from wasmer import Instance

def run_wasm(*args):
    wasm_bytes = open("$(funcname).wasm", 'rb').read()
    instance = Instance(wasm_bytes)
    res = instance.exports.$(funcname)(*args)
    return res
"""

pyfile = string(funcname) * ".py"
open(pyfile,"w") do f
    print(f, pycode)
end

pushfirst!(PyVector(pyimport("sys")."path"), @__DIR__)
pymodule = pyimport(funcname)
pymodule.run_wasm(3.) # 3. + 3. should be 6.0

こんな感じです。IR生成のExampleは

にあります。

ぬくもりの自動化について

手で IR 書くのたぶんしんどいので Mjolnir.jl を使って構築する議論がされています。が、完全ではないようです。Mjolnir ってMjölnirのことで、銀河英雄伝説のトールハンマー的な位置づけらしいです。これもマクロ・コンパイラー周りのツールになっているようです。

https://github.com/MikeInnes/WebAssembly.jl/issues/27

Appenix

私はWebAssemblyもコンパイラーも専門領域じゃないのでこれ以上深入りできないのですが、

JuliaCon 2020 | Advanced Metaprogramming Tools | Mike Innes

をみるともっと発展的なことが述べられています。

IRTools.jl 自体が FluxML 配下にあるので 機械学習フレームワークFluxのなかでつかいたいモチベーションがあるはずです。上記のYouTube動画でもコンパイラーとDeepLearningレイヤーの話も軽く触れられています。

ここらへんの領域に興味がある人はぜひ取り組んでみてはいかがでしょうか?