Commit b2881df5 authored by one's avatar one
Browse files

Update gfx936 and gfx938 support for kokkos

parent c9dcc243
......@@ -199,10 +199,15 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
when="+rocm",
)
hygongpu_arch_list = [
"gfx936",
"gfx938",
]
amdgpu_arch_map = {
"gfx936": "amd_gfx90a",
"gfx938": "amd_gfx942",
"gfx900": "vega900",
"gfx906": "vega906",
"gfx936": "vega906",
"gfx908": "vega908",
"gfx90a": "vega90A",
"gfx940": "amd_gfx940",
......@@ -397,6 +402,7 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
if kokkos_microarch_name:
spack_microarches.append(kokkos_microarch_name)
hygon_offload_arch = None
if spec.satisfies("+rocm"):
amdgpu_target = spec.variants["amdgpu_target"].value
if amdgpu_target != "none":
......@@ -405,6 +411,8 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
spack_microarches.append(self.amdgpu_apu_arch_map[amdgpu_target])
else:
spack_microarches.append(self.amdgpu_arch_map[amdgpu_target])
if amdgpu_target in self.hygongpu_arch_list:
hygon_offload_arch = amdgpu_target
else:
# Note that conflict declarations should prevent
# choosing an unsupported AMD GPU target
......@@ -415,6 +423,10 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
for arch in spack_microarches:
options.append(self.define("Kokkos_ARCH_" + arch.upper(), True))
if hygon_offload_arch:
offload_arch = "--offload-arch={0}".format(hygon_offload_arch)
options.append(self.define("Kokkos_IMPL_AMDGPU_FLAGS", offload_arch))
options.append(self.define("Kokkos_IMPL_AMDGPU_LINK", offload_arch))
self.append_args("ENABLE", self.devices_variants.keys(), options)
self.append_args("ENABLE", self.options_variants.keys(), options)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment