如果数组是在 GPU 中定义的，为什么 np.einsum() 可以在 GPU 上运行？

Question

我最近意识到，只要在 GPU 中定义了变量（或数组），numpy 函数的运行速度就与 cupys 一样快。我找不到监控 numpy 函数是否在 GPU 上实际执行的方法。如果有人有类似经历，请分享您的答案。谢谢。

我正在使用谷歌 colabpro GPU 运行时。

下面是两组代码：

import numpy as np
import cupy as cp

第一组：

a = np.random.randn(500, 500, 500)
b = np.random.randn(500, 500, 500)   
start_time = time.time()
for i in range(1): 
  c = np.einsum('ijk,ikm->ijm', a, b)        
end_time = time.time()
print('forwar gpu time')
print(end_time - start_time)

forwar gpu time
55.88586902618408

第 2 组：

a = cp.random.randn(500, 500, 500)   # change to cupy
b = cp.random.randn(500, 500, 500)   # change to cupy
start_time = time.time()
for i in range(1): 
  c = np.einsum('ijk,ikm->ijm', a, b)    # remain numpy    
end_time = time.time()
print('forwar gpu time')
print(end_time - start_time)

forwar gpu time
0.0009937286376953125

Answer 1

让我们看一下einsumfunc.py中的numpy.einsum：

@array_function_dispatch(_einsum_dispatcher, module='numpy')
def einsum(*operands, out=None, optimize=False, **kwargs):

嗯，那个装饰器看起来很有前途；您确实在描述参数类型的函数dispatching。让我们看一下overrides.py中的array_function_dispatch

def array_function_dispatch(dispatcher, module=None, verify=True,
                            docs_from_dispatcher=False):
    """Decorator for adding dispatch with the __array_function__ protocol.
    See NEP-18 for example usage.

这将我们带到更具可读性的NEP-18：“我们提出__array_function__ 协议，以允许NumPy 函数的参数定义该函数如何对它们进行操作......”因此NumPy 函数检查__array_function__或__array_ufunc__ (NEP-13) 的论点。

CuPy 在core.pyx（Cython 文件）中定义了cupy.ndarray.__array_function__。它在类似组织的 CuPy 模块中查找 NumPy 函数的名称，并在参数上调用它。所以当你打电话给np.einsum(...)时，它最终还是找到了cp.einsum(...)：

cdef class ndarray:
...
...
...
    def __array_function__(self, func, types, args, kwargs):
        try:
            module = functools.reduce(
                getattr, func.__module__.split('.')[1:], cupy)
            cupy_func = getattr(module, func.__name__)
        except AttributeError:
            return NotImplemented
        if cupy_func is func:
            # avoid NumPy func
            return NotImplemented
        for t in types:
            if t not in _HANDLED_TYPES:
                return NotImplemented
        return cupy_func(*args, **kwargs)