RubyとApache Arrowの開発に参加している須藤です。RubyのMemoryViewの使い方がわからないという話を聞いてsumをするサンプルを作ったので紹介します。

MemoryView

RubyのMemoryViewは、ざっくり言うと、同じ型の値が連続して並んでいるデータをゼロコピーで交換するための仕組みです。まじめな説明はMemoryView: Ruby 3.0 から導入される数値配列のライブラリ間共有のための仕組みなどを参照してください。

MemoryViewを使うと、異なるライブラリー間で効率よくデータを交換したり、データを高速に処理したりできます。今回はデータを高速に処理する方の使い方を紹介します。

Rubyはすべてがオブジェクトなので1もオブジェクトです。CRubyは1はCの1としては表現していません。そのため、1 + 2は「Rubyの1をCの1に変換したもの」と「Rubyの2をCの2に変換したもの」をCの+で足し算しています。そして、足した結果であるCの3をRubyの3に変換します。

MemoryViewで扱うデータは1をCの1として表現しています。そのため、RubyとCの変換をすることなく1 + 2を計算できます。このやり方を紹介します。

MemoryViewでsum

MemoryViewはCのAPIなので拡張ライブラリー内でMemoryViewを使うことになります。今回は[1, 2, 3]のような数値の配列データのsumを求める処理を実装します。[1, 2, 3].sumと同じ挙動をするものです。

まず、必要なヘッダーファイルをincludeします。

#include <ruby.h>
#include <ruby/memory_view.h>

拡張ライブラリーの初期化関数を作ります。今回はMemoryViewSample.#sumを作ります。

void
Init_memory_view_sample(void)
{
  VALUE rb_mMemoryViewSample = rb_define_module("MemoryViewSample");
  rb_define_module_function(rb_mMemoryViewSample, "sum", mvs_sum, 1);
}

それでは、本体のmvs_sum()を実装します。

まず、rb_memory_view_get()で対象オブジェクトからMemoryViewをエクスポートします。エクスポートできない場合はfalseが返るので、戻り値で結果を判断できます。

static VALUE
mvs_sum(VALUE module, VALUE numbers) {
  rb_memory_view_t view = {0};
  if (!rb_memory_view_get(numbers, &view, 0)) {
    rb_raise(rb_eArgError, "Unable to get a memory view: %+"PRIsVALUE, numbers);
  }
  // ...
}

今回はint64_tの配列データのみを対象にします。rb_memory_view_t::formatにデータのフォーマットがArray#pack/String#unpack互換の文字列で表現されているので、そこを確認することで型を判断できます。なお、int64_tは"q"になります。

  if (!view.format) {
    rb_memory_view_release(&view);
    rb_raise(rb_eArgError, "Unknown format: %+"PRIsVALUE, numbers);
  }
  if (strcmp(view.format, "q") != 0) {
    rb_memory_view_release(&view);
    rb_raise(rb_eArgError, "int64_t format is only supported: <%s>", view.format);
  }

これで期待するデータが入っているMemoryViewであることを確認できたので（本当は次元数とかもチェックしないといけないけどね）、あとはこのデータを使ってsumを計算するだけです。データはrb_memory_view_t::dataに入っているので、それをconst int64_t *にキャストして各要素を取り出します。要素数はrb_memory_view_t::byte_size（総データサイズ）とrb_memory_view_t::item_size（各要素のデータサイズ）を使えば計算できます。

  size_t n = view.byte_size / view.item_size;
  const int64_t *raw_numbers = (const int64_t *)(view.data);
  int64_t sum = 0;
  size_t i;
  for (i = 0; i < n; i++) {
    sum += raw_numbers[i];
  }

sumの計算にRubyとCの変換が入っていませんね。これにより高速にデータ処理できます。

最後に計算結果をRubyのオブジェクトとして返しておしまいです。

  rb_memory_view_release(&view);
  return LL2NUM(sum);

計測

それでは、本当にMemoryViewを使うと高速になるのか計測してみましょう。

ここでは、Red Arrow（Apache ArrowのRuby実装）を使ってMemoryView対応オブジェクトを作ります。Red ArrowはMemoryViewに対応しているのです。

require "benchmark"
require "arrow"
require "memory_view_sample"

raw_numbers = 128000.times.to_a
arrow_numbers = Arrow::Int64Array.new(raw_numbers)

GC.disable

Benchmark.bmbm do |x|
  x.report("Array#sum") do
    raw_numbers.sum
  end
  x.report("MemoryViewSample.sum") do
    MemoryViewSample.sum(arrow_numbers)
  end
end

手元のマシンでは次のような結果になりました。Array#sumよりも2倍くらい速いですね。ちなみに、Array#sumはEnumerable#sumを使っていなくて、Array用に最適化された実装になっています。それよりも速いです。

Rehearsal --------------------------------------------------------
Array#sum              0.000120   0.000013   0.000133 (  0.000130)
MemoryViewSample.sum   0.000133   0.000015   0.000148 (  0.000148)
----------------------------------------------- total: 0.000281sec

                           user     system      total        real
Array#sum              0.000131   0.000000   0.000131 (  0.000127)
MemoryViewSample.sum   0.000067   0.000000   0.000067 (  0.000063)

速さの秘密

RubyとCの変換がなくなっただけで2倍近くも差がつくものでしょうか。つくかもしれないしつかないかもしれない。。。どうなんだろう。実は、今回のサンプルにはもう一つ工夫がしてありました。sumの実装は次のようにシンプルな実装で、ここに高速化のための特別な秘密はありません。

  size_t n = view.byte_size / view.item_size;
  const int64_t *raw_numbers = (const int64_t *)(view.data);
  int64_t sum = 0;
  size_t i;
  for (i = 0; i < n; i++) {
    sum += raw_numbers[i];
  }

秘密はビルドオプションの方にあります。extconf.rbを次のようにして-O3 -march=nativeでビルドしていました。

require "mkmf"
$CFLAGS << " -O3 -march=native"
create_makefile("memory_view_sample")

このオプションを指定すると、対象のマシンで最速になるようにビルドしようとしてくれます。たとえば、今回の実装のように明示的にSIMDを使っていなくても、SIMDを使える場合はSIMDを使ったバイナリーを生成してくれます。

一方、次のように-O0を指定すると最適化してくれません。-O0の場合は、手元では次のような結果になりました。Array#sumより遅かったです。

Rehearsal --------------------------------------------------------
Array#sum              0.000123   0.000015   0.000138 (  0.000134)
MemoryViewSample.sum   0.000231   0.000027   0.000258 (  0.000259)
----------------------------------------------- total: 0.000396sec

                           user     system      total        real
Array#sum              0.000121   0.000014   0.000135 (  0.000131)
MemoryViewSample.sum   0.000180   0.000000   0.000180 (  0.000176)

Apache Arrowのsum実装との比較

実はApache Arrowは高速なデータ処理機能も含まれています。今回の実装とApache Arrow実装の速度も比べてみましょう。

Benchmark.bmbm do |x|
  x.report("Array#sum") do
    raw_numbers.sum
  end
  x.report("MemoryViewSample.sum") do
    MemoryViewSample.sum(arrow_numbers)
  end
  x.report("Arrow::Int64Array#sum") do
    arrow_numbers.sum
  end
end

手元での結果はこんな感じでArray#sumより少し速いくらいでした。。。マジかよ。。。sumの処理じゃないところのオーバーヘッドが大きいのでしょう。。。あとでなんとかしよう。。。

Rehearsal ---------------------------------------------------------
Array#sum               0.000124   0.000010   0.000134 (  0.000131)
MemoryViewSample.sum    0.000122   0.000010   0.000132 (  0.000132)
Arrow::Int64Array#sum   0.002768   0.000000   0.002768 (  0.002769)
------------------------------------------------ total: 0.003034sec

                            user     system      total        real
Array#sum               0.000128   0.000000   0.000128 (  0.000125)
MemoryViewSample.sum    0.000064   0.000000   0.000064 (  0.000062)
Arrow::Int64Array#sum   0.000114   0.000000   0.000114 (  0.000111)

まとめ

RubyのMemoryViewを使ってデータを高速に処理するサンプルを紹介しました。ファイル一式は https://gitlab.com/ktou/memory-view-sample にあります。

個別に質問されたやつだったのですが、個別に回答するだけなのはもったいないという気持ちになったので、ここで供養しました。

ツイートする

フォローする