mrubyのJITのPArrayモジュール

はじめに

twitterでRuby版の深層学習のライブラリがあったらいいねーというツイートがあって、私もあったらいいねーと思ったので深層学習ってどうやるかよく知らないのですが、mrubyのJITでそううのができる枠組みを作ることにしました。
今あるニューラルネットのプログラムをみてみるとベクトルの計算速度がカギになりそうなのでまずはベクトルのライブラリを作ることにしました。ベクトルのライブラリというとRubyではNArrayとかいろいろあるわけですが、mrubyでは知らないしmrubyのJITに対応したものは無いです。
mrubyのJITはJITコンパイラでmrubyの通常の実行も機械語で行うのでベクトルライブラリも機械語を生成するようにすることでインピーダンスミスマッチの無い高速なパフォーマンスが期待できるわけですが、期待というのは裏切るためにありますからね、どうなることか。

PArrayは当然実験システムです。興味のある方はこちらからソースコードが見えます。
https://github.com/miura1729/mruby/tree/parray

PArrayの定義は主にここです
https://github.com/miura1729/mruby/tree/parray/mrbgems/mruby-parray

PArray

　ライブラリの名前はPArrayとしました。Pはparallelの略でSIMD命令の使用を前提としてとにかく高速化することを目指しています。GPUも対応したのですが、Cygwinな環境ではそもそもGPUをいじる環境がない、ような気がします。
PArrayモジュールの中で色々データ構造を定義する予定ですが、今の所4要素のベクターのPVector4しかないです。4要素なのは3Dグラフィックスを考慮(具体的にはaoベンチ)してです。3要素でもいいのですが、SIMD命令を無駄なく使いたいから4要素です。
まだ、途中なのでほとんど作っていません。こんな感じで使います。

parray3.rb

a = PArray::PVector4[1.9, 2.1, 3.9, 4.9]
b = PArray::PVector4[3.2, 2.2, 23.2, 4.1]
1000000.times do |i|
  c = a + b
  d = c + b
  c.move
  d.move
end

この足し算はJITコンパイル時には機械語にインライン展開されます。一応SSEのSIMD命令を使って行います。他のオーバヘッドが大きいからそんなに速いわけではないのですが…
生成されるコードはこんな感じ。4要素の配列の足し算を一度に2要素ずつ行っています。

  movupd(xmm0, ptr [ebx]);
  movupd(xmm1, ptr [reg_tmp1]);
  addpd(xmm0, xmm1);
  movupd(ptr [reg_tmp0], xmm0);

  movupd(xmm0, ptr [ebx + 16]);
  movupd(xmm1, ptr [reg_tmp1 + 16]);
  addpd(xmm0, xmm1);
  movupd(ptr [reg_tmp0 + 16], xmm0);

ただ、せっかくSIMD命令使っているのにオブジェクトのアロケーションで速度を落としたくない。そこで、昔作ったmruby-mmmを引っ張り出しています。mruby-mmmは使わなくなったオブジェクトをキャッシュしておいて必要な時に再利用するライブラリで詳しくはここを見てください。
http://d.hatena.ne.jp/miura1729/20140525/1401018927
サンプルプログラム中のmoveはオブジェクトはもう使わないという宣言でC++のSTD::moveから来ています。

どのくらい速くなるのか、このままでは何も比べるものがありませんから同等のRubyプログラムを用意します。

parray2.rb

class Array
  def add(other)
    Array[
    self[0] + other[0],
    self[1] + other[1],
    self[2] + other[2],
    self[3] + other[3]]
  end

a = Array[1.9, 2.1, 3.9, 4.9]
b = Array[3.2, 2.2, 23.2, 4.1]

1000000.times do 
  c =  a.add(b)
  c =  a.add(b)
end

mrubyのJITでの測定

PArray版
$ time bin/mruby benchmark/parray3.rb

real 0m0.172s
user 0m0.156s
sys 0m0.000s

RubyのArray版
$ time bin/mruby benchmark/parray2.rb

real 0m0.811s
user 0m0.795s
sys 0m0.000s

RubyのArray版　Ruby2.2.2p95での測定
$ ruby -v
ruby 2.2.2p95 (2015-04-13 revision 50295) [i386-cygwin]

$ time ruby benchmark/parray2.rb

real 0m2.122s
user 0m1.981s
sys 0m0.124s

RubyのArray版　オリジナルのmrubyでの測定
$ time ../../src/mruby/bin/mruby benchmark/parray2.rb

real 0m5.694s
user 0m5.491s
sys 0m0.015s

このようにmrubyのJITで比べても配列でベクターを定義するのの5倍くらい速いし、CRubyの15倍くらい速いし、オリジナルのmrubyの30倍強速いということが分かります。イテレータのオーバヘッドとかあるからベクトル計算そのものはもうちょっと速いでしょう。

PArrayの内部構造

PArrayは、こんな感じの定義になります。

#
# PArray - Fast vector/matrix library for number cranching
#
module PArray
  class PVector
  end

  class PVector4
    include MMM

    def initialize(a, b, c, d)
      self[0] = a
      self[1] = b
      self[2] = c
      self[3] = d
    end

    def inspect
      "#<Vec4 [#{self[0]}, #{self[1]}, #{self[2]}, #{self[3]}]>"
    end

    def +(other)
      PVector4[self[0] + other[0], self[1] + other[1], self[2] + other[2], self[3] + other[3]]
    end

    def -(other)
      PVector4[self[0] - other[0], self[1] - other[1], self[2] - other[2], self[3] - other[3]]
    end
  end
end

PVector4っていうのは4要素のベクトルで今の所これの足し算だけを作っています。ちょっと見るとなんてことないRubyのプログラムですが、よく見るとself[0]とかやっていて変です。これはCでオブジェクトのアロケータをカスタマイズしていて配列としてアロケーションしているためです。そして、クラスをPVector4に差し替えるみたいなことをしています。