PyCUDAをいじってみる

とりあえず

サンプルコードをいじってみる。

import pycuda.autoinit
import pycuda.driver as drv
import numpy
import time
from pycuda.compiler import SourceModule
mod = SourceModule("""
__global__ void square(float* p)
{
  const int i = blockDim.x*blockIdx.x + threadIdx.x;
  p[i] *= p[i];
}
""")

square = mod.get_function("square")

size = 4096
threads = 512
assert threads <= 512 and size % threads == 0

lst = [i for i in xrange(size)]
p = numpy.array(lst, dtype=numpy.float32)

# cputime
s = time.time()
p**2
e = time.time()
cputime= e - s

# gputime
s = time.time()
square(drv.InOut(p), block=(threads,1,1), grid=(size/threads,1))
e = time.time()
gputime= e - s

print "%.10f %.10f" % (cputime, gputime)
print "win", ["cpu", "gpu"][cputime > gputime]

実行結果

$ python x.py
0.0000278950 0.0003490448
win cpu

C++のテンプレートを試してみる

C++のテンプレートは今のところ使えない？

import pycuda.autoinit
import pycuda.driver as drv
import numpy

from pycuda.compiler import SourceModule
mod = SourceModule("""
template<class Type>
__global__ void square(Type* p)
{
  const int i = blockDim.x*blockIdx.x + threadIdx.x;
  p[i] *= p[i];
}
""",
# デフォルトではソースコードをextern "C"で囲うため
# error: this declaration may not have extern "C" linkage
# などとエラーが出力されるので以下を指定
no_extern_c=True,)

square = mod.get_function("square")

size = 256

p = numpy.array([i for i in xrange(size)], dtype=numpy.float32)

square(drv.InOut(p), block=(size,1,1))

実行結果

$ python test.py
Traceback (most recent call last):
  File "test.py", line 17, in <module>
    square = mod.get_function("square")
  File "/usr/local/lib/python2.6/dist-packages/pycuda-0.93-py2.6-linux-x86_64.egg/pycuda/compiler.py", line 208, in get_function
    func = self.module.get_function(name)
pycuda._driver.LogicError: cuModuleGetFunction failed: not found

2つのバイト列の平均を取る

これでwaveファイルの合成を高速化できたりしないかななどと考えながら。

import pycuda.autoinit
import pycuda.driver as drv
import numpy
import time

from pycuda.compiler import SourceModule
mod = SourceModule("""
__global__ void bytesadd(unsigned char* dst, unsigned char* p, unsigned char* q)
{
  const int i = blockDim.x*blockIdx.x + threadIdx.x;
  dst[i] = (unsigned char)((short)p[i] + q[i]) / 2;
}
""")

bytesadd = mod.get_function("bytesadd")

def main():
    threads = 512
    for i in range(2**14, 2**22, 8192):
        size = i
        if size/ threads >= 65536:
            break

        lst = [j for j in xrange(size)]
        p = numpy.array(lst, dtype=numpy.uint8)
        q = numpy.array(lst, dtype=numpy.uint8)
        dst = numpy.zeros(size)

        # cputime
        s = time.time()
        dst = (p+q)/2
        e = time.time()
        cputime = e - s

        # gputime
        s = time.time()
        bytesadd(drv.Out(dst), drv.In(p), drv.In(q), block=(threads,1,1), grid=(size/threads,1))
        e = time.time()
        gputime = e - s
        
        print ",".join(map(str,[size, gputime, cputime]))
main()

実行結果

プロットしてみた。

PyCUDAのドキュメントにはテンプレートエンジンを使ってループを直接アンロールしたり、組み込み変数を即値としてコードに埋め込むメタプログラミングなどが紹介されていて面白い。
http://documen.tician.de/pycuda/metaprog.html