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

[Alex MacLean 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)
+    if not supported_sms:
+        raise RuntimeError("No SM architecture values found in ptxas help output")
+    return supported_sms
+
+
+def ptxas_supports_address_size_32(ptxas_executable):
+    result = subprocess.run(
+        [ptxas_executable, "-m 32"],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    if "is not defined for option 'machine'" in result.stderr:
+        return False
+    if "Missing .version directive at start of file" in result.stderr:
+        return True
+    raise RuntimeError("Unexpected ptxas output: {}".format(result.stderr))
 
-        for known_version in ptxas_known_versions:
-            if version_int(known_version) <= version_int(version):
-                major, minor = known_version
-                config.available_features.add("ptxas-{}.{}".format(major, minor))
 
+def enable_ptxas(ptxas_executable):
     config.available_features.add("ptxas")
     tools.extend(
         [
             ToolSubst("%ptxas", ptxas_executable),
-            ToolSubst("%ptxas-verify", "{} -arch=sm_60 -c -".format(ptxas_executable)),
+            ToolSubst("%ptxas-verify", f"{ptxas_executable} -c -"),
         ]
     )
 
+    major_version, minor_version = ptxas_version(ptxas_executable)
+    config.available_features.add("ptxas-{}.{}".format(major_version, minor_version))
----------------
AlexMaclean wrote:

Nit: use an `f` string here as well

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