R
for

forってfunctionでwrapすると速くなんの?

昨日偉そうなことを呟いてしまったので、実際どれくらいRのforを最適化できるだろうかと今朝、電車の中で遊んでいた。

とりあえずcolSumsと比較するぞー!
と思い、テストコードを書いてmicrobenchmarkする。

for高速化の基本として、メモリの動的確保を避け、結果を格納するオブジェクト(y)を事前に用意する。
更に関数呼び出しや代入は少なければ少ないほどいいので、結果を格納するオブジェクト(y)をインデックスとしても利用することにした。
しかしここで不思議な事態が発生。
なんとforをfunctionに入れると速くなるのだ。

2017/2/1追記: どうやらループ回数が少ない時は、特にJITコンパイルのオーバーヘッドが、計算時間の差に現れるようです(後述)。

set.seed(123)
x <- matrix(runif(10000), 10, 1000)

microbenchmark::microbenchmark(
  `naked` = {
    y <- 1:ncol(x)
    for(i in y) y[i] <- sum(x[, i])
    y
  },
  `wrapped` = {
    (function(x) {
      y <- 1:ncol(x)
      for(i in y) y[i] <- sum(x[, i])
      y
    })(x)
  }
)

expr min lq mean median uq max neval
naked 1492.253 1617.5290 1901.4042 1706.717 1777.8755 4297.369 100
wrapped 431.556 490.3475 555.0379 510.185 533.5315 2340.449 100

image.png

左が剥き出しで右が関数にラップしたもの。
速度が全然違う。
なじぇ……?

しかも行列を1000列から100000列に増やすと差が縮むっぽい(ループ回数100倍)。

expr min lq mean median uq max neval
naked 45.29335 50.87810 53.68207 51.23251 52.36761 122.4935 100
wrapped 45.42338 50.00987 51.31302 50.44030 51.08073 61.9420 100

image.png

ループを更に10倍する

image.png

剥き出しのほうが速いことが多いけれど、ラップするほうが速度が安定している……?

勿論両者の結果は同じ

y <- 1:ncol(x)
for(i in y) y[i] <- sum(x[, i])
y2 <- (function(x) {
      y <- 1:ncol(x)
      for(i in y) y[i] <- sum(x[, i])
      y
    })(x)
identical(y, y2)
# [1] TRUE

2017/2/1追記

R3.4.0以降では、JITコンパイルが導入されました(compiler::enableJIT(3))。
ループ回数が少ない時は、JITコンパイルのオーバーヘッドが、計算時間の差に現れるようです。
Twitter上で考察にご協力頂いた、yutannihilationさん、TobakuCptlsmさんに感謝です。
JITを無効化するにはcompiler::enableJIT(0)を実行します。
元に戻すにはcompiler::enableJIT(3)

compiler::enableJIT(0)
set.seed(123)
x <- matrix(runif(10000), 10, 1000)

microbenchmark::microbenchmark(
  `naked` = {
    y <- 1:ncol(x)
    for(i in y) y[i] <- sum(x[, i])
    y
  },
  `wrapped` = {
    (function(x) {
      y <- 1:ncol(x)
      for(i in y) y[i] <- sum(x[, i])
      y
    })(x)
  }
)

expr min lq mean median uq max neval
naked 889.575 911.6020 930.1113 924.7345 936.240 1084.229 100
wrapped 770.046 796.0295 843.1757 811.8900 822.234 3968.241 100

image.png

それでも若干wrappedの方が速いのが気持ち悪いですが……。