exawind: updates to package to allow mixed device (#44159)

* exawind: updates to package to allow mixed device

* Style.

* Remove ninja variants.

* Add conflict for amr-wind+hypre with mixed device.

* Relax amr-wind~hypre requirement.

* Move runtime variables to nalu-wind.

* Update suggestions.

* Remove umpire.

2 files changed, 74 insertions, 39 deletions

diff --git a/var/spack/repos/builtin/packages/exawind/package.py b/var/spack/repos/builtin/packages/exawind/package.py
index 8ea3720414..2db941c664 100644
--- a/var/spack/repos/builtin/packages/exawind/package.py
+++ b/var/spack/repos/builtin/packages/exawind/package.py
@@ -16,59 +16,84 @@ class Exawind(CMakePackage, CudaPackage, ROCmPackage):
 
     tags = ["ecp", "ecp-apps"]
 
-    # Testing is currently always enabled, but should be optional in the future
-    # to avoid cloning the mesh submodule
-    version("master", branch="main", submodules=True)
-    version("1.0.0", tag="v1.0.0", submodules=True)
     license("Apache-2.0")
 
-    variant("openfast", default=False, description="Enable OpenFAST integration")
-    variant("hypre", default=True, description="Enable hypre solver")
-    variant("stk_simd", default=False, description="Enable SIMD in STK")
-    variant("umpire", default=False, description="Enable Umpire")
-    variant("tiny_profile", default=False, description="Turn on AMR-wind with tiny profile")
+    version("master", branch="main", submodules=True, preferred=True)
+    version("1.0.0", tag="v1.0.0", submodules=True)
+
+    variant("amr_wind_gpu", default=False, description="Enable AMR-Wind on the GPU")
+    variant("nalu_wind_gpu", default=False, description="Enable Nalu-Wind on the GPU")
     variant("sycl", default=False, description="Enable SYCL backend for AMR-Wind")
     variant("gpu-aware-mpi", default=False, description="gpu-aware-mpi")
 
-    conflicts("amr-wind+hypre", when="+sycl")
-
     for arch in CudaPackage.cuda_arch_values:
-        depends_on("amr-wind+cuda cuda_arch=%s" % arch, when="+cuda cuda_arch=%s" % arch)
-        depends_on("nalu-wind+cuda cuda_arch=%s" % arch, when="+cuda cuda_arch=%s" % arch)
+        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="+rocm amdgpu_target=%s" % arch)
-        depends_on("nalu-wind+rocm amdgpu_target=%s" % arch, when="+rocm amdgpu_target=%s" % arch)
-
-    depends_on("nalu-wind+tioga")
-    depends_on("amr-wind+netcdf+mpi")
-    depends_on("tioga~nodegid")
+        depends_on(
+            "amr-wind+rocm amdgpu_target=%s" % arch,
+            when="+amr_wind_gpu+rocm amdgpu_target=%s" % arch,
+        )
+        depends_on(
+            "nalu-wind+rocm amdgpu_target=%s" % arch,
+            when="+nalu_wind_gpu+rocm amdgpu_target=%s" % arch,
+        )
+        depends_on(
+            "trilinos+rocm amdgpu_target=%s" % arch,
+            when="+nalu_wind_gpu+rocm amdgpu_target=%s" % arch,
+        )
+
+    depends_on("nalu-wind+hypre+fsi+openfast+tioga")
+    depends_on("amr-wind+netcdf+mpi+tiny_profile")
+    depends_on("trilinos")
     depends_on("yaml-cpp@0.6:")
-    depends_on("nalu-wind+openfast", when="+openfast")
-    depends_on("amr-wind+hypre", when="+hypre~sycl")
-    depends_on("amr-wind~hypre", when="~hypre")
-    depends_on("nalu-wind+hypre", when="+hypre")
-    depends_on("nalu-wind~hypre", when="~hypre")
-    depends_on("amr-wind+sycl", when="+sycl")
-    depends_on("nalu-wind+umpire", when="+umpire")
-    depends_on("amr-wind+umpire", when="+umpire")
-    depends_on("amr-wind+tiny_profile", when="+tiny_profile")
+    depends_on("tioga~nodegid")
+    depends_on("openfast+cxx@2.6.0:")
+    depends_on("amr-wind+sycl", when="+amr_wind_gpu+sycl")
+    depends_on("kokkos-nvcc-wrapper", type="build", when="+cuda")
+    depends_on("mpi")
     depends_on("nalu-wind+gpu-aware-mpi", when="+gpu-aware-mpi")
     depends_on("amr-wind+gpu-aware-mpi", when="+gpu-aware-mpi")
     depends_on("nalu-wind@2.0.0:", when="@1.0.0:")
     depends_on("amr-wind@0.9.0:", when="@1.0.0:")
     depends_on("tioga@1.0.0:", when="@1.0.0:")
 
+    with when("~amr_wind_gpu~nalu_wind_gpu"):
+        conflicts("+cuda")
+        conflicts("+rocm")
+        conflicts("+sycl")
+    with when("~nalu_wind_gpu"):
+        conflicts("^nalu-wind+cuda")
+        conflicts("^nalu-wind+rocm")
+    with when("~amr_wind_gpu"):
+        conflicts("^amr-wind+cuda")
+        conflicts("^amr-wind+rocm")
+        conflicts("^amr-wind+sycl")
+    conflicts("+amr_wind_gpu", when="~cuda~rocm~sycl")
+    conflicts("+nalu_wind_gpu", when="~cuda~rocm")
+    conflicts("+nalu_wind_gpu", when="+sycl")
+    conflicts("^amr-wind+hypre", when="~amr_wind_gpu+nalu_wind_gpu")
+    conflicts("^amr-wind+hypre", when="+amr_wind_gpu~nalu_wind_gpu")
+    conflicts("+sycl", when="+cuda")
+    conflicts("+rocm", when="+cuda")
+    conflicts("+sycl", when="+rocm")
+
     def cmake_args(self):
         spec = self.spec
 
         args = [self.define("MPI_HOME", spec["mpi"].prefix)]
 
-        if "+umpire" in self.spec:
-            args.append(self.define_from_variant("EXAWIND_ENABLE_UMPIRE", "umpire"))
-            args.append(self.define("UMPIRE_DIR", self.spec["umpire"].prefix))
-
         if spec.satisfies("+cuda"):
+            args.append(self.define("CMAKE_CXX_COMPILER", spec["mpi"].mpicxx))
+            args.append(self.define("CMAKE_C_COMPILER", spec["mpi"].mpicc))
             args.append(self.define("EXAWIND_ENABLE_CUDA", True))
             args.append(self.define("CUDAToolkit_ROOT", self.spec["cuda"].prefix))
             args.append(self.define("EXAWIND_CUDA_ARCH", self.spec.variants["cuda_arch"].value))
@@ -77,6 +102,7 @@ class Exawind(CMakePackage, CudaPackage, ROCmPackage):
             targets = self.spec.variants["amdgpu_target"].value
             args.append(self.define("EXAWIND_ENABLE_ROCM", True))
             args.append(self.define("CMAKE_CXX_COMPILER", self.spec["hip"].hipcc))
+            # Optimization to only build one specific target architecture:
             args.append(self.define("CMAKE_HIP_ARCHITECTURES", ";".join(str(x) for x in targets)))
             args.append(self.define("AMDGPU_TARGETS", ";".join(str(x) for x in targets)))
             args.append(self.define("GPU_TARGETS", ";".join(str(x) for x in targets)))
@@ -84,14 +110,17 @@ class Exawind(CMakePackage, CudaPackage, ROCmPackage):
         if spec.satisfies("^amr-wind+hdf5"):
             args.append(self.define("H5Z_ZFP_USE_STATIC_LIBS", True))
 
-        if spec.satisfies("^amr-wind+ascent"):
-            args.append(self.define("CMAKE_EXE_LINKER_FLAGS", self.compiler.openmp_flag))
-
         return args
 
     def setup_build_environment(self, env):
-        if "~stk_simd" in self.spec:
-            env.append_flags("CXXFLAGS", "-DUSE_STK_SIMD_NONE")
+        env.append_flags("CXXFLAGS", "-DUSE_STK_SIMD_NONE")
+        if "+rocm+amr_wind_gpu~nalu_wind_gpu" in self.spec:
+            # Manually turn off device self.defines to solve Kokkos issues in Nalu-Wind headers
+            env.append_flags("CXXFLAGS", "-U__HIP_DEVICE_COMPILE__ -DDESUL_HIP_RDC")
+        if "+cuda" in self.spec:
+            env.set("OMPI_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
+            env.set("MPICH_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
+            env.set("MPICXX_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
         if "+rocm" in self.spec:
             env.set("OMPI_CXX", self.spec["hip"].hipcc)
             env.set("MPICH_CXX", self.spec["hip"].hipcc)
diff --git a/var/spack/repos/builtin/packages/nalu-wind/package.py b/var/spack/repos/builtin/packages/nalu-wind/package.py
index 18efa6d6fe..17429964d0 100644
--- a/var/spack/repos/builtin/packages/nalu-wind/package.py
+++ b/var/spack/repos/builtin/packages/nalu-wind/package.py
@@ -118,12 +118,18 @@ class NaluWind(CMakePackage, CudaPackage, ROCmPackage):
     conflicts("^trilinos+cuda", when="~cuda")
     conflicts("^trilinos+rocm", when="~rocm")
 
+    def setup_dependent_run_environment(self, env, dependent_spec):
+        spec = self.spec
+        if spec.satisfies("+cuda") or spec.satisfies("+rocm"):
+            env.set("CUDA_LAUNCH_BLOCKING", "1")
+            env.set("CUDA_MANAGED_FORCE_DEVICE_ALLOC", "1")
+            env.set("HIP_LAUNCH_BLOCKING", "1")
+            env.set("HIP_MANAGED_FORCE_DEVICE_ALLOC", "1")
+
     def setup_build_environment(self, env):
         spec = self.spec
         env.append_flags("CXXFLAGS", "-DUSE_STK_SIMD_NONE")
         if spec.satisfies("+cuda"):
-            env.set("CUDA_LAUNCH_BLOCKING", "1")
-            env.set("CUDA_MANAGED_FORCE_DEVICE_ALLOC", "1")
             env.set("OMPI_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
             env.set("MPICH_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
             env.set("MPICXX_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)