GPU-Aware Communication

charm4py/channel.py

@@ -1,4 +1,6 @@
 from .threads import LocalFuture
+from .charm import charm
+import time
 
 
 class Channel(object):
@@ -52,20 +54,55 @@ def ready(self):
     def waitReady(self, f):
         self.wait_ready = f
 
-    def send(self, *msg):
+    def send(self, *msg, **kwargs):
         if not self.established:
             self.established_fut = LocalFuture()
             self.established_fut.get()
             self.setEstablished()
-        self.remote._channelRecv__(self.remote_port, self.send_seqno, *msg)
+        self.remote._channelRecv__(self.remote_port, self.send_seqno, *msg, **kwargs)
         self.send_seqno = (self.send_seqno + 1) % CHAN_BUF_SIZE
 
-    def recv(self):
+    def recv(self, *post_buffers, post_addresses = None, post_sizes = None, stream_ptrs = None):
         if self.recv_seqno in self.data:
             ret = self.data.pop(self.recv_seqno)
         else:
             self.recv_fut = LocalFuture()
             ret = self.recv_fut.get()
             self.recv_fut = None
         self.recv_seqno = (self.recv_seqno + 1) % CHAN_BUF_SIZE
+
+        if post_buffers:
+            if isinstance(ret, tuple):
+                gpu_recv_bufs = ret[-1]
+                ret = ret[0:-1]
+                if len(ret) == 1:
+                    ret = ret[0]
+            else:
+                gpu_recv_bufs = ret
+
+            assert len(post_buffers) == len(gpu_recv_bufs)
+
+            recv_future = charm.getGPUDirectData(post_buffers, gpu_recv_bufs, stream_ptrs)
+            recv_future.get()
+        elif post_addresses is not None:
+            if isinstance(ret, tuple):
+                gpu_recv_bufs = ret[-1]
+                ret = ret[0:-1]
+                if len(ret) == 1:
+                    ret = ret[0]
+            else:
+                gpu_recv_bufs = ret
+
+            assert len(post_addresses) == len(gpu_recv_bufs)
+            assert post_sizes
+            recv_future = charm.getGPUDirectDataFromAddresses(post_addresses, post_sizes, gpu_recv_bufs, stream_ptrs)
+            recv_future.get()
+
+
         return ret
+
+
+
+
+
+

charm4py/chare.py

@@ -461,6 +461,9 @@ def proxy_entry_method(proxy, *args, **kwargs):
             for i in range(num_args, argcount):
                 argname = argnames[i]
                 # first look for argument in kwargs
+                # TODO: Should stream_ptrs be skipped?
+                if argname in {'stream_ptrs', 'src_ptrs', 'src_sizes'}:
+                    continue
                 if argname in kwargs:
                     args.append(kwargs[argname])
                 else:
@@ -485,8 +488,28 @@ def proxy_entry_method(proxy, *args, **kwargs):
             gid = proxy.gid
             if Options.local_msg_optim and (elemIdx == charm._myPe) and (len(args) > 0):
                 destObj = charm.groups[gid]
-            msg = charm.packMsg(destObj, args, header)
-            charm.CkGroupSend(gid, elemIdx, ep, msg)
+            should_pack_gpu = True
+            if 'src_ptrs' in kwargs:
+                should_pack_gpu = False
+            msg = charm.packMsg(destObj, args, header, pack_gpu=should_pack_gpu)
+            if msg[1] or not should_pack_gpu:
+                if 'stream_ptrs' in kwargs:
+                    stream_ptrs = kwargs['stream_ptrs']
+                else:
+                    stream_ptrs = None
+                if should_pack_gpu:
+                    charm.CkGroupSendWithDeviceData(gid, elemIdx, ep,
+                                                    msg, stream_ptrs
+                                                    )
+                else:
+                    charm.CkGroupSendWithDeviceDataFromPointers(gid, elemIdx, ep,
+                                                                msg, kwargs['src_ptrs'],
+                                                                kwargs['src_sizes'],
+                                                                stream_ptrs
+                                                                )
+
+            else:
+                charm.CkGroupSend(gid, elemIdx, ep, msg)
         else:
             root, sid = proxy.section
             header[b'sid'] = sid
@@ -721,7 +744,10 @@ def proxy_entry_method(proxy, *args, **kwargs):
             for i in range(num_args, argcount):
                 argname = argnames[i]
                 # first look for argument in kwargs
-                if argname in kwargs:
+                # TODO: Should stream_ptrs be skipped?
+                if argname in {'stream_ptrs', 'src_ptrs', 'src_sizes'}:
+                    continue
+                if argname in kwargs and argname:
                     args.append(kwargs[argname])
                 else:
                     # if not there, see if there is a default value
@@ -741,15 +767,39 @@ def proxy_entry_method(proxy, *args, **kwargs):
             if elemIdx == ():
                 header[b'bcast'] = True
         if not proxy.issec or elemIdx != ():
+            # TODO: Check that this is channel proxy method?
             destObj = None
             aid = proxy.aid
             if Options.local_msg_optim and (len(args) > 0):
                 array = charm.arrays[aid]
                 if elemIdx in array:
                     destObj = array[elemIdx]
-            msg = charm.packMsg(destObj, args, header)
-            charm.CkArraySend(aid, elemIdx, ep, msg)
+            should_pack_gpu = True
+            if 'src_ptrs' in kwargs:
+                should_pack_gpu = False
+            msg = charm.packMsg(destObj, args, header, pack_gpu = should_pack_gpu)
+            if msg[1] or not should_pack_gpu:
+                if 'stream_ptrs' in kwargs:
+                    stream_ptrs = kwargs['stream_ptrs']
+                else:
+                    stream_ptrs = None
+                if should_pack_gpu:
+                    charm.CkArraySendWithDeviceData(aid, elemIdx, ep,
+                                                    msg, stream_ptrs
+                                                    )
+                else:
+                    charm.CkArraySendWithDeviceDataFromPointers(aid, elemIdx, ep,
+                                                                msg, kwargs['src_ptrs'],
+                                                                kwargs['src_sizes'],
+                                                                stream_ptrs
+                                                                )
+
+
+
+            else:
+                charm.CkArraySend(aid, elemIdx, ep, msg)
         else:
+            # TODO: Error if trying to send ZC data
             root, sid = proxy.section
             header[b'sid'] = sid
             if Options.local_msg_optim and root == charm._myPe:

charm4py/charm.py

@@ -28,6 +28,8 @@
 from . import reduction
 from . import wait
 import array
+import numpy as np
+import greenlet
 try:
     import numpy
 except ImportError:
@@ -40,6 +42,14 @@ class NumpyDummy:
 def SECTION_ALL(obj):
     return 0
 
+def getDeviceDataInfo(devArray):
+    return devArray.__cuda_array_interface__['data']
+
+def getDeviceDataAddress(devArray):
+    return getDeviceDataInfo(devArray)[0]
+
+def getDeviceDataSizeInBytes(devArray):
+    return devArray.nbytes
 
 class Options(object):
 
@@ -106,11 +116,13 @@ def __init__(self):
         self.options.interactive.verbose = 1
         self.options.interactive.broadcast_imports = True
 
+        '''
         if 'OMPI_COMM_WORLD_SIZE' in os.environ:
             # this is needed for OpenMPI, see:
             # https://svn.open-mpi.org/trac/ompi/wiki/Linkers
             import ctypes
             self.__libmpi__ = ctypes.CDLL('libmpi.so', mode=ctypes.RTLD_GLOBAL)
+        '''
         self.lib = load_charm_library(self)
         self.ReducerType = self.lib.ReducerType
         self.CkContributeToChare = self.lib.CkContributeToChare
@@ -120,6 +132,13 @@ def __init__(self):
         self.CkChareSend = self.lib.CkChareSend
         self.CkGroupSend = self.lib.CkGroupSend
         self.CkArraySend = self.lib.CkArraySend
+        self.CkArraySendWithDeviceData = self.lib.CkArraySendWithDeviceData
+        self.CkGroupSendWithDeviceData = self.lib.CkGroupSendWithDeviceData
+        self.CkArraySendWithDeviceDataFromPointersArray = self.lib.CkArraySendWithDeviceDataFromPointersArray
+        self.CkArraySendWithDeviceDataFromPointersOther = self.lib.CkArraySendWithDeviceDataFromPointersOther
+        self.CkGroupSendWithDeviceDataFromPointersArray = self.lib.CkGroupSendWithDeviceDataFromPointersArray
+        self.CkGroupSendWithDeviceDataFromPointersOther = self.lib.CkGroupSendWithDeviceDataFromPointersOther
+        self.CkCudaEnabled = self.lib.CkCudaEnabled
         self.reducers = reduction.ReducerContainer(self)
         self.redMgr = reduction.ReductionManager(self, self.reducers)
         self.mainchareRegistered = False
@@ -305,6 +324,21 @@ def recvArrayMsg(self, aid, index, ep, msg, dcopy_start):
                 self.arrays[aid][index] = obj
                 em.run(obj, header, args)  # now call the user's array element __init__
 
+    def recvGPUDirectArrayMsg(self, aid, index, ep,
+                              devBuf_ptrs, msg, dcopy_start
+                              ):
+        obj = self.arrays[aid][index]
+        header, args = self.unpackMsg(msg, dcopy_start, obj)
+        args.append(devBuf_ptrs)
+
+        self.invokeEntryMethod(obj, ep, header, args)
+
+    def recvGPUDirectGroupMsg(self, gid, ep, devBuf_ptrs, msg, dcopy_start):
+        obj = self.groups[gid]
+        header, args = self.unpackMsg(msg, dcopy_start, obj)
+        args.append(devBuf_ptrs)
+        self.invokeEntryMethod(obj, ep, header, args)
+
     def recvArrayBcast(self, aid, indexes, ep, msg, dcopy_start):
         header, args = self.unpackMsg(msg, dcopy_start, None)
         array = self.arrays[aid]
@@ -332,6 +366,70 @@ def unpackMsg(self, msg, dcopy_start, dest_obj):
 
         return header, args
 
+    def getGPUDirectData(self, post_buffers, remote_bufs, stream_ptrs):
+        # this future will only be satisfied when all buffers have been received
+        return_fut = self.Future()
+        post_buf_data = [getDeviceDataAddress(buf) for buf in post_buffers]
+        post_buf_sizes = [getDeviceDataSizeInBytes(buf) for buf in post_buffers]
+        if not stream_ptrs:
+            stream_ptrs = [0] * len(post_buffers)
+        self.lib.getGPUDirectData(post_buf_data, post_buf_sizes, remote_bufs, stream_ptrs, return_fut)
+        return return_fut
+
+    def getGPUDirectDataFromAddresses(self, post_buf_ptrs, post_buf_sizes, remote_bufs, stream_ptrs):
+        # this future will only be satisfied when all buffers have been received
+        return_fut = self.Future()
+        if not stream_ptrs:
+            stream_ptrs = array.array('L', [0] * len(post_buf_ptrs))
+        self.lib.getGPUDirectDataFromAddresses(post_buf_ptrs, post_buf_sizes, remote_bufs, stream_ptrs, return_fut)
+        return return_fut
+
+    def CkArraySendWithDeviceDataFromPointers(self, array_id, index, ep,
+                                              msg, gpu_src_ptrs,
+                                              gpu_src_sizes,
+                                              stream_ptrs
+                                              ):
+        if isinstance(gpu_src_ptrs, array.array):
+            assert isinstance(gpu_src_sizes, array.array), \
+                "GPU source pointers and sizes must be of the same type."
+            self.CkArraySendWithDeviceDataFromPointersArray(array_id, index, ep,
+                                                            msg, gpu_src_ptrs,
+                                                            gpu_src_sizes,
+                                                            stream_ptrs,
+                                                            len(gpu_src_ptrs)
+                                                            )
+        else:
+            self.CkArraySendWithDeviceDataFromPointersOther(array_id, index, ep,
+                                                            msg, gpu_src_ptrs,
+                                                            gpu_src_sizes,
+                                                            stream_ptrs,
+                                                            len(gpu_src_ptrs)
+                                                            )
+
+    def CkGroupSendWithDeviceDataFromPointers(self, gid, elemIdx, ep,
+                                              msg, gpu_src_ptrs, gpu_src_sizes,
+                                              stream_ptrs):
+        if isinstance(gpu_src_ptrs, array.array):
+            assert isinstance(gpu_src_sizes, array.array), \
+                "GPU source pointers and sizes must be of the same type."
+            self.CkGroupSendWithDeviceDataFromPointersArray(gid, elemIdx, ep, msg,
+                                                            gpu_src_ptrs,
+                                                            gpu_src_sizes,
+                                                            stream_ptrs,
+                                                            len(gpu_src_ptrs)
+                                                            )
+        else:
+            self.CkGroupSendWithDeviceDataFromPointersOther(gid, elemIdx, ep, msg,
+                                                            gpu_src_ptrs,
+                                                            gpu_src_sizes,
+                                                            stream_ptrs,
+                                                            len(gpu_src_ptrs)
+                                                            )
+
+    # deposit value of one of the futures that was created on this PE
+    def _future_deposit_result(self, fid, result=None):
+        self.threadMgr.depositFuture(fid, result)
+
     def packMsg(self, destObj, msgArgs, header):
         """Prepares a message for sending, given arguments to an entry method invocation.
 
@@ -815,6 +913,7 @@ def triggerCallable(self, tag):
     def iwait(self, objs):
         n = len(objs)
         f = LocalFuture()
+
         for obj in objs:
             if obj.ready():
                 n -= 1
@@ -826,6 +925,43 @@ def iwait(self, objs):
             n -= 1
             yield obj
 
+    def iwait_map(self, func, objs):
+        n = len(objs)
+        f = LocalFuture()
+        remaining_grs = [n]
+
+        def map_func(remaining, obj):
+            gr = greenlet.getcurrent()
+            gr.notify = gr.parent.notify
+            gr.obj = gr.parent.obj
+            gr.fu = 1
+            func(obj)
+            remaining[0] -= 1
+
+        def gr_func():
+            return map_func(remaining_grs, obj)
+
+        for obj in objs:
+            if obj.ready():
+                new_gr = greenlet.greenlet(gr_func)
+                n -= 1
+                obj = new_gr.switch()
+                while obj:
+                    new_gr = greenlet.greenlet(gr_func)
+                    n -= 1
+                    obj = new_gr.switch()
+            else:
+                obj.waitReady(f)
+        while n > 0:
+            obj = self.threadMgr.pauseThread()
+            while obj:
+                new_gr = greenlet.greenlet(gr_func)
+                n -= 1
+                obj = new_gr.switch()
+
+        while remaining_grs[0]:
+            self.threadMgr.pauseThread()
+
     def wait(self, objs):
         for o in self.iwait(objs):
             pass
@@ -1156,6 +1292,5 @@ def rebuildNumpyArray(data, shape, dt):
     a.shape = shape
     return a.copy()
 
-
 charm = Charm()
 readonlies = __ReadOnlies()