From 40886599ceb71c9b7a487fe5fe0973b80bbd173d Mon Sep 17 00:00:00 2001 From: Satish Balay Date: Thu, 16 Dec 2021 16:47:31 -0600 Subject: petsc: enable +rocm+kokkos (#27779) Also fix kokkos-kernels when kokkos is built with +rocm --- .../builtin/packages/kokkos-kernels/package.py | 7 ++++-- var/spack/repos/builtin/packages/petsc/package.py | 25 +++++++++++++++++++++- 2 files changed, 29 insertions(+), 3 deletions(-) (limited to 'var') diff --git a/var/spack/repos/builtin/packages/kokkos-kernels/package.py b/var/spack/repos/builtin/packages/kokkos-kernels/package.py index c41a1607aa..8b046dd89f 100644 --- a/var/spack/repos/builtin/packages/kokkos-kernels/package.py +++ b/var/spack/repos/builtin/packages/kokkos-kernels/package.py @@ -102,8 +102,11 @@ class KokkosKernels(CMakePackage, CudaPackage): options.append("-DSpack_WORKAROUND=On") options.append("-DKokkos_ROOT=%s" % spec["kokkos"].prefix) - # Compiler weirdness due to nvcc_wrapper - options.append("-DCMAKE_CXX_COMPILER=%s" % spec["kokkos"].kokkos_cxx) + if spec.satisfies('^kokkos+rocm'): + options.append("-DCMAKE_CXX_COMPILER=%s" % spec['hip'].hipcc) + else: + # Compiler weirdness due to nvcc_wrapper + options.append("-DCMAKE_CXX_COMPILER=%s" % spec["kokkos"].kokkos_cxx) if self.run_tests: options.append("-DKokkosKernels_ENABLE_TESTS=ON") diff --git a/var/spack/repos/builtin/packages/petsc/package.py b/var/spack/repos/builtin/packages/petsc/package.py index f9e0255741..8c60a31675 100644 --- a/var/spack/repos/builtin/packages/petsc/package.py +++ b/var/spack/repos/builtin/packages/petsc/package.py @@ -221,6 +221,12 @@ class Petsc(Package, CudaPackage, ROCmPackage): depends_on('hip', when='+rocm') depends_on('hipblas', when='+rocm') depends_on('hipsparse', when='+rocm') + depends_on('rocsparse', when='+rocm') + depends_on('rocsolver', when='+rocm') + depends_on('rocblas', when='+rocm') + depends_on('rocrand', when='+rocm') + depends_on('rocthrust', when='+rocm') + depends_on('rocprim', when='+rocm') # Build dependencies depends_on('python@2.6:2.8', type='build', when='@:3.10') @@ -321,6 +327,7 @@ class Petsc(Package, CudaPackage, ROCmPackage): depends_on('kokkos-kernels', when='+kokkos') depends_on('kokkos+cuda+wrapper+cuda_lambda', when='+kokkos +cuda') depends_on('kokkos-kernels+cuda', when='+kokkos +cuda') + depends_on('kokkos+rocm', when='+kokkos +rocm') # Using the following tarballs # * petsc-3.12 (and older) - includes docs @@ -414,7 +421,7 @@ class Petsc(Package, CudaPackage, ROCmPackage): # if not (useinc || uselib): usedir - i.e (False, False) for library in ( ('cuda', 'cuda', False, False), - ('hip', 'hip', False, False), + ('hip', 'hip', True, False), 'metis', 'hypre', 'parmetis', @@ -496,6 +503,22 @@ class Petsc(Package, CudaPackage, ROCmPackage): else: options.append('CUDAFLAGS=-gencode arch=compute_{0},code=sm_{0}' .format(cuda_arch[0])) + if '+rocm' in spec: + if not spec.satisfies('amdgpu_target=none'): + hip_arch = spec.variants['amdgpu_target'].value + options.append('--with-hip-arch={0}'.format(hip_arch[0])) + hip_pkgs = ['hipsparse', 'hipblas', 'rocsparse', 'rocsolver', 'rocblas'] + hip_ipkgs = hip_pkgs + ['rocthrust', 'rocprim'] + hip_lpkgs = hip_pkgs + ['rocrand'] + hip_inc = '' + hip_lib = '' + for pkg in hip_ipkgs: + hip_inc += spec[pkg].headers.include_flags + ' ' + for pkg in hip_lpkgs: + hip_lib += spec[pkg].libs.joined() + ' ' + options.append('HIPPPFLAGS=%s' % hip_inc) + options.append('with-hip-lib=%s -L%s -lamdhip64' % + (hip_lib, spec['hip'].prefix.lib)) if 'superlu-dist' in spec: if spec.satisfies('@3.10.3:3.15'): -- cgit v1.2.3-70-g09d2