Scalaでの行列演算の計算速度

2次元配列の乗算をいろいろな方法で書いてみて、計算速度を比較してみました。

背景

機械学習の勉強にニューラルネットワークを自作しているのですが、ニューラルネットワークの処理時間のほとんどを行列の積の計算が占めることを実感できたので、計算を高速化するための検討がこの記事です。

機械学習の実用目的であれば適切なライブラリを使ってGPUなども使えばいいのでしょうが、あくまで勉強目的です。

備考

速度の測定は、JITによる高速化後の安定した状態を測定するために、最初に10回実行してから、そのあとの10回を測定し、その平均をとりました。実際には最初の数回で速度は安定していました。そのあとの速度はどれだけ回数を回しても速度は安定しているのですが、JVM起動ごとに大きく速度が異なり、実行時間に数倍もの差があるときもありました。1回のJVM起動で2つの処理を繰り返して速度比較するのですが、2つとも速くなったり遅いこともあれば、片方だけが極端に遅いこともあり、速度比較が難しかったです。理由はよくわかっていません。なのでこの記事では処理速度の数字は明確には示さずに、JVMを何度も動かしてどの程度どっちが速いことが多かったか、ということだけ書きます。

2次元配列の2つの軸のうちどちらを行列の列とし、どちらを行とするのかは解釈次第ですが、以下のサンプルコードでは、行列の積と考えると、2つの行列で行と列の解釈が逆になっています。行列の積 $ X Y $ ではなく、片方を転置した $ X Y^T $ を計算しているともみなせます。ニューラルネットワークの計算においては、こっちのほうが個人的には解釈しやすかったので、こうしています。

A. 最初の素直なコード

2次元配列を IndexedSeq[IndexedSeq[Double]] で表現して map で演算する素直な方法。

サンプルデータの作成（このコードの実行時間は測定の対象外です）

val a1 = (0 until 100).map { _ =>
  (0 until 100).map { _ =>
    Random.nextDouble();
  }
}
val a2 = (0 until 100).map { _ =>
  (0 until 100).map { _ =>
    Random.nextDouble();
  }
}

乗算処理本体（実行時間の測定対象）

    val z = (0 until 100).map { i =>
      (0 until 100).map { j =>
        (0 until 100).map { k =>
          a1(i)(k) * a2(j)(k);
        }.sum;
      }
    }

B. 入力側は IndexedSeq の代わりに Array

IndexedSeq[IndexedSeq[Double]] の代わりにJavaの配列である Array[Array[Double]] を使う方法。演算時は手抜きで IndexedSeq の map を使ってから toArray でJavaの配列に変換しています。

サンプルデータの作成（このコードの実行時間は測定の対象外です）

val b1 = (0 until 100).map { _ =>
  (0 until 100).map { _ =>
    Random.nextDouble();
  }.toArray;
}.toArray;
val b2 = (0 until 100).map { _ =>
  (0 until 100).map { _ =>
    Random.nextDouble();
  }.toArray;
}.toArray;

乗算処理本体（実行時間の測定対象）

    val z = (0 until 100).map { i =>
      (0 until 100).map { j =>
        (0 until 100).map { k =>
          b1(i)(k) * b2(j)(k);
        }.sum;
      }.toArray;
    }.toArray;

AとBを比較すると、1割から3割ほどBのほうが速いことが多かったです（たまにAのほうが速かった）。

Bのほうがいちいち toArray で変換する処理があるにも関わらずそれでもBのほうが速いということは、インデックスでの要素へのアクセスが IndexedSeq よりも Array のほうが速いということのようです。次のCの比較でもそれが裏付けられます。

C. 配列アクセスの回数削減

配列の外側のインデックスへのアクセスを外側にまとめることで、アクセス回数を減らした方法。

    val z = (0 until 100).map { i =>
      val b1i = b1(i);
      (0 until 100).map { j =>
        val b2j = b2(j);
        (0 until 100).map { k =>
          b1i(k) * b2j(k);
        }.sum;
      }.toArray;
    }.toArray;

BとCを比較して、どっちが速いとは言えず、ほとんど同じでした。

Cのようにソースコードに明示的に書かなくてもBのままでもコンパイラまたは実行時にJVMがいい具合に勝手に最適化してくれているということでしょうか。

ちなみにAの IndexedSeq に対してB->Cの改良と同じことをしたところ、A->Bと同じくらい速くなりました。このことからBの最後に書いた通り、インデックスでの要素へのアクセスが Array はほとんど時間がかからず、 IndexedSeq は一定の時間がかかるということになります。

D. 計算結果を直接配列に書き込み

計算結果を IndexedSeq ではなく初めから Array にして構築する方法。

    val z = new Array[Array[Double]](100);
    (0 until 100).foreach { i =>
      z(i) = new Array[Double](100);
      (0 until 100).foreach { j =>
        z(i)(j) = (0 until 100).map { k =>
          b1(i)(k) * b2(j)(k);
        }.sum;
      }
    }

BとDを比較して、Dのほうが0～2割ほどDのほうが速いことが多かったです（たまにBのほうが速かった）。

単純に toArray の処理を省略できているからでしょうか。

E. sum を使わずに直接加算

合計の計算を IndexedSeq の sum ではなく、直接自力で加算していく方法。

    val z = new Array[Array[Double]](100);
    (0 until 100).foreach { i =>
      z(i) = new Array[Double](100);
      (0 until 100).foreach { j =>
        var s: Double = 0.0;
        (0 until 100).foreach { k =>
          s = s + b1(i)(k) * b2(j)(k);
        }
        z(i)(j) = s;
      }
    }

DとEを比較して、Eが圧倒的に速く、Eの処理時間がDの20分の1から30分の1でした。 IndexedSeq の map して sum だと、 IndexedSeq の中に Double のオブジェクトを保存するためのboxingが必要になり、boxingにかかる処理時間がほとんどということでしょうか。

F. foreach の代わりに while

一番内側のループを foreach ではなくJavaっぽい while にする方法。

    val z = new Array[Array[Double]](100);
    (0 until 100).foreach { i =>
      z(i) = new Array[Double](100);
      (0 until 100).foreach { j =>
        var s: Double = 0.0;
        var k = 0;
        while (k < 100) {
          s = s + b1(i)(k) * b2(j)(k);
          k = k + 1;
        }
        z(i)(j) = s;
      }
    }

EとFを比較して、どっちが速いとは言えず、ほとんど同じでした。

foreach のほうが高階関数の呼び出しになりオーバーヘッドが大きいイメージがあったのですが、意外とそうでもないようです。

G. 配列の配列の代わりに配列

2次元配列を Array[Array[Double]] の代わりに1次元の Array[Double] を使う方法。

val g1: Array[Double] = (0 until 100).flatMap { _ =>
  (0 until 100).map { _ =>
    Random.nextDouble();
  }
}.toArray;
val g2: Array[Double] = (0 until 100).flatMap { _ =>
  (0 until 100).map { _ =>
    Random.nextDouble();
  }
}.toArray;

    val z = new Array[Double](100 * 100);
    (0 until 100).foreach { i =>
      (0 until 100).foreach { j =>
        var s: Double = 0.0;
        var k = 0;
        while (k < 100) {
          s = s + g1(100 * i + k) * g2(100 * j + k);
          k = k + 1;
        }
        z(100 * i + j) = s;
      }
    }

EとGを比較して、Eのほうが1割ぐらい速かったです。

Array[Array[Double]] は配列オブジェクトの配列となって1次元の配列に比べてオーバーヘッドが多少あるかなと思ってたのですが、気にしなくてもいいようです。

H. インデックス計算の重複を共通化

Gを改良して、インデックスの計算を一部キャッシュしておく方法。

    val z = new Array[Double](100 * 100);
    (0 until 100).foreach { i =>
      val i2 = 100 * i;
      val i3 = 100 * i;
      (0 until 100).foreach { j =>
        val j2 = 100 * j;
        var s: Double = 0.0;
        var k = 0;
        while (k < 100) {
          s = s + g1(i2 + k) * g2(j2 + k);
          k = k + 1;
        }
        z(i3 + j) = s;
      }
    }

EとHを比較して、Gのときと同じく、Eのほうが1割ぐらい速かったです。

インデックスの計算は気にすることではなく、配列オブジェクトの配列というオーバーヘッドはやはり気にしなくていいようです。

I. ループ展開

Fを改良して、一番内側のループの中を4倍に展開し、ループの回数を4分の1にする方法。（ループアンローリング）

    val z = new Array[Array[Double]](100);
    (0 until 100).foreach { i =>
      z(i) = new Array[Double](100);
      (0 until 100).foreach { j =>
        var s: Double = 0.0;
        var k = 0;
        while (k < 100) {
          s = s + b1(i)(k    ) * b2(j)(k    );
          s = s + b1(i)(k + 1) * b2(j)(k + 1);
          s = s + b1(i)(k + 2) * b2(j)(k + 2);
          s = s + b1(i)(k + 3) * b2(j)(k + 3);
          k = k + 4;
        }
        z(i)(j) = s;
      }
    }

EとIを比較して、どっちが速いとは言えず、ほとんど同じでした。

Iのほうがジャンプの回数が減る分速くなるかなと思ってましたが、このぐらいのループ展開は実行時にJVMが勝手にやってるのかもしれません。

J. ループ内の依存性解消

Iを改良して、ループ展開に合わせて合計を保存する変数も4つに分割する方法。

    val z = new Array[Array[Double]](100);
    (0 until 100).foreach { i =>
      z(i) = new Array[Double](100);
      (0 until 100).foreach { j =>
        var s0: Double = 0.0;
        var s1: Double = 0.0;
        var s2: Double = 0.0;
        var s3: Double = 0.0;
        var k = 0;
        while (k < 100) {
          s0 = s0 + b1(i)(k    ) * b2(j)(k    );
          s1 = s1 + b1(i)(k + 1) * b2(j)(k + 1);
          s2 = s2 + b1(i)(k + 2) * b2(j)(k + 2);
          s3 = s3 + b1(i)(k + 3) * b2(j)(k + 3);
          k = k + 4;
        }
        z(i)(j) = s0 + s1 + s2 + s3;
      }
    }

EとJを比較して、Iのほうが3割から6割ぐらい速かったです。

ループの中の4つの演算に依存性がなくなったので、CPUのパイプライン処理が効率化されたのかもしれません。もしくはさらにSIMDの命令セットが使われているのかもしれません。JVMがどの程度SIMDを使いこなしているのかはわかりませんが。

K. Double の代わりに Float

Double の代わりに Float を使用する方法。

val k1: Array[Array[Float]] = (0 until 100).map { _ =>
  (0 until 100).map { _ =>
    Random.nextFloat();
  }.toArray;
}.toArray;
val k2: Array[Array[Float]] = (0 until 100).map { _ =>
  (0 until 100).map { _ =>
    Random.nextFloat();
  }.toArray;
}.toArray;

    val z = new Array[Array[Float]](100);
    (0 until 100).foreach { i =>
      z(i) = new Array[Float](100);
      (0 until 100).foreach { j =>
        var s0: Float = 0.0f;
        var s1: Float = 0.0f;
        var s2: Float = 0.0f;
        var s3: Float = 0.0f;
        var k = 0;
        while (k < 100) {
          s0 = s0 + k1(i)(k    ) * k2(j)(k    );
          s1 = s1 + k1(i)(k + 1) * k2(j)(k + 1);
          s2 = s2 + k1(i)(k + 2) * k2(j)(k + 2);
          s3 = s3 + k1(i)(k + 3) * k2(j)(k + 3);
          k = k + 4;
        }
        z(i)(j) = s0 + s1 + s2 + s3;
      }
    }

JとKを比較して、どっちが速いとは言えず、ほとんど同じでした。

ニューラルネットワークの計算では高い精度は不要なので、floatでも十分なのですが、JVMでCPUで計算している限りはfloatでもdoubleでも処理速度にあまり差はないのかもしれません。よくわかってないです。

ちなみにintでもあまり差はありませんでした。

まとめ

A -> B -> D -> E -> J と改良していき、おおざっぱに100倍ぐらい速くなりました。

• B. 入力側は IndexedSeq の代わりに Array
• D. 計算結果を直接 Array に書き込み
• E. sum を使わずに直接加算
• I., J. ループ展開してさらにループ内の依存性解消

計算にかかる時間は、Aは50ms～200msぐらい、Jは0.5msから2msぐらいのイメージでした。