Ruby製の自作WASM Runtimeの実行時間を43%削った話

Ruby Advent Calendar 2024の20日目の記事です。

qiita.com

今年後半からWarditeという名前のPure Ruby WASM Runtimeを作り始めたのですが、

github.com

udzura.hatenablog.jp

その内部の話とかは今日は置いといて（東京Ruby会議12前夜祭で話しちゃおっかな）、今回、Warditeを作る上で計測しつつ高速化した話をします。と言いつつ、めちゃくちゃ基本的なことをしただけです。

なんとかgrayscale.wasmが完走したが遅い

以前grayscale処理を完走させた話をしました*1。

udzura.hatenablog.jp

しかし、最初の実装では、YJIT有効でも、 150x150 のPNGファイルのグレースケール処理に14.43sが必要でした（コマンドを起動して、パースして...の処理全体での計測）。

bundle exec ruby --yjit tmp/grayscale-cmd.rb --wasm-file  --source  --dest     14.43s user 0.06s system 99% cpu 14.533 total

ひとまずボトルネックを探りたいな〜、このレベルで遅いと何か基本的なことを見逃しているのかも？と思いました。

ruby-profでシュッと計測

まずは定番ではありますがruby-profを用いてみました。このgemは導入が簡単で、Gemfileに以下を記述して bundle installで導入できます。

gem "ruby-prof", require: false

その後計測スクリプトに以下のように追記するだけです。

require 'ruby-prof'
profile = RubyProf::Profile.new

# 関数呼び出しの前後で計測する
profile.start
offset = instance.runtime.grayscale($options.width, $options.height, start, data_url.size)
result = profile.stop

# ...処理の最後に出力
printer = RubyProf::GraphPrinter.new(result)
printer.print(STDOUT, {})

ということで、これで計測した結果の抜粋を掲載します*2。

------------------------------------------------------------------------------------------------------------------------------------------------------
                     54.539      9.845      0.000     44.69413069318/13069318     Kernel#loop
  73.63%  13.29%     54.539      9.845      0.000     44.694         13069318     Wardite::Runtime#eval_insn     /Users/udzura/ghq/github.com/udzura/wardite/lib/wardi
te.rb:420
                     19.493      0.024      0.000     19.469      95886/95886     Wardite::Runtime#fetch_ops_while_end
                     15.638      6.483      0.000      9.156  5225913/5225913     <Module::Wardite::Evaluator>#i32_eval_insn
                      3.330      1.238      0.000      2.093  1155055/1155055     Wardite::Runtime#do_branch
                      0.773      0.773      0.000      0.00013069318/13069318     Wardite::Op#namespace
                      0.757      0.757      0.000      0.00012631671/73174992     Array#[]
                      0.749      0.749      0.000      0.00015275900/53340046     BasicObject#!
                      0.542      0.542      0.000      0.00010109073/23362733     Wardite::Runtime#stack

#fetch_ops_while_end という関数がまあまあ支配的、という結果が出ました。

で、この #fetch_ops_while_end が何だったかというと、もともと、WebAssemblyの :block/:loop/:if といった命令を評価した際に、対応する :end の位置を命令列をたどって計算するというメソッドでした。

      when :if
        block = insn.operand[0]
        raise EvalError, "if op without block" if !block.is_a?(Block)
        cond = stack.pop 
        raise EvalError, "cond not found" if !cond.is_a?(I32)
        next_pc = fetch_ops_while_end(frame.body, frame.pc)

github.com

例えばifから抜ける際にそのendまでジャンプする必要があるので、その位置を計算してそこに飛ばす必要があるわけです。

問題はこのendの位置を、毎回 :block/:loop/:if に入るたびにいちいち計算していたということでした。この状況ではブロックやifのネストが深かったりする時、とんでもなく遅くなることが考えられます。また、何度も呼ばれる関数にこれらの命令があると無駄に毎回計算することになるはず。

この対応するendの位置は、パースして命令列が決まってしまえば変わることがないものです。そこで、一度命令列をパースしてから、もう一度命令列を検査して、 :block/:loop/:if ごとに対応するendの位置を事前計算してキャッシュするように変更しました。

    def revisit!
      @ops.each_with_index do |op, idx|
        case op.code
        when :block
          next_pc = fetch_ops_while_end(idx)
          op.meta[:end_pos] = next_pc

        when :loop
          next_pc = fetch_ops_while_end(idx)
          op.meta[:end_pos] = next_pc

        when :if
          next_pc = fetch_ops_while_end(idx)
          else_pc = fetch_ops_while_else_or_end(idx)
          op.meta[:end_pos] = next_pc
          op.meta[:else_pos] = else_pc
        end
      end
    end

詳しくはP/Rを参照してください。

github.com

これで毎回の再計算を防ぐことができました。

再計測

上記の対応を入れたところ、計測結果が8秒台前半になりました。

bundle exec ruby --yjit tmp/grayscale-cmd.rb --wasm-file  --source  --dest     8.17s user 0.07s system 99% cpu 8.300 total

1 - 8.17/14.43 = 0.4338 ということで素朴に解釈すれば実行時間の43%削減を実現できたということになります。

その後、もう一箇所怪しかった、 関数に入る際にローカル変数を毎回ゼロから作っていた箇所を、事前生成するように変えた ところ、どうやら安定して8秒を切るようになりました。

bundle exec ruby --yjit examples/grayscale.rb --wasm-file  --source  --dest    7.84s user 0.06s system 99% cpu 7.962 total

とはいえ、ねえ...

まだまだ普通に遅いと思うのでもっと頑張らねばというところですね。

ちなみに、以下は perf を用いて支配的な関数をreportした結果の抜粋です。

# Children      Self  Command  Shared Object          Symbol 
# ........  ........  .......  .....................  ....................................................
#
   100.00%     0.00%  ruby     ruby                   [.] _start
            |
            ---_start
               __libc_start_main
               0xffffbd7173fc
               main
               ruby_run_node
               rb_ec_exec_node
               rb_vm_exec
               |          
               |--98.93%--vm_exec_core
               |          |          
               |          |--9.67%--0xffffbe152a58
               |          |          |          
               |          |          |--7.26%--rb_vm_set_ivar_id
               |          |          |          |          
               |          |          |          |--2.86%--rb_shape_get_iv_index
               |          |          |          |          
               |          |          |          |--1.91%--rb_shape_get_shape_by_id@plt
               |          |          |          |          
               |          |          |           --0.94%--rb_shape_get_shape_by_id
               |          |          |          
               |          |          |--0.81%--rb_shape_get_shape_by_id@plt
               |          |          |          
               |          |           --0.67%--rb_shape_get_shape_by_id
               |          |          
               |          |--7.32%--0xffffbe1525cc
               |          |          |          
               |          |           --6.62%--rb_class_new_instance_pass_kw
               |          |                     |          
               |          |                     |--2.52%--vm_call0_cc
               |          |                     |          |          
               |          |                     |          |--1.51%--vm_call0_body
               |          |                     |          |          |          
               |          |                     |          |           --0.81%--vm_call_iseq_setup
               |          |                     |          |          
               |          |                     |           --0.67%--rb_vm_exec
               |          |                     |          
               |          |                     |--2.22%--rb_class_allocate_instance
               |          |                     |          |          
               |          |                     |           --2.15%--newobj_alloc
               |          |                     |                     |          
               |          |                     |                      --1.58%--gc_continue
               |          |                     |                                |          
               |          |                     |                                 --1.41%--gc_sweep_step
               |          |                     |                                           |          
               |          |                     |                                            --0.94%--obj_free
               |          |                     |          
               |          |                      --0.50%--rb_call0
               |          |          
               |          |--5.78%--0xffffbe14a65c
               |          |          |          
               |          |           --4.90%--rb_class_new_instance_pass_kw
               |          |                     |          
               |          |                     |--1.95%--rb_class_allocate_instance
               |          |                     |          newobj_alloc
               |          |                     |          |          
               |          |                     |           --1.38%--gc_continue
               |          |                     |                     |          
               |          |                     |                      --1.24%--gc_sweep_step
               |          |                     |                                |          
               |          |                     |                                 --0.74%--obj_free
               |          |                     |          
               |          |                      --1.81%--vm_call0_cc
               |          |                                |          
               |          |                                |--0.94%--vm_call0_body
               |          |                                |          |          
               |          |                                |           --0.57%--vm_call_iseq_setup
               |          |                                |          
               |          |                                 --0.74%--rb_vm_exec
...

正直めちゃくちゃ効果がありそうな箇所は見えてはいないのですが、まず、 rb_vm_set_ivar_id() 、おそらくインスタンス変数へのセットが上に来ています。あと目立つのは rb_class_new_instance_pass_kw() など、オブジェクトを作っているところでしょうか。

ところでこのgrayscale処理の中身を考えると、大きくいうとbase64のencode/decodeと（dataUrlを扱うため）、PNGを処理するためのdeflateの処理の時間が支配的だと思われます*3。

つまり、数値計算が多いはずです。

実際、WASMの内部の数値型に対応するオブジェクト（I32, I64, F32, F64）が、一回の処理でどれくらい生成されるかを計測してみました。

ちなみに計測には TracePoint を使いました。

$COUNTER = {}

TracePoint.trace(:call) do |tp|
  # こういう名前のメソッドを作って初期化しているので、名前を追えばわかる
  if %i(I32 I64 F32 F64).include?(tp.method_id)
    $COUNTER[tp.method_id] ||= 0
    $COUNTER[tp.method_id] += 1
  end
end

END {
  pp $COUNTER
}

結果は、150x150 のPNGの処理で以下のような感じでした。

{:I32=>18847048, :I64=>1710695, :F32=>247501, :F64=>2}

1880万回も I32 クラスのオブジェクトが生成されているということがわかります。そう考えると、 I32 クラスなどを作るときには、当然 値をインスタンス変数にセット しているので、 rb_vm_set_ivar_id() が上がってくることにも関係していそうです。

このあたりの箇所について、オブジェクトを介さずにRubyの Integer で直接計算処理を扱うなどでチューニングができるかもしれません。そのためには内部設計に大きく手を加えることになりそうなのですが。

---

ということで、Wardite高速化坂を登り始めたばかりだからな... と締めておきます。

環境など

最後に検証した環境情報を残しておきます。

• Apple M3 Pro MacOS 14.4 (perfは、この上にmultipassでUbuntu 24.04のインスタンスを作成)
• ruby 3.3.4 w/ YJIT (3.3.6 だともっと早い？)
• WASMは rustc 1.80.1 で生成