Allocate a GPUArray on one GPU in one process and copy it to some other GPU in another process using IPC handles.

p2p_pycuda_ipc_demo.py

#!/usr/bin/env python

"""
Allocate a GPUArray on one GPU in one process and copy it to
some other GPU in another process using IPC handles.

Notes
-----
Requires that the two GPUs support peer-to-peer data transfers.
"""

import multiprocessing as mp
import numpy as np
import zmq

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

def proc1():
    sock = zmq.Context().socket(zmq.REQ)
    sock.connect('ipc://p2p_pycuda_ipc_demo')

    drv.init()
    dev = drv.Device(0)
    ctx = dev.make_context()

    x_gpu = gpuarray.to_gpu(np.random.rand(8))
    h = drv.mem_get_ipc_handle(x_gpu.ptr)
    sock.send_pyobj((x_gpu.shape, x_gpu.dtype, h))
    sock.recv_pyobj()

    ctx.detach()

def proc2():
    sock = zmq.Context().socket(zmq.REP)
    sock.bind('ipc://p2p_pycuda_ipc_demo')

    drv.init()
    dev = drv.Device(1)
    ctx = dev.make_context()

    shape, dtype, h = sock.recv_pyobj()
    sock.send_pyobj('')

    x_gpu = gpuarray.GPUArray(shape, dtype, gpudata=drv.IPCMemoryHandle(h))
    print 'x_gpu: ', x_gpu
    y_gpu = gpuarray.zeros_like(x_gpu)
    drv.memcpy_peer(y_gpu.ptr, x_gpu.ptr,
                    x_gpu.dtype.itemsize*x_gpu.size, ctx, ctx)
    print 'y_gpu:', y_gpu

    ctx.detach()

if __name__ == '__main__':
    p1 = mp.Process(target=proc1)
    p2 = mp.Process(target=proc2)

    p1.start()
    p2.start()