From 360dbe41f7b7a380651db3735ea6e0a7cf1b1920 Mon Sep 17 00:00:00 2001
From: Cédric Chevalier <cedric.chevalier@cea.fr>
Date: Tue, 29 Oct 2024 13:28:12 +0100
Subject: kokkos: async malloc (#46464)
---
var/spack/repos/builtin/packages/kokkos/package.py | 13 +++++++------
1 file 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)
--
cgit v1.2.3-70-g09d2