[llvm] [lit] Refactor available `ptxas` features (PR #154439)

[Alex MacLean via llvm-commits]

================
@@ -294,80 +294,146 @@ def get_asan_rtlib():
 )
 
 
-# Find (major, minor) version of ptxas
 def ptxas_version(ptxas):
-    ptxas_cmd = subprocess.Popen([ptxas, "--version"], stdout=subprocess.PIPE)
-    ptxas_out = ptxas_cmd.stdout.read().decode("ascii")
-    ptxas_cmd.wait()
-    match = re.search(r"release (\d+)\.(\d+)", ptxas_out)
-    if match:
-        return (int(match.group(1)), int(match.group(2)))
-    print("couldn't determine ptxas version")
-    return None
-
-
-# Enable %ptxas and %ptxas-verify tools.
-# %ptxas-verify defaults to sm_60 architecture. It can be overriden
-# by specifying required one, for instance: %ptxas-verify -arch=sm_80.
-def enable_ptxas(ptxas_executable):
-    version = ptxas_version(ptxas_executable)
-    if version:
-        # ptxas is supposed to be backward compatible with previous
-        # versions, so add a feature for every known version prior to
-        # the current one.
-        ptxas_known_versions = [
-            (9, 0),
-            (9, 1),
-            (9, 2),
-            (10, 0),
-            (10, 1),
-            (10, 2),
-            (11, 0),
-            (11, 1),
-            (11, 2),
-            (11, 3),
-            (11, 4),
-            (11, 5),
-            (11, 6),
-            (11, 7),
-            (11, 8),
-            (12, 0),
-            (12, 1),
-            (12, 2),
-            (12, 3),
-            (12, 4),
-            (12, 5),
-            (12, 6),
-            (12, 8),
-        ]
+    output = subprocess.check_output([ptxas, "--version"], text=True)
+    match = re.search(r"release (\d+)\.(\d+)", output)
+    if not match:
+        raise RuntimeError("Couldn't determine ptxas version")
+    return int(match.group(1)), int(match.group(2))
+
+
+def ptxas_supported_isa_versions(ptxas):
+    result = subprocess.run(
+        [ptxas, "--list-version"],
+        capture_output=True,
+        text=True,
+    )
+    versions = []
+    for line in result.stdout.splitlines():
+        match = re.match(r"(\d+)\.(\d+)", line)
+        if match:
+            versions.append((int(match.group(1)), int(match.group(2))))
+    return versions
+
+
+def ptxas_add_isa_features(major_version, minor_version):
+    supported_isa_versions = ptxas_supported_isa_versions(ptxas_executable)
+    if supported_isa_versions:
+        for major_version, minor_version in supported_isa_versions:
+            config.available_features.add(f"ptxas-isa-{major_version}.{minor_version}")
+        return
+    if major_version >= 13:
+        raise RuntimeError(
+            f"ptxas {ptxas_executable} does not support ISA version listing"
+        )
+    # Use a more pythonic approach: define a list of (major, minor, feature) tuples and iterate.
+    isa_features = [
+        (12, 9, "ptxas-isa-8.8"),
+        (12, 8, "ptxas-isa-8.7"),
+        (12, 7, "ptxas-isa-8.6"),
+        (12, 6, "ptxas-isa-8.5"),
+        (12, 5, "ptxas-isa-8.5"),
+        (12, 4, "ptxas-isa-8.4"),
+        (12, 3, "ptxas-isa-8.3"),
+        (12, 2, "ptxas-isa-8.2"),
+        (12, 1, "ptxas-isa-8.1"),
+        (12, 0, "ptxas-isa-8.0"),
+        (11, 8, "ptxas-isa-7.8"),
+        (11, 7, "ptxas-isa-7.7"),
+        (11, 6, "ptxas-isa-7.6"),
+        (11, 5, "ptxas-isa-7.5"),
+        (11, 4, "ptxas-isa-7.4"),
+        (11, 3, "ptxas-isa-7.3"),
+        (11, 2, "ptxas-isa-7.2"),
+        (11, 1, "ptxas-isa-7.1"),
+        (11, 0, "ptxas-isa-7.0"),
+        (10, 2, "ptxas-isa-6.5"),
+        (10, 1, "ptxas-isa-6.4"),
+        (10, 0, "ptxas-isa-6.3"),
+        (9, 2, "ptxas-isa-6.2"),
+        (9, 1, "ptxas-isa-6.1"),
+        (9, 0, "ptxas-isa-6.0"),
+        (8, 0, "ptxas-isa-5.0"),
+        (7, 5, "ptxas-isa-4.3"),
+        (7, 0, "ptxas-isa-4.2"),
+        (6, 5, "ptxas-isa-4.1"),
+        (6, 0, "ptxas-isa-4.0"),
+        (5, 5, "ptxas-isa-3.2"),
+        (5, 0, "ptxas-isa-3.1"),
+        (4, 1, "ptxas-isa-3.0"),
+        (4, 0, "ptxas-isa-2.3"),
+        (3, 2, "ptxas-isa-2.2"),
+        (3, 1, "ptxas-isa-2.1"),
+        (3, 0, "ptxas-isa-2.0"),
+        (3, 0, "ptxas-isa-1.5"),
+        (2, 2, "ptxas-isa-1.4"),
+        (2, 1, "ptxas-isa-1.3"),
+        (2, 0, "ptxas-isa-1.2"),
+        (1, 1, "ptxas-isa-1.1"),
+        (1, 0, "ptxas-isa-1.0"),
+    ]
+    for major, minor, feature in isa_features:
+        if major_version >= major and minor_version >= minor:
+            config.available_features.add(feature)
----------------
AlexMaclean wrote:

This seems incorrect. 

https://github.com/llvm/llvm-project/pull/154439