Update: Ginkgo ROCm setup and smoke tests (#23280)

Add Ginkgo ROCmPackage support and improve tests.

3 files changed, 95 insertions, 21 deletions

diff --git a/var/spack/repos/builtin/packages/ginkgo/package.py b/var/spack/repos/builtin/packages/ginkgo/package.py
index 9ab3115860..a5af280996 100644
--- a/var/spack/repos/builtin/packages/ginkgo/package.py
+++ b/var/spack/repos/builtin/packages/ginkgo/package.py
@@ -7,7 +7,7 @@ from spack import *
 import sys
 
 
-class Ginkgo(CMakePackage, CudaPackage):
+class Ginkgo(CMakePackage, CudaPackage, ROCmPackage):
     """High-performance linear algebra library for manycore systems,
     with a focus on sparse solution of linear systems."""
 
@@ -28,29 +28,36 @@ class Ginkgo(CMakePackage, CudaPackage):
     variant('full_optimizations', default=False, description='Compile with all optimizations')
     variant('openmp', default=sys.platform != 'darwin',  description='Build with OpenMP')
     variant('develtools', default=False, description='Compile with develtools enabled')
+    variant('hwloc', default=False, description='Enable HWLOC support')
     variant('build_type', default='Release',
             description='The build type to build',
             values=('Debug', 'Release'))
-    variant('hip', default=False, description='Compile Ginkgo with HIP support')
 
     depends_on('cmake@3.9:', type='build')
     depends_on('cuda@9:',    when='+cuda')
 
-    depends_on('hip',       when='+hip')
-    depends_on('hipsparse', type="link", when='+hip')
-    depends_on('hipblas',   type="link", when='+hip')
-    depends_on('rocrand',   type="link", when='@develop+hip')
-    depends_on('rocthrust',   type="build", when='+hip')
+    depends_on('rocthrust',  type="build", when='+rocm')
+    depends_on('hipsparse',  type="link", when='+rocm')
+    depends_on('hipblas',    type="link", when='+rocm')
+    depends_on('rocrand',    type="link", when='+rocm')
+    depends_on('hwloc@2.1:', type="link", when='+hwloc')
 
-    # Somehow, these dependencies not propagated by the HIP stack?
-    depends_on('rocm-device-libs',   type="link", when='+hip')
-    depends_on('comgr',   type="link", when='+hip')
+    depends_on('googletest', type="test")
+    depends_on('numactl',    type="test", when="+hwloc")
 
     conflicts('%gcc@:5.2.9')
-    conflicts("+hip", when="@:1.1.1")
-    # The HIP packages from spack doen't seem to work well with CUDA
-    # backend for now, so disable HIP with CUDA backend.
-    conflicts("+cuda", when="+hip")
+    conflicts("+rocm", when="@:1.1.1")
+    conflicts("+cuda", when="+rocm")
+
+    # ROCm 4.1.0 breaks platform settings which breaks Ginkgo's HIP support.
+    conflicts("^hip@4.1.0:", when="@:1.3.0")
+    conflicts("^hip@4.1.0:", when="@master")
+    conflicts("^hipblas@4.1.0:", when="@:1.3.0")
+    conflicts("^hipblas@4.1.0:", when="@master")
+    conflicts("^hipsparse@4.1.0:", when="@:1.3.0")
+    conflicts("^hipsparse@4.1.0:", when="@master")
+    conflicts("^rocthrust@4.1.0:", when="@:1.3.0")
+    conflicts("^rocthrust@4.1.0:", when="@master")
 
     def cmake_args(self):
         # Check that the have the correct C++ standard is available
@@ -68,21 +75,34 @@ class Ginkgo(CMakePackage, CudaPackage):
         spec = self.spec
         args = [
             '-DGINKGO_BUILD_CUDA=%s' % ('ON' if '+cuda' in spec else 'OFF'),
+            '-DGINKGO_BUILD_HIP=%s' % ('ON' if '+rocm' in spec else 'OFF'),
             '-DGINKGO_BUILD_OMP=%s' % ('ON' if '+openmp' in spec else 'OFF'),
             '-DBUILD_SHARED_LIBS=%s' % ('ON' if '+shared' in spec else 'OFF'),
             '-DGINKGO_JACOBI_FULL_OPTIMIZATIONS=%s' % (
                 'ON' if '+full_optimizations' in spec else 'OFF'),
+            '-DGINKGO_BUILD_HWLOC=%s' % ('ON' if '+hwloc' in spec else 'OFF'),
             '-DGINKGO_DEVEL_TOOLS=%s' % (
                 'ON' if '+develtools' in spec else 'OFF'),
-            '-DGINKGO_BUILD_HIP=%s' % ('ON' if '+hip' in spec else 'OFF'),
             # As we are not exposing benchmarks, examples, tests nor doc
             # as part of the installation, disable building them altogether.
             '-DGINKGO_BUILD_BENCHMARKS=OFF',
             '-DGINKGO_BUILD_DOC=OFF',
             '-DGINKGO_BUILD_EXAMPLES=OFF',
-            '-DGINKGO_BUILD_TESTS=OFF'
+            '-DGINKGO_BUILD_TESTS=%s' % ('ON' if self.run_tests else 'OFF'),
+            # Let spack handle the RPATH
+            '-DGINKGO_INSTALL_RPATH=OFF'
         ]
-        if '+hip' in spec:
+
+        if self.run_tests:
+            args.append('-DGINKGO_USE_EXTERNAL_GTEST=ON')
+
+        if '+cuda' in spec:
+            archs = spec.variants['cuda_arch'].value
+            if archs != 'none':
+                arch_str = ";".join(archs)
+                args.append('-DGINKGO_CUDA_ARCHITECTURES={0}'.format(arch_str))
+
+        if '+rocm' in spec:
             args.append('-DHIP_PATH={0}'. format(spec['hip'].prefix))
             args.append('-DHIP_CLANG_PATH={0}/bin'.
                         format(spec['llvm-amdgpu'].prefix))
@@ -92,11 +112,43 @@ class Ginkgo(CMakePackage, CudaPackage):
                         format(spec['hipsparse'].prefix))
             args.append('-DHIPBLAS_PATH={0}'.
                         format(spec['hipblas'].prefix))
+            args.append('-DHIPRAND_PATH={0}/hiprand'.
+                        format(spec['rocrand'].prefix))
+            args.append('-DROCRAND_PATH={0}/rocrand'.
+                        format(spec['rocrand'].prefix))
+            archs = self.spec.variants['amdgpu_target'].value
+            if archs != 'none':
+                arch_str = ";".join(archs)
+                args.append(
+                    '-DGINKGO_HIP_AMDGPU={0}'.format(arch_str)
+                )
         return args
 
     @run_after('install')
-    @on_package_attributes(run_tests=True)
-    def test_install(self):
-        """Perform smoke tests on the installed package."""
+    def setup_build_tests(self):
+        """Build and install the smoke tests."""
+        # For now only develop and next releases support this scheme.
+        if not self.spec.satisfies('@develop') and not self.spec.satisfies('@1.4.0:'):
+            return
         with working_dir(self.build_directory):
             make("test_install")
+        smoke_test_path = join_path(self.build_directory, 'test_install')
+        with working_dir(smoke_test_path):
+            make("install")
+
+    def test(self):
+        """Run the smoke tests."""
+        # For now only develop and next releases support this scheme.
+        if not self.spec.satisfies('@develop') and not self.spec.satisfies('@1.4.0:'):
+            print("SKIPPED: smoke tests not supported with this Ginkgo version.")
+            return
+        files = [('test_install', [r'REFERENCE',
+                                   r'correctly detected and is complete']),
+                 ('test_install_cuda', [r'CUDA',
+                                        r'correctly detected and is complete']),
+                 ('test_install_hip', [r'HIP',
+                                       r'correctly detected and is complete'])]
+        smoke_test_path = join_path(self.prefix, 'smoke_tests')
+        for f, expected in files:
+            self.run_test(f, [], expected, skip_missing=True, installed=True,
+                          work_dir=smoke_test_path)
diff --git a/var/spack/repos/builtin/packages/hipsparse/package.py b/var/spack/repos/builtin/packages/hipsparse/package.py
index abdd53144f..43d09c9adf 100644
--- a/var/spack/repos/builtin/packages/hipsparse/package.py
+++ b/var/spack/repos/builtin/packages/hipsparse/package.py
@@ -30,7 +30,7 @@ class Hipsparse(CMakePackage):
     for ver in ['3.5.0', '3.7.0', '3.8.0', '3.9.0', '3.10.0', '4.0.0', '4.1.0']:
         depends_on('rocm-cmake@' + ver, type='build', when='@' + ver)
         depends_on('rocm-device-libs@' + ver, type='build', when='@' + ver)
-        depends_on('rocsparse@' + ver, type='build', when='@' + ver)
+        depends_on('rocsparse@' + ver, type='link', when='@' + ver)
         depends_on('hip@' + ver, when='@' + ver)
         depends_on('comgr@' + ver, type='build', when='@' + ver)
         depends_on('hsa-rocr-dev@' + ver, type='link', when='@' + ver)
diff --git a/var/spack/repos/builtin/packages/rocrand/package.py b/var/spack/repos/builtin/packages/rocrand/package.py
index 85054c23ec..d2b821179c 100644
--- a/var/spack/repos/builtin/packages/rocrand/package.py
+++ b/var/spack/repos/builtin/packages/rocrand/package.py
@@ -3,6 +3,8 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+import os
+import glob
 
 from spack import *
 
@@ -41,6 +43,26 @@ class Rocrand(CMakePackage):
     def setup_build_environment(self, env):
         env.set('CXX', self.spec['hip'].hipcc)
 
+    @run_after('install')
+    def fix_library_locations(self):
+        """Fix the rocRAND and hipRAND libraries location"""
+        # rocRAND installs librocrand.so* and libhiprand.so* to rocrand/lib and
+        # hiprand/lib, respectively. This confuses spack's RPATH management. We
+        # fix it by adding a symlink to the libraries.
+        hiprand_lib_path = join_path(self.prefix, 'hiprand', 'lib')
+        rocrand_lib_path = join_path(self.prefix, 'rocrand', 'lib')
+        mkdirp(self.prefix.lib)
+        with working_dir(hiprand_lib_path):
+            hiprand_libs = glob.glob('*.so*')
+            for lib in hiprand_libs:
+                os.symlink(join_path(hiprand_lib_path, lib),
+                           join_path(self.prefix.lib, lib))
+        with working_dir(rocrand_lib_path):
+            rocrand_libs = glob.glob('*.so*')
+            for lib in rocrand_libs:
+                os.symlink(join_path(rocrand_lib_path, lib),
+                           join_path(self.prefix.lib, lib))
+
     def cmake_args(self):
         args = ['-DBUILD_BENCHMARK=OFF',
                 '-DBUILD_TEST=OFF']