arborx: Enable use of Kokkos from Trilinos for cuda/rocm. (#45220)

This commit is contained in:
Victor Brunini 2024-08-20 15:48:42 -07:00 committed by GitHub
parent d714a9b223
commit c29652580a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -81,14 +81,17 @@ class Arborx(CMakePackage, CudaPackage, ROCmPackage):
depends_on("kokkos+%s" % backend.lower(), when="~trilinos+%s" % backend.lower()) depends_on("kokkos+%s" % backend.lower(), when="~trilinos+%s" % backend.lower())
for arch in CudaPackage.cuda_arch_values: for arch in CudaPackage.cuda_arch_values:
cuda_dep = "+cuda cuda_arch={0}".format(arch) cuda_dep = f"+cuda cuda_arch={arch}"
depends_on("kokkos {0}".format(cuda_dep), when=cuda_dep) depends_on(f"kokkos {cuda_dep}", when=f"~trilinos {cuda_dep}")
depends_on(f"trilinos {cuda_dep}", when=f"+trilinos {cuda_dep}")
for arch in ROCmPackage.amdgpu_targets: for arch in ROCmPackage.amdgpu_targets:
rocm_dep = "+rocm amdgpu_target={0}".format(arch) rocm_dep = f"+rocm amdgpu_target={arch}"
depends_on("kokkos {0}".format(rocm_dep), when=rocm_dep) depends_on(f"kokkos {rocm_dep}", when=f"~trilinos {rocm_dep}")
depends_on(f"trilinos {rocm_dep}", when=f"+trilinos {rocm_dep}")
conflicts("+cuda", when="cuda_arch=none") conflicts("+cuda", when="cuda_arch=none")
conflicts("^kokkos", when="+trilinos")
depends_on("kokkos+cuda_lambda", when="~trilinos+cuda") depends_on("kokkos+cuda_lambda", when="~trilinos+cuda")
# Trilinos/Kokkos # Trilinos/Kokkos
@ -104,20 +107,22 @@ class Arborx(CMakePackage, CudaPackage, ROCmPackage):
depends_on("trilinos@14.4.0:", when="@1.6:+trilinos") depends_on("trilinos@14.4.0:", when="@1.6:+trilinos")
patch("trilinos14.0-kokkos-major-version.patch", when="@1.4+trilinos ^trilinos@14.0.0") patch("trilinos14.0-kokkos-major-version.patch", when="@1.4+trilinos ^trilinos@14.0.0")
conflicts("~serial", when="+trilinos") conflicts("~serial", when="+trilinos")
conflicts("+cuda", when="+trilinos")
def cmake_args(self): def cmake_args(self):
spec = self.spec spec = self.spec
if "~trilinos" in spec:
kokkos_spec = spec["kokkos"]
else:
kokkos_spec = spec["trilinos"]
options = [ options = [
"-DKokkos_ROOT=%s" f"-DKokkos_ROOT={kokkos_spec.prefix}",
% (spec["kokkos"].prefix if "~trilinos" in spec else spec["trilinos"].prefix),
self.define_from_variant("ARBORX_ENABLE_MPI", "mpi"), self.define_from_variant("ARBORX_ENABLE_MPI", "mpi"),
] ]
if spec.satisfies("+cuda"): if spec.satisfies("+cuda"):
# Only Kokkos allows '+cuda' for now options.append(f"-DCMAKE_CXX_COMPILER={kokkos_spec.kokkos_cxx}")
options.append("-DCMAKE_CXX_COMPILER=%s" % spec["kokkos"].kokkos_cxx)
if spec.satisfies("+rocm"): if spec.satisfies("+rocm"):
options.append("-DCMAKE_CXX_COMPILER=%s" % spec["hip"].hipcc) options.append("-DCMAKE_CXX_COMPILER=%s" % spec["hip"].hipcc)