CODEWITHSUNDEEP
pythongpugoogle-colaboratory
Yifei Li
Yifei Li

Reputation: 33

why np.einsum() can run on GPU if the arrays are defined in GPU?

I recently realize that as long as variables (or arrays) are defined in GPU, numpy functions run as fast as cupys. I couldn't find out a way to monitor whether numpy functions are actually executed on GPU. if someone had similar experience, please share your answer. Thanks.

I am using google colabpro GPU runtime.

below are two sets of code:

import numpy as np
import cupy as cp

set 1:

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

set 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

Upvotes: 2

Views: 487

Answers (1)

BatWannaBe
BatWannaBe

Reputation: 4510

Let's look at a bit of numpy.einsum in einsumfunc.py:

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

Well, that decorator looks promising; you are indeed describing a function dispatching on argument types. Let's look at a bit of array_function_dispatch in overrides.py

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.

And that takes us to the far more readable NEP-18: "We propose the __array_function__ protocol, to allow arguments of NumPy functions to define how that function operates on them..." So NumPy functions check __array_function__ or __array_ufunc__ (NEP-13) of its arguments.

CuPy defines cupy.ndarray.__array_function__ in core.pyx (a Cython file). It looks for the NumPy function's name in the similarly organized CuPy module and calls it on the arguments. So when you called np.einsum(...), it ended up finding cp.einsum(...) anyway:

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)

Upvotes: 5

Related Questions