Python/CFFIでCUDAしてみる

ふと思い立って、CUDAでCPUからGPUにデータを転送する時の速度を測ってみた。

普通にCUDA SDKのサンプルで測定しても良いが、PythonのCFFI で実行してみよう。

コードはこんな感じで書ける。

import time
import sys
from cffi import FFI
ffi=FFI()

ffi.cdef(r"""
typedef unsigned int cudaError_t;
cudaError_t cudaGetLastError();
char* cudaGetErrorString(cudaError_t error);

cudaError_t cudaDeviceSynchronize();
cudaError_t cudaMalloc(void **p, size_t s);
cudaError_t cudaFree(void *p);
cudaError_t cudaMallocHost(void **p, size_t s);
cudaError_t cudaFreeHost(void *p);
cudaError_t cudaMemcpy(void *dst, const void *src, size_t count, unsigned int kind);

""")

cudaMemcpyHostToHost          =   0
cudaMemcpyHostToDevice        =   1
cudaMemcpyDeviceToHost        =   2
cudaMemcpyDeviceToDevice      =   3
cudaMemcpyDefault             =   4

cuda = ffi.dlopen("/usr/local/cuda/lib64/libcudart.so")

def check():
    err = cuda.cudaGetLastError()
    if err:
        print(ffi.string(cuda.cudaGetErrorString(ret)))
    assert err == 0

def run(host, dev, n, size):
    for i in range(n):
        ret = cuda.cudaMemcpy(dev[0], host[0], size, cudaMemcpyHostToDevice)
        check()

size = int(sys.argv[1])

pinned = ffi.new("void **")
cuda.cudaMallocHost(pinned, size)
check()

ffi.memmove(pinned[0], b'a'*size, size)

dev = ffi.new("void **")
cuda.cudaMalloc(dev, size)
check()

run(pinned, dev, 1, 1) # warm up

buf = ffi.new("char []", b'a'*size)
host = ffi.new("char **", buf)

N = 1000
for s in range(1, size, size//20):
    f = time.time()
    run(pinned, dev, N, s)
    p = time.time()-f

    f = time.time()
    run(host, dev, N, s)
    n= time.time()-f

    print(f'{s}, {p:0.5f}, {n:0.5f}')

AWSのp2.xlargeとp3.2xlargeで実行してみた。

p2.xlarge
Tesla K80
E5-2686 v4 @ 2.30GHz

p3.2xlarge   
Tesla V100-SXM2-16GB
E5-2686 v4 @ 2.30GHz

ざっくり、Pinnedメモリのほうが15%くらい速い。また、転送速度に関してはp2.xlargeでもp3.2xlargeでも大差はないようだ。

ちなみに、手元にあったi7-6700 CPU @ 3.40GHz/GeForce GTX 1070 と E5-1650 v4 @ 3.60GHz/GeForce GTX 1080 Ti のデスクトップは、EC2のインスタンスより2割程度高速だった。

また、~100KBぐらいまでの転送量を見てみると、50KBぐらいまではPinnedメモリより普通のメモリのほうが速いようだ。

これは、データの転送よりも、転送の前処理・後処理に時間がかかってしまっているのかもしれない。ストリームを使って並列化するなどすればまた違ってくるかもしれないので、あとで実験してみたい。