kokkos et al. : don't monkeypatch spec in callbacks (#48916)

Currently, a few packages using kokkos rely on
kokkos itself monkeypatching its own spec to
provide some attribute.

In this commit we change this attribute to be
defined on the package, and never be monkeypatched.
This commit is contained in:
Massimiliano Culpo 2025-02-12 07:27:49 +01:00 committed by GitHub
parent 46f5b192ef
commit af89bdf632
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
15 changed files with 41 additions and 48 deletions

View File

@ -110,22 +110,15 @@ class Arborx(CMakePackage, CudaPackage, ROCmPackage):
conflicts("~serial", when="+trilinos")
def cmake_args(self):
spec = self.spec
if "+trilinos" in spec:
kokkos_spec = spec["trilinos"]
else:
kokkos_spec = spec["kokkos"]
kokkos_pkg = self["trilinos"] if self.spec.satisfies("+trilinos") else self["kokkos"]
options = [
f"-DKokkos_ROOT={kokkos_spec.prefix}",
self.define("Kokkos_ROOT", kokkos_pkg.prefix),
self.define_from_variant("ARBORX_ENABLE_MPI", "mpi"),
]
if spec.satisfies("+cuda"):
options.append(f"-DCMAKE_CXX_COMPILER={kokkos_spec.kokkos_cxx}")
if spec.satisfies("+rocm"):
options.append("-DCMAKE_CXX_COMPILER=%s" % spec["hip"].hipcc)
if self.spec.satisfies("+cuda"):
options.append(self.define("CMAKE_CXX_COMPILER", kokkos_pkg.kokkos_cxx))
if self.spec.satisfies("+rocm"):
options.append(self.define("CMAKE_CXX_COMPILER", self.spec["hip"].hipcc))
return options

View File

@ -66,7 +66,7 @@ def cmake_args(self):
[
"-DKokkosCore_PREFIX={0}".format(kokkos.prefix),
"-DKokkosKernels_PREFIX={0}".format(kokkos_kernels.prefix),
"-DCMAKE_CXX_COMPILER:STRING={0}".format(spec["kokkos"].kokkos_cxx),
"-DCMAKE_CXX_COMPILER:STRING={0}".format(self["kokkos"].kokkos_cxx),
# Compadre_USE_PYTHON is OFF by default
"-DCompadre_USE_PYTHON=OFF",
]

View File

@ -552,7 +552,7 @@ def cmake_args(self):
)
# Make sure we use the same compiler that Trilinos uses
if spec.satisfies("+trilinos"):
options.extend([self.define("CMAKE_CXX_COMPILER", spec["trilinos"].kokkos_cxx)])
options.extend([self.define("CMAKE_CXX_COMPILER", self["trilinos"].kokkos_cxx)])
# Complex support
options.append(self.define_from_variant("DEAL_II_WITH_COMPLEX_VALUES", "complex"))

View File

@ -121,9 +121,9 @@ def setup_build_environment(self, env):
# Manually turn off device self.defines to solve Kokkos issues in Nalu-Wind headers
env.append_flags("CXXFLAGS", "-U__HIP_DEVICE_COMPILE__ -DDESUL_HIP_RDC")
if self.spec.satisfies("+cuda"):
env.set("OMPI_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("OMPI_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
if self.spec.satisfies("+rocm"):
env.set("OMPI_CXX", self.spec["hip"].hipcc)
env.set("MPICH_CXX", self.spec["hip"].hipcc)

View File

@ -147,7 +147,7 @@ def cmake_args(self):
# CMake pulled in via find_package(Legion) won't work without this
options.append(self.define("HIP_PATH", "{0}/hip".format(spec["hip"].prefix)))
elif self.spec.satisfies("^kokkos"):
options.append(self.define("CMAKE_CXX_COMPILER", self.spec["kokkos"].kokkos_cxx))
options.append(self.define("CMAKE_CXX_COMPILER", self["kokkos"].kokkos_cxx))
else:
# kept for supporing version prior to 2.2
options = [

View File

@ -267,7 +267,7 @@ def cmake_args(self):
options.append("-DCMAKE_CXX_COMPILER=%s" % spec["hip"].hipcc)
else:
# Compiler weirdness due to nvcc_wrapper
options.append("-DCMAKE_CXX_COMPILER=%s" % spec["kokkos"].kokkos_cxx)
options.append("-DCMAKE_CXX_COMPILER=%s" % self["kokkos"].kokkos_cxx)
if self.run_tests:
options.append("-DKokkosKernels_ENABLE_TESTS=ON")

View File

@ -69,6 +69,6 @@ def setup_dependent_build_environment(self, env, dependent_spec):
env.set("OMPI_CXX", wrapper)
env.set("MPICXX_CXX", wrapper) # HPE MPT
def setup_dependent_package(self, module, dependent_spec):
wrapper = join_path(self.prefix.bin, "nvcc_wrapper")
self.spec.kokkos_cxx = wrapper
@property
def kokkos_cxx(self) -> str:
return join_path(self.prefix.bin, "nvcc_wrapper")

View File

@ -408,11 +408,12 @@ def append_args(self, cmake_prefix, cmake_options, spack_options):
if option:
spack_options.append(option)
def setup_dependent_package(self, module, dependent_spec):
try:
self.spec.kokkos_cxx = self.spec["kokkos-nvcc-wrapper"].kokkos_cxx
except Exception:
self.spec.kokkos_cxx = spack_cxx
@property
def kokkos_cxx(self) -> str:
if self.spec.satisfies("+wrapper"):
return self["kokkos-nvcc-wrapper"].kokkos_cxx
# Assumes build-time globals have been set already
return spack_cxx
def cmake_args(self):
spec = self.spec
@ -474,9 +475,7 @@ def cmake_args(self):
options.append(self.define(tpl + "_DIR", spec[tpl].prefix))
if self.spec.satisfies("+wrapper"):
options.append(
self.define("CMAKE_CXX_COMPILER", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
)
options.append(self.define("CMAKE_CXX_COMPILER", self.kokkos_cxx))
elif "+rocm" in self.spec:
if "+cmake_lang" in self.spec:
options.append(

View File

@ -406,7 +406,7 @@ def cmake_args(self):
if spec.satisfies("+kokkos"):
# default is off.
options.append("-DLegion_USE_Kokkos=ON")
os.environ["KOKKOS_CXX_COMPILER"] = spec["kokkos"].kokkos_cxx
os.environ["KOKKOS_CXX_COMPILER"] = self["kokkos"].kokkos_cxx
if spec.satisfies("+cuda+cuda_unsupported_compiler ^kokkos%clang +cuda"):
# Keep CMake CUDA compiler detection happy
options.append(

View File

@ -140,9 +140,9 @@ def setup_build_environment(self, env):
spec = self.spec
env.append_flags("CXXFLAGS", "-DUSE_STK_SIMD_NONE")
if spec.satisfies("+cuda"):
env.set("OMPI_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("OMPI_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
if spec.satisfies("+rocm"):
env.append_flags("CXXFLAGS", "-fgpu-rdc")

View File

@ -116,7 +116,7 @@ def cmake_args(self):
if "+cuda%gcc" in self.spec:
options += [
self.define(
"CMAKE_CXX_COMPILER", "{0}".format(self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
"CMAKE_CXX_COMPILER", "{0}".format(self["kokkos-nvcc-wrapper"].kokkos_cxx)
)
]

View File

@ -61,7 +61,7 @@ def set_cmake_from_variants(self):
with open("cmake_opts.txt", "w") as f:
f.write("KokkosCore_PREFIX:PATH=%s\n" % spec["kokkos"].prefix)
f.write("KokkosKernels_PREFIX:PATH=%s\n" % spec["kokkos-kernels"].prefix)
f.write("CMAKE_CXX_COMPILER:STRING={0}\n".format(spec["kokkos"].kokkos_cxx))
f.write("CMAKE_CXX_COMPILER:STRING={0}\n".format(self["kokkos"].kokkos_cxx))
if spec.variants["debug"].value == "0":
f.write(
"CMAKE_CXX_FLAGS:STRING=%s\n"

View File

@ -184,7 +184,7 @@ def cmake_args(self):
]
if "+kokkos+cuda" in self.spec:
args.append(self.define("CMAKE_CXX_COMPILER", self.spec["kokkos"].kokkos_cxx))
args.append(self.define("CMAKE_CXX_COMPILER", self["kokkos"].kokkos_cxx))
return args

View File

@ -98,5 +98,5 @@ def cmake_args(self):
args.append(self.define("CMAKE_CXX_COMPILER", self.spec["hip"].hipcc))
args.append(self.define("CMAKE_C_COMPILER", self.spec["hip"].hipcc))
if self.spec.satisfies("^kokkos+cuda"):
args.append(self.define("CMAKE_CXX_COMPILER", self.spec["kokkos"].kokkos_cxx))
args.append(self.define("CMAKE_CXX_COMPILER", self["kokkos"].kokkos_cxx))
return args

View File

@ -589,21 +589,22 @@ def setup_dependent_run_environment(self, env, dependent_spec):
# in case the dependent app also run a CUDA backend via Trilinos
env.set("CUDA_LAUNCH_BLOCKING", "1")
def setup_dependent_package(self, module, dependent_spec):
if "+wrapper" in self.spec:
self.spec.kokkos_cxx = self.spec["kokkos-nvcc-wrapper"].kokkos_cxx
else:
self.spec.kokkos_cxx = spack_cxx
@property
def kokkos_cxx(self) -> str:
if self.spec.satisfies("+wrapper"):
return self["kokkos-nvcc-wrapper"].kokkos_cxx
# Assumes build-time globals have been set already
return spack_cxx
def setup_build_environment(self, env):
spec = self.spec
if "+cuda" in spec and "+wrapper" in spec:
if "+mpi" in spec:
env.set("OMPI_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("OMPI_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
else:
env.set("CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
if "+rocm" in spec:
if "+mpi" in spec: