如何正确使用CuPy流

问题描述

我目前正在尝试找出如何有效使用CuPy流的方法。以下代码通过重复的矩阵乘法计算矩阵幂。我希望下面的代码将其大部分时间花费在同步行上，但似乎将其大部分时间花费在matmul行上。这是CuPy中的错误，还是我在滥用CuPy流？

#!/usr/bin/env python

"""
stream_example.py
Inefficiently calculates a matrix power through repeated matrix multiplication.  
"""

import numpy as np
import cupy
import sys
import time

def main(N,power):
    compute_stream = cupy.cuda.stream.Stream(non_blocking=True)

    with compute_stream:
        d_mat = cupy.random.randn(N*N,dtype=cupy.float64).reshape(N,N)
        d_ret = d_mat

        cupy.matmul(d_ret,d_mat)

        start_time = time.time()
        for i in range(power - 1):
            d_ret = cupy.matmul(d_ret,d_mat)
        end_time = time.time()
        print(f"Time spent on cupy.matmul for loop: {end_time - start_time}")

        start_time = time.time()
        compute_stream.synchronize()
        end_time = time.time()
        print(f"Time spent compute_stream.synchronize(): {end_time - start_time}")

if __name__ == "__main__":
    main(int(sys.argv[1]),int(sys.argv[2]))

结果表明，大部分时间都花在循环的重复乘法中，而不是stream.synchronize（）。 cupy.matmul（）不能异步使用吗？

$ python3 stream_example.py 16384 1024
Time spent on cupy.matmul for loop: 2.667935609817505
Time spent compute_stream.synchronize(): 4.2438507080078125e-05

解决方法

似乎可以解决此问题并添加以下内容。我将绿色复选标记保留给可以提出一种不太hacky解决方案的人：

import cupy_backends.cuda.libs.cublas
from cupy.cuda import device
handle = device.get_cublas_handle()
...
cupy_backends.cuda.libs.cublas.setStream(handle,compute_stream.ptr)

$ python3 stream_example.py 16384 4
Time spent on cupy.matmul for loop: 0.007548093795776367
Time spent compute_stream.synchronize(): 5.099333047866821

cuda cuda cupy python