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") conflicts("~serial", when="+trilinos")
def cmake_args(self): def cmake_args(self):
spec = self.spec kokkos_pkg = self["trilinos"] if self.spec.satisfies("+trilinos") else self["kokkos"]
if "+trilinos" in spec:
kokkos_spec = spec["trilinos"]
else:
kokkos_spec = spec["kokkos"]
options = [ options = [
f"-DKokkos_ROOT={kokkos_spec.prefix}", self.define("Kokkos_ROOT", kokkos_pkg.prefix),
self.define_from_variant("ARBORX_ENABLE_MPI", "mpi"), self.define_from_variant("ARBORX_ENABLE_MPI", "mpi"),
] ]
if self.spec.satisfies("+cuda"):
if spec.satisfies("+cuda"): options.append(self.define("CMAKE_CXX_COMPILER", kokkos_pkg.kokkos_cxx))
options.append(f"-DCMAKE_CXX_COMPILER={kokkos_spec.kokkos_cxx}") if self.spec.satisfies("+rocm"):
if spec.satisfies("+rocm"): options.append(self.define("CMAKE_CXX_COMPILER", self.spec["hip"].hipcc))
options.append("-DCMAKE_CXX_COMPILER=%s" % spec["hip"].hipcc)
return options return options

View File

@ -66,7 +66,7 @@ def cmake_args(self):
[ [
"-DKokkosCore_PREFIX={0}".format(kokkos.prefix), "-DKokkosCore_PREFIX={0}".format(kokkos.prefix),
"-DKokkosKernels_PREFIX={0}".format(kokkos_kernels.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 # Compadre_USE_PYTHON is OFF by default
"-DCompadre_USE_PYTHON=OFF", "-DCompadre_USE_PYTHON=OFF",
] ]

View File

@ -552,7 +552,7 @@ def cmake_args(self):
) )
# Make sure we use the same compiler that Trilinos uses # Make sure we use the same compiler that Trilinos uses
if spec.satisfies("+trilinos"): 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 # Complex support
options.append(self.define_from_variant("DEAL_II_WITH_COMPLEX_VALUES", "complex")) 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 # 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") env.append_flags("CXXFLAGS", "-U__HIP_DEVICE_COMPILE__ -DDESUL_HIP_RDC")
if self.spec.satisfies("+cuda"): if self.spec.satisfies("+cuda"):
env.set("OMPI_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx) env.set("OMPI_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx) env.set("MPICH_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx) env.set("MPICXX_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
if self.spec.satisfies("+rocm"): if self.spec.satisfies("+rocm"):
env.set("OMPI_CXX", self.spec["hip"].hipcc) env.set("OMPI_CXX", self.spec["hip"].hipcc)
env.set("MPICH_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 # CMake pulled in via find_package(Legion) won't work without this
options.append(self.define("HIP_PATH", "{0}/hip".format(spec["hip"].prefix))) options.append(self.define("HIP_PATH", "{0}/hip".format(spec["hip"].prefix)))
elif self.spec.satisfies("^kokkos"): 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: else:
# kept for supporing version prior to 2.2 # kept for supporing version prior to 2.2
options = [ options = [

View File

@ -267,7 +267,7 @@ def cmake_args(self):
options.append("-DCMAKE_CXX_COMPILER=%s" % spec["hip"].hipcc) options.append("-DCMAKE_CXX_COMPILER=%s" % spec["hip"].hipcc)
else: else:
# Compiler weirdness due to nvcc_wrapper # 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: if self.run_tests:
options.append("-DKokkosKernels_ENABLE_TESTS=ON") 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("OMPI_CXX", wrapper)
env.set("MPICXX_CXX", wrapper) # HPE MPT env.set("MPICXX_CXX", wrapper) # HPE MPT
def setup_dependent_package(self, module, dependent_spec): @property
wrapper = join_path(self.prefix.bin, "nvcc_wrapper") def kokkos_cxx(self) -> str:
self.spec.kokkos_cxx = wrapper 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: if option:
spack_options.append(option) spack_options.append(option)
def setup_dependent_package(self, module, dependent_spec): @property
try: def kokkos_cxx(self) -> str:
self.spec.kokkos_cxx = self.spec["kokkos-nvcc-wrapper"].kokkos_cxx if self.spec.satisfies("+wrapper"):
except Exception: return self["kokkos-nvcc-wrapper"].kokkos_cxx
self.spec.kokkos_cxx = spack_cxx # Assumes build-time globals have been set already
return spack_cxx
def cmake_args(self): def cmake_args(self):
spec = self.spec spec = self.spec
@ -474,9 +475,7 @@ def cmake_args(self):
options.append(self.define(tpl + "_DIR", spec[tpl].prefix)) options.append(self.define(tpl + "_DIR", spec[tpl].prefix))
if self.spec.satisfies("+wrapper"): if self.spec.satisfies("+wrapper"):
options.append( options.append(self.define("CMAKE_CXX_COMPILER", self.kokkos_cxx))
self.define("CMAKE_CXX_COMPILER", self.spec["kokkos-nvcc-wrapper"].kokkos_cxx)
)
elif "+rocm" in self.spec: elif "+rocm" in self.spec:
if "+cmake_lang" in self.spec: if "+cmake_lang" in self.spec:
options.append( options.append(

View File

@ -406,7 +406,7 @@ def cmake_args(self):
if spec.satisfies("+kokkos"): if spec.satisfies("+kokkos"):
# default is off. # default is off.
options.append("-DLegion_USE_Kokkos=ON") 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"): if spec.satisfies("+cuda+cuda_unsupported_compiler ^kokkos%clang +cuda"):
# Keep CMake CUDA compiler detection happy # Keep CMake CUDA compiler detection happy
options.append( options.append(

View File

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

View File

@ -116,7 +116,7 @@ def cmake_args(self):
if "+cuda%gcc" in self.spec: if "+cuda%gcc" in self.spec:
options += [ options += [
self.define( 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: with open("cmake_opts.txt", "w") as f:
f.write("KokkosCore_PREFIX:PATH=%s\n" % spec["kokkos"].prefix) f.write("KokkosCore_PREFIX:PATH=%s\n" % spec["kokkos"].prefix)
f.write("KokkosKernels_PREFIX:PATH=%s\n" % spec["kokkos-kernels"].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": if spec.variants["debug"].value == "0":
f.write( f.write(
"CMAKE_CXX_FLAGS:STRING=%s\n" "CMAKE_CXX_FLAGS:STRING=%s\n"

View File

@ -184,7 +184,7 @@ def cmake_args(self):
] ]
if "+kokkos+cuda" in self.spec: 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 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_CXX_COMPILER", self.spec["hip"].hipcc))
args.append(self.define("CMAKE_C_COMPILER", self.spec["hip"].hipcc)) args.append(self.define("CMAKE_C_COMPILER", self.spec["hip"].hipcc))
if self.spec.satisfies("^kokkos+cuda"): 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 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 # in case the dependent app also run a CUDA backend via Trilinos
env.set("CUDA_LAUNCH_BLOCKING", "1") env.set("CUDA_LAUNCH_BLOCKING", "1")
def setup_dependent_package(self, module, dependent_spec): @property
if "+wrapper" in self.spec: def kokkos_cxx(self) -> str:
self.spec.kokkos_cxx = self.spec["kokkos-nvcc-wrapper"].kokkos_cxx if self.spec.satisfies("+wrapper"):
else: return self["kokkos-nvcc-wrapper"].kokkos_cxx
self.spec.kokkos_cxx = spack_cxx # Assumes build-time globals have been set already
return spack_cxx
def setup_build_environment(self, env): def setup_build_environment(self, env):
spec = self.spec spec = self.spec
if "+cuda" in spec and "+wrapper" in spec: if "+cuda" in spec and "+wrapper" in spec:
if "+mpi" in spec: if "+mpi" in spec:
env.set("OMPI_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx) env.set("OMPI_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx) env.set("MPICH_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx) env.set("MPICXX_CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
else: else:
env.set("CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx) env.set("CXX", self["kokkos-nvcc-wrapper"].kokkos_cxx)
if "+rocm" in spec: if "+rocm" in spec:
if "+mpi" in spec: if "+mpi" in spec: