PyCuda/Examples/MultipleThreads

   1 # Derived from a test case by Chris Heuser
   2 # Also see FAQ about PyCUDA and threads.
   3 
   4 
   5 import pycuda
   6 import pycuda.driver as cuda
   7 from pycuda.compiler import SourceModule
   8 import threading
   9 import numpy
  10 
  11 class GPUThread(threading.Thread):
  12     def __init__(self, number, some_array):
  13         threading.Thread.__init__(self)
  14 
  15         self.number = number
  16         self.some_array = some_array
  17 
  18     def run(self):
  19         self.dev = cuda.Device(self.number)
  20         self.ctx = self.dev.make_context()
  21 
  22         self.array_gpu = cuda.mem_alloc(some_array.nbytes)
  23         cuda.memcpy_htod(self.array_gpu, some_array)
  24 
  25         test_kernel(self.array_gpu)
  26         print "successful exit from thread %d" % self.number
  27         self.ctx.pop()
  28 
  29         del self.array_gpu
  30         del self.ctx
  31 
  32 def test_kernel(input_array_gpu):
  33     mod = SourceModule("""
  34         __global__ void f(float * out, float * in)
  35         {
  36             int idx = threadIdx.x;
  37             out[idx] = in[idx] + 6;
  38         }
  39         """)
  40     func = mod.get_function("f")
  41 
  42     output_array = numpy.zeros((1,512))
  43     output_array_gpu = cuda.mem_alloc(output_array.nbytes)
  44 
  45     func(output_array_gpu,
  46           input_array_gpu,
  47           block=(512,1,1))
  48     cuda.memcpy_dtoh(output_array, output_array_gpu)
  49 
  50     return output_array
  51 
  52 cuda.init()
  53 some_array = numpy.ones((1,512), dtype=numpy.float32)
  54 num = cuda.Device.count()
  55 
  56 gpu_thread_list = []
  57 for i in range(num):
  58     gpu_thread = GPUThread(i, some_array)
  59     gpu_thread.start()
  60     gpu_thread_list.append(gpu_thread)