|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +import importlib.util |
| 8 | +import json |
| 9 | +import shlex |
| 10 | +import shutil |
| 11 | +import subprocess |
| 12 | +import sys |
| 13 | +import tempfile |
| 14 | +import unittest |
| 15 | +from pathlib import Path |
| 16 | + |
| 17 | +ROOT = Path(__file__).resolve().parents[3] |
| 18 | + |
| 19 | + |
| 20 | +def _compile_flags(entry: dict) -> set: |
| 21 | + """The flags of one compile command, without the compiler, source and outputs.""" |
| 22 | + flags = set() |
| 23 | + arguments = iter(shlex.split(entry["command"])[1:]) |
| 24 | + for argument in arguments: |
| 25 | + if argument in ("-o", "-c", "-MT", "-MF"): |
| 26 | + next(arguments) |
| 27 | + elif argument == "-isystem": |
| 28 | + flags.add(f"-isystem {next(arguments)}") |
| 29 | + elif argument != "-MD": |
| 30 | + flags.add(argument) |
| 31 | + return flags |
| 32 | + |
| 33 | + |
| 34 | +def _copied_from_torch(flag: str) -> bool: |
| 35 | + """Whether a flag comes from what find_package_torch copies from torch. |
| 36 | +
|
| 37 | + A CUDA build of torch also gives torch the CUDA toolkit include dir and torch_cuda's |
| 38 | + definitions, which torch_cpu_library correctly does not carry. |
| 39 | + """ |
| 40 | + return ( |
| 41 | + flag == "-DET_TORCH_CXX_FLAGS_PROBE" |
| 42 | + or flag.startswith("-std=") |
| 43 | + or (flag.startswith("-isystem ") and flag.endswith("torch/csrc/api/include")) |
| 44 | + ) |
| 45 | + |
| 46 | + |
| 47 | +@unittest.skipIf(sys.platform == "win32", "Visual Studio writes no compile commands") |
| 48 | +class TestTorchCpuLibraryCMake(unittest.TestCase): |
| 49 | + def test_torch_cpu_library_gets_torch_settings(self): |
| 50 | + cmake = shutil.which("cmake") |
| 51 | + if cmake is None or importlib.util.find_spec("torch") is None: |
| 52 | + self.skipTest("CMake and torch are required") |
| 53 | + |
| 54 | + found_first = ( |
| 55 | + "add_torch_to_cmake_prefix_path()\nfind_package(Torch CONFIG REQUIRED)" |
| 56 | + ) |
| 57 | + for find_torch in ("", found_first): |
| 58 | + with self.subTest(found_first=bool(find_torch)): |
| 59 | + self._check(cmake, find_torch) |
| 60 | + |
| 61 | + def _check(self, cmake: str, find_torch: str) -> None: |
| 62 | + with tempfile.TemporaryDirectory() as directory: |
| 63 | + source = Path(directory) |
| 64 | + build = source / "build" |
| 65 | + (source / "probe.cpp").write_text("int probe() { return 0; }\n") |
| 66 | + (source / "CMakeLists.txt").write_text( |
| 67 | + f"""cmake_minimum_required(VERSION 3.24) |
| 68 | +project(TorchCpuLibraryProbe LANGUAGES CXX) |
| 69 | +include("${{EXECUTORCH_ROOT}}/tools/cmake/Utils.cmake") |
| 70 | +{find_torch} |
| 71 | +find_package_torch() |
| 72 | +add_library(links_torch OBJECT probe.cpp) |
| 73 | +target_link_libraries(links_torch PRIVATE torch) |
| 74 | +add_library(links_torch_cpu_library OBJECT probe.cpp) |
| 75 | +target_link_libraries(links_torch_cpu_library PRIVATE torch_cpu_library) |
| 76 | +add_library(links_it_publicly STATIC probe.cpp) |
| 77 | +target_link_libraries(links_it_publicly PUBLIC torch_cpu_library) |
| 78 | +add_library(consumer OBJECT probe.cpp) |
| 79 | +target_link_libraries(consumer PRIVATE links_it_publicly) |
| 80 | +""" |
| 81 | + ) |
| 82 | + result = subprocess.run( |
| 83 | + [ |
| 84 | + cmake, |
| 85 | + "-S", |
| 86 | + str(source), |
| 87 | + "-B", |
| 88 | + str(build), |
| 89 | + f"-DEXECUTORCH_ROOT={ROOT}", |
| 90 | + f"-DPYTHON_EXECUTABLE={sys.executable}", |
| 91 | + "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON", |
| 92 | + # Recent torch sets no TORCH_CXX_FLAGS, so set one to check it travels too. |
| 93 | + "-DTORCH_CXX_FLAGS=-DET_TORCH_CXX_FLAGS_PROBE", |
| 94 | + ], |
| 95 | + capture_output=True, |
| 96 | + text=True, |
| 97 | + ) |
| 98 | + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) |
| 99 | + |
| 100 | + flags = {} |
| 101 | + for entry in json.loads((build / "compile_commands.json").read_text()): |
| 102 | + target = Path(entry["output"]).parent.name.removesuffix(".dir") |
| 103 | + flags[target] = _compile_flags(entry) |
| 104 | + expected = { |
| 105 | + flag for flag in flags["links_torch"] if _copied_from_torch(flag) |
| 106 | + } |
| 107 | + self.assertIn("-DET_TORCH_CXX_FLAGS_PROBE", expected) |
| 108 | + self.assertTrue(any(flag.startswith("-isystem ") for flag in expected)) |
| 109 | + for target in ("links_torch_cpu_library", "links_it_publicly", "consumer"): |
| 110 | + self.assertEqual(expected - flags[target], set(), target) |
| 111 | + |
| 112 | + def test_find_package_torch_without_torch_cpu_library(self): |
| 113 | + """A parent project or a static libtorch can define torch alone.""" |
| 114 | + cmake = shutil.which("cmake") |
| 115 | + if cmake is None: |
| 116 | + self.skipTest("CMake is required") |
| 117 | + with tempfile.TemporaryDirectory() as directory: |
| 118 | + source = Path(directory) |
| 119 | + (source / "CMakeLists.txt").write_text( |
| 120 | + """cmake_minimum_required(VERSION 3.24) |
| 121 | +project(TorchWithoutTorchCpuLibrary LANGUAGES NONE) |
| 122 | +include("${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake") |
| 123 | +add_library(torch INTERFACE) |
| 124 | +find_package_torch() |
| 125 | +""" |
| 126 | + ) |
| 127 | + result = subprocess.run( |
| 128 | + [ |
| 129 | + cmake, |
| 130 | + "-S", |
| 131 | + str(source), |
| 132 | + "-B", |
| 133 | + str(source / "build"), |
| 134 | + f"-DEXECUTORCH_ROOT={ROOT}", |
| 135 | + ], |
| 136 | + capture_output=True, |
| 137 | + text=True, |
| 138 | + ) |
| 139 | + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) |
| 140 | + |
| 141 | + |
| 142 | +if __name__ == "__main__": |
| 143 | + unittest.main() |
0 commit comments