kokkos: allow using new gfx942_apu arch (#48609)

Add an apu variant that promotes GPU architectures to their APU
equivalent. Right now this is just gfx942 -> gfx942_apu.
This commit is contained in:
Richard Berger 2025-04-04 23:29:29 -06:00 committed by GitHub
parent d00b05b71e
commit ff82ba24e9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -271,16 +271,26 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
"gfx1030": "navi1030",
"gfx1100": "navi1100",
}
amdgpu_apu_arch_map = {"gfx942": "amd_gfx942_apu"}
amd_support_conflict_msg = (
"{0} is not supported; "
"Kokkos supports the following AMD GPU targets: " + ", ".join(amdgpu_arch_map.keys())
)
amd_apu_support_conflict_msg = (
"{0} is not supported; "
"Kokkos supports the following AMD GPU targets with unified memory: "
+ ", ".join(amdgpu_apu_arch_map.keys())
)
for arch in ROCmPackage.amdgpu_targets:
if arch not in amdgpu_arch_map:
conflicts(
"+rocm",
when="amdgpu_target={0}".format(arch),
msg=amd_support_conflict_msg.format(arch),
"+rocm", when=f"amdgpu_target={arch}", msg=amd_support_conflict_msg.format(arch)
)
if arch not in amdgpu_apu_arch_map:
conflicts(
"+rocm+apu",
when=f"amdgpu_target={arch}",
msg=amd_apu_support_conflict_msg.format(arch),
)
intel_gpu_arches = (
@ -298,6 +308,7 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
values=("none",) + intel_gpu_arches,
description="Intel GPU architecture",
)
variant("apu", default=False, description="Enable APU support", when="@4.5: +rocm")
for dev, (dflt, desc) in devices_variants.items():
variant(dev, default=dflt, description=desc)
@ -456,7 +467,10 @@ def cmake_args(self):
for amdgpu_target in spec.variants["amdgpu_target"].value:
if amdgpu_target != "none":
if amdgpu_target in self.amdgpu_arch_map:
spack_microarches.append(self.amdgpu_arch_map[amdgpu_target])
if spec.satisfies("+apu") and amdgpu_target in self.amdgpu_apu_arch_map:
spack_microarches.append(self.amdgpu_apu_arch_map[amdgpu_target])
else:
spack_microarches.append(self.amdgpu_arch_map[amdgpu_target])
else:
# Note that conflict declarations should prevent
# choosing an unsupported AMD GPU target