Add ROCm support for Spock (#60)

configs/spock/packages.yaml

@@ -12,14 +12,24 @@ packages:
     version: [4.3.0]
     buildable: false
     externals:
-      - spec: "hip@4.3.0%gcc@9.3.0"
-        modules:
-          - rocm/4.3.0
+      - spec: hip@4.3.0
+        prefix: /opt/rocm-4.3.0/hip
+  hsa-rocr-dev:
+    version: [4.3.0]
+    buildable: false
+    externals:
+      - spec: hsa-rocr-dev@4.3.0
+        prefix: /opt/rocm-4.3.0/hsa
+  llvm-amdgpu:
+    version: [4.3.0]
+    buildable: false
+    externals:
+      - spec: llvm-amdgpu@4.3.0
+        prefix: /opt/rocm-4.3.0/llvm
   all:
     compiler:
       - gcc@9.3.0
     providers:
-      mpi:
-        - cray-mpich@8.1.10
+      mpi: [cray-mpich@8.1.10]
       blas: [openblas]
       lapack: [openblas]

env-templates/exawind_spock.yaml

@@ -0,0 +1,8 @@
+spack:
+  include:
+  - include.yaml
+  concretization: separately
+  view: false
+  specs:
+    - 'exawind~hypre+amr_wind_gpu~nalu_wind_gpu+rocm amdgpu_target=gfx908 ^nalu-wind+hypre ^amr-wind~hypre'
+    - 'exawind+hypre~amr_wind_gpu~nalu_wind_gpu'

repos/exawind/packages/amr-wind/package.py

@@ -3,7 +3,7 @@
 import os
 from shutil import copyfile
 
-class AmrWind(bAmrWind):
+class AmrWind(bAmrWind, ROCmPackage):
 
     depends_on('hypre+unified-memory', when='+hypre+cuda')
 
@@ -15,6 +15,12 @@ def cmake_args(self):
         if '+cuda' in spec:
             options.append(define('BUILD_SHARED_LIBS', False))
 
+        if '+rocm' in self.spec:
+            targets = self.spec.variants['amdgpu_target'].value
+            options.append('-DCMAKE_CXX_COMPILER={0}'.format(self.spec['hip'].hipcc))
+            options.append('-DAMR_WIND_ENABLE_ROCM=ON')
+            options.append('-DAMReX_AMD_ARCH=' + ';'.join(str(x) for x in targets))
+
         if spec['mpi'].name == 'openmpi':
             options.append(define('MPIEXEC_PREFLAGS', '--oversubscribe'))
 

repos/exawind/packages/exawind/package.py

@@ -4,7 +4,7 @@
 import os
 
 
-class Exawind(CMakePackage, CudaPackage):
+class Exawind(CMakePackage, CudaPackage, ROCmPackage):
     """Multi-application driver for Exawind project."""
 
     homepage = "https://github.com/Exawind/exawind-driver"
@@ -17,9 +17,9 @@ class Exawind(CMakePackage, CudaPackage):
     # Testing is currently always enabled, but should be optional in the future
     # to avoid cloning the mesh submodule
     version('master', branch='main', submodules=True)
+
     variant('asan', default=False,
             description='turn on address sanitizer')
-
     variant('openfast', default=False,
             description='Enable OpenFAST integration')
     variant('hypre', default=True,
@@ -29,15 +29,18 @@ class Exawind(CMakePackage, CudaPackage):
     variant('nalu_wind_gpu', default=False,
             description='Enable Nalu-Wind on the GPU')
 
-    conflicts('+amr_wind_gpu', when='~cuda')
-    conflicts('+nalu_wind_gpu', when='~cuda')
+    conflicts('+amr_wind_gpu', when='~cuda~rocm')
+    conflicts('+nalu_wind_gpu', when='~cuda~rocm')
     conflicts('+amr_wind_gpu~nalu_wind_gpu', when='^amr-wind+hypre')
 
     for arch in CudaPackage.cuda_arch_values:
         depends_on('amr-wind+cuda cuda_arch=%s' % arch, when='+amr_wind_gpu+cuda cuda_arch=%s' % arch)
         depends_on('nalu-wind+cuda cuda_arch=%s' % arch, when='+nalu_wind_gpu+cuda cuda_arch=%s' % arch)
         depends_on('trilinos+cuda cuda_arch=%s' % arch, when='+nalu_wind_gpu+cuda cuda_arch=%s' % arch)
 
+    for arch in ROCmPackage.amdgpu_targets:
+        depends_on('amr-wind+rocm amdgpu_target=%s' % arch, when='+amr_wind_gpu+rocm amdgpu_target=%s' % arch)
+
     depends_on('nalu-wind+tioga')
     depends_on('amr-wind+netcdf+mpi')
     depends_on('tioga+shared~nodegid')
@@ -66,6 +69,10 @@ def cmake_args(self):
             args.append(define('EXAWIND_ENABLE_CUDA', True))
             args.append(define('CUDAToolkit_ROOT', self.spec['cuda'].prefix))
 
+        if spec.satisfies('+rocm'):
+            args.append(define('EXAWIND_ENABLE_ROCM', True))
+            args.append('-DCMAKE_CXX_COMPILER={0}'.format(self.spec['hip'].hipcc))
+
         return args
 
     def setup_build_environment(self, env):

scripts/install-exawind.sh

@@ -28,6 +28,8 @@ fi
 
 if [ "${SPACK_MANAGER_MACHINE}" == 'eagle' ] || [ "${SPACK_MANAGER_MACHINE}" == 'summit' ]; then
   cmd "spack manager create-env -y ${SPACK_MANAGER}/env-templates/exawind_matrix.yaml -d ${SPACK_MANAGER}/environments/exawind"
+elif [ "${SPACK_MANAGER_MACHINE}" == 'spock' ]; then
+  cmd "spack manager create-env -y ${SPACK_MANAGER}/env-templates/exawind_spock.yaml -d ${SPACK_MANAGER}/environments/exawind"
 else
   cmd "spack manager create-env -s exawind+hypre -d ${SPACK_MANAGER}/environments/exawind"
 fi

spack-scripting/scripting/cmd/manager_cmds/create_env.py

@@ -86,7 +86,6 @@ def create_env(parser, args):
         envcmd.env_activate(activate_args)
 
 
-
 def add_command(parser, command_dict):
     sub_parser = parser.add_parser('create-env', help='convenience script'
                                    ' for setting up a spack environment')