如何正确使用CuPy Streams

Question

我目前正在尝试弄清楚如何有效地使用CuPy streams。下面的代码通过重复的矩阵乘法计算矩阵的幂。我希望下面的代码将大部分时间花在同步行上，但它似乎大部分时间都花在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]))

结果表明，大部分时间都花在重复的for循环乘法上，而不是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

Answer 1

它看起来像是添加了以下解决这个问题的方法。我将为那些能想出不那么麻烦的解决方案的人保留绿色的复选标记：

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