lebedov · April 16, 2012 18:41 · Apr 16, 2012
diff --git a/pycuda_req_rep_demo.py b/pycuda_req_rep_demo.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+
+"""
+Pass data between processes started through the multiprocessing module
+using pyzmq and process them with PyCUDA
+"""
+
+import numpy as np
+import zmq
+import multiprocessing as mp
+
+gpu = 0
+
+def worker():
+    import pycuda.driver as drv
+    import pycuda.gpuarray as gpuarray
+    import atexit
+
+    # Initialize device:
+    drv.init()
+    dev = drv.Device(gpu)
+    ctx = dev.make_context()
+    atexit.register(ctx.pop)
+
+    print dev.name()
+
+    context = zmq.Context()
+    socket = context.socket(zmq.REP)
+    socket.connect("tcp://localhost:5555")
+
+    # Process data sent to worker until a quit signal is transmitted:
+    while True:
+        data = socket.recv_pyobj()
+        print "Worker %i: %s" % (gpu, data)
+        if data == 'quit':
+            break
+
+        # Do something with the data on the GPU:
+        data_gpu = gpuarray.to_gpu(data)
+        result_gpu = -data_gpu
+
+        socket.send_pyobj(result_gpu.get())
+
+def master():
+
+    # Data to send to worker:
+    data_list = map(lambda x: np.random.rand(4, 4), xrange(4))
+
+    context = zmq.Context()
+    socket = context.socket(zmq.REQ)
+    socket.bind("tcp://*:5555")
+
+    # Send data out for processing and get back the results:
+    for i in xrange(len(data_list)):
+        socket.send_pyobj(data_list[i])
+        result = socket.recv_pyobj()
+        print "Master: ", result
+    socket.send_pyobj('quit') 
+
+if __name__ == '__main__':    
+    worker = mp.Process(target=worker)
+    worker.start()
+
+    master()
+
+
No results found