如何正确使用 CuPy 流

Posted 2023-02-16

【中文标题】如何正确使用 CuPy 流

【英文标题】：How to Properly use CuPy Streams

【发布时间】：2020-10-28 20:47:48

【问题描述】：

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

【问题讨论】：

问题是 with 构造不会导致在所选流上发出 matmul 操作。相反，它们被发布到空流。 （对于给定的操作，少于 3 秒的时间是不可信的 - 您的应用程序花费的时间要长得多。）您可以通过将流同步转换为 cupy.cuda.Device().synchronize()

来说服自己

【参考方案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

【讨论】：

由于某种原因，此解决方法似乎也需要更长的时间才能运行？

运行时间较长，因为您没有正确测量上一次运行的执行时间。 matmul 启动是/仍然是异步的。所以〜3s不是前一个案例的有效时间。

这是有道理的。我认为这是导致挂起的 slurm 问题。我可能只是同步了错误的流。

是的，我同意。我将把谈话转到：github.com/cupy/cupy/issues/4199

为了节省未来读者的点击次数，这似乎已经解决了cupy 中的“已知”问题，该问题应该在cupy 8.1 中修复