Openmpi cuda support (#3725)

* Update cuda to latest ga2 release

* Add CUDA support to openmpi

* Use spec's lib directories for cuda

* flake8 compliance

2 files changed, 14 insertions, 2 deletions

diff --git a/var/spack/repos/builtin/packages/cuda/package.py b/var/spack/repos/builtin/packages/cuda/package.py
index bf5a64fa59..eabb5a846c 100644
--- a/var/spack/repos/builtin/packages/cuda/package.py
+++ b/var/spack/repos/builtin/packages/cuda/package.py
@@ -38,6 +38,8 @@ class Cuda(Package):
 
     homepage = "http://www.nvidia.com/object/cuda_home_new.html"
 
+    version('8.0.61', '33e1bd980e91af4e55f3ef835c103f9b', expand=False,
+            url="https://developer.nvidia.com/compute/cuda/8.0/Prod2/local_installers/cuda_8.0.61_375.26_linux-run")
     version('8.0.44', '6dca912f9b7e2b7569b0074a41713640', expand=False,
             url="https://developer.nvidia.com/compute/cuda/8.0/prod/local_installers/cuda_8.0.44_linux-run")
     version('7.5.18', '4b3bcecf0dfc35928a0898793cf3e4c6', expand=False,
diff --git a/var/spack/repos/builtin/packages/openmpi/package.py b/var/spack/repos/builtin/packages/openmpi/package.py
index 7c07c515fa..63579efe0e 100644
--- a/var/spack/repos/builtin/packages/openmpi/package.py
+++ b/var/spack/repos/builtin/packages/openmpi/package.py
@@ -105,14 +105,14 @@ class Openmpi(AutotoolsPackage):
     variant('vt', default=True, description='Build VampirTrace support')
     variant('thread_multiple', default=False,
             description='Enable MPI_THREAD_MULTIPLE support')
-
-    # TODO: support for CUDA is missing
+    variant('cuda', default=False, description='Enable CUDA support')
 
     provides('mpi@:2.2', when='@1.6.5')
     provides('mpi@:3.0', when='@1.7.5:')
     provides('mpi@:3.1', when='@2.0.0:')
 
     depends_on('hwloc')
+    depends_on('hwloc +cuda', when='+cuda')
     depends_on('jdk', when='+java')
     depends_on('sqlite', when='+sqlite3')
 
@@ -255,6 +255,16 @@ class Openmpi(AutotoolsPackage):
             else:
                 config_args.append('--disable-mpi-thread-multiple')
 
+        # CUDA support
+        if spec.satisfies('@1.6:'):
+            if '+cuda' in spec:
+                config_args.append('--with-cuda={0}'.format(
+                    spec['cuda'].prefix))
+                config_args.append('--with-cuda-libdir={0}'.format(
+                    spec['cuda'].libs.directories))
+            else:
+                config_args.append('--without-cuda')
+
         return config_args
 
     @run_after('install')