Check the shipped device code covers the GPUs a wheel claims

.ci/scripts/wheel/test_cpp_sdk.py

@@ -30,6 +30,7 @@
 import sys
 import tempfile
 from pathlib import Path
+from typing import Optional
 
 # Registry entry points. A second definer of any of these means a second
 # process-wide registry.
@@ -358,10 +359,7 @@ def test_shipped_libraries_load() -> None:
             # list instead would hide a genuinely under-linked symbol on the same wheel.
             # The leading-underscore forms matter too: nvcc emits host stubs such as
             # __cudaRegisterFatBinary for every compiled .cu file.
-            and not (
-                skip_undefined
-                and re.search(_CUDA_SYMBOL, line)
-            )
+            and not (skip_undefined and re.search(_CUDA_SYMBOL, line))
         ]
         if undefined:
             unresolved[str(library.relative_to(package_dir))] = undefined[:5]
@@ -1100,17 +1098,27 @@ def test_documented_example_compiles(work_dir: Path) -> None:
 
     build_dir = work_dir / "documented-build"
     configure = subprocess.run(
-        ["cmake", "-S", str(source_dir), "-B", str(build_dir),
-         f"-DCMAKE_PREFIX_PATH={_installed_package_dir()}"],
-        capture_output=True, text=True, check=False,
+        [
+            "cmake",
+            "-S",
+            str(source_dir),
+            "-B",
+            str(build_dir),
+            f"-DCMAKE_PREFIX_PATH={_installed_package_dir()}",
+        ],
+        capture_output=True,
+        text=True,
+        check=False,
     )
     assert configure.returncode == 0, (
         "the documented example does not configure against the installed wheel: "
         f"{(configure.stderr or configure.stdout).strip()[-500:]}"
     )
     build = subprocess.run(
         ["cmake", "--build", str(build_dir)],
-        capture_output=True, text=True, check=False,
+        capture_output=True,
+        text=True,
+        check=False,
     )
     assert build.returncode == 0, (
         "the documented example does not compile against the installed wheel: "
@@ -1162,6 +1170,149 @@ def test_python_extension_links_shared_runtime() -> None:
         )
 
 
+def _requested_cuda_architectures() -> list:
+    """GPU architectures this wheel was built to support, from the build environment.
+
+    Empty when the build did not name any, which is the case for a CPU wheel and for a local
+    build that used detection.
+    """
+    raw = os.environ.get("CMAKE_CUDA_ARCHITECTURES", "")
+    if not raw:
+        match = re.search(
+            r"-DCMAKE_CUDA_ARCHITECTURES=([^\s]+)", os.environ.get("CMAKE_ARGS", "")
+        )
+        raw = match.group(1) if match else ""
+    if not raw:
+        return []
+    found = []
+    for entry in raw.replace(",", ";").split(";"):
+        number = re.match(r"(\d+)", entry.strip())
+        # A "-virtual" entry asks for a portable format rather than compiled code for that
+        # architecture, so it is not expected to appear as device code.
+        if number and "virtual" not in entry:
+            found.append(number.group(1))
+    return found
+
+
+def _cuobjdump() -> Optional[str]:
+    """Path to the CUDA object inspector, or None when it cannot be found.
+
+    It ships with the CUDA toolkit rather than the base system, and is not on the default PATH
+    on a typical CUDA machine, so the toolkit's own directory is searched too. Without this the
+    device-code audit finds no tool and has nothing to inspect.
+    """
+    found = shutil.which("cuobjdump")
+    if found:
+        return found
+    for root in (
+        os.environ.get("CUDA_HOME"),
+        os.environ.get("CUDA_PATH"),
+        "/usr/local/cuda",
+    ):
+        if not root:
+            continue
+        candidate = Path(root) / "bin" / "cuobjdump"
+        if candidate.is_file():
+            return str(candidate)
+    return None
+
+
+def _is_accelerator_row() -> bool:
+    """Whether this build is an accelerator row rather than a CPU wheel.
+
+    Taken from the build variable that names the row's CUDA train, which is the same signal
+    the rest of the wheel build uses. A CPU row leaves it unset or names no CUDA train.
+    """
+    train = (
+        os.environ.get("CU_VERSION") or os.environ.get("DESIRED_CUDA") or ""
+    ).strip()
+    return bool(train) and train.lower() not in {"cpu", "none"}
+
+
+def test_device_code_covers_claimed_architectures() -> None:
+    """Every GPU architecture the row claims must be present as compiled device code.
+
+    A wheel whose device code covers only the build machine's GPU installs on all the
+    hardware the row promises and then fails when a model runs. That failure appears late and
+    looks like a model problem rather than a packaging one, so it is caught here.
+
+    On an accelerator row this fails closed. A missing architecture claim, a missing
+    inspection tool, a missing library, or a failed inspection each leave the shipped device
+    code unverified, so each blocks rather than passes.
+    """
+    accelerator_row = _is_accelerator_row()
+    claimed = _requested_cuda_architectures()
+
+    if not claimed:
+        assert not accelerator_row, (
+            "this is an accelerator row but the build named no GPU architectures, so its "
+            "device code is whatever the build machine's GPU happened to be; the row's "
+            "architecture list has to reach the build"
+        )
+        print("- this is a CPU wheel, so there is no device code to audit")
+        return
+
+    inspector = _cuobjdump()
+    assert inspector or not accelerator_row, (
+        "this is an accelerator row but cuobjdump is not available, so the shipped device "
+        "code cannot be audited"
+    )
+    if not inspector:
+        print(
+            "- cuobjdump is not available and this is not an accelerator row, skipping"
+        )
+        return
+
+    package_dir = _installed_package_dir()
+    # Every shipped library is inspected rather than a subset chosen by file name. A
+    # library carrying device code under an unexpected name would otherwise be skipped,
+    # which is the failure this check exists to catch.
+    libraries = _shipped_shared_objects(package_dir)
+    assert libraries, f"no shared libraries found under {package_dir}"
+
+    # Each library is checked on its own. Unioning architectures across libraries would let
+    # one library's device code stand in for another's, which is the case worth catching.
+    audited = 0
+    for library in libraries:
+        result = subprocess.run(
+            [inspector, "--list-elf", str(library)],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        combined = result.stdout + result.stderr
+        # A library with no device code is a legitimate case, for example the delegate itself,
+        # which links the runtime but carries no kernels. cuobjdump reports that with a
+        # non-zero status, so it has to be told apart from a real inspection failure.
+        if "does not contain device code" in combined:
+            continue
+        assert result.returncode == 0 or not accelerator_row, (
+            f"could not inspect {library.name} on an accelerator row: "
+            f"{result.stderr.strip()[:200]}"
+        )
+        if result.returncode != 0:
+            continue
+        present = set(re.findall(r"sm_(\d+)", result.stdout))
+        if not present:
+            continue
+        missing = sorted(set(claimed) - present, key=int)
+        assert not missing, (
+            f"{library.name} claims GPU architectures {sorted(claimed, key=int)} but only "
+            f"carries device code for {sorted(present, key=int)}; a model would fail on "
+            f"hardware needing {missing}"
+        )
+        audited += 1
+
+    assert (
+        audited or not accelerator_row
+    ), "this is an accelerator row but no library carried device code to audit"
+    print(
+        f"\u2713 device code covers every claimed GPU architecture "
+        f"{sorted(claimed, key=int)} in {audited} "
+        f"librar{'y' if audited == 1 else 'ies'}"
+    )
+
+
 def run_tests(work_dir: Path) -> None:
     test_shipped_libraries_load()
     test_shipped_libraries_resolve_without_build_tree()
@@ -1175,6 +1326,7 @@ def run_tests(work_dir: Path) -> None:
     test_single_kernel_registration()
     test_single_xnnpack_delegate()
     test_single_cuda_delegate()
+    test_device_code_covers_claimed_architectures()
     test_cpp_consumer(work_dir)
     test_documented_example_compiles(work_dir)
     test_component_targets_link(work_dir)