有没有一种在调用之前编译(即缓存)cupy RawKernel的方法?

时间:2019-05-26 21:49:57

标签: cupy

我正在编写一个处理大量图像的python应用程序。 应用程序的计算速度很重要,因此我试图通过编写Cupy内核来最大程度地减少执行时间。

为简单起见,假设我下面有一个粗糙的原始内核。

import cupy as cp


add_kernel = cp.RawKernel(r'''
extern "C" __global__
void add_one(float* dimg, float* y) {
    int j = threadIdx.x;
    int i = blockIdx.x;
    int k = blockDim.x;
    int tid = k*i+j;

    y[tid] = dimg[tid] + 1;
}
''', 'add_one')


if __name__ == '__main__':
    h, w = 192, 256
    dimg_cp = cp.zeros(shape=(h, w), dtype=cp.float32)
    y = cp.zeros(shape=(h, w), dtype=cp.float32)
    add_kernel((h,), (w,), (dimg_cp, y))
    print(y)

在这里,“ add_kernel”仅复制一个输入矩阵并将一个矩阵添加到复制矩阵的每个元素中,然后将其返回。效果很好,但我相信可以在执行速度方面进一步优化代码。

根据link,当首次调用内核(即未缓存)时,将会产生编译开销。

我想避免这个编译时间。 所以我想问一下在第一次调用内核之前是否有一种编译cp.RawKernel的方法吗?

谢谢。

1 个答案:

答案 0 :(得分:1)

当前没有不调用内核就可以预编译内核的明确方法。一种简单的解决方案是只需少量输入就可以调用它一次。请注意,编译后的内核也被缓存到文件中,因此开销仅在环境中首次执行脚本时存在。