kokkos: async malloc (#46464)

1 files changed, 7 insertions, 6 deletions

diff --git a/var/spack/repos/builtin/packages/kokkos/package.py b/var/spack/repos/builtin/packages/kokkos/package.py
index ca9791fab1..01ee5235e9 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 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
         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)