kokkos: async malloc (#46464)

This commit is contained in:
Cédric Chevalier 2024-10-29 13:28:12 +01:00 committed by GitHub
parent ea1aa0714b
commit 360dbe41f7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -227,6 +227,10 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
conflicts("+cuda", when="cxxstd=17 ^cuda@:10")
conflicts("+cuda", when="cxxstd=20 ^cuda@:11")
# Expose a way to disable CudaMallocAsync that can cause problems
# with some MPI such as cray-mpich
variant("alloc_async", default=False, description="Use CudaMallocAsync", when="@4.2: +cuda")
# SYCL and OpenMPTarget require C++17 or higher
for cxxstdver in cxxstds[: cxxstds.index("17")]:
conflicts(
@ -371,12 +375,9 @@ def cmake_args(self):
if self.spec.satisfies("%oneapi") or self.spec.satisfies("%intel"):
options.append(self.define("CMAKE_CXX_FLAGS", "-fp-model=precise"))
# Kokkos 4.2.00+ changed the default to Kokkos_ENABLE_IMPL_CUDA_MALLOC_ASYNC=on
# which breaks GPU-aware with Cray-MPICH
# See https://github.com/kokkos/kokkos/pull/6402
# TODO: disable this once Cray-MPICH is fixed
if self.spec.satisfies("@4.2.00:") and self.spec.satisfies("^[virtuals=mpi] cray-mpich"):
options.append(self.define("Kokkos_ENABLE_IMPL_CUDA_MALLOC_ASYNC", False))
options.append(
self.define_from_variant("Kokkos_ENABLE_IMPL_CUDA_MALLOC_ASYNC", "alloc_async")
)
# Remove duplicate options
return lang.dedupe(options)