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