Lbann update (#6751)

* Added a package for the MDAnalysis toolkit.

* Updated the LBANN recipe to use the revised cmake environment being
rolled out.

* Flake8 errors

1 files changed, 59 insertions, 2 deletions

diff --git a/var/spack/repos/builtin/packages/lbann/package.py b/var/spack/repos/builtin/packages/lbann/package.py
index 5e65566cb2..8cfd6e4c62 100644
--- a/var/spack/repos/builtin/packages/lbann/package.py
+++ b/var/spack/repos/builtin/packages/lbann/package.py
@@ -34,12 +34,17 @@ class Lbann(CMakePackage):
     url      = "https://github.com/LLNL/lbann/archive/v0.91.tar.gz"
 
     version('develop', git='https://github.com/LLNL/lbann.git', branch="develop")
+    version('0.93', '1913a25a53d4025fa04c16f14afdaa55')
+    version('0.92', 'c0eb1595a7c74640e96f280beb497564')
     version('0.91', '83b0ec9cd0b7625d41dfb06d2abd4134')
 
     variant('gpu', default=False, description='Builds with support for GPUs via CUDA and cuDNN')
+    variant('nccl', default=False, description='Builds with support for NCCL communication lib')
     variant('opencv', default=True, description='Builds with support for image processing routines with OpenCV')
     variant('seq_init', default=False, description='Force serial initialization of weight matrices.')
-    variant('dtype', default=4, description='Size (bits) of floating point representation for weights')
+    variant('dtype', default='float',
+            description='Type for floating point representation of weights',
+            values=('float', 'double'))
     variant('build_type', default='Release',
             description='The build type to build',
             values=('Debug', 'Release'))
@@ -55,7 +60,55 @@ class Lbann(CMakePackage):
     depends_on('opencv@3.2.0: +openmp +core +highgui +imgproc +jpeg +png +tiff +zlib ~eigen', when='+opencv')
     depends_on('protobuf@3.0.2:')
     depends_on('cnpy')
+    depends_on('nccl', when='+gpu +nccl')
 
+    @when('@0.94:')
+    def cmake_args(self):
+        spec = self.spec
+        # Environment variables
+        CPPFLAGS = []
+        CPPFLAGS.append('-DLBANN_SET_EL_RNG')
+
+        args = [
+            '-DCMAKE_INSTALL_MESSAGE=LAZY',
+            '-DCMAKE_CXX_FLAGS=%s' % ' '.join(CPPFLAGS),
+            '-DLBANN_WITH_TOPO_AWARE:BOOL=%s' % ('+gpu +nccl' in spec),
+            '-DLBANN_SEQUENTIAL_INITIALIZATION:BOOL=%s' %
+            ('+seq_init' in spec),
+            '-DLBANN_WITH_TBINF=OFF',
+            '-DLBANN_WITH_VTUNE=OFF',
+            '-DElemental_DIR={0}/CMake/elemental'.format(
+                spec['elemental'].prefix),
+            '-DCNPY_DIR={0}'.format(spec['cnpy'].prefix),
+            '-DLBANN_DATATYPE={0}'.format(spec.variants['dtype'].value),
+            '-DLBANN_VERBOSE=0',
+            '-DLBANN_VERSION=spack']
+
+        if '+opencv' in spec:
+            args.extend(['-DOpenCV_DIR:STRING={0}'.format(
+                spec['opencv'].prefix)])
+
+        if '+gpu' in spec:
+            args.extend([
+                '-DLBANN_WITH_CUDA:BOOL=%s' % ('+gpu' in spec),
+                '-DLBANN_WITH_SOFTMAX_CUDA:BOOL=%s' % ('+gpu' in spec),
+                '-DCUDA_TOOLKIT_ROOT_DIR={0}'.format(
+                    spec['cuda'].prefix)])
+            args.extend([
+                '-DLBANN_WITH_CUDNN:BOOL=%s' % ('+gpu' in spec),
+                '-DcuDNN_DIR={0}'.format(
+                    spec['cudnn'].prefix)])
+            args.extend(['-DCUB_DIR={0}'.format(
+                spec['cub'].prefix)])
+            if '+nccl' in spec:
+                args.extend([
+                    '-DLBANN_WITH_NCCL:BOOL=%s' % ('+gpu +nccl' in spec),
+                    '-DNCCL_DIR={0}'.format(
+                        spec['nccl'].prefix)])
+
+        return args
+
+    @when('@:0.93')
     def cmake_args(self):
         spec = self.spec
         # Environment variables
@@ -76,11 +129,15 @@ class Lbann(CMakePackage):
             '-DELEMENTAL_MATH_LIBS={0}'.format(
                 spec['elemental'].libs),
             '-DSEQ_INIT:BOOL=%s' % ('+seq_init' in spec),
-            '-DDATATYPE={0}'.format(int(spec.variants['dtype'].value)),
             '-DVERBOSE=0',
             '-DLBANN_HOME=.',
             '-DLBANN_VER=spack']
 
+        if spec.variants['dtype'].value == 'float':
+            args.extend(['-DDATATYPE=4'])
+        elif spec.variants['dtype'].value == 'double':
+            args.extend(['-DDATATYPE=8'])
+
         if '+opencv' in spec:
             args.extend(['-DOpenCV_DIR:STRING={0}'.format(
                 spec['opencv'].prefix)])