Added example script

GPUSimulators/CudaContext.py

@@ -44,8 +44,12 @@ Class which keeps track of the CUDA context and some helper functions
 """
 class CudaContext(object):
     
-    def __init__(self, blocking=False, use_cache=True, autotuning=True):
-        self.blocking = blocking
+    def __init__(self, device=None, context_flags=None, use_cache=True, autotuning=True):
+        """
+        Create a new CUDA context
+        Set device to an id or pci_bus_id to select a specific GPU
+        Set context_flags to cuda.ctx_flags.SCHED_BLOCKING_SYNC for a blocking context
+        """
         self.use_cache = use_cache
         self.logger =  logging.getLogger(__name__)
         self.modules = {}
@@ -60,17 +64,19 @@ class CudaContext(object):
         #Print some info about CUDA
         self.logger.info("CUDA version %s", str(cuda.get_version()))
         self.logger.info("Driver version %s",  str(cuda.get_driver_version()))
-
-        self.cuda_device = cuda.Device(0)
-        self.logger.info("Using '%s' GPU", self.cuda_device.name())
+        
+        if device is None:
+            device = 0
+        
+        self.cuda_device = cuda.Device(device)
+        self.logger.info("Using device %d/%d '%s' (%s) GPU", device, cuda.Device.count(), self.cuda_device.name(), self.cuda_device.pci_bus_id())
         self.logger.debug(" => compute capability: %s", str(self.cuda_device.compute_capability()))
 
         # Create the CUDA context
-        if (self.blocking):
-            self.cuda_context = self.cuda_device.make_context(flags=cuda.ctx_flags.SCHED_BLOCKING_SYNC)
-            self.logger.warning("Using blocking context")
-        else:
-            self.cuda_context = self.cuda_device.make_context(flags=cuda.ctx_flags.SCHED_AUTO)
+        if context_flags is None:
+            context_flags=cuda.ctx_flags.SCHED_AUTO
+            
+        self.cuda_context = self.cuda_device.make_context(flags=context_flags)
             
         free, total = cuda.mem_get_info()
         self.logger.debug(" => memory: %d / %d MB available", int(free/(1024*1024)), int(total/(1024*1024)))

GPUSimulators/MPISimulator.py

@@ -143,6 +143,55 @@ class MPIGrid(object):
             out_data = np.empty([self.comm.size] + list(data.shape), dtype=data.dtype)
         self.comm.Gather(data, out_data, root)
         return out_data
+        
+    def getLocalRank(self):
+        """
+        Returns the local rank on this node for this MPI process
+        """
+        
+        # This function has been adapted from 
+        # https://github.com/SheffieldML/PyDeepGP/blob/master/deepgp/util/parallel.py
+        # by Zhenwen Dai released under BSD 3-Clause "New" or "Revised" License:
+        # 
+        # Copyright (c) 2016, Zhenwen Dai
+        # All rights reserved.
+        # 
+        # Redistribution and use in source and binary forms, with or without
+        # modification, are permitted provided that the following conditions are met:
+        # 
+        # * Redistributions of source code must retain the above copyright notice, this
+        #   list of conditions and the following disclaimer.
+        # 
+        # * Redistributions in binary form must reproduce the above copyright notice,
+        #   this list of conditions and the following disclaimer in the documentation
+        #   and/or other materials provided with the distribution.
+        # 
+        # * Neither the name of DGP nor the names of its
+        #   contributors may be used to endorse or promote products derived from
+        #   this software without specific prior written permission.
+        # 
+        # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+        # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+        # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+        # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+        # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+        # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+        # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+        # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+        # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+        # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+        
+        #Get this ranks unique (physical) node name
+        node_name = MPI.Get_processor_name()
+        
+        #Gather the list of all node names on all nodes
+        node_names = self.comm.allgather(node_name)
+                
+        #Loop over all node names up until our rank
+        #and count how many duplicates of our nodename we find
+        local_rank = len([0 for name in node_names[:self.comm.rank] if name==node_name])
+        
+        return local_rank
 
 
 class MPISimulator(Simulator.BaseSimulator):
@@ -233,9 +282,7 @@ class MPISimulator(Simulator.BaseSimulator):
         self.out_n = np.empty_like(self.in_n)
         self.out_s = np.empty_like(self.in_s)
         
-        self.logger.debug("Simlator rank {:d} has neighbors {:s}".format(self.grid.comm.rank, str([self.north, self.south, self.east, self.west])))
-        
-        self.logger.debug("Simlator rank {:d} initialized ".format(self.grid.comm.rank))
+        self.logger.debug("Simlator rank {:d} initialized on {:s}".format(self.grid.comm.rank, MPI.Get_processor_name()))
     
         
     def substep(self, dt, step_number):