import pycuda.autoinit
import pycuda.driver as drv
import numpy

from pycuda.compiler import SourceModule
mod = SourceModule("""
__global__ void multiplicar(float *dest, float *a, float *b)
{
  const int i = threadIdx.x;
  dest[i] = a[i] * b[i];
}
""")

multiplicar = mod.get_function("multiplicar")

a = numpy.random.randn(400).astype(numpy.float32)
b = numpy.random.randn(400).astype(numpy.float32)

dest = numpy.zeros_like(a)

print dest

multiplicar(
        drv.Out(dest), drv.In(a), drv.In(b),
        block=(400,1,1), grid=(1,1))

print dest

print dest-a*b

"""
Explanation: ¿Por qué PyCUDA?
Hasta ahora hemos visto que si bien CUDA no es un lenguaje imposible de aprender, puede llegar a ser un dolor de cabeza el tener muchos apuntadores y manejar la memoria de un modo tan rudimentario.
Sin embargo hay alternativas que nos permiten trabajar en entornos más agradables, un ejemplo de ellos es PyCUDA creado con Andreas Klöckner. Básicamente PyCUDA se encarga de mapear todo CUDA dentro de Python. 
Por poner un ejemplo, un código simple sería el siguiente
End of explanation
"""


import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule

"""
Explanation: Al correr este programa vamos a obtener un montón de ceros; algo no muy interesante. Sin embargo detrás de escenas sí pasó algo interesante.

PyCUDA compiló el código fuente y lo cargó a la tarjeta.
Se asignó memoria automáticamente, además de copiar las cosas de CPU a GPU y de vuelta.
Por último la limpieza (liberación de memoria) se hace sola.

Útil ¿cierto?
Usando PyCUDA
Para empezar debemos importar e incializar PyCUDA
End of explanation
"""


import numpy
a = numpy.random.randn(4,4)

"""
Explanation: Transferir datos
El siguiente paso es transferir datos al GPU. Principalmente arreglos de numpy. Por ejemplo, tomemos un arreglo de números aleatorios de $4 \times 4$
End of explanation
"""


a = a.astype(numpy.float32)

"""
Explanation: sin embargo nuestro arreglo a consiste en números de doble precisión, dado que no todos los GPU de NVIDIA cuentan con doble precisión es que hacemos lo siguiente
End of explanation
"""


a_gpu = cuda.mem_alloc(a.nbytes)

"""
Explanation: finalmente, necesitmos un arreglo hacia el cuál transferir la información, así que deberíamos guardar la memoria en el dispositivo:
End of explanation
"""


cuda.memcpy_htod(a_gpu, a)

"""
Explanation: como último paso, necesitamos tranferir los datos al GPU
End of explanation
"""


mod = SourceModule("""
  __global__ void duplicar(float *a)
  {
    int idx = threadIdx.x + threadIdx.y*4;
    a[idx] *= 2;
  }
  """)

"""
Explanation: Ejecutando kernels
Durante este capítulo nos centraremos en un ejemplo muy simple. Escribir un código para duplicar cada una de las entradas en un arreglo, en seguida escribimos el kernel en CUDA C, y se lo otorgamos al constructor de pycuda.compiler.SourceModule
End of explanation
"""


mod

func = mod.get_function("duplicar")
func(a_gpu, block=(4,4,1))

func

"""
Explanation: Si no hay errores, el código ahora ha sido compilado y cargado en el dispositivo. Encontramos una referencia a nuestra pycuda.driver.Function y la llamamos, especificando a_gpu como el argumento, y un tamaño de bloque de $4\times 4$:
End of explanation
"""


a_duplicado = numpy.empty_like(a)
cuda.memcpy_dtoh(a_duplicado, a_gpu)
print a_duplicado
print a

print(type(a))
print(type(a_gpu))
print(type(a_duplicado))

a_duplicado

"""
Explanation: Finalmente recogemos la información del GPU y la mostramos con el a original
End of explanation
"""


