| # Copyright (c) Meta Platforms, Inc. and affiliates. |
| # Copyright 2024 Arm Limited and/or its 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. |
| |
| # Part of this code is from pybind11 cmake_example: |
| # https://github.com/pybind/cmake_example/blob/master/setup.py so attach the |
| # license below. |
| |
| # Copyright (c) 2016 The Pybind Development Team, All rights reserved. |
| # |
| # Redistribution and use in source and binary forms, with or without |
| # modification, are permitted provided that the following conditions are met: |
| # |
| # 1. Redistributions of source code must retain the above copyright notice, this |
| # list of conditions and the following disclaimer. |
| # |
| # 2. Redistributions in binary form must reproduce the above copyright notice, |
| # this list of conditions and the following disclaimer in the documentation |
| # and/or other materials provided with the distribution. |
| # |
| # 3. Neither the name of the copyright holder nor the names of its contributors |
| # may be used to endorse or promote products derived from this software |
| # without specific prior written permission. |
| # |
| # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND |
| # ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED |
| # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE |
| # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE |
| # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL |
| # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR |
| # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER |
| # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, |
| # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| # |
| # You are under no obligation whatsoever to provide any bug fixes, patches, or |
| # upgrades to the features, functionality or performance of the source code |
| # ("Enhancements") to anyone; however, if you choose to make your Enhancements |
| # available either publicly, or directly to the author of this software, without |
| # imposing a separate written license agreement for such Enhancements, then you |
| # hereby grant the following license: a non-exclusive, royalty-free perpetual |
| # license to install, use, modify, prepare derivative works, incorporate into |
| # other computer software, distribute, and sublicense such enhancements or |
| # derivative works thereof, in binary and source code form. |
| |
| import contextlib |
| |
| # Import this before distutils so that setuptools can intercept the distuils |
| # imports. |
| import importlib.util |
| import logging |
| import os |
| import re |
| import shutil |
| import site |
| import subprocess |
| import sys |
| from distutils import log # type: ignore[import-not-found] |
| from distutils.sysconfig import get_python_lib # type: ignore[import-not-found] |
| from pathlib import Path |
| from typing import List, Optional |
| |
| # Clean dynamic import using importlib |
| _install_utils_path = Path(__file__).parent / "install_utils.py" |
| _spec = importlib.util.spec_from_file_location("install_utils", _install_utils_path) |
| if _spec is None: |
| raise ImportError(f"Could not create module spec for {_install_utils_path}") |
| install_utils = importlib.util.module_from_spec(_spec) |
| if _spec.loader is None: |
| raise ImportError(f"Module spec has no loader for {_install_utils_path}") |
| _spec.loader.exec_module(install_utils) |
| |
| from setuptools import Extension, setup |
| from setuptools.command.build import build |
| from setuptools.command.build_ext import build_ext |
| from setuptools.command.build_py import build_py |
| |
| logging.basicConfig( |
| level=logging.INFO, |
| format="%(asctime)s [%(levelname)s] %(message)s", |
| ) |
| |
| try: |
| from tools.cmake.cmake_cache import CMakeCache |
| except ImportError: |
| sys.path.insert( |
| 0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "tools", "cmake") |
| ) |
| from cmake_cache import CMakeCache # type: ignore[no-redef, import-not-found] |
| |
| |
| def _is_macos() -> bool: |
| return sys.platform == "darwin" |
| |
| |
| def _is_windows() -> bool: |
| return sys.platform == "win32" |
| |
| |
| class Version: |
| """Static strings that describe the version of the pip package.""" |
| |
| # Cached values returned by the properties. |
| __root_dir_attr: Optional[str] = None |
| __string_attr: Optional[str] = None |
| __git_hash_attr: Optional[str] = None |
| |
| @classmethod |
| def _root_dir(cls) -> str: |
| """The path to the root of the git repo.""" |
| if cls.__root_dir_attr is None: |
| # This setup.py file lives in the root of the repo. |
| cls.__root_dir_attr = str(Path(__file__).parent.resolve()) |
| return str(cls.__root_dir_attr) |
| |
| @classmethod |
| def git_hash(cls) -> Optional[str]: |
| """The current git hash, if known.""" |
| if cls.__git_hash_attr is None: |
| import subprocess |
| |
| try: |
| cls.__git_hash_attr = ( |
| subprocess.check_output( |
| ["git", "rev-parse", "HEAD"], cwd=cls._root_dir() |
| ) |
| .decode("ascii") |
| .strip() |
| ) |
| except subprocess.CalledProcessError: |
| cls.__git_hash_attr = "" # Non-None but empty. |
| # A non-None but empty value indicates that we don't know it. |
| return cls.__git_hash_attr if cls.__git_hash_attr else None |
| |
| @classmethod |
| def string(cls) -> str: |
| """The version string.""" |
| if cls.__string_attr is None: |
| # If set, BUILD_VERSION should override any local version |
| # information. CI will use this to manage, e.g., release vs. nightly |
| # versions. |
| version = os.getenv("BUILD_VERSION", "").strip() |
| if not version: |
| # Otherwise, read the version from a local file and add the git |
| # commit if available. |
| version = ( |
| open(os.path.join(cls._root_dir(), "version.txt")).read().strip() |
| ) |
| if cls.git_hash(): |
| version += "+" + cls.git_hash()[:7] # type: ignore[index] |
| cls.__string_attr = version |
| return cls.__string_attr |
| |
| @classmethod |
| def write_to_python_file(cls, path: str) -> None: |
| """Creates a file similar to PyTorch core's `torch/version.py`.""" |
| |
| lines = [ |
| "from typing import Optional", |
| '__all__ = ["__version__", "git_version"]', |
| f'__version__ = "{cls.string()}"', |
| # A string or None. |
| f"git_version: Optional[str] = {repr(cls.git_hash())}", |
| ] |
| with open(path, "w") as fp: |
| fp.write("\n".join(lines) + "\n") |
| |
| |
| # The build type is determined by the DEBUG environment variable. If DEBUG is |
| # set to a non-empty value, the build type is Debug. Otherwise, the build type |
| # is Release. |
| def get_build_type(is_debug=None) -> str: |
| debug = int(os.environ.get("DEBUG", 0) or 0) if is_debug is None else is_debug |
| return "Debug" if debug else "Release" |
| |
| |
| def get_dynamic_lib_name(name: str) -> str: |
| if _is_windows(): |
| return f"{name}.dll" |
| elif _is_macos(): |
| return f"lib{name}.dylib" |
| else: |
| return f"lib{name}.so" |
| |
| |
| def get_executable_name(name: str) -> str: |
| if _is_windows(): |
| return name + ".exe" |
| else: |
| return name |
| |
| |
| class _BaseExtension(Extension): |
| """A base class that maps an abstract source to an abstract destination.""" |
| |
| def __init__( |
| self, |
| src: str, |
| dst: str, |
| name: str, |
| dependent_cmake_flags: List[str], |
| ): |
| # Source path; semantics defined by the subclass. |
| self.src: str = src |
| |
| # Destination path relative to a namespace defined elsewhere. If this ends |
| # in "/", it is treated as a directory. If this is "", it is treated as the |
| # root of the namespace. |
| # Destination path; semantics defined by the subclass. |
| self.dst: str = dst |
| |
| # Other parts of setuptools expects .name to exist. For actual extensions |
| # this can be the module path, but otherwise it should be somehing unique |
| # that doesn't look like a module path. |
| self.name: str = name |
| |
| self.dependent_cmake_flags = dependent_cmake_flags |
| self.cmake_cache: Optional[CMakeCache] = None |
| |
| super().__init__(name=self.name, sources=[]) |
| |
| def _get_build_dir(self, installer: "InstallerBuildExt") -> Path: |
| # Share the cmake-out location with CustomBuild. |
| build_cmd = installer.get_finalized_command("build") |
| if "%CMAKE_CACHE_DIR%" in self.src: |
| if not hasattr(build_cmd, "cmake_cache_dir"): |
| raise RuntimeError( |
| f"Extension {self.name} has a src {self.src} that contains" |
| " %CMAKE_CACHE_DIR% but CMake does not run in the `build` " |
| "command. Please double check if the command is correct." |
| ) |
| else: |
| return Path(build_cmd.cmake_cache_dir) |
| else: |
| # If the src path doesn't contain %CMAKE_CACHE_DIR% placeholder, |
| # try to find it under the current directory. |
| return Path(".") |
| |
| def is_cmake_artifact_used(self, installer: "InstallerBuildExt") -> bool: |
| cache_path = str(self._get_build_dir(installer) / "CMakeCache.txt") |
| if not os.path.exists(cache_path): |
| # If this is not a CMake folder, then assume it's used. |
| return True |
| elif self.cmake_cache is None: |
| self.cmake_cache = CMakeCache(cache_path=cache_path) |
| |
| return all( |
| self.cmake_cache.is_enabled(flag) for flag in self.dependent_cmake_flags |
| ) |
| |
| def src_path(self, installer: "InstallerBuildExt") -> Path: |
| """Returns the path to the source file, resolving globs. |
| |
| Args: |
| installer: The InstallerBuildExt instance that is installing the |
| file. |
| """ |
| build_dir = self._get_build_dir(installer) |
| |
| src_path = self.src.replace("%CMAKE_CACHE_DIR%/", "") |
| |
| cfg = get_build_type(installer.debug) |
| |
| if os.name == "nt": |
| # Replace %BUILD_TYPE% with the current build type. |
| src_path = src_path.replace("%BUILD_TYPE%", cfg) |
| else: |
| # Remove %BUILD_TYPE% from the path. |
| src_path = src_path.replace("/%BUILD_TYPE%", "") |
| |
| # Construct the full source path, resolving globs. If there are no glob |
| # pattern characters, this will just ensure that the source file exists. |
| srcs = tuple(build_dir.glob(src_path)) |
| if len(srcs) != 1: |
| raise ValueError( |
| f"Expecting exactly 1 file matching {self.src} in {build_dir}, " |
| f"found {repr(srcs)}. Resolved src pattern: {src_path}." |
| ) |
| return srcs[0] |
| |
| def inplace_dir(self, installer: "InstallerBuildExt") -> Path: |
| """Returns the path of this file to be installed to, under inplace mode. |
| |
| It will be a relative path to the project root directory. For more info |
| related to inplace/editable mode, please checkout this doc: |
| https://setuptools.pypa.io/en/latest/userguide/development_mode.html |
| """ |
| raise NotImplementedError() |
| |
| |
| class BuiltFile(_BaseExtension): |
| """An extension that installs a single file that was built by cmake. |
| |
| This isn't technically a `build_ext` style python extension, but there's no |
| dedicated command for installing arbitrary data. It's convenient to use |
| this, though, because it lets us manage the files to install as entries in |
| `ext_modules`. |
| """ |
| |
| def __init__( |
| self, |
| src_dir: str, |
| src_name: str, |
| dst: str, |
| dependent_cmake_flags: List[str], |
| is_executable: bool = False, |
| is_dynamic_lib: bool = False, |
| ): |
| """Initializes a BuiltFile. |
| |
| Args: |
| src_dir: The directory of the file to install, relative to the cmake-out |
| directory. A placeholder %BUILD_TYPE% will be replaced with the build |
| type for multi-config generators (like Visual Studio) where the build |
| output is in a subdirectory named after the build type. For single- |
| config generators (like Makefile Generators or Ninja), this placeholder |
| will be removed. |
| src_name: The name of the file to install |
| dst: The path to install to, relative to the root of the pip |
| package. If dst ends in "/", it is treated as a directory. |
| Otherwise it is treated as a filename. |
| is_executable: If True, the file is an executable. This is used to |
| determine the destination filename for executable. |
| is_dynamic_lib: If True, the file is a dynamic library. This is used |
| to determine the destination filename for dynamic library. |
| """ |
| if is_executable and is_dynamic_lib: |
| raise ValueError("is_executable and is_dynamic_lib cannot be both True.") |
| if is_executable: |
| src_name = get_executable_name(src_name) |
| elif is_dynamic_lib: |
| src_name = get_dynamic_lib_name(src_name) |
| src = os.path.join(src_dir, src_name) |
| # This is not a real extension, so use a unique name that doesn't look |
| # like a module path. Some of setuptools's autodiscovery will look for |
| # extension names with prefixes that match certain module paths. |
| super().__init__( |
| src=src, |
| dst=dst, |
| name=f"@EXECUTORCH_BuiltFile_{src}:{dst}", |
| dependent_cmake_flags=dependent_cmake_flags, |
| ) |
| |
| def dst_path(self, installer: "InstallerBuildExt") -> Path: |
| """Returns the path to the destination file. |
| |
| Args: |
| installer: The InstallerBuildExt instance that is installing the |
| file. |
| """ |
| dst_root = Path(installer.build_lib).resolve() |
| |
| if self.dst.endswith("/"): |
| # Destination looks like a directory. Use the basename of the source |
| # file for its final component. |
| return dst_root / Path(self.dst) / self.src_path(installer).name |
| else: |
| # Destination looks like a file. |
| return dst_root / Path(self.dst) |
| |
| def inplace_dir(self, installer: "InstallerBuildExt") -> Path: |
| """For a `BuiltFile`, we use self.dst as its inplace directory path. |
| Need to handle directory vs file. |
| """ |
| # HACK: get rid of the leading "executorch" in ext.dst. |
| # This is because we don't have a root level "executorch" module. |
| package_dir = self.dst.removeprefix("executorch/") |
| # If dst is a file, use it's directory |
| if not package_dir.endswith("/"): |
| package_dir = os.path.dirname(package_dir) |
| return Path(package_dir) |
| |
| |
| class BuiltExtension(_BaseExtension): |
| """An extension that installs a python extension that was built by cmake.""" |
| |
| def __init__( |
| self, |
| src: str, |
| modpath: str, |
| dependent_cmake_flags: List[str], |
| src_dir: Optional[str] = None, |
| ): |
| """Initializes a BuiltExtension. |
| |
| Args: |
| src_dir: The directory of the file to install, relative to the cmake-out |
| directory. A placeholder %BUILD_TYPE% will be replaced with the build |
| type for multi-config generators (like Visual Studio) where the build |
| output is in a subdirectory named after the build type. For single- |
| config generators (like Makefile Generators or Ninja), this placeholder |
| will be removed. |
| src_name: The name of the file to install. If the path ends in `.so`, |
| modpath: The dotted path of the python module that maps to the |
| extension. |
| """ |
| assert ( |
| "/" not in modpath |
| ), f"modpath must be a dotted python module path: saw '{modpath}'" |
| full_src = src |
| if src_dir is None and _is_windows(): |
| src_dir = "%BUILD_TYPE%/" |
| if src_dir is not None: |
| full_src = os.path.join(src_dir, src) |
| self.dependent_cmake_flags = dependent_cmake_flags |
| # This is a real extension, so use the modpath as the name. |
| super().__init__( |
| src=f"%CMAKE_CACHE_DIR%/{full_src}", |
| dst=modpath, |
| name=modpath, |
| dependent_cmake_flags=self.dependent_cmake_flags, |
| ) |
| |
| def src_path(self, installer: "InstallerBuildExt") -> Path: |
| """Returns the path to the source file, resolving globs. |
| |
| Args: |
| installer: The InstallerBuildExt instance that is installing the |
| file. |
| """ |
| try: |
| return super().src_path(installer) |
| except ValueError: |
| # Probably couldn't find the file. If the path ends with .so, try |
| # looking for a .dylib file instead, in case we're running on macos. |
| if self.src.endswith(".so"): |
| dylib_src = re.sub(r"\.so$", ".dylib", self.src) |
| return BuiltExtension( |
| src=dylib_src, |
| modpath=self.dst, |
| dependent_cmake_flags=self.dependent_cmake_flags, |
| ).src_path(installer) |
| else: |
| raise |
| |
| def dst_path(self, installer: "InstallerBuildExt") -> Path: |
| """Returns the path to the destination file. |
| |
| Args: |
| installer: The InstallerBuildExt instance that is installing the |
| file. |
| """ |
| # Our destination is a dotted module path. get_ext_fullpath() returns |
| # the relative path to the .so/.dylib/etc. file that maps to the module |
| # path: that's the file we're creating. |
| return Path(installer.get_ext_fullpath(self.dst)) |
| |
| def inplace_dir(self, installer: "InstallerBuildExt") -> Path: |
| """For BuiltExtension, deduce inplace dir path from extension name.""" |
| build_py = installer.get_finalized_command("build_py") |
| modpath = self.name.split(".") |
| package = ".".join(modpath[:-1]) |
| package_dir = os.path.abspath(build_py.get_package_dir(package)) |
| |
| return Path(package_dir) |
| |
| |
| class InstallerBuildExt(build_ext): |
| """Installs files that were built by cmake.""" |
| |
| def __init__(self, *args, **kwargs): |
| self._ran_build = False |
| super().__init__(*args, **kwargs) |
| |
| def run(self): |
| # Run the build command first in editable mode. Since `build` command |
| # will also trigger `build_ext` command, only run this once. |
| if self._ran_build: |
| return |
| |
| if self.editable_mode: |
| self._ran_build = True |
| self.run_command("build") |
| super().run() |
| |
| def copy_extensions_to_source(self) -> None: |
| """For each extension in `ext_modules`, we need to copy the extension |
| file from the build directory to the correct location in the local |
| directory. |
| |
| This should only be triggered when inplace mode (editable mode) is enabled. |
| |
| Args: |
| |
| Returns: |
| """ |
| for ext in self.extensions: |
| if not ext.is_cmake_artifact_used(self): |
| continue |
| |
| package_dir = ext.inplace_dir(self) |
| |
| # Ensure that the destination directory exists. |
| self.mkpath(os.fspath(package_dir)) |
| |
| regular_file = ext.src_path(self) |
| inplace_file = os.path.join( |
| package_dir, os.path.basename(ext.dst_path(self)) |
| ) |
| |
| # Always copy, even if source is older than destination, to ensure |
| # that the right extensions for the current Python/platform are |
| # used. |
| if os.path.exists(regular_file) or not ext.optional: |
| self.copy_file(regular_file, inplace_file, level=self.verbose) |
| |
| if ext._needs_stub: |
| inplace_stub = self._get_equivalent_stub(ext, inplace_file) |
| self._write_stub_file(inplace_stub, ext, compile=True) |
| # Always compile stub and remove the original (leave the cache behind) |
| # (this behaviour was observed in previous iterations of the code) |
| |
| # TODO(dbort): Depend on the "build" command to ensure it runs first |
| |
| def build_extension(self, ext: _BaseExtension) -> None: |
| if not ext.is_cmake_artifact_used(self): |
| return |
| |
| src_file: Path = ext.src_path(self) |
| dst_file: Path = ext.dst_path(self) |
| |
| # Ensure that the destination directory exists. |
| if not dst_file.parent.exists(): |
| self.mkpath(os.fspath(dst_file.parent)) |
| |
| # Copy the file. |
| self.copy_file(os.fspath(src_file), os.fspath(dst_file)) |
| |
| # Ensure that the destination file is writable, even if the source was |
| # not. build_py does this by passing preserve_mode=False to copy_file, |
| # but that would clobber the X bit on any executables. TODO(dbort): This |
| # probably won't work on Windows. |
| if not os.access(src_file, os.W_OK): |
| # Make the file writable. This should respect the umask. |
| os.chmod(src_file, os.stat(src_file).st_mode | 0o222) |
| |
| |
| class CustomBuildPy(build_py): |
| """Copies platform-independent files from the source tree into the output |
| package directory. |
| |
| Override it so we can copy some files to locations that don't match their |
| original relative locations. |
| |
| Standard setuptools features like package_data and MANIFEST.in can only |
| include or exclude a file in the source tree; they don't have a way to map |
| a file to a different relative location under the output package directory. |
| """ |
| |
| def run(self): |
| # Copy python files to the output directory. This set of files is |
| # defined by the py_module list and package_data patterns. |
| build_py.run(self) |
| |
| # dst_root is the root of the `executorch` module in the output package |
| # directory. build_lib is the platform-independent root of the output |
| # package, and will look like `pip-out/lib`. It can contain multiple |
| # python packages, so be sure to copy the files into the `executorch` |
| # package subdirectory. |
| if self.editable_mode: |
| # In editable mode, the package directory is the original source directory |
| dst_root = self.get_package_dir("executorch") |
| else: |
| dst_root = os.path.join(self.build_lib, "executorch") |
| # Create the version file. |
| Version.write_to_python_file(os.path.join(dst_root, "version.py")) |
| |
| # Manually copy files into the output package directory. These are |
| # typically python "resource" files that will live alongside the python |
| # code that uses them. |
| src_to_dst = [ |
| # TODO(dbort): See if we can add a custom pyproject.toml section for |
| # these, instead of hard-coding them here. See |
| # https://setuptools.pypa.io/en/latest/userguide/extension.html |
| ("schema/scalar_type.fbs", "exir/_serialize/scalar_type.fbs"), |
| ("schema/program.fbs", "exir/_serialize/program.fbs"), |
| ( |
| "devtools/bundled_program/schema/bundled_program_schema.fbs", |
| "devtools/bundled_program/serialize/bundled_program_schema.fbs", |
| ), |
| ( |
| "devtools/bundled_program/schema/scalar_type.fbs", |
| "devtools/bundled_program/serialize/scalar_type.fbs", |
| ), |
| # Install executorch-wheel-config.cmake to pip package. |
| ( |
| "tools/cmake/executorch-wheel-config.cmake", |
| "share/cmake/executorch-config.cmake", |
| ), |
| ] |
| # Copy all the necessary headers into include/executorch/ so that they can |
| # be found in the pip package. This is the subset of headers that are |
| # essential for building custom ops extensions. |
| # TODO: Use cmake to gather the headers instead of hard-coding them here. |
| # For example: |
| # https://discourse.cmake.org/t/installing-headers-the-modern-way-regurgitated-and-revisited/3238/3 |
| for include_dir in [ |
| "runtime/core/", |
| "runtime/kernel/", |
| "runtime/platform/", |
| "extension/kernel_util/", |
| "extension/tensor/", |
| "extension/threadpool/", |
| ]: |
| src_list = Path(include_dir).rglob("*.h") |
| for src in src_list: |
| src_to_dst.append( |
| (str(src), os.path.join("include/executorch", str(src))) |
| ) |
| for src, dst in src_to_dst: |
| dst = os.path.join(dst_root, dst) |
| |
| # When modifying the filesystem, use the self.* methods defined by |
| # Command to benefit from the same logging and dry_run logic as |
| # setuptools. |
| |
| # Ensure that the destination directory exists. |
| self.mkpath(os.path.dirname(dst)) |
| # Follow the example of the base build_py class by not preserving |
| # the mode. This ensures that the output file is read/write even if |
| # the input file is read-only. |
| self.copy_file(src, dst, preserve_mode=False) |
| |
| |
| class Buck2EnvironmentFixer(contextlib.AbstractContextManager): |
| """Removes HOME from the environment when running as root. |
| |
| This script is sometimes run as root in docker containers. buck2 doesn't |
| allow running as root unless $HOME is owned by root or is not set. |
| |
| TODO(pytorch/test-infra#5091): Remove this once the CI jobs stop running as |
| root. |
| """ |
| |
| def __init__(self): |
| self.saved_env = {} |
| |
| def __enter__(self): |
| if os.name != "nt" and os.geteuid() == 0 and "HOME" in os.environ: |
| log.info("temporarily unsetting HOME while running as root") |
| self.saved_env["HOME"] = os.environ.pop("HOME") |
| return self |
| |
| def __exit__(self, *args, **kwargs): |
| if "HOME" in self.saved_env: |
| log.info("restored HOME") |
| os.environ["HOME"] = self.saved_env["HOME"] |
| |
| |
| # TODO(dbort): For editable wheels, may need to update get_source_files(), |
| # get_outputs(), and get_output_mapping() to satisfy |
| # https://setuptools.pypa.io/en/latest/userguide/extension.html#setuptools.command.build.SubCommand.get_output_mapping |
| |
| |
| class CustomBuild(build): |
| def initialize_options(self): |
| super().initialize_options() |
| # The default build_base directory is called "build", but we have a |
| # top-level directory with that name. Setting build_base in setup() |
| # doesn't affect this, so override the core build command. |
| # |
| # See build.initialize_options() in |
| # setuptools/_distutils/command/build.py for the default. |
| self.build_base = "pip-out" |
| |
| def run(self): # noqa C901 |
| self.dump_options() |
| cmake_build_type = get_build_type(self.debug) |
| # get_python_lib() typically returns the path to site-packages, where |
| # all pip packages in the environment are installed. |
| cmake_prefix_path = os.environ.get("CMAKE_PREFIX_PATH", get_python_lib()) |
| # Put the cmake cache under the temp directory, like |
| # "pip-out/temp.<plat>/cmake-out". |
| pip_build_dir = os.path.join( |
| os.path.dirname(os.path.abspath(__file__)), self.build_temp |
| ) |
| cmake_cache_dir = os.path.join(pip_build_dir, "cmake-out") |
| self.mkpath(cmake_cache_dir) |
| |
| cmake_configuration_args = [ |
| f"-DPYTHON_EXECUTABLE={sys.executable}", |
| # Let cmake calls like `find_package(Torch)` find cmake config files |
| # like `TorchConfig.cmake` that are provided by pip packages. |
| f"-DCMAKE_PREFIX_PATH={cmake_prefix_path}", |
| f"-DCMAKE_BUILD_TYPE={cmake_build_type}", |
| ] |
| |
| # Use ClangCL on Windows. |
| if _is_windows(): |
| cmake_configuration_args += ["-T ClangCL"] |
| |
| # Allow adding extra cmake args through the environment. Used by some |
| # tests and demos to expand the set of targets included in the pip |
| # package. |
| cmake_configuration_args += [ |
| item for item in re.split(r"\s+", os.environ.get("CMAKE_ARGS", "")) if item |
| ] |
| |
| # Check if CUDA is available, and if so, enable building the CUDA |
| # backend by default. |
| if install_utils.is_cuda_available() and install_utils.is_cmake_option_on( |
| cmake_configuration_args, "EXECUTORCH_BUILD_CUDA", default=True |
| ): |
| cmake_configuration_args += ["-DEXECUTORCH_BUILD_CUDA=ON"] |
| |
| # Check if QNN SDK is available (via QNN_SDK_ROOT env var), and if so, |
| # enable building the Qualcomm backend by default. |
| qnn_sdk_root = os.environ.get("QNN_SDK_ROOT", "").strip() |
| if qnn_sdk_root and install_utils.is_cmake_option_on( |
| cmake_configuration_args, "EXECUTORCH_BUILD_QNN", default=True |
| ): |
| cmake_configuration_args += [ |
| "-DEXECUTORCH_BUILD_QNN=ON", |
| f"-DQNN_SDK_ROOT={qnn_sdk_root}", |
| ] |
| |
| with Buck2EnvironmentFixer(): |
| # Generate the cmake cache from scratch to ensure that the cache state |
| # is predictable. |
| if os.path.exists(cmake_cache_dir): |
| log.info(f"clearing {cmake_cache_dir}") |
| shutil.rmtree(cmake_cache_dir) |
| |
| subprocess.run( |
| [ |
| "cmake", |
| *cmake_configuration_args, |
| "--preset", |
| "pybind", |
| "-B", |
| cmake_cache_dir, |
| ], |
| check=True, |
| ) |
| |
| cmake_cache = CMakeCache( |
| cache_path=os.path.join(cmake_cache_dir, "CMakeCache.txt") |
| ) |
| cmake_build_args = [ |
| # Default build parallelism based on number of cores, but allow |
| # overriding through the environment. |
| "-j{parallelism}".format( |
| parallelism=os.environ.get( |
| "CMAKE_BUILD_PARALLEL_LEVEL", os.cpu_count() - 1 |
| ) |
| ), |
| # CMAKE_BUILD_TYPE variable specifies the build type (configuration) for |
| # single-configuration generators (e.g., Makefile Generators or Ninja). |
| # For multi-config generators (like Visual Studio), CMAKE_BUILD_TYPE |
| # isn’t directly applicable. |
| # During the build step, --config specifies the configuration to build |
| # for multi-config generators. |
| f"--config={cmake_build_type}", |
| ] |
| |
| # Allow adding extra build args through the environment. Used by some |
| # tests and demos to expand the set of targets included in the pip |
| # package. |
| cmake_build_args += [ |
| item |
| for item in re.split(r"\s+", os.environ.get("CMAKE_BUILD_ARGS", "")) |
| if item |
| ] |
| |
| if cmake_cache.is_enabled("EXECUTORCH_BUILD_PYBIND"): |
| cmake_build_args += ["--target", "portable_lib"] |
| cmake_build_args += ["--target", "data_loader"] |
| cmake_build_args += ["--target", "selective_build"] |
| |
| if cmake_cache.is_enabled("EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER"): |
| cmake_build_args += ["--target", "_llm_runner"] |
| |
| if cmake_cache.is_enabled("EXECUTORCH_BUILD_CUDA"): |
| cmake_build_args += ["--target", "aoti_cuda_backend"] |
| cmake_build_args += ["--target", "aoti_common_shims_slim"] |
| |
| if cmake_cache.is_enabled("EXECUTORCH_BUILD_EXTENSION_MODULE"): |
| cmake_build_args += ["--target", "extension_module"] |
| |
| if cmake_cache.is_enabled("EXECUTORCH_BUILD_EXTENSION_TRAINING"): |
| cmake_build_args += ["--target", "_training_lib"] |
| |
| if cmake_cache.is_enabled("EXECUTORCH_BUILD_COREML"): |
| cmake_build_args += ["--target", "executorchcoreml"] |
| |
| if cmake_cache.is_enabled("EXECUTORCH_BUILD_KERNELS_LLM_AOT"): |
| cmake_build_args += ["--target", "custom_ops_aot_lib"] |
| cmake_build_args += ["--target", "quantized_ops_aot_lib"] |
| |
| if cmake_cache.is_enabled("EXECUTORCH_BUILD_QNN"): |
| cmake_build_args += ["--target", "qnn_executorch_backend"] |
| cmake_build_args += ["--target", "PyQnnManagerAdaptor"] |
| |
| # Set PYTHONPATH to the location of the pip package. |
| os.environ["PYTHONPATH"] = ( |
| site.getsitepackages()[0] + ";" + os.environ.get("PYTHONPATH", "") |
| ) |
| # Build the system. |
| self.spawn(["cmake", "--build", cmake_cache_dir, *cmake_build_args]) |
| # Share the cmake-out location with _BaseExtension. |
| self.cmake_cache_dir = cmake_cache_dir |
| # Finally, run the underlying subcommands like build_py, build_ext. |
| build.run(self) |
| |
| |
| setup( |
| version=Version.string(), |
| cmdclass={ |
| "build": CustomBuild, |
| "build_ext": InstallerBuildExt, |
| "build_py": CustomBuildPy, |
| }, |
| # Note that setuptools uses the presence of ext_modules as the main signal |
| # that a wheel is platform-specific. If we install any platform-specific |
| # files, this list must be non-empty. Therefore, we should always install |
| # platform-specific files using InstallerBuildExt. |
| ext_modules=[ |
| BuiltFile( |
| src_dir="%CMAKE_CACHE_DIR%/third-party/flatc_ep/bin/", |
| src_name="flatc", |
| dst="executorch/data/bin/", |
| is_executable=True, |
| dependent_cmake_flags=[], |
| ), |
| BuiltFile( |
| src_dir="tools/wheel", |
| src_name="pip_data_bin_init.py.in", |
| dst="executorch/data/bin/__init__.py", |
| dependent_cmake_flags=[], |
| ), |
| # Install the prebuilt pybindings extension wrapper for the runtime, |
| # portable kernels, and a selection of backends. This lets users |
| # load and execute .pte files from python. |
| BuiltExtension( |
| src="_portable_lib.cp*" if _is_windows() else "_portable_lib.*", |
| modpath="executorch.extension.pybindings._portable_lib", |
| dependent_cmake_flags=["EXECUTORCH_BUILD_PYBIND"], |
| ), |
| # Install the data_loader pybindings extension which provides the |
| # PyDataLoader type for external pybinding extensions. |
| BuiltExtension( |
| src="data_loader.cp*" if _is_windows() else "data_loader.*", |
| modpath="executorch.extension.pybindings.data_loader", |
| dependent_cmake_flags=["EXECUTORCH_BUILD_PYBIND"], |
| ), |
| BuiltExtension( |
| src="extension/training/_training_lib.*", # @lint-ignore https://github.com/pytorch/executorch/blob/cb3eba0d7f630bc8cec0a9cc1df8ae2f17af3f7a/scripts/lint_xrefs.sh |
| modpath="executorch.extension.training.pybindings._training_lib", |
| dependent_cmake_flags=["EXECUTORCH_BUILD_EXTENSION_TRAINING"], |
| ), |
| BuiltExtension( |
| src_dir="%CMAKE_CACHE_DIR%/codegen/tools/%BUILD_TYPE%/", |
| src="selective_build.cp*" if _is_windows() else "selective_build.*", |
| modpath="executorch.codegen.tools.selective_build", |
| dependent_cmake_flags=["EXECUTORCH_BUILD_PYBIND"], |
| ), |
| BuiltExtension( |
| src="extension/llm/runner/_llm_runner.*", # @lint-ignore https://github.com/pytorch/executorch/blob/cb3eba0d7f630bc8cec0a9cc1df8ae2f17af3f7a/scripts/lint_xrefs.sh |
| modpath="executorch.extension.llm.runner._llm_runner", |
| dependent_cmake_flags=["EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER"], |
| ), |
| BuiltExtension( |
| src="executorchcoreml.*", |
| src_dir="backends/apple/coreml", |
| modpath="executorch.backends.apple.coreml.executorchcoreml", |
| dependent_cmake_flags=["EXECUTORCH_BUILD_COREML"], |
| ), |
| BuiltFile( |
| src_dir="%CMAKE_CACHE_DIR%/extension/llm/custom_ops/%BUILD_TYPE%/", |
| src_name="custom_ops_aot_lib", |
| dst="executorch/extension/llm/custom_ops/", |
| is_dynamic_lib=True, |
| dependent_cmake_flags=["EXECUTORCH_BUILD_KERNELS_LLM_AOT"], |
| ), |
| BuiltFile( |
| src_dir="%CMAKE_CACHE_DIR%/kernels/quantized/%BUILD_TYPE%/", |
| src_name="quantized_ops_aot_lib", |
| dst="executorch/kernels/quantized/", |
| is_dynamic_lib=True, |
| dependent_cmake_flags=["EXECUTORCH_BUILD_KERNELS_LLM_AOT"], |
| ), |
| BuiltFile( |
| src_dir="backends/cuda/runtime/", |
| src_name="aoti_cuda_shims.lib", |
| dst="executorch/data/lib/", |
| dependent_cmake_flags=[], |
| ), |
| BuiltFile( |
| src_dir="%CMAKE_CACHE_DIR%/backends/cuda/%BUILD_TYPE%/", |
| src_name="aoti_cuda_shims", |
| dst="executorch/backends/cuda/", |
| is_dynamic_lib=True, |
| dependent_cmake_flags=["EXECUTORCH_BUILD_CUDA"], |
| ), |
| BuiltFile( |
| src_dir="%CMAKE_CACHE_DIR%/backends/qualcomm/%BUILD_TYPE%/", |
| src_name="qnn_executorch_backend", |
| dst="executorch/backends/qualcomm/", |
| is_dynamic_lib=True, |
| dependent_cmake_flags=["EXECUTORCH_BUILD_QNN"], |
| ), |
| BuiltExtension( |
| src_dir="%CMAKE_CACHE_DIR%/backends/qualcomm/%BUILD_TYPE%/", |
| src="PyQnnManagerAdaptor.*", |
| modpath="executorch.backends.qualcomm.python.PyQnnManagerAdaptor", |
| dependent_cmake_flags=["EXECUTORCH_BUILD_QNN"], |
| ), |
| ], |
| ) |
