diff --git a/var/spack/repos/builtin/packages/kokkos/package.py b/var/spack/repos/builtin/packages/kokkos/package.py index ca9791fab14..01ee5235e9c 100644 --- a/var/spack/repos/builtin/packages/kokkos/package.py +++ b/var/spack/repos/builtin/packages/kokkos/package.py @@ -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)