From 77b2b77e6c8892a60c3b1cfc21b2c81022bf99a5 Mon Sep 17 00:00:00 2001 From: Jacob Faibussowitsch Date: Mon, 21 Jul 2025 13:35:27 -0400 Subject: [PATCH] Fixup CUDA architecture defaults to make them into a list --- src/aedifix/packages/cuda.py | 7 +++++-- tests/packages/test_cuda.py | 2 ++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/aedifix/packages/cuda.py b/src/aedifix/packages/cuda.py index adf215b..9ef58b4 100644 --- a/src/aedifix/packages/cuda.py +++ b/src/aedifix/packages/cuda.py @@ -4,6 +4,7 @@ from __future__ import annotations import os +import re import shutil from argparse import Action, ArgumentParser, Namespace from pathlib import Path @@ -55,7 +56,7 @@ def map_cuda_arch_names(in_arch: str) -> list[str]: # TODO(jfaibussowit): rubin? } arch = [] - for sub_arch in in_arch.split(","): + for sub_arch in re.split(r"[;,]", in_arch): # support Turing, TURING, and, if the user is feeling spicy, tUrInG sub_arch_lo = sub_arch.strip().casefold() if not sub_arch_lo: @@ -131,7 +132,9 @@ class CUDA(Package): spec=ArgSpec( dest="cuda_arch", required=False, - default=_guess_cuda_architecture(), + default=CudaArchAction.map_cuda_arch_names( + _guess_cuda_architecture() + ), action=CudaArchAction, help=( "Specify the target GPU architecture. Available choices are: " diff --git a/tests/packages/test_cuda.py b/tests/packages/test_cuda.py index afc4b39..1d0ef6f 100644 --- a/tests/packages/test_cuda.py +++ b/tests/packages/test_cuda.py @@ -18,6 +18,7 @@ ("turing,hopper", ["75", "90"]), ("volta,60,all-major", ["70", "60", "all-major"]), ("60,,80", ["60", "80"]), + ("50-real;120-real;121", ["50-real", "120-real", "121"]), ) @@ -33,6 +34,7 @@ class TestCUDA: ("env_var", "env_value", "expected"), [ ("CUDAARCHS", "volta", "volta"), + ("CUDAARCHS", "volta;ampere", "volta;ampere"), ("CMAKE_CUDA_ARCHITECTURES", "75", "75"), ("", "", "all-major"), ],