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

[Justin Fargnoli via llvm-commits]

================
@@ -294,80 +294,82 @@ 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),
-        ]
-
-        def version_int(ver):
-            return ver[0] * 100 + ver[1]
-
-        # ignore ptxas if its version is below the minimum supported
-        # version
-        min_version = ptxas_known_versions[0]
-        if version_int(version) < version_int(min_version):
-            print(
-                "Warning: ptxas version {}.{} is not supported".format(
-                    version[0], version[1]
-                )
-            )
-            return
+    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,
+        check=True,
+    )
+    versions = []
+    for line in result.stdout.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        match = re.match(r"(\d+)\.(\d+)", line)
+        if match:
+            versions.append((int(match.group(1)), int(match.group(2))))
+    return versions
+
+
+def ptxas_supported_sms(ptxas_executable):
+    result = subprocess.run(
+        [ptxas_executable, "--help"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    supported_sms = re.findall(r"'sm_(\d+(?:[af]?))'", result.stdout)
----------------
justinfargnoli wrote:

2102b1b2bffbdb32f64ab68a9e023d034eeaa3fc should address this

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