# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # GPU Computing with [CuPy](https://cupy.chainer.org/)

# ## CuPy in a Nutshell
#
# * CuPy is a matrix library accelerated with [NVIDIA CUDA](https://developer.nvidia.com/cuda-zone). Using [CUDA-related GPU-accelerated libraries](https://developer.nvidia.com/gpu-accelerated-libraries) (cuBLAS, cuDNN, cuRand, cuSolver, cuSPARSE, cuFFT, NCCL) it allows to take full advantage of the computing power of GPUs.
# * CuPy's core component is the `cupy.ndarray` class which is highly compatible with `numpy.ndarray`.
# * It supports most of the high level operations of `NumPy` arrays and allows to write user defined kernels to execute on the GPU.
# * CuPy produces GPU kernels optimized for the shapes and dtypes of given arguments on the fly.

import numpy as np
import cupy as cp
import matplotlib.pyplot as plt
# cp.cuda.set_allocator(None) # You can disable CuPy memory pool with this 
# %load_ext Cython

# +
from contextlib import contextmanager
from timeit import default_timer

@contextmanager
def cupy_timer():
    start = cp.cuda.Event()
    end = cp.cuda.Event()
    start.record()
    yield
    end.record()
    end.synchronize()
    elapsed_time = cp.cuda.get_elapsed_time(start, end)
    print(f'Elapsed time: {elapsed_time} ms')
    

@contextmanager
def cpu_timer():
    start = default_timer()
    yield
    end = default_timer()
    print(f'Elapsed time: {(end - start) * 1000} ms')


# -

# ## Creating CuPy arrays directly on the GPU

# +
A = cp.zeros((1000, 1000))
print(f'Array A is of type: {type(A)}')

B = cp.array([[1, 2, 3], 
              [4, 5, 6]])

print(B)
# -

B = cp.random.random((1000, 1000))
mu = B.mean()
print(f'The mean value is: {mu}')
print(f'The type of mu is: {type(mu)}')

# ## Transferring Arrays from Host to Device

A_cpu = np.random.random((1000, 1000))
A_gpu = cp.array(A_cpu)
mu_cpu = A_cpu.mean()
mu_gpu = A_gpu.mean()
print(f'Mean using CPU: {mu_cpu}')
print(f'Mean using GPU: {mu_gpu}')

# ## Tranferring CuPy arrays to Host
#
# Tranferring of CuPy arrays to host is performed using the `cupy.ndarrray.get` method or the `cupy.asnumpy` function. 

x_gpu = cp.ones((1000, 1000))
x_cpu = x_gpu.get()
print(f'Type of CPU array is: {type(x_cpu)}')

y_gpu = cp.random.random((1000, 1000))
y_cpu = cp.asnumpy(y_gpu)
print(f'Type of CPU array is: {type(y_cpu)}')

# ## Typical CuPy Matrix Operations

# ### Matrix-Matrix Multiplication

# +
x_cpu = np.random.random((1000, 1000))
y_cpu = np.random.random((1000, 1000))
x_gpu = cp.array(x_cpu)
y_gpu = cp.array(y_cpu)

with cpu_timer():
    z_cpu = x_cpu @ y_cpu

with cupy_timer():
    z_gpu = x_gpu @ y_gpu
# -

# ## Linear System Solving NumPy vs CuPy

# +
A_cpu = np.random.random((10000, 10000))
b_cpu = np.random.random(10000)
A_gpu = cp.array(A_cpu)
b_gpu = cp.array(b_cpu)
with cpu_timer():
    np.linalg.solve(A_cpu, b_cpu)
    
with cupy_timer():
    cp.linalg.solve(A_gpu, b_gpu)


# -

# ## Euclidean distance matrix
#
# $
#     d_e(\mathbf x, \mathbf y) =
#     \begin{bmatrix}
#     \sum_{i=1}^n (x_{1i}-y_{1i})^2 & \sum_{i=1}^n(x_{1i}-y_{2i})^2 & \cdots & \sum_{i=1}^n (x_{1i}-y_{ni})^2 \\  
#     \sum_{i=1}^n(x_{2i}-y_{1i})^2 & \sum_{i=1}^n(x_{2i}-y_{2i})^2 & \cdots & \sum_{i=1}^n(x_{2i}-y_{ni})^2 \\  
#     \vdots & \vdots & \ddots & \vdots \\
#     \sum_{i=1}^n(x_{ni}-y_{1i})^2 & \sum_{i=1}^n(x_{ni}-y_{2i})^2 & \cdots & \sum_{i=1}^n(x_{ni}-y_{ni})^2 \\  
#     \end{bmatrix}
# $
#
# ## Vectorization friendly summation 
# $ 
# \sum_{k=1}^n \left(x_{ik}-y_{jk}\right)^2 = \left(\vec{x_i} - \vec {y_j}\right)\cdot \left(\vec{x_i} - \vec{y_j}\right)=\vec{x_i} \cdot \vec{x_i} + \vec{y_j} \cdot \vec{y_j} -2\vec{x_i}\cdot \vec{y_j}$
#

def euclidean_distance_cpu(x, y):
    x2 = np.einsum('ij,ij->i', x, x)[:, np.newaxis]
    y2 = np.einsum('ij,ij->i', y, y)[np.newaxis, :]
    xy = x @ y.T
    return np.abs(x2 + y2 - 2.0 * xy)


def euclidean_distance_gpu(x, y):
    x2 = cp.einsum('ij,ij->i', x, x)[:, cp.newaxis]
    y2 = cp.einsum('ij,ij->i', y, y)[cp.newaxis, :]
    xy = x @ y.T
    return cp.abs(x2 + y2 - 2.0 * xy)


# +
x_cpu = np.random.random((5000, 5000))
y_cpu = np.random.random((5000, 5000))
x_gpu = cp.array(x_cpu)
y_gpu = cp.array(y_cpu)

with cpu_timer():
    eu_cpu = euclidean_distance_cpu(x_cpu, y_cpu)

with cupy_timer():
    eu_gpu = euclidean_distance_gpu(x_gpu, y_gpu)
    
assert(np.allclose(eu_cpu, eu_gpu.get()))


# -

# ## Make function work for both CuPy/NumPy 

def euclidean_distance(x, y):
    array_mod = cp.get_array_module(x)
    x2 = array_mod.einsum('ij,ij->i', x, x)[:, array_mod.newaxis]
    y2 = array_mod.einsum('ij,ij->i', y, y)[array_mod.newaxis, :]
    xy = x @ y.T
    return array_mod.abs(x2 + y2 - 2.0 * xy)


# +
with cpu_timer():
    eu_cpu = euclidean_distance(x_cpu, y_cpu)

with cupy_timer():
    eu_gpu = euclidean_distance(x_gpu, y_gpu)
    
assert(np.allclose(eu_cpu, eu_gpu.get()))
# -

# ## Calculating/Plotting [Julia Sets](https://en.wikipedia.org/wiki/Julia_set)

# + language="cython"
#
# # distutils: extra_compile_args = -fopenmp -march=native
# # distutils: extra_link_args = -fopenmp
# from cython cimport boundscheck, wraparound
# from cython.parallel cimport prange
#
# @boundscheck(False)
# @wraparound(False)
# def julia_set_cython(const double[:, :] X, const double[:, :] Y,
#                      const double cx, const double cy,
#                      const int iter_max, const double radius2, 
#                      int[:, :] julia):
#     cdef:
#         int i, j, k, nx, ny
#         double x, y
#     ny = X.shape[0]
#     nx = Y.shape[1]
#     for i in prange(ny, nogil=True, schedule='static'):
#         for j in range(nx):
#             x = X[i, j]
#             y = Y[i, j]
#             k = 0
#             while x * x + y * y < radius2 and k < iter_max:
#                 x, y = x * x - y * y + cx, 2.0 * x * y + cy
#                 k = k + 1
#                 
#             julia[i, j] = k 

# +
X, Y = np.meshgrid(np.linspace(-2.0 , 2.0, 5000), np.linspace(-2.0, 2.0, 5000))
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111)
julia = np.zeros_like(X, dtype=np.int32)
c = -0.9 + 0.22143j
radius2 = 4.0
with cpu_timer():
    julia_set_cython(X, Y, c.real, c.imag, 100, radius2, julia)

ax.set_aspect('equal')
ax.imshow(julia, extent=[-2, 2, -2, 2]);
# -

# ### Elementwise Kernel
#
# Using the `cupy.ElementwiseKernel` class it is easy to create GPU kernels by defining the computation that is going to be applied to each element of the input array(s).

julia_kernel = cp.ElementwiseKernel('float64 X, float64 Y, float64 cx, float64 cy, int32 itermax, float64 radius2',
                                    'int32 julia',
                                    f'''julia = 0;
                                    double x = X, y = Y;
                                    double xtemp;
                                    int nit = 0;
                                    while(x * x + y * y < radius2 && nit < itermax) {{
                                        xtemp = x * x - y * y + cx;
                                        y = 2.0 * x * y + cy;
                                        x = xtemp;
                                        nit += 1;
                                    }}
                                    julia = nit;''', 'julia_kernel')

X, Y = cp.meshgrid(cp.linspace(-2.0 , 2.0, 5000), cp.linspace(-2.0, 2.0, 5000))
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111)
c = -0.9 + 0.22143j
with cupy_timer():
    julia_gpu = julia_kernel(X, Y, c.real, c.imag, 100, 4.0)
    julia_array = julia_gpu.get()
ax.set_aspect('equal')
ax.imshow(julia_array, extent=[-2, 2, -2, 2]);

# ### Raw Kernel
#
# It is possible to create raw CUDA kernels using the `cupy.RawKernel` class but the block and grid dimensions are not handled automatically.

julia_rawkernel = cp.RawKernel(r'''
        extern "C" 
        __global__ void julia_rawkernel(const double* X, const double* Y, const double cx,
                              const double cy, const int itermax, const int nx,
                              const int ny, const double radius, int* julia)  {
        int tidx = blockDim.x * blockIdx.x + threadIdx.x;
        int tidy = blockDim.y * blockIdx.y + threadIdx.y;
        int niter = 0;
        double tempx;
        if(tidx < nx && tidy < ny) 
        {
            int tid = tidy * nx + tidx;
            double x = X[tid], y = Y[tid];
            while((x * x + y * y) < (radius * radius) && niter < itermax) {
                tempx = x * x - y * y + cx;
                y = 2.0 * x * y + cy;
                x = tempx;
                niter +=1 ;
            }
            julia[tid] = niter;
        }
    }''', 'julia_rawkernel')

fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111)
c = -0.9 + 0.22143j
threads_x = 32
threads_y = 32
x_points = 5000
y_points = 5000
x_blocks = (x_points - 1) // threads_x + 1 
y_blocks = (y_points - 1) // threads_y + 1
X_gpu, Y_gpu = cp.meshgrid(cp.linspace(-2.0 , 2.0, x_points), cp.linspace(-2.0, 2.0, y_points))
with cupy_timer():
    julia_gpu = cp.zeros_like(X_gpu, dtype=cp.int32)
    julia_rawkernel((x_blocks, y_blocks), (threads_x, threads_y), (X_gpu, Y_gpu, c.real, c.imag, 100, 
                    x_points, y_points, 2.0, julia_gpu))
    julia_array = julia_gpu.get()
ax.imshow(julia_array, extent=[-2, 2, -2, 2]);

# ### Raw Kernel JIT (experimental)

# +
from cupyx import jit

@jit.rawkernel()
def julia_kernel(X, Y, cx, cy, itermax, nx, ny, radius, julia):
    tidx = jit.blockDim.x * jit.blockIdx.x + jit.threadIdx.x
    tidy = jit.blockDim.y * jit.blockIdx.y + jit.threadIdx.y
    x = X[tidy, tidx]
    y = Y[tidy, tidx]   
    niter = 0
    if tidx < nx and tidy < ny:
        while (x * x + y * y) < radius * radius and niter < itermax:
            x, y = (x * x - y * y + cx), (2.0 * x * y + cy)
            niter += 1
            
        julia[tidy, tidx] = niter


# +
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111)
c = cp.complex64(-0.9 + 0.22143j)
threads_x = 32
threads_y = 32
x_points = 5000
y_points = 5000
x_blocks = (x_points - 1) // threads_x + 1 
y_blocks = (y_points - 1) // threads_y + 1
X_gpu, Y_gpu = cp.meshgrid(cp.linspace(-2.0 , 2.0, x_points), cp.linspace(-2.0, 2.0, y_points))
radius = cp.float64(2.0)
with cupy_timer():
    julia_gpu = cp.zeros_like(X_gpu, dtype=cp.int32)
    julia_kernel[(x_blocks, y_blocks), (threads_x, threads_y)](
        X_gpu, Y_gpu, c.real, c.imag, 100, x_points, y_points, radius, julia_gpu
    )
    julia_array = julia_gpu.get()
    
ax.imshow(julia_array, extent=[-2, 2, -2, 2]);
# -

# ### Inspecting the source code of the kernel

print(julia_kernel.cached_code)


# ## Kernel fusion

# +
def sigmoid_cpu(x):
    return 1.0 / (1.0 + np.exp(-x))

sigmoid_kernel = cp.ElementwiseKernel(
    'float64 X',
    'float64 sigmoid',
    'sigmoid = 1.0 / (1.0 + exp(-X));', 
    'sigmoid_kernel')

@cp.fuse
def sigmoid_fused(x):
    return 1.0 / (1.0 + cp.exp(-x))

def sigmoid_gpu(x):
    return 1.0 / (1.0 + cp.exp(-x))


# +
x_cpu = np.linspace(-10.0, 10.0, 10_000_000)
x_gpu = cp.array(x_cpu)

print('Sigmoid Cpu: ', end='')
with cpu_timer():
    s_cpu = sigmoid_cpu(x_cpu)
    
print('Sigmoid Gpu', end='')
with cupy_timer():
    s_gpu = sigmoid_gpu(x_gpu)
    
print('Sigmoid Kernel', end='')
with cupy_timer():
    s_gpu_kernel = sigmoid_kernel(x_gpu)
    
print('Sigmoid Fused', end='')
with cupy_timer():
    s_gpu_fused = sigmoid_fused(x_gpu)
    
assert(np.allclose(s_cpu, s_gpu.get()))
assert(cp.allclose(s_cpu, s_gpu_kernel))
assert(cp.allclose(s_cpu, s_gpu_fused))

plt.close('all')
fig = plt.figure(figsize=(20, 10))
ax = fig.add_subplot(111)
ax.plot(x_cpu, s_cpu, '--');
ax.set_title('Sigmoid function', )
ax.set_xlabel('x')
ax.set_ylabel('y(x)')
ax.grid('Both');
# -

# ## CuPy Memory Pools 
#
# In order to improve performance, CuPy uses a memory pool for memory allocations by default.

# !nvidia-smi #

# %whos ndarray

mempool = cp.get_default_memory_pool()

used_gbs = mempool.used_bytes() / 1024 ** 3
total_gbs = mempool.total_bytes() / 1024 ** 3
print(f'Memory pool uses: {used_gbs:.3f} out of {total_gbs:.3f} GB')

mempool.free_all_blocks()

StarcoderdataJupyterScriptsDedupFiltered