diff --git a/.ci/scripts/wheel/test_cpp_sdk.py b/.ci/scripts/wheel/test_cpp_sdk.py index 61265d6435d..9b13fda0a4e 100644 --- a/.ci/scripts/wheel/test_cpp_sdk.py +++ b/.ci/scripts/wheel/test_cpp_sdk.py @@ -30,6 +30,7 @@ import sys import tempfile from pathlib import Path +from typing import List, Optional # Registry entry points. A second definer of any of these means a second # process-wide registry. @@ -603,10 +604,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] @@ -1344,9 +1342,17 @@ 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: " @@ -1354,7 +1360,9 @@ def test_documented_example_compiles(work_dir: Path) -> None: ) 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: " @@ -1406,6 +1414,157 @@ 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_names: List[str] = [] + 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_names.append(library.name) + + # A count is not enough. The library that carries the compiled kernels has to be the one + # audited, or a build that produced only forward-compatible intermediate code there would + # still pass on the strength of some other library's coverage. + assert audited_names or not accelerator_row, ( + "this is an accelerator row but no library carried device code to audit" + ) + if accelerator_row: + assert any("shim" in name for name in audited_names), ( + "no shim library carried device code; the row claims " + f"{sorted(claimed, key=int)} but only {sorted(audited_names)} were audited" + ) + print( + f"\u2713 device code covers every claimed GPU architecture " + f"{sorted(claimed, key=int)} in {len(audited_names)} " + f"librar{'y' if len(audited_names) == 1 else 'ies'}" + ) + + def run_tests(work_dir: Path) -> None: test_shipped_libraries_load() test_shipped_libraries_resolve_without_build_tree() @@ -1419,6 +1578,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)