Merge pull request #83 from helmholtz-analytics/features/82-cuda-awar…

…e-mpi

Added CUDA-aware MPI checks

heat/core/communication.py

@@ -1,9 +1,23 @@
 from mpi4py import MPI
 import abc
+import subprocess
 import torch
 
 from .stride_tricks import sanitize_axis
 
+# check whether OpenMPI support CUDA-aware MPI
+try:
+    buffer = subprocess.check_output(['mpirun', '--help'])
+
+    # OpenMPI
+    if buffer.startswith(b'mpirun (Open MPI)'):
+        buffer = subprocess.check_output(['ompi_info', '--parsable', '--all'])
+        CUDA_AWARE_MPI = b'mpi_built_with_cuda_support:value:true' in buffer
+    else:
+        CUDA_AWARE_MPI = False
+except FileNotFoundError:
+    CUDA_AWARE_MPI = False
+
 
 class Communication(metaclass=abc.ABCMeta):
     @staticmethod
@@ -93,11 +107,12 @@ def chunk(self, shape, split):
     def as_buffer(obj):
         if isinstance(obj, tensor.tensor):
             obj = obj._tensor__array
-
         if not isinstance(obj, torch.Tensor):
             return obj
 
-        return MPI.memory.fromaddress(obj.cpu().data_ptr(), obj.element_size() * torch.numel(obj))
+        pointer = obj.data_ptr() if CUDA_AWARE_MPI else obj.cpu().data_ptr()
+
+        return MPI.memory.fromaddress(pointer, obj.element_size() * torch.numel(obj))
 
     def convert_tensors(self, a_callable):
         def wrapped(*args, **kwargs):

heat/core/tests/test_communicator.py → heat/core/tests/test_communication.py

@@ -54,3 +54,6 @@ def test_mpi_communicator(self):
 
         self.assertIsInstance(chunks, tuple)
         self.assertEqual(len(chunks), len(self.data.shape))
+
+    def test_cuda_aware_mpi(self):
+        self.assertTrue(hasattr(ht.communication, 'CUDA_AWARE_MPI'))