cupy のメモリアロケータ

cupy の中の人の投稿によると、cupy.cuda.set_allocator() で
メモリプールを使い回すようにすると高速になると書いてある。

実際そのとおりなのだが、投稿にある test() を２回連続で
回すと、アロケータをセットした方が圧倒的に遅くなった。
ソースコードを見ていないが、どこかに実行結果をキャッシュしている？

numpy: 0:00:00.004890 sec
cupy: 0:00:00.000179 sec
cupy(memory pool): 0:00:00.000456 sec

なお、記事では触れていないが、いわいる「バーンイン」を
行わないで GPU 回りのメソッドを呼ぶと初期化設定のため、
最初の１回だけは非常に時間がかかることに注意。

使用したテストコードは以下の通り。

# -*- coding: utf-8 -*-
import numpy as np
import cupy as cp
import datetime

def test(xp):
    x = xp.arange(1000000).reshape(1000, -1)
    return x.T * 2

def test_with_time(title, xp):
    t1 = datetime.datetime.now()
    test(xp)
    t2 = datetime.datetime.now()
    print(title, t2 - t1, "sec")

# バーンイン
#cp.arange(1)
test(cp)
 
test_with_time("numpy:", np)
test_with_time("cupy:", cp)
cp.cuda.set_allocator(cp.cuda.MemoryPool().malloc)
test_with_time("cupy(memory pool):", cp)