PyCuda/Examples/Demo

   1 # Sample source code from the Tutorial Introduction in the documentation.
   2 
   3 import pycuda.driver as cuda
   4 import pycuda.autoinit
   5 from pycuda.compiler import SourceModule
   6 
   7 import numpy
   8 a = numpy.random.randn(4,4)
   9 
  10 a = a.astype(numpy.float32)
  11 
  12 a_gpu = cuda.mem_alloc(a.size * a.dtype.itemsize)
  13 
  14 cuda.memcpy_htod(a_gpu, a)
  15 
  16 mod = SourceModule("""
  17     __global__ void doublify(float *a)
  18     {
  19       int idx = threadIdx.x + threadIdx.y*4;
  20       a[idx] *= 2;
  21     }
  22     """)
  23 
  24 func = mod.get_function("doublify")
  25 func(a_gpu, block=(4,4,1))
  26 
  27 a_doubled = numpy.empty_like(a)
  28 cuda.memcpy_dtoh(a_doubled, a_gpu)
  29 print "original array:"
  30 print a
  31 print "doubled with kernel:"
  32 print a_doubled
  33 
  34 # alternate kernel invocation -------------------------------------------------
  35 
  36 func(cuda.InOut(a), block=(4, 4, 1))
  37 print "doubled with InOut:"
  38 print a
  39 
  40 # part 2 ----------------------------------------------------------------------
  41 
  42 import pycuda.gpuarray as gpuarray
  43 a_gpu = gpuarray.to_gpu(numpy.random.randn(4,4).astype(numpy.float32))
  44 a_doubled = (2*a_gpu).get()
  45 
  46 print "original array:"
  47 print a_gpu
  48 print "doubled with gpuarray:"
  49 print a_doubled