Merge pull request #270 from abergeron/send_buffer

Add a way to use the cuda IPC functionality to send and receive buffers across processes.

Author: nouiz

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

@@ -789,7 +789,7 @@ def from_gpudata(size_t data, offset, dtype, shape, GpuContext context=None,
     :type shape: iterable of ints
     :param context: context of the gpudata
     :type context: GpuContext
-    :param strides: strides for the results
+    :param strides: strides for the results (C contiguous if not specified)
     :type strides: iterable of ints
     :param writable: is the data writable?
     :type writeable: bool
@@ -839,7 +839,7 @@ def from_gpudata(size_t data, offset, dtype, shape, GpuContext context=None,
         else:
             size = gpuarray_get_elsize(typecode)
             for i in range(nd-1, -1, -1):
-                strides[i] = size
+                cstrides[i] = size
                 size *= cdims[i]
 
         return pygpu_fromgpudata(<gpudata *>data, offset, typecode, nd, cdims,
@@ -1424,6 +1424,33 @@ 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.
+
+    :param c: context
+    :param hpy: binary handle data received
+    :param l: size of the referred memory block
+
+    """
+    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 +1588,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 != b'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/ext_cuda.h

@@ -16,10 +16,12 @@ static void (*cuda_exit)(gpucontext *);
 static gpucontext *(*cuda_make_ctx)(CUcontext, int);
 static CUstream (*cuda_get_stream)(void *);
 static gpudata *(*cuda_make_buf)(void *, CUdeviceptr, size_t);
-static CUdeviceptr (*cuda_get_ptr)(gpudata *);
 static size_t (*cuda_get_sz)(gpudata *);
 static int (*cuda_wait)(gpudata *, int);
 static int (*cuda_record)(gpudata *, int);
+static CUipcMemHandle (*cuda_get_ipc_handle)(gpudata *d);
+static gpudata *(*cuda_open_ipc_handle)(gpucontext *c, CUipcMemHandle h,
+                                        size_t sz);
 
 static void setup_ext_cuda(void) {
   // The casts are necessary to reassure C++ compilers
@@ -28,10 +30,11 @@ static void setup_ext_cuda(void) {
   cuda_make_ctx = (gpucontext *(*)(CUcontext, int))gpuarray_get_extension("cuda_make_ctx");
   cuda_get_stream = (CUstream (*)(void *))gpuarray_get_extension("cuda_get_stream");
   cuda_make_buf = (gpudata *(*)(void *, CUdeviceptr, size_t))gpuarray_get_extension("cuda_make_buf");
-  cuda_get_ptr = (CUdeviceptr (*)(gpudata *))gpuarray_get_extension("cuda_get_ptr");
   cuda_get_sz = (size_t (*)(gpudata *))gpuarray_get_extension("cuda_get_sz");
   cuda_wait = (int (*)(gpudata *, int))gpuarray_get_extension("cuda_wait");
   cuda_record = (int (*)(gpudata *, int))gpuarray_get_extension("cuda_record");
+  cuda_get_ipc_handle = (CUipcMemHandle (*)(gpudata *))gpuarray_get_extension("cuda_get_ipc_handle");
+  cuda_open_ipc_handle = (gpudata *(*)(gpucontext *c, CUipcMemHandle h, size_t sz))gpuarray_get_extension("cuda_open_ipc_handle");
 }
 
 #ifdef __cplusplus

src/gpuarray/extension.h

@@ -19,6 +19,10 @@ extern "C" {
 #define GPUARRAY_CUDA_WAIT_READ  0x10000 /* CUDA_WAIT_READ */
 #define GPUARRAY_CUDA_WAIT_WRITE 0x20000 /* CUDA_WAIT_WRITE */
 
+typedef struct _GpuArrayIpcMemHandle {
+  char priv[64];
+} GpuArrayIpcMemHandle;
+
 /**
  * Obtain a function pointer for an extension.
  *

src/gpuarray_buffer_cuda.c

@@ -16,9 +16,15 @@
 #include "gpuarray/buffer.h"
 #include "gpuarray/util.h"
 #include "gpuarray/error.h"
-#include "gpuarray/extension.h"
 #include "gpuarray/buffer_blas.h"
 
+#include "gpuarray/extension.h"
+
+STATIC_ASSERT(DONTFREE == GPUARRAY_CUDA_CTX_NOFREE, cuda_nofree_eq);
+STATIC_ASSERT(CUDA_WAIT_READ == GPUARRAY_CUDA_WAIT_READ, cuda_wait_read_eq);
+STATIC_ASSERT(CUDA_WAIT_WRITE == GPUARRAY_CUDA_WAIT_WRITE, cuda_wait_write_eq);
+STATIC_ASSERT(sizeof(GpuArrayIpcMemHandle) == sizeof(CUipcMemHandle), cuda_ipcmem_eq);
+
 /* Allocations will be made in blocks of at least this size */
 #define BLOCK_SIZE (4 * 1024 * 1024)
 
@@ -555,6 +561,32 @@ static gpudata *cuda_alloc(gpucontext *c, size_t size, void *data, int flags,
   return res;
 }
 
+int cuda_get_ipc_handle(gpudata *d, GpuArrayIpcMemHandle *h) {
+  ASSERT_BUF(d);
+  cuda_enter(d->ctx);
+  CUDA_EXIT_ON_ERROR(d->ctx,
+                     cuIpcGetMemHandle((CUipcMemHandle *)h, d->ptr));
+  cuda_exit(d->ctx);
+  return GA_NO_ERROR;
+}
+
+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);
+  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);
+    if (d != NULL)
+      d->flags |= CUDA_IPC_MEMORY;
+  }
+  cuda_exit(ctx);
+  return d;
+}
+
 static void cuda_retain(gpudata *d) {
   ASSERT_BUF(d);
   d->refcnt++;
@@ -580,6 +612,9 @@ static void cuda_free(gpudata *d) {
     if (d->flags & DONTFREE) {
       /* This is the path for "external" buffers */
       deallocate(d);
+    } else if (d->flags & CUDA_IPC_MEMORY) {
+      cuIpcCloseMemHandle(d->ptr);
+      deallocate(d);
     } else if (ctx->flags & GA_CTX_DISABLE_ALLOCATION_CACHE) {
       /* Just free the pointer */
       cuMemFree(d->ptr);
@@ -1354,12 +1389,16 @@ static int cuda_sync(gpudata *b) {
 
   ASSERT_BUF(b);
   cuda_enter(ctx);
-  ctx->err = cuEventSynchronize(b->wev);
-  if (ctx->err != CUDA_SUCCESS)
-    err = GA_IMPL_ERROR;
-  ctx->err = cuEventSynchronize(b->rev);
-  if (ctx->err != CUDA_SUCCESS)
-    err = GA_IMPL_ERROR;
+  if (ctx->flags & GA_CTX_SINGLE_STREAM) {
+    cuStreamSynchronize(ctx->s);
+  } else {
+    ctx->err = cuEventSynchronize(b->wev);
+    if (ctx->err != CUDA_SUCCESS)
+      err = GA_IMPL_ERROR;
+    ctx->err = cuEventSynchronize(b->rev);
+    if (ctx->err != CUDA_SUCCESS)
+      err = GA_IMPL_ERROR;
+  }
   cuda_exit(ctx);
   return err;
 }

src/gpuarray_extension.c

@@ -16,6 +16,8 @@ extern void *cuda_make_buf(void);
 extern void *cuda_get_sz(void);
 extern void *cuda_wait(void);
 extern void *cuda_record(void);
+extern void *cuda_get_ipc_handle(void);
+extern void *cuda_open_ipc_handle(void);
 #endif
 #ifdef WITH_OPENCL
 extern void *cl_make_ctx(void);
@@ -34,6 +36,8 @@ static ext ext_list[] = {
   {"cuda_get_sz", cuda_get_sz},
   {"cuda_wait", cuda_wait},
   {"cuda_record", cuda_record},
+  {"cuda_get_ipc_handle", cuda_get_ipc_handle},
+  {"cuda_open_ipc_handle", cuda_open_ipc_handle},
 #endif
 #ifdef WITH_OPENCL
   {"cl_make_ctx", cl_make_ctx},

src/private_cuda.h

@@ -123,7 +123,6 @@ struct _gpudata {
 
 GPUARRAY_LOCAL gpudata *cuda_make_buf(cuda_context *c, CUdeviceptr p,
                                       size_t sz);
-GPUARRAY_LOCAL CUdeviceptr cuda_get_ptr(gpudata *g);
 GPUARRAY_LOCAL size_t cuda_get_sz(gpudata *g);
 GPUARRAY_LOCAL int cuda_wait(gpudata *, int);
 GPUARRAY_LOCAL int cuda_record(gpudata *, int);
@@ -135,8 +134,9 @@ GPUARRAY_LOCAL int cuda_record(gpudata *, int);
 
 #define CUDA_WAIT_ALL   (CUDA_WAIT_READ|CUDA_WAIT_WRITE)
 
-#define CUDA_HEAD_ALLOC 0x40000
-#define CUDA_MAPPED_PTR 0x80000
+#define CUDA_IPC_MEMORY 0x100000
+#define CUDA_HEAD_ALLOC 0x200000
+#define CUDA_MAPPED_PTR 0x400000
 
 struct _gpukernel {
   cuda_context *ctx; /* Keep the context first */