From ea8489fce266d69f2fbe314c1385956b1a342e12 Mon Sep 17 00:00:00 2001 From: Rasmus Larsen Date: Mon, 29 Jan 2024 19:52:31 +0100 Subject: [PATCH] ROCm: Allow setting compilation target (#2581) --- setup.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/setup.py b/setup.py index b552cb67..8fad433f 100644 --- a/setup.py +++ b/setup.py @@ -287,11 +287,15 @@ if _is_cuda(): }, )) elif _is_hip(): - amd_arch = get_amdgpu_offload_arch() - if amd_arch not in ROCM_SUPPORTED_ARCHS: - raise RuntimeError( - f"Only the following arch is supported: {ROCM_SUPPORTED_ARCHS}" - f"amdgpu_arch_found: {amd_arch}") + amd_archs = os.getenv("GPU_ARCHS") + if amd_archs is None: + amd_archs = get_amdgpu_offload_arch() + for arch in amd_archs.split(";"): + if arch not in ROCM_SUPPORTED_ARCHS: + raise RuntimeError( + f"Only the following arch is supported: {ROCM_SUPPORTED_ARCHS}" + f"amdgpu_arch_found: {arch}") + NVCC_FLAGS += [f"--offload-arch={arch}"] elif _is_neuron(): neuronxcc_version = get_neuronxcc_version()