CUDA wheel: check every dependency, and give torch_cpu_library targets torch's compile settings back

.ci/scripts/tests/test_torch_cpu_library_cmake.py

@@ -0,0 +1,143 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import importlib.util
+import json
+import shlex
+import shutil
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[3]
+
+
+def _compile_flags(entry: dict) -> set:
+    """The flags of one compile command, without the compiler, source and outputs."""
+    flags = set()
+    arguments = iter(shlex.split(entry["command"])[1:])
+    for argument in arguments:
+        if argument in ("-o", "-c", "-MT", "-MF"):
+            next(arguments)
+        elif argument == "-isystem":
+            flags.add(f"-isystem {next(arguments)}")
+        elif argument != "-MD":
+            flags.add(argument)
+    return flags
+
+
+def _copied_from_torch(flag: str) -> bool:
+    """Whether a flag comes from what find_package_torch copies from torch.
+
+    A CUDA build of torch also gives torch the CUDA toolkit include dir and torch_cuda's
+    definitions, which torch_cpu_library correctly does not carry.
+    """
+    return (
+        flag == "-DET_TORCH_CXX_FLAGS_PROBE"
+        or flag.startswith("-std=")
+        or (flag.startswith("-isystem ") and flag.endswith("torch/csrc/api/include"))
+    )
+
+
+@unittest.skipIf(sys.platform == "win32", "Visual Studio writes no compile commands")
+class TestTorchCpuLibraryCMake(unittest.TestCase):
+    def test_torch_cpu_library_gets_torch_settings(self):
+        cmake = shutil.which("cmake")
+        if cmake is None or importlib.util.find_spec("torch") is None:
+            self.skipTest("CMake and torch are required")
+
+        found_first = (
+            "add_torch_to_cmake_prefix_path()\nfind_package(Torch CONFIG REQUIRED)"
+        )
+        for find_torch in ("", found_first):
+            with self.subTest(found_first=bool(find_torch)):
+                self._check(cmake, find_torch)
+
+    def _check(self, cmake: str, find_torch: str) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            source = Path(directory)
+            build = source / "build"
+            (source / "probe.cpp").write_text("int probe() { return 0; }\n")
+            (source / "CMakeLists.txt").write_text(
+                f"""cmake_minimum_required(VERSION 3.24)
+project(TorchCpuLibraryProbe LANGUAGES CXX)
+include("${{EXECUTORCH_ROOT}}/tools/cmake/Utils.cmake")
+{find_torch}
+find_package_torch()
+add_library(links_torch OBJECT probe.cpp)
+target_link_libraries(links_torch PRIVATE torch)
+add_library(links_torch_cpu_library OBJECT probe.cpp)
+target_link_libraries(links_torch_cpu_library PRIVATE torch_cpu_library)
+add_library(links_it_publicly STATIC probe.cpp)
+target_link_libraries(links_it_publicly PUBLIC torch_cpu_library)
+add_library(consumer OBJECT probe.cpp)
+target_link_libraries(consumer PRIVATE links_it_publicly)
+"""
+            )
+            result = subprocess.run(
+                [
+                    cmake,
+                    "-S",
+                    str(source),
+                    "-B",
+                    str(build),
+                    f"-DEXECUTORCH_ROOT={ROOT}",
+                    f"-DPYTHON_EXECUTABLE={sys.executable}",
+                    "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
+                    # Recent torch sets no TORCH_CXX_FLAGS, so set one to check it travels too.
+                    "-DTORCH_CXX_FLAGS=-DET_TORCH_CXX_FLAGS_PROBE",
+                ],
+                capture_output=True,
+                text=True,
+            )
+            self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
+
+            flags = {}
+            for entry in json.loads((build / "compile_commands.json").read_text()):
+                target = Path(entry["output"]).parent.name.removesuffix(".dir")
+                flags[target] = _compile_flags(entry)
+            expected = {
+                flag for flag in flags["links_torch"] if _copied_from_torch(flag)
+            }
+            self.assertIn("-DET_TORCH_CXX_FLAGS_PROBE", expected)
+            self.assertTrue(any(flag.startswith("-isystem ") for flag in expected))
+            for target in ("links_torch_cpu_library", "links_it_publicly", "consumer"):
+                self.assertEqual(expected - flags[target], set(), target)
+
+    def test_find_package_torch_without_torch_cpu_library(self):
+        """A parent project or a static libtorch can define torch alone."""
+        cmake = shutil.which("cmake")
+        if cmake is None:
+            self.skipTest("CMake is required")
+        with tempfile.TemporaryDirectory() as directory:
+            source = Path(directory)
+            (source / "CMakeLists.txt").write_text(
+                """cmake_minimum_required(VERSION 3.24)
+project(TorchWithoutTorchCpuLibrary LANGUAGES NONE)
+include("${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake")
+add_library(torch INTERFACE)
+find_package_torch()
+"""
+            )
+            result = subprocess.run(
+                [
+                    cmake,
+                    "-S",
+                    str(source),
+                    "-B",
+                    str(source / "build"),
+                    f"-DEXECUTORCH_ROOT={ROOT}",
+                ],
+                capture_output=True,
+                text=True,
+            )
+            self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
+
+
+if __name__ == "__main__":
+    unittest.main()

.ci/scripts/wheel/test_cuda_linux.py

@@ -143,25 +143,79 @@ def test_cuda_libraries_resolve_relatively() -> None:
         print(f"✓ {name} resolves the CUDA runtime relatively ({relative[0]})")
 
 
-def test_no_library_needs_cuda_torch() -> None:
-    """No shipped library may record a dependency on PyTorch's own CUDA libraries.
+# The torch libraries a shipped library may link. A CUDA build of torch adds libtorch_cuda,
+# libc10_cuda and more, which a CPU-only torch does not have.
+_ALLOWED_TORCH_LIBRARIES = frozenset(
+    {"libtorch.so", "libtorch_cpu.so", "libtorch_python.so", "libc10.so"}
+)
+
+# The system libraries the manylinux policy lets a wheel record, and the loader.
+# libatomic is one of them; the aarch64 build links it for sentencepiece.
+_SYSTEM_LIBRARIES = frozenset(
+    {
+        "libatomic.so.1",
+        "libc.so.6",
+        "libm.so.6",
+        "libdl.so.2",
+        "libpthread.so.0",
+        "librt.so.1",
+        "libstdc++.so.6",
+        "libgcc_s.so.1",
+        "ld-linux-x86-64.so.2",
+        "ld-linux-aarch64.so.1",
+    }
+)
+
 
-    The CUDA delegate needs only the CUDA runtime. A library that names libtorch_cuda or
-    libc10_cuda cannot load next to a CPU-only torch, so the whole wheel then fails to import
-    there even though nothing in it calls into those libraries.
+def _declared_requirement_libraries() -> Set[str]:
+    """The shared library file names installed by the packages a plain install requires."""
+    import importlib.metadata as metadata
+
+    from packaging.requirements import Requirement
+
+    names: Set[str] = set()
+    for line in metadata.requires("executorch") or []:
+        requirement = Requirement(line)
+        # Skip extras and requirements for other platforms, which a plain install does not get.
+        if requirement.marker and not requirement.marker.evaluate({"extra": ""}):
+            continue
+        try:
+            files = metadata.files(requirement.name) or []
+        except metadata.PackageNotFoundError:
+            continue
+        names.update(Path(file).name for file in files if ".so" in Path(file).name)
+    return names
+
+
+def test_recorded_dependencies_are_allowed() -> None:
+    """Every library a shipped library records must be on an allow list.
+
+    The list is: shipped in the wheel, a torch library a CPU-only torch also has, the system
+    C and C++ runtime, or installed by a package a plain install requires. A name outside
+    it, such as libtorch_cuda or libnvrtc, may be missing next to a CPU-only torch, and
+    then the library that records it fails to load there.
     """
     package_dir = _package_dir()
-    cuda_torch = {"libtorch_cuda.so", "libc10_cuda.so"}
+    libraries = test_shared_libraries._shipped_shared_objects(package_dir)
+    allowed = (
+        {library.name for library in libraries}
+        | _ALLOWED_TORCH_LIBRARIES
+        | _SYSTEM_LIBRARIES
+        | _declared_requirement_libraries()
+    )
     offenders = {}
-    for library in test_shared_libraries._shipped_shared_objects(package_dir):
-        needed = test_shared_libraries._recorded_dependencies(library) & cuda_torch
-        if needed:
-            offenders[str(library.relative_to(package_dir))] = sorted(needed)
+    for library in libraries:
+        unexpected = test_shared_libraries._recorded_dependencies(library) - allowed
+        if unexpected:
+            offenders[str(library.relative_to(package_dir))] = sorted(unexpected)
     assert not offenders, (
-        f"these libraries depend on PyTorch's CUDA libraries, so the wheel cannot be imported "
-        f"with a CPU-only torch: {offenders}"
+        f"these libraries record dependencies outside the allow list, so they may fail to "
+        f"load next to a CPU-only torch: {offenders}"
+    )
+    print(
+        f"✓ every dependency recorded by the {len(libraries)} shipped libraries is on the "
+        f"allow list"
     )
-    print("✓ no shipped library depends on PyTorch's CUDA libraries")
 
 
 def _row_architectures() -> list[str]:
@@ -607,7 +661,7 @@ def test_a_model_runs_through_the_delegate() -> None:
     test_cuda_libraries_are_shipped()
     test_cuda_runtime_is_declared()
     test_cuda_libraries_resolve_relatively()
-    test_no_library_needs_cuda_torch()
+    test_recorded_dependencies_are_allowed()
     test_device_code_covers_the_row()
     test_portable_device_code_is_present()
     test_the_delegate_registers()

CMakeLists.txt

@@ -1439,7 +1439,7 @@ if(EXECUTORCH_BUILD_PYBIND)
       TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib"
     )
     add_library(
-      aten_bridge
+      aten_bridge STATIC
       ${CMAKE_CURRENT_SOURCE_DIR}/extension/aten_util/aten_bridge.cpp
     )
     set_target_properties(aten_bridge PROPERTIES CXX_STANDARD 20)
@@ -1450,9 +1450,8 @@ if(EXECUTORCH_BUILD_PYBIND)
       aten_bridge PUBLIC C10_USING_CUSTOM_GENERATED_MACROS
     )
     target_compile_options(aten_bridge PUBLIC ${_pybind_compile_options})
-    # torch_cpu_library is the CPU half of torch. Linking torch would also force
-    # libtorch_cuda into DT_NEEDED in a CUDA build, and a CPU-only torch could
-    # then not load this library.
+    # aten_bridge is static, so _training_lib gets this link. torch would record
+    # libtorch_cuda with a CUDA build of PyTorch; torch_cpu_library does not.
     if(EXECUTORCH_BUILD_SHARED)
       target_link_libraries(
         aten_bridge PRIVATE torch_cpu_library c10 executorch_shared

extension/llm/custom_ops/CMakeLists.txt

@@ -230,9 +230,8 @@ if(EXECUTORCH_BUILD_KERNELS_LLM_AOT)
   endif()
   executorch_target_link_shared_runtime(custom_ops_aot_lib)
 
-  # torch_cpu_library is the CPU half of torch. Linking torch would also force
-  # libtorch_cuda into DT_NEEDED in a CUDA build, and a CPU-only torch could
-  # then not load this library.
+  # torch would record libtorch_cuda with a CUDA build of PyTorch, so a CPU-only
+  # torch could not load this. torch_cpu_library does not.
   target_link_libraries(
     custom_ops_aot_lib PUBLIC cpublas torch_cpu_library extension_tensor
                               extension_threadpool

extension/llm/runner/CMakeLists.txt

@@ -122,9 +122,8 @@ if(EXECUTORCH_BUILD_PYBIND)
   )
   # Link with the extension_llm_runner library and its dependencies.
   #
-  # torch_cpu_library is the CPU half of torch. Linking torch would also force
-  # libtorch_cuda into DT_NEEDED in a CUDA build, and a CPU-only torch could
-  # then not load this library.
+  # torch would record libtorch_cuda with a CUDA build of PyTorch, so a CPU-only
+  # torch could not load this. torch_cpu_library does not.
   target_link_libraries(
     _llm_runner PRIVATE extension_llm_runner tokenizers::tokenizers
                         portable_lib ${TORCH_PYTHON_LIBRARY} torch_cpu_library

extension/training/CMakeLists.txt

@@ -53,9 +53,8 @@ if(EXECUTORCH_BUILD_PYBIND)
   # is resolved from it rather than from the whole-archive-forcing static
   # `executorch` target, so this module shares the one backend registry.
   #
-  # torch_cpu_library is the CPU half of torch. Linking torch would also force
-  # libtorch_cuda into DT_NEEDED in a CUDA build, and a CPU-only torch could
-  # then not load this library.
+  # torch would record libtorch_cuda with a CUDA build of PyTorch, so a CPU-only
+  # torch could not load this. torch_cpu_library does not.
   if(EXECUTORCH_BUILD_SHARED)
     set(_pybind_training_dep_libs ${TORCH_PYTHON_LIBRARY} etdump aten_bridge
                                   torch_cpu_library extension_training

tools/cmake/Codegen.cmake

@@ -334,9 +334,8 @@ function(gen_custom_ops_aot_lib)
   set_target_properties(${GEN_LIB_NAME} PROPERTIES CXX_STANDARD 20)
   target_compile_definitions(${GEN_LIB_NAME} PRIVATE USE_ATEN_LIB=1)
   include_directories(${TORCH_INCLUDE_DIRS})
-  # torch_cpu_library is the CPU half of torch. Linking torch would also force
-  # libtorch_cuda into DT_NEEDED in a CUDA build, and a CPU-only torch could
-  # then not load this library.
+  # torch would record libtorch_cuda with a CUDA build of PyTorch, so a CPU-only
+  # torch could not load this. torch_cpu_library does not.
   target_link_libraries(${GEN_LIB_NAME} PRIVATE torch_cpu_library)
 
   executorch_target_link_options_shared_lib(${GEN_LIB_NAME})

tools/cmake/Utils.cmake

@@ -278,6 +278,36 @@ macro(find_package_torch)
     add_torch_to_cmake_prefix_path()
     find_package(Torch CONFIG REQUIRED)
   endif()
+  # TorchConfig gives its include dirs, compile flags and C++ standard to torch
+  # only. Our targets link torch_cpu_library instead, so copy them there once,
+  # also when the project found torch before calling this. This affects every
+  # target in the build that links torch_cpu_library.
+  if(TARGET torch_cpu_library)
+    get_target_property(
+      _torch_settings_copied torch_cpu_library EXECUTORCH_TORCH_SETTINGS_COPIED
+    )
+    if(NOT _torch_settings_copied)
+      foreach(_torch_property
+              INTERFACE_INCLUDE_DIRECTORIES INTERFACE_COMPILE_OPTIONS
+              INTERFACE_COMPILE_FEATURES
+      )
+        get_target_property(_torch_value torch ${_torch_property})
+        if(_torch_value)
+          set_property(
+            TARGET torch_cpu_library
+            APPEND
+            PROPERTY ${_torch_property} ${_torch_value}
+          )
+        endif()
+      endforeach()
+      set_property(
+        TARGET torch_cpu_library PROPERTY EXECUTORCH_TORCH_SETTINGS_COPIED TRUE
+      )
+    endif()
+    unset(_torch_settings_copied)
+    unset(_torch_property)
+    unset(_torch_value)
+  endif()
 endmacro()
 
 # Modify ${targetName}'s INTERFACE_INCLUDE_DIRECTORIES by wrapping each entry in