如何正确使用CuPy Streams

Question

我目前正在尝试弄清楚如何有效地使用 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]))

结果表明，大部分时间都花在重复乘法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

Answer 2

你不能像那样对你的代码进行基准测试。 Cupy 在幕后做了各种各样的事情，这些事情可能会扰乱你的测量。您需要使用cupy分析器，请参阅：https://docs.cupy.dev/en/stable/user_guide/performance.html#benchmarking

我猜你的代码是在调用同步后测量 end_time 之后进行实际同步，因为该调用仅在当前 cuda 流上安排同步，而不实际执行它。