Add python interface for ipc memory handles.

Author: abergeron

pygpu/gpuarray.pxd

@@ -196,6 +196,9 @@ cdef extern from "gpuarray/array.h":
 
 cdef extern from "gpuarray/extension.h":
     void *gpuarray_get_extension(const char *)
+    ctypedef struct GpuArrayIpcMemHandle:
+        pass
+
     cdef int GPUARRAY_CUDA_CTX_NOFREE
 
 cdef type get_exc(int errcode)

pygpu/gpuarray.pyx

@@ -1424,6 +1424,28 @@ def _concatenate(list al, unsigned int axis, int restype, object cls,
     finally:
         PyMem_Free(als)
 
+cdef int (*cuda_get_ipc_handle)(gpudata *, GpuArrayIpcMemHandle *)
+cdef gpudata *(*cuda_open_ipc_handle)(gpucontext *, GpuArrayIpcMemHandle *, size_t)
+
+cuda_get_ipc_handle = <int (*)(gpudata *, GpuArrayIpcMemHandle *)>gpuarray_get_extension("cuda_get_ipc_handle")
+cuda_open_ipc_handle = <gpudata *(*)(gpucontext *, GpuArrayIpcMemHandle *, size_t)>gpuarray_get_extension("cuda_open_ipc_handle")
+
+def open_ipc_handle(GpuContext c, bytes hpy, size_t l):
+    """
+    Open an IPC handle to get a new GpuArray from it.
+    """
+    cdef char *b
+    cdef GpuArrayIpcMemHandle h
+    cdef gpudata *d
+
+    b = hpy
+    memcpy(&h, b, sizeof(h))
+
+    d = cuda_open_ipc_handle(c.ctx, &h, l)
+    if d is NULL:
+        raise GpuArrayException, "could not open handle"
+    return <size_t>d
+
 cdef class GpuArray:
     """
     Device array
@@ -1561,6 +1583,19 @@ cdef class GpuArray:
             raise ValueError, "GpuArray and Numpy array do not have the same size in bytes"
         array_read(np.PyArray_DATA(dst), sz, self)
 
+    def get_ipc_handle(self):
+        cdef GpuArrayIpcMemHandle h
+        cdef int err
+        if cuda_get_ipc_handle is NULL:
+            raise SystemError, "Could not get necessary extension"
+        if self.context.kind == 'cuda':
+            raise ValueError, "Only works for cuda contexts"
+        err = cuda_get_ipc_handle(self.ga.data, &h)
+        if err != GA_NO_ERROR:
+            raise get_exc(err), GpuArray_error(&self.ga, err)
+        res = <bytes>(<char *>&h)[:sizeof(h)]
+        return res
+
     def __array__(self):
         """
         __array__()

src/gpuarray_buffer_cuda.c

@@ -560,27 +560,22 @@ static gpudata *cuda_alloc(gpucontext *c, size_t size, void *data, int flags,
   return res;
 }
 
-GpuArrayIpcMemHandle cuda_get_ipc_handle(gpudata *d) {
-  CUipcMemHandle h = {{0}};
-
+int cuda_get_ipc_handle(gpudata *d, GpuArrayIpcMemHandle *h) {
   ASSERT_BUF(d);
   cuda_enter(d->ctx);
-  d->ctx->err = cuIpcGetMemHandle(&h, d->ptr);
+  CUDA_EXIT_ON_ERROR(d->ctx,
+                     cuIpcGetMemHandle((CUipcMemHandle *)h, d->ptr));
   cuda_exit(d->ctx);
-  /* You need to do this stupid dance because direct casting of
-   * structures is not allowed */
-  return *((GpuArrayIpcMemHandle *)&h);
+  return GA_NO_ERROR;
 }
 
-gpudata *cuda_open_ipc_handle(gpucontext *c, GpuArrayIpcMemHandle h, size_t sz) {
+gpudata *cuda_open_ipc_handle(gpucontext *c, GpuArrayIpcMemHandle *h, size_t sz) {
   CUdeviceptr p;
   cuda_context *ctx = (cuda_context *)c;
   gpudata *d = NULL;
 
   cuda_enter(ctx);
-  /* You need to do this stupid dance because direct casting of
-   * structures is not allowed */
-  ctx->err = cuIpcOpenMemHandle(&p, *((CUipcMemHandle *)&h),
+  ctx->err = cuIpcOpenMemHandle(&p, *((CUipcMemHandle *)h),
                                 CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS);
   if (ctx->err == CUDA_SUCCESS) {
     d = cuda_make_buf(ctx, p, sz);