Demonstrate how to pass IPC handles to GPU data between processes in Python

ipc_demo.py

#!/usr/bin/env python

"""
Demonstrate how to pass IPC handles to GPU data between processes in Python.
"""

import ctypes
import numpy as np
import multiprocessing as mp
import zmq

import pycuda.driver as drv
import pycuda.gpuarray as gpuarray

N = 8
dtype = np.float64

def func1():
    drv.init()
    dev = drv.Device(0)
    ctx_gpu = dev.make_context()

    ctx = zmq.Context()
    sock = ctx.socket(zmq.REQ)
    sock.connect('tcp://localhost:6000')

    x = np.asarray(np.random.rand(N), dtype)
    print 'orig: ', x

    x_gpu = gpuarray.to_gpu(x)
    h = drv.mem_get_ipc_handle(x_gpu.ptr)

    sock.send_pyobj((h, x_gpu.shape, x_gpu.dtype))
    sock.recv()

    ctx_gpu.pop()

def func2():
    drv.init()
    dev = drv.Device(0)
    ctx_gpu = dev.make_context()

    ctx = zmq.Context()
    sock = ctx.socket(zmq.REP)
    sock.bind('tcp://*:6000')

    h, s, d = sock.recv_pyobj()
    sock.send('')

    x_ptr = drv.IPCMemoryHandle(h)
    x_gpu = gpuarray.GPUArray(s, d, gpudata=x_ptr)
    print 'gpu:  ', x_gpu.get()

    ctx_gpu.pop()

if __name__ == '__main__':
    p1 = mp.Process(target=func1)
    p2 = mp.Process(target=func2)

    p1.start()
    p2.start()