Cuda + Python

cupy.py

import numpy as np
import cupy as cp

print(cp.cuda.is_available())

# Limitando o tamanho de memória a ser usado
mempool = cp.get_default_memory_pool()
with cp.cuda.Device(0):
    mempool.set_limit(size=1024**3) # 1GB

# Alocando arrays na CPU e na GPU
x_cpu = np.array([1,2,3])
l2_cpu = np.linalg.norm(x_cpu)

x_gpu = cp.array([1,2,3,])
l2_gpu = cp.linalg.norm(x_gpu)

# Somando matriz na GPU
l1_matriz = cp.array([1,2,4,8,16,32,64])
l2_matriz = cp.array([128,256,512,1024,2048,3072,4096])
l3_matriz = l1_matriz + l2_matriz
l4_matriz = l3_matriz * 1024
print(l4_matriz)

# Definindo um kernel personalizado
diferenca_do_quadrado = cp.ElementwiseKernel(
    'float32 x, float32 y',
    'float32 z',
    'z = (x - y) * (x - y)',
    'squared_diff'
)

x = cp.arange(1024, dtype=np.float32).reshape(64, 16)
y = cp.arange(16, dtype=np.float32)

print(diferenca_do_quadrado(x, y))