xref: /aosp_15_r20/external/pytorch/setup.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# Welcome to the PyTorch setup.py.
2# Environment variables you are probably interested in:
3#
4#   DEBUG
5#     build with -O0 and -g (debug symbols)
6#
7#   REL_WITH_DEB_INFO
8#     build with optimizations and -g (debug symbols)
9#
10#   USE_CUSTOM_DEBINFO="path/to/file1.cpp;path/to/file2.cpp"
11#     build with debug info only for specified files
12#
13#   MAX_JOBS
14#     maximum number of compile jobs we should use to compile your code
15#
16#   USE_CUDA=0
17#     disables CUDA build
18#
19#   CFLAGS
20#     flags to apply to both C and C++ files to be compiled (a quirk of setup.py
21#     which we have faithfully adhered to in our build system is that CFLAGS
22#     also applies to C++ files (unless CXXFLAGS is set), in contrast to the
23#     default behavior of autogoo and cmake build systems.)
24#
25#   CC
26#     the C/C++ compiler to use
27#
28# Environment variables for feature toggles:
29#
30#   DEBUG_CUDA=1
31#     if used in conjunction with DEBUG or REL_WITH_DEB_INFO, will also
32#     build CUDA kernels with -lineinfo --source-in-ptx.  Note that
33#     on CUDA 12 this may cause nvcc to OOM, so this is disabled by default.
34
35#   USE_CUDNN=0
36#     disables the cuDNN build
37#
38#   USE_CUSPARSELT=0
39#     disables the cuSPARSELt build
40#
41#   USE_CUDSS=0
42#     disables the cuDSS build
43#
44#   USE_CUFILE=0
45#     disables the cuFile build
46#
47#   USE_FBGEMM=0
48#     disables the FBGEMM build
49#
50#   USE_KINETO=0
51#     disables usage of libkineto library for profiling
52#
53#   USE_NUMPY=0
54#     disables the NumPy build
55#
56#   BUILD_TEST=0
57#     disables the test build
58#
59#   USE_MKLDNN=0
60#     disables use of MKLDNN
61#
62#   USE_MKLDNN_ACL
63#     enables use of Compute Library backend for MKLDNN on Arm;
64#     USE_MKLDNN must be explicitly enabled.
65#
66#   MKLDNN_CPU_RUNTIME
67#     MKL-DNN threading mode: TBB or OMP (default)
68#
69#   USE_STATIC_MKL
70#     Prefer to link with MKL statically - Unix only
71#   USE_ITT=0
72#     disable use of Intel(R) VTune Profiler's ITT functionality
73#
74#   USE_NNPACK=0
75#     disables NNPACK build
76#
77#   USE_DISTRIBUTED=0
78#     disables distributed (c10d, gloo, mpi, etc.) build
79#
80#   USE_TENSORPIPE=0
81#     disables distributed Tensorpipe backend build
82#
83#   USE_GLOO=0
84#     disables distributed gloo backend build
85#
86#   USE_MPI=0
87#     disables distributed MPI backend build
88#
89#   USE_SYSTEM_NCCL=0
90#     disables use of system-wide nccl (we will use our submoduled
91#     copy in third_party/nccl)
92#
93#   USE_OPENMP=0
94#     disables use of OpenMP for parallelization
95#
96#   USE_FLASH_ATTENTION=0
97#     disables building flash attention for scaled dot product attention
98#
99#   USE_MEM_EFF_ATTENTION=0
100#    disables building memory efficient attention for scaled dot product attention
101#
102#   BUILD_BINARY
103#     enables the additional binaries/ build
104#
105#   ATEN_AVX512_256=TRUE
106#     ATen AVX2 kernels can use 32 ymm registers, instead of the default 16.
107#     This option can be used if AVX512 doesn't perform well on a machine.
108#     The FBGEMM library also uses AVX512_256 kernels on Xeon D processors,
109#     but it also has some (optimized) assembly code.
110#
111#   PYTORCH_BUILD_VERSION
112#   PYTORCH_BUILD_NUMBER
113#     specify the version of PyTorch, rather than the hard-coded version
114#     in this file; used when we're building binaries for distribution
115#
116#   TORCH_CUDA_ARCH_LIST
117#     specify which CUDA architectures to build for.
118#     ie `TORCH_CUDA_ARCH_LIST="6.0;7.0"`
119#     These are not CUDA versions, instead, they specify what
120#     classes of NVIDIA hardware we should generate PTX for.
121#
122#   PYTORCH_ROCM_ARCH
123#     specify which AMD GPU targets to build for.
124#     ie `PYTORCH_ROCM_ARCH="gfx900;gfx906"`
125#
126#   ONNX_NAMESPACE
127#     specify a namespace for ONNX built here rather than the hard-coded
128#     one in this file; needed to build with other frameworks that share ONNX.
129#
130#   BLAS
131#     BLAS to be used by Caffe2. Can be MKL, Eigen, ATLAS, FlexiBLAS, or OpenBLAS. If set
132#     then the build will fail if the requested BLAS is not found, otherwise
133#     the BLAS will be chosen based on what is found on your system.
134#
135#   MKL_THREADING
136#     MKL threading mode: SEQ, TBB or OMP (default)
137#
138#   USE_ROCM_KERNEL_ASSERT=1
139#     Enable kernel assert in ROCm platform
140#
141# Environment variables we respect (these environment variables are
142# conventional and are often understood/set by other software.)
143#
144#   CUDA_HOME (Linux/OS X)
145#   CUDA_PATH (Windows)
146#     specify where CUDA is installed; usually /usr/local/cuda or
147#     /usr/local/cuda-x.y
148#   CUDAHOSTCXX
149#     specify a different compiler than the system one to use as the CUDA
150#     host compiler for nvcc.
151#
152#   CUDA_NVCC_EXECUTABLE
153#     Specify a NVCC to use. This is used in our CI to point to a cached nvcc
154#
155#   CUDNN_LIB_DIR
156#   CUDNN_INCLUDE_DIR
157#   CUDNN_LIBRARY
158#     specify where cuDNN is installed
159#
160#   MIOPEN_LIB_DIR
161#   MIOPEN_INCLUDE_DIR
162#   MIOPEN_LIBRARY
163#     specify where MIOpen is installed
164#
165#   NCCL_ROOT
166#   NCCL_LIB_DIR
167#   NCCL_INCLUDE_DIR
168#     specify where nccl is installed
169#
170#   ACL_ROOT_DIR
171#     specify where Compute Library is installed
172#
173#   LIBRARY_PATH
174#   LD_LIBRARY_PATH
175#     we will search for libraries in these paths
176#
177#   ATEN_THREADING
178#     ATen parallel backend to use for intra- and inter-op parallelism
179#     possible values:
180#       OMP - use OpenMP for intra-op and native backend for inter-op tasks
181#       NATIVE - use native thread pool for both intra- and inter-op tasks
182#
183#   USE_SYSTEM_LIBS (work in progress)
184#      Use system-provided libraries to satisfy the build dependencies.
185#      When turned on, the following cmake variables will be toggled as well:
186#        USE_SYSTEM_CPUINFO=ON USE_SYSTEM_SLEEF=ON BUILD_CUSTOM_PROTOBUF=OFF
187#
188#   USE_MIMALLOC
189#      Static link mimalloc into C10, and use mimalloc in alloc_cpu & alloc_free.
190#      By default, It is only enabled on Windows.
191#
192#   USE_PRIORITIZED_TEXT_FOR_LD
193#      Uses prioritized text form cmake/prioritized_text.txt for LD
194#
195#   BUILD_LIBTORCH_WHL
196#      Builds libtorch.so and its dependencies as a wheel
197#
198#   BUILD_PYTHON_ONLY
199#      Builds pytorch as a wheel using libtorch.so from a seperate wheel
200
201import os
202import sys
203
204
205if sys.platform == "win32" and sys.maxsize.bit_length() == 31:
206    print(
207        "32-bit Windows Python runtime is not supported. Please switch to 64-bit Python."
208    )
209    sys.exit(-1)
210
211import platform
212
213
214BUILD_LIBTORCH_WHL = os.getenv("BUILD_LIBTORCH_WHL", "0") == "1"
215BUILD_PYTHON_ONLY = os.getenv("BUILD_PYTHON_ONLY", "0") == "1"
216
217python_min_version = (3, 8, 0)
218python_min_version_str = ".".join(map(str, python_min_version))
219if sys.version_info < python_min_version:
220    print(
221        f"You are using Python {platform.python_version()}. Python >={python_min_version_str} is required."
222    )
223    sys.exit(-1)
224
225import filecmp
226import glob
227import importlib
228import importlib.util
229import json
230import shutil
231import subprocess
232import sysconfig
233import time
234from collections import defaultdict
235
236import setuptools.command.build_ext
237import setuptools.command.install
238import setuptools.command.sdist
239from setuptools import Extension, find_packages, setup
240from setuptools.dist import Distribution
241from tools.build_pytorch_libs import build_caffe2
242from tools.generate_torch_version import get_torch_version
243from tools.setup_helpers.cmake import CMake
244from tools.setup_helpers.env import build_type, IS_DARWIN, IS_LINUX, IS_WINDOWS
245from tools.setup_helpers.generate_linker_script import gen_linker_script
246
247
248def _get_package_path(package_name):
249    spec = importlib.util.find_spec(package_name)
250    if spec:
251        # The package might be a namespace package, so get_data may fail
252        try:
253            loader = spec.loader
254            if loader is not None:
255                file_path = loader.get_filename()  # type: ignore[attr-defined]
256                return os.path.dirname(file_path)
257        except AttributeError:
258            pass
259    return None
260
261
262# set up appropriate env variables
263if BUILD_LIBTORCH_WHL:
264    # Set up environment variables for ONLY building libtorch.so and not libtorch_python.so
265    # functorch is not supported without python
266    os.environ["BUILD_FUNCTORCH"] = "OFF"
267
268
269if BUILD_PYTHON_ONLY:
270    os.environ["BUILD_LIBTORCHLESS"] = "ON"
271    os.environ["LIBTORCH_LIB_PATH"] = f"{_get_package_path('torch')}/lib"
272
273################################################################################
274# Parameters parsed from environment
275################################################################################
276
277VERBOSE_SCRIPT = True
278RUN_BUILD_DEPS = True
279# see if the user passed a quiet flag to setup.py arguments and respect
280# that in our parts of the build
281EMIT_BUILD_WARNING = False
282RERUN_CMAKE = False
283CMAKE_ONLY = False
284filtered_args = []
285for i, arg in enumerate(sys.argv):
286    if arg == "--cmake":
287        RERUN_CMAKE = True
288        continue
289    if arg == "--cmake-only":
290        # Stop once cmake terminates. Leave users a chance to adjust build
291        # options.
292        CMAKE_ONLY = True
293        continue
294    if arg == "rebuild" or arg == "build":
295        arg = "build"  # rebuild is gone, make it build
296        EMIT_BUILD_WARNING = True
297    if arg == "--":
298        filtered_args += sys.argv[i:]
299        break
300    if arg == "-q" or arg == "--quiet":
301        VERBOSE_SCRIPT = False
302    if arg in ["clean", "egg_info", "sdist"]:
303        RUN_BUILD_DEPS = False
304    filtered_args.append(arg)
305sys.argv = filtered_args
306
307if VERBOSE_SCRIPT:
308
309    def report(*args):
310        print(*args)
311
312else:
313
314    def report(*args):
315        pass
316
317    # Make distutils respect --quiet too
318    setuptools.distutils.log.warn = report
319
320# Constant known variables used throughout this file
321cwd = os.path.dirname(os.path.abspath(__file__))
322lib_path = os.path.join(cwd, "torch", "lib")
323third_party_path = os.path.join(cwd, "third_party")
324caffe2_build_dir = os.path.join(cwd, "build")
325
326# CMAKE: full path to python library
327if IS_WINDOWS:
328    cmake_python_library = "{}/libs/python{}.lib".format(
329        sysconfig.get_config_var("prefix"), sysconfig.get_config_var("VERSION")
330    )
331    # Fix virtualenv builds
332    if not os.path.exists(cmake_python_library):
333        cmake_python_library = "{}/libs/python{}.lib".format(
334            sys.base_prefix, sysconfig.get_config_var("VERSION")
335        )
336else:
337    cmake_python_library = "{}/{}".format(
338        sysconfig.get_config_var("LIBDIR"), sysconfig.get_config_var("INSTSONAME")
339    )
340cmake_python_include_dir = sysconfig.get_path("include")
341
342
343################################################################################
344# Version, create_version_file, and package_name
345################################################################################
346
347package_name = os.getenv("TORCH_PACKAGE_NAME", "torch")
348LIBTORCH_PKG_NAME = os.getenv("LIBTORCH_PACKAGE_NAME", "torch_no_python")
349if BUILD_LIBTORCH_WHL:
350    package_name = LIBTORCH_PKG_NAME
351
352
353package_type = os.getenv("PACKAGE_TYPE", "wheel")
354version = get_torch_version()
355report(f"Building wheel {package_name}-{version}")
356
357cmake = CMake()
358
359
360def get_submodule_folders():
361    git_modules_path = os.path.join(cwd, ".gitmodules")
362    default_modules_path = [
363        os.path.join(third_party_path, name)
364        for name in [
365            "gloo",
366            "cpuinfo",
367            "onnx",
368            "fbgemm",
369            "cutlass",
370        ]
371    ]
372    if not os.path.exists(git_modules_path):
373        return default_modules_path
374    with open(git_modules_path) as f:
375        return [
376            os.path.join(cwd, line.split("=", 1)[1].strip())
377            for line in f
378            if line.strip().startswith("path")
379        ]
380
381
382def check_submodules():
383    def check_for_files(folder, files):
384        if not any(os.path.exists(os.path.join(folder, f)) for f in files):
385            report("Could not find any of {} in {}".format(", ".join(files), folder))
386            report("Did you run 'git submodule update --init --recursive'?")
387            sys.exit(1)
388
389    def not_exists_or_empty(folder):
390        return not os.path.exists(folder) or (
391            os.path.isdir(folder) and len(os.listdir(folder)) == 0
392        )
393
394    if bool(os.getenv("USE_SYSTEM_LIBS", False)):
395        return
396    folders = get_submodule_folders()
397    # If none of the submodule folders exists, try to initialize them
398    if all(not_exists_or_empty(folder) for folder in folders):
399        try:
400            print(" --- Trying to initialize submodules")
401            start = time.time()
402            subprocess.check_call(
403                ["git", "submodule", "update", "--init", "--recursive"], cwd=cwd
404            )
405            end = time.time()
406            print(f" --- Submodule initialization took {end - start:.2f} sec")
407        except Exception:
408            print(" --- Submodule initalization failed")
409            print("Please run:\n\tgit submodule update --init --recursive")
410            sys.exit(1)
411    for folder in folders:
412        check_for_files(
413            folder,
414            [
415                "CMakeLists.txt",
416                "Makefile",
417                "setup.py",
418                "LICENSE",
419                "LICENSE.md",
420                "LICENSE.txt",
421            ],
422        )
423    check_for_files(
424        os.path.join(third_party_path, "fbgemm", "third_party", "asmjit"),
425        ["CMakeLists.txt"],
426    )
427    check_for_files(
428        os.path.join(third_party_path, "onnx", "third_party", "benchmark"),
429        ["CMakeLists.txt"],
430    )
431
432
433# Windows has very bad support for symbolic links.
434# Instead of using symlinks, we're going to copy files over
435def mirror_files_into_torchgen():
436    # (new_path, orig_path)
437    # Directories are OK and are recursively mirrored.
438    paths = [
439        (
440            "torchgen/packaged/ATen/native/native_functions.yaml",
441            "aten/src/ATen/native/native_functions.yaml",
442        ),
443        ("torchgen/packaged/ATen/native/tags.yaml", "aten/src/ATen/native/tags.yaml"),
444        ("torchgen/packaged/ATen/templates", "aten/src/ATen/templates"),
445        ("torchgen/packaged/autograd", "tools/autograd"),
446        ("torchgen/packaged/autograd/templates", "tools/autograd/templates"),
447    ]
448    for new_path, orig_path in paths:
449        # Create the dirs involved in new_path if they don't exist
450        if not os.path.exists(new_path):
451            os.makedirs(os.path.dirname(new_path), exist_ok=True)
452
453        # Copy the files from the orig location to the new location
454        if os.path.isfile(orig_path):
455            shutil.copyfile(orig_path, new_path)
456            continue
457        if os.path.isdir(orig_path):
458            if os.path.exists(new_path):
459                # copytree fails if the tree exists already, so remove it.
460                shutil.rmtree(new_path)
461            shutil.copytree(orig_path, new_path)
462            continue
463        raise RuntimeError("Check the file paths in `mirror_files_into_torchgen()`")
464
465
466# all the work we need to do _before_ setup runs
467def build_deps():
468    report("-- Building version " + version)
469
470    check_submodules()
471    check_pydep("yaml", "pyyaml")
472    build_python = not BUILD_LIBTORCH_WHL
473    build_caffe2(
474        version=version,
475        cmake_python_library=cmake_python_library,
476        build_python=build_python,
477        rerun_cmake=RERUN_CMAKE,
478        cmake_only=CMAKE_ONLY,
479        cmake=cmake,
480    )
481
482    if CMAKE_ONLY:
483        report(
484            'Finished running cmake. Run "ccmake build" or '
485            '"cmake-gui build" to adjust build options and '
486            '"python setup.py install" to build.'
487        )
488        sys.exit()
489
490    # Use copies instead of symbolic files.
491    # Windows has very poor support for them.
492    sym_files = [
493        "tools/shared/_utils_internal.py",
494        "torch/utils/benchmark/utils/valgrind_wrapper/callgrind.h",
495        "torch/utils/benchmark/utils/valgrind_wrapper/valgrind.h",
496    ]
497    orig_files = [
498        "torch/_utils_internal.py",
499        "third_party/valgrind-headers/callgrind.h",
500        "third_party/valgrind-headers/valgrind.h",
501    ]
502    for sym_file, orig_file in zip(sym_files, orig_files):
503        same = False
504        if os.path.exists(sym_file):
505            if filecmp.cmp(sym_file, orig_file):
506                same = True
507            else:
508                os.remove(sym_file)
509        if not same:
510            shutil.copyfile(orig_file, sym_file)
511
512
513################################################################################
514# Building dependent libraries
515################################################################################
516
517missing_pydep = """
518Missing build dependency: Unable to `import {importname}`.
519Please install it via `conda install {module}` or `pip install {module}`
520""".strip()
521
522
523def check_pydep(importname, module):
524    try:
525        importlib.import_module(importname)
526    except ImportError as e:
527        raise RuntimeError(
528            missing_pydep.format(importname=importname, module=module)
529        ) from e
530
531
532class build_ext(setuptools.command.build_ext.build_ext):
533    def _embed_libomp(self):
534        # Copy libiomp5.dylib/libomp.dylib inside the wheel package on MacOS
535        lib_dir = os.path.join(self.build_lib, "torch", "lib")
536        libtorch_cpu_path = os.path.join(lib_dir, "libtorch_cpu.dylib")
537        if not os.path.exists(libtorch_cpu_path):
538            return
539        # Parse libtorch_cpu load commands
540        otool_cmds = (
541            subprocess.check_output(["otool", "-l", libtorch_cpu_path])
542            .decode("utf-8")
543            .split("\n")
544        )
545        rpaths, libs = [], []
546        for idx, line in enumerate(otool_cmds):
547            if line.strip() == "cmd LC_LOAD_DYLIB":
548                lib_name = otool_cmds[idx + 2].strip()
549                assert lib_name.startswith("name ")
550                libs.append(lib_name.split(" ", 1)[1].rsplit("(", 1)[0][:-1])
551
552            if line.strip() == "cmd LC_RPATH":
553                rpath = otool_cmds[idx + 2].strip()
554                assert rpath.startswith("path ")
555                rpaths.append(rpath.split(" ", 1)[1].rsplit("(", 1)[0][:-1])
556
557        omp_lib_name = (
558            "libomp.dylib" if os.uname().machine == "arm64" else "libiomp5.dylib"
559        )
560        omp_rpath_lib_path = os.path.join("@rpath", omp_lib_name)
561        if omp_rpath_lib_path not in libs:
562            return
563
564        # Copy libomp/libiomp5 from rpath locations
565        for rpath in rpaths:
566            source_lib = os.path.join(rpath, omp_lib_name)
567            if not os.path.exists(source_lib):
568                continue
569            target_lib = os.path.join(self.build_lib, "torch", "lib", omp_lib_name)
570            self.copy_file(source_lib, target_lib)
571            # Delete old rpath and add @loader_lib to the rpath
572            # This should prevent delocate from attempting to package another instance
573            # of OpenMP library in torch wheel as well as loading two libomp.dylib into
574            # the address space, as libraries are cached by their unresolved names
575            subprocess.check_call(
576                [
577                    "install_name_tool",
578                    "-rpath",
579                    rpath,
580                    "@loader_path",
581                    libtorch_cpu_path,
582                ]
583            )
584            break
585
586        # Copy omp.h from OpenMP_C_FLAGS and copy it into include folder
587        omp_cflags = get_cmake_cache_vars()["OpenMP_C_FLAGS"]
588        if not omp_cflags:
589            return
590        for include_dir in [f[2:] for f in omp_cflags.split(" ") if f.startswith("-I")]:
591            omp_h = os.path.join(include_dir, "omp.h")
592            if not os.path.exists(omp_h):
593                continue
594            target_omp_h = os.path.join(self.build_lib, "torch", "include", "omp.h")
595            self.copy_file(omp_h, target_omp_h)
596            break
597
598    def run(self):
599        # Report build options. This is run after the build completes so # `CMakeCache.txt` exists and we can get an
600        # accurate report on what is used and what is not.
601        cmake_cache_vars = defaultdict(lambda: False, cmake.get_cmake_cache_variables())
602        if cmake_cache_vars["USE_NUMPY"]:
603            report("-- Building with NumPy bindings")
604        else:
605            report("-- NumPy not found")
606        if cmake_cache_vars["USE_CUDNN"]:
607            report(
608                "-- Detected cuDNN at "
609                + cmake_cache_vars["CUDNN_LIBRARY"]
610                + ", "
611                + cmake_cache_vars["CUDNN_INCLUDE_DIR"]
612            )
613        else:
614            report("-- Not using cuDNN")
615        if cmake_cache_vars["USE_CUDA"]:
616            report("-- Detected CUDA at " + cmake_cache_vars["CUDA_TOOLKIT_ROOT_DIR"])
617        else:
618            report("-- Not using CUDA")
619        if cmake_cache_vars["USE_XPU"]:
620            report("-- Detected XPU runtime at " + cmake_cache_vars["SYCL_LIBRARY_DIR"])
621        else:
622            report("-- Not using XPU")
623        if cmake_cache_vars["USE_MKLDNN"]:
624            report("-- Using MKLDNN")
625            if cmake_cache_vars["USE_MKLDNN_ACL"]:
626                report("-- Using Compute Library for the Arm architecture with MKLDNN")
627            else:
628                report(
629                    "-- Not using Compute Library for the Arm architecture with MKLDNN"
630                )
631            if cmake_cache_vars["USE_MKLDNN_CBLAS"]:
632                report("-- Using CBLAS in MKLDNN")
633            else:
634                report("-- Not using CBLAS in MKLDNN")
635        else:
636            report("-- Not using MKLDNN")
637        if cmake_cache_vars["USE_NCCL"] and cmake_cache_vars["USE_SYSTEM_NCCL"]:
638            report(
639                "-- Using system provided NCCL library at {}, {}".format(
640                    cmake_cache_vars["NCCL_LIBRARIES"],
641                    cmake_cache_vars["NCCL_INCLUDE_DIRS"],
642                )
643            )
644        elif cmake_cache_vars["USE_NCCL"]:
645            report("-- Building NCCL library")
646        else:
647            report("-- Not using NCCL")
648        if cmake_cache_vars["USE_DISTRIBUTED"]:
649            if IS_WINDOWS:
650                report("-- Building without distributed package")
651            else:
652                report("-- Building with distributed package: ")
653                report(
654                    "  -- USE_TENSORPIPE={}".format(cmake_cache_vars["USE_TENSORPIPE"])
655                )
656                report("  -- USE_GLOO={}".format(cmake_cache_vars["USE_GLOO"]))
657                report("  -- USE_MPI={}".format(cmake_cache_vars["USE_OPENMPI"]))
658        else:
659            report("-- Building without distributed package")
660        if cmake_cache_vars["STATIC_DISPATCH_BACKEND"]:
661            report(
662                "-- Using static dispatch with backend {}".format(
663                    cmake_cache_vars["STATIC_DISPATCH_BACKEND"]
664                )
665            )
666        if cmake_cache_vars["USE_LIGHTWEIGHT_DISPATCH"]:
667            report("-- Using lightweight dispatch")
668        if cmake_cache_vars["BUILD_EXECUTORCH"]:
669            report("-- Building Executorch")
670
671        if cmake_cache_vars["USE_ITT"]:
672            report("-- Using ITT")
673        else:
674            report("-- Not using ITT")
675
676        # Do not use clang to compile extensions if `-fstack-clash-protection` is defined
677        # in system CFLAGS
678        c_flags = str(os.getenv("CFLAGS", ""))
679        if (
680            IS_LINUX
681            and "-fstack-clash-protection" in c_flags
682            and "clang" in os.environ.get("CC", "")
683        ):
684            os.environ["CC"] = str(os.environ["CC"])
685
686        # It's an old-style class in Python 2.7...
687        setuptools.command.build_ext.build_ext.run(self)
688
689        if IS_DARWIN and package_type != "conda":
690            self._embed_libomp()
691
692        # Copy the essential export library to compile C++ extensions.
693        if IS_WINDOWS:
694            build_temp = self.build_temp
695
696            ext_filename = self.get_ext_filename("_C")
697            lib_filename = ".".join(ext_filename.split(".")[:-1]) + ".lib"
698
699            export_lib = os.path.join(
700                build_temp, "torch", "csrc", lib_filename
701            ).replace("\\", "/")
702
703            build_lib = self.build_lib
704
705            target_lib = os.path.join(build_lib, "torch", "lib", "_C.lib").replace(
706                "\\", "/"
707            )
708
709            # Create "torch/lib" directory if not exists.
710            # (It is not created yet in "develop" mode.)
711            target_dir = os.path.dirname(target_lib)
712            if not os.path.exists(target_dir):
713                os.makedirs(target_dir)
714
715            self.copy_file(export_lib, target_lib)
716
717    def build_extensions(self):
718        self.create_compile_commands()
719        # The caffe2 extensions are created in
720        # tmp_install/lib/pythonM.m/site-packages/caffe2/python/
721        # and need to be copied to build/lib.linux.... , which will be a
722        # platform dependent build folder created by the "build" command of
723        # setuptools. Only the contents of this folder are installed in the
724        # "install" command by default.
725        # We only make this copy for Caffe2's pybind extensions
726        caffe2_pybind_exts = [
727            "caffe2.python.caffe2_pybind11_state",
728            "caffe2.python.caffe2_pybind11_state_gpu",
729            "caffe2.python.caffe2_pybind11_state_hip",
730        ]
731        if BUILD_LIBTORCH_WHL:
732            caffe2_pybind_exts = []
733        i = 0
734        while i < len(self.extensions):
735            ext = self.extensions[i]
736            if ext.name not in caffe2_pybind_exts:
737                i += 1
738                continue
739            fullname = self.get_ext_fullname(ext.name)
740            filename = self.get_ext_filename(fullname)
741            report(f"\nCopying extension {ext.name}")
742
743            relative_site_packages = (
744                sysconfig.get_path("purelib")
745                .replace(sysconfig.get_path("data"), "")
746                .lstrip(os.path.sep)
747            )
748            src = os.path.join("torch", relative_site_packages, filename)
749            if not os.path.exists(src):
750                report(f"{src} does not exist")
751                del self.extensions[i]
752            else:
753                dst = os.path.join(os.path.realpath(self.build_lib), filename)
754                report(f"Copying {ext.name} from {src} to {dst}")
755                dst_dir = os.path.dirname(dst)
756                if not os.path.exists(dst_dir):
757                    os.makedirs(dst_dir)
758                self.copy_file(src, dst)
759                i += 1
760
761        # Copy functorch extension
762        for i, ext in enumerate(self.extensions):
763            if ext.name != "functorch._C":
764                continue
765            fullname = self.get_ext_fullname(ext.name)
766            filename = self.get_ext_filename(fullname)
767            fileext = os.path.splitext(filename)[1]
768            src = os.path.join(os.path.dirname(filename), "functorch" + fileext)
769            dst = os.path.join(os.path.realpath(self.build_lib), filename)
770            if os.path.exists(src):
771                report(f"Copying {ext.name} from {src} to {dst}")
772                dst_dir = os.path.dirname(dst)
773                if not os.path.exists(dst_dir):
774                    os.makedirs(dst_dir)
775                self.copy_file(src, dst)
776
777        setuptools.command.build_ext.build_ext.build_extensions(self)
778
779    def get_outputs(self):
780        outputs = setuptools.command.build_ext.build_ext.get_outputs(self)
781        outputs.append(os.path.join(self.build_lib, "caffe2"))
782        report(f"setup.py::get_outputs returning {outputs}")
783        return outputs
784
785    def create_compile_commands(self):
786        def load(filename):
787            with open(filename) as f:
788                return json.load(f)
789
790        ninja_files = glob.glob("build/*compile_commands.json")
791        cmake_files = glob.glob("torch/lib/build/*/compile_commands.json")
792        all_commands = [entry for f in ninja_files + cmake_files for entry in load(f)]
793
794        # cquery does not like c++ compiles that start with gcc.
795        # It forgets to include the c++ header directories.
796        # We can work around this by replacing the gcc calls that python
797        # setup.py generates with g++ calls instead
798        for command in all_commands:
799            if command["command"].startswith("gcc "):
800                command["command"] = "g++ " + command["command"][4:]
801
802        new_contents = json.dumps(all_commands, indent=2)
803        contents = ""
804        if os.path.exists("compile_commands.json"):
805            with open("compile_commands.json") as f:
806                contents = f.read()
807        if contents != new_contents:
808            with open("compile_commands.json", "w") as f:
809                f.write(new_contents)
810
811
812class concat_license_files:
813    """Merge LICENSE and LICENSES_BUNDLED.txt as a context manager
814
815    LICENSE is the main PyTorch license, LICENSES_BUNDLED.txt is auto-generated
816    from all the licenses found in ./third_party/. We concatenate them so there
817    is a single license file in the sdist and wheels with all of the necessary
818    licensing info.
819    """
820
821    def __init__(self, include_files=False):
822        self.f1 = "LICENSE"
823        self.f2 = "third_party/LICENSES_BUNDLED.txt"
824        self.include_files = include_files
825
826    def __enter__(self):
827        """Concatenate files"""
828
829        old_path = sys.path
830        sys.path.append(third_party_path)
831        try:
832            from build_bundled import create_bundled
833        finally:
834            sys.path = old_path
835
836        with open(self.f1) as f1:
837            self.bsd_text = f1.read()
838
839        with open(self.f1, "a") as f1:
840            f1.write("\n\n")
841            create_bundled(
842                os.path.relpath(third_party_path), f1, include_files=self.include_files
843            )
844
845    def __exit__(self, exception_type, exception_value, traceback):
846        """Restore content of f1"""
847        with open(self.f1, "w") as f:
848            f.write(self.bsd_text)
849
850
851try:
852    from wheel.bdist_wheel import bdist_wheel
853except ImportError:
854    # This is useful when wheel is not installed and bdist_wheel is not
855    # specified on the command line. If it _is_ specified, parsing the command
856    # line will fail before wheel_concatenate is needed
857    wheel_concatenate = None
858else:
859    # Need to create the proper LICENSE.txt for the wheel
860    class wheel_concatenate(bdist_wheel):
861        """check submodules on sdist to prevent incomplete tarballs"""
862
863        def run(self):
864            with concat_license_files(include_files=True):
865                super().run()
866
867        def write_wheelfile(self, *args, **kwargs):
868            super().write_wheelfile(*args, **kwargs)
869
870            if BUILD_LIBTORCH_WHL:
871                # Remove extraneneous files in the libtorch wheel
872                for root, dirs, files in os.walk(self.bdist_dir):
873                    for file in files:
874                        if file.endswith((".a", ".so")) and os.path.isfile(
875                            os.path.join(self.bdist_dir, file)
876                        ):
877                            os.remove(os.path.join(root, file))
878                        elif file.endswith(".py"):
879                            os.remove(os.path.join(root, file))
880                # need an __init__.py file otherwise we wouldn't have a package
881                open(os.path.join(self.bdist_dir, "torch", "__init__.py"), "w").close()
882
883
884class install(setuptools.command.install.install):
885    def run(self):
886        super().run()
887
888
889class clean(setuptools.Command):
890    user_options = []
891
892    def initialize_options(self):
893        pass
894
895    def finalize_options(self):
896        pass
897
898    def run(self):
899        import glob
900        import re
901
902        with open(".gitignore") as f:
903            ignores = f.read()
904            pat = re.compile(r"^#( BEGIN NOT-CLEAN-FILES )?")
905            for wildcard in filter(None, ignores.split("\n")):
906                match = pat.match(wildcard)
907                if match:
908                    if match.group(1):
909                        # Marker is found and stop reading .gitignore.
910                        break
911                    # Ignore lines which begin with '#'.
912                else:
913                    # Don't remove absolute paths from the system
914                    wildcard = wildcard.lstrip("./")
915
916                    for filename in glob.glob(wildcard):
917                        try:
918                            os.remove(filename)
919                        except OSError:
920                            shutil.rmtree(filename, ignore_errors=True)
921
922
923class sdist(setuptools.command.sdist.sdist):
924    def run(self):
925        with concat_license_files():
926            super().run()
927
928
929def get_cmake_cache_vars():
930    try:
931        return defaultdict(lambda: False, cmake.get_cmake_cache_variables())
932    except FileNotFoundError:
933        # CMakeCache.txt does not exist. Probably running "python setup.py clean" over a clean directory.
934        return defaultdict(lambda: False)
935
936
937def configure_extension_build():
938    r"""Configures extension build options according to system environment and user's choice.
939
940    Returns:
941      The input to parameters ext_modules, cmdclass, packages, and entry_points as required in setuptools.setup.
942    """
943
944    cmake_cache_vars = get_cmake_cache_vars()
945
946    ################################################################################
947    # Configure compile flags
948    ################################################################################
949
950    library_dirs = []
951    extra_install_requires = []
952
953    if IS_WINDOWS:
954        # /NODEFAULTLIB makes sure we only link to DLL runtime
955        # and matches the flags set for protobuf and ONNX
956        extra_link_args = ["/NODEFAULTLIB:LIBCMT.LIB"]
957        # /MD links against DLL runtime
958        # and matches the flags set for protobuf and ONNX
959        # /EHsc is about standard C++ exception handling
960        extra_compile_args = ["/MD", "/FS", "/EHsc"]
961    else:
962        extra_link_args = []
963        extra_compile_args = [
964            "-Wall",
965            "-Wextra",
966            "-Wno-strict-overflow",
967            "-Wno-unused-parameter",
968            "-Wno-missing-field-initializers",
969            "-Wno-unknown-pragmas",
970            # Python 2.6 requires -fno-strict-aliasing, see
971            # http://legacy.python.org/dev/peps/pep-3123/
972            # We also depend on it in our code (even Python 3).
973            "-fno-strict-aliasing",
974        ]
975
976    library_dirs.append(lib_path)
977
978    main_compile_args = []
979    main_libraries = ["torch_python"]
980
981    main_link_args = []
982    main_sources = ["torch/csrc/stub.c"]
983
984    if BUILD_LIBTORCH_WHL:
985        main_libraries = ["torch"]
986        main_sources = []
987
988    if build_type.is_debug():
989        if IS_WINDOWS:
990            extra_compile_args.append("/Z7")
991            extra_link_args.append("/DEBUG:FULL")
992        else:
993            extra_compile_args += ["-O0", "-g"]
994            extra_link_args += ["-O0", "-g"]
995
996    if build_type.is_rel_with_deb_info():
997        if IS_WINDOWS:
998            extra_compile_args.append("/Z7")
999            extra_link_args.append("/DEBUG:FULL")
1000        else:
1001            extra_compile_args += ["-g"]
1002            extra_link_args += ["-g"]
1003
1004    # pypi cuda package that requires installation of cuda runtime, cudnn and cublas
1005    # should be included in all wheels uploaded to pypi
1006    pytorch_extra_install_requirements = os.getenv(
1007        "PYTORCH_EXTRA_INSTALL_REQUIREMENTS", ""
1008    )
1009    if pytorch_extra_install_requirements:
1010        report(
1011            f"pytorch_extra_install_requirements: {pytorch_extra_install_requirements}"
1012        )
1013        extra_install_requires += pytorch_extra_install_requirements.split("|")
1014
1015    # Cross-compile for M1
1016    if IS_DARWIN:
1017        macos_target_arch = os.getenv("CMAKE_OSX_ARCHITECTURES", "")
1018        if macos_target_arch in ["arm64", "x86_64"]:
1019            macos_sysroot_path = os.getenv("CMAKE_OSX_SYSROOT")
1020            if macos_sysroot_path is None:
1021                macos_sysroot_path = (
1022                    subprocess.check_output(
1023                        ["xcrun", "--show-sdk-path", "--sdk", "macosx"]
1024                    )
1025                    .decode("utf-8")
1026                    .strip()
1027                )
1028            extra_compile_args += [
1029                "-arch",
1030                macos_target_arch,
1031                "-isysroot",
1032                macos_sysroot_path,
1033            ]
1034            extra_link_args += ["-arch", macos_target_arch]
1035
1036    def make_relative_rpath_args(path):
1037        if IS_DARWIN:
1038            return ["-Wl,-rpath,@loader_path/" + path]
1039        elif IS_WINDOWS:
1040            return []
1041        else:
1042            return ["-Wl,-rpath,$ORIGIN/" + path]
1043
1044    ################################################################################
1045    # Declare extensions and package
1046    ################################################################################
1047
1048    extensions = []
1049    excludes = ["tools", "tools.*"]
1050    if not cmake_cache_vars["BUILD_CAFFE2"]:
1051        excludes.extend(["caffe2", "caffe2.*"])
1052    if not cmake_cache_vars["BUILD_FUNCTORCH"]:
1053        excludes.extend(["functorch", "functorch.*"])
1054    packages = find_packages(exclude=excludes)
1055    C = Extension(
1056        "torch._C",
1057        libraries=main_libraries,
1058        sources=main_sources,
1059        language="c",
1060        extra_compile_args=main_compile_args + extra_compile_args,
1061        include_dirs=[],
1062        library_dirs=library_dirs,
1063        extra_link_args=extra_link_args
1064        + main_link_args
1065        + make_relative_rpath_args("lib"),
1066    )
1067    extensions.append(C)
1068
1069    # These extensions are built by cmake and copied manually in build_extensions()
1070    # inside the build_ext implementation
1071    if cmake_cache_vars["BUILD_CAFFE2"]:
1072        extensions.append(
1073            Extension(name="caffe2.python.caffe2_pybind11_state", sources=[]),
1074        )
1075        if cmake_cache_vars["USE_CUDA"]:
1076            extensions.append(
1077                Extension(name="caffe2.python.caffe2_pybind11_state_gpu", sources=[]),
1078            )
1079        if cmake_cache_vars["USE_ROCM"]:
1080            extensions.append(
1081                Extension(name="caffe2.python.caffe2_pybind11_state_hip", sources=[]),
1082            )
1083    if cmake_cache_vars["BUILD_FUNCTORCH"]:
1084        extensions.append(
1085            Extension(name="functorch._C", sources=[]),
1086        )
1087
1088    cmdclass = {
1089        "bdist_wheel": wheel_concatenate,
1090        "build_ext": build_ext,
1091        "clean": clean,
1092        "install": install,
1093        "sdist": sdist,
1094    }
1095
1096    entry_points = {
1097        "console_scripts": [
1098            "convert-caffe2-to-onnx = caffe2.python.onnx.bin.conversion:caffe2_to_onnx",
1099            "convert-onnx-to-caffe2 = caffe2.python.onnx.bin.conversion:onnx_to_caffe2",
1100            "torchrun = torch.distributed.run:main",
1101        ],
1102        "torchrun.logs_specs": [
1103            "default = torch.distributed.elastic.multiprocessing:DefaultLogsSpecs",
1104        ],
1105    }
1106
1107    if cmake_cache_vars["USE_DISTRIBUTED"]:
1108        # Only enable fr_trace command if distributed is enabled
1109        entry_points["console_scripts"].append(
1110            "torchfrtrace = tools.flight_recorder.fr_trace:main",
1111        )
1112    return extensions, cmdclass, packages, entry_points, extra_install_requires
1113
1114
1115# post run, warnings, printed at the end to make them more visible
1116build_update_message = """
1117    It is no longer necessary to use the 'build' or 'rebuild' targets
1118
1119    To install:
1120      $ python setup.py install
1121    To develop locally:
1122      $ python setup.py develop
1123    To force cmake to re-generate native build files (off by default):
1124      $ python setup.py develop --cmake
1125"""
1126
1127
1128def print_box(msg):
1129    lines = msg.split("\n")
1130    size = max(len(l) + 1 for l in lines)
1131    print("-" * (size + 2))
1132    for l in lines:
1133        print("|{}{}|".format(l, " " * (size - len(l))))
1134    print("-" * (size + 2))
1135
1136
1137def main():
1138    if BUILD_LIBTORCH_WHL and BUILD_PYTHON_ONLY:
1139        raise RuntimeError(
1140            "Conflict: 'BUILD_LIBTORCH_WHL' and 'BUILD_PYTHON_ONLY' can't both be 1. Set one to 0 and rerun."
1141        )
1142    install_requires = [
1143        "filelock",
1144        "typing-extensions>=4.8.0",
1145        'setuptools ; python_version >= "3.12"',
1146        'sympy==1.12.1 ; python_version == "3.8"',
1147        'sympy==1.13.1 ; python_version >= "3.9"',
1148        "networkx",
1149        "jinja2",
1150        "fsspec",
1151    ]
1152
1153    if BUILD_PYTHON_ONLY:
1154        install_requires.append(f"{LIBTORCH_PKG_NAME}=={get_torch_version()}")
1155
1156    use_prioritized_text = str(os.getenv("USE_PRIORITIZED_TEXT_FOR_LD", ""))
1157    if (
1158        use_prioritized_text == ""
1159        and platform.system() == "Linux"
1160        and platform.processor() == "aarch64"
1161    ):
1162        print_box(
1163            """
1164            WARNING: we strongly recommend enabling linker script optimization for ARM + CUDA.
1165            To do so please export USE_PRIORITIZED_TEXT_FOR_LD=1
1166            """
1167        )
1168    if use_prioritized_text == "1" or use_prioritized_text == "True":
1169        gen_linker_script(
1170            filein="cmake/prioritized_text.txt", fout="cmake/linker_script.ld"
1171        )
1172        linker_script_path = os.path.abspath("cmake/linker_script.ld")
1173        os.environ["LDFLAGS"] = os.getenv("LDFLAGS", "") + f" -T{linker_script_path}"
1174        os.environ["CFLAGS"] = (
1175            os.getenv("CFLAGS", "") + " -ffunction-sections -fdata-sections"
1176        )
1177        os.environ["CXXFLAGS"] = (
1178            os.getenv("CXXFLAGS", "") + " -ffunction-sections -fdata-sections"
1179        )
1180
1181    # Parse the command line and check the arguments before we proceed with
1182    # building deps and setup. We need to set values so `--help` works.
1183    dist = Distribution()
1184    dist.script_name = os.path.basename(sys.argv[0])
1185    dist.script_args = sys.argv[1:]
1186    try:
1187        dist.parse_command_line()
1188    except setuptools.distutils.errors.DistutilsArgError as e:
1189        print(e)
1190        sys.exit(1)
1191
1192    mirror_files_into_torchgen()
1193    if RUN_BUILD_DEPS:
1194        build_deps()
1195
1196    (
1197        extensions,
1198        cmdclass,
1199        packages,
1200        entry_points,
1201        extra_install_requires,
1202    ) = configure_extension_build()
1203    install_requires += extra_install_requires
1204
1205    extras_require = {
1206        "optree": ["optree>=0.12.0"],
1207        "opt-einsum": ["opt-einsum>=3.3"],
1208    }
1209
1210    # Read in README.md for our long_description
1211    with open(os.path.join(cwd, "README.md"), encoding="utf-8") as f:
1212        long_description = f.read()
1213
1214    version_range_max = max(sys.version_info[1], 12) + 1
1215    torch_package_data = [
1216        "py.typed",
1217        "bin/*",
1218        "test/*",
1219        "*.pyi",
1220        "_C/*.pyi",
1221        "cuda/*.pyi",
1222        "fx/*.pyi",
1223        "optim/*.pyi",
1224        "autograd/*.pyi",
1225        "jit/*.pyi",
1226        "nn/*.pyi",
1227        "nn/modules/*.pyi",
1228        "nn/parallel/*.pyi",
1229        "utils/data/*.pyi",
1230        "utils/data/datapipes/*.pyi",
1231        "lib/*.pdb",
1232        "lib/*shm*",
1233        "lib/torch_shm_manager",
1234        "lib/*.h",
1235        "include/*.h",
1236        "include/ATen/*.h",
1237        "include/ATen/cpu/*.h",
1238        "include/ATen/cpu/vec/vec256/*.h",
1239        "include/ATen/cpu/vec/vec256/vsx/*.h",
1240        "include/ATen/cpu/vec/vec256/zarch/*.h",
1241        "include/ATen/cpu/vec/vec512/*.h",
1242        "include/ATen/cpu/vec/*.h",
1243        "include/ATen/core/*.h",
1244        "include/ATen/cuda/*.cuh",
1245        "include/ATen/cuda/*.h",
1246        "include/ATen/cuda/detail/*.cuh",
1247        "include/ATen/cuda/detail/*.h",
1248        "include/ATen/cuda/tunable/*.h",
1249        "include/ATen/cudnn/*.h",
1250        "include/ATen/functorch/*.h",
1251        "include/ATen/ops/*.h",
1252        "include/ATen/hip/*.cuh",
1253        "include/ATen/hip/*.h",
1254        "include/ATen/hip/detail/*.cuh",
1255        "include/ATen/hip/detail/*.h",
1256        "include/ATen/hip/impl/*.h",
1257        "include/ATen/hip/tunable/*.h",
1258        "include/ATen/mps/*.h",
1259        "include/ATen/miopen/*.h",
1260        "include/ATen/detail/*.h",
1261        "include/ATen/native/*.h",
1262        "include/ATen/native/cpu/*.h",
1263        "include/ATen/native/cuda/*.h",
1264        "include/ATen/native/cuda/*.cuh",
1265        "include/ATen/native/hip/*.h",
1266        "include/ATen/native/hip/*.cuh",
1267        "include/ATen/native/mps/*.h",
1268        "include/ATen/native/nested/*.h",
1269        "include/ATen/native/quantized/*.h",
1270        "include/ATen/native/quantized/cpu/*.h",
1271        "include/ATen/native/transformers/*.h",
1272        "include/ATen/native/sparse/*.h",
1273        "include/ATen/native/utils/*.h",
1274        "include/ATen/quantized/*.h",
1275        "include/ATen/xpu/*.h",
1276        "include/ATen/xpu/detail/*.h",
1277        "include/caffe2/serialize/*.h",
1278        "include/c10/*.h",
1279        "include/c10/macros/*.h",
1280        "include/c10/core/*.h",
1281        "include/ATen/core/boxing/*.h",
1282        "include/ATen/core/boxing/impl/*.h",
1283        "include/ATen/core/dispatch/*.h",
1284        "include/ATen/core/op_registration/*.h",
1285        "include/c10/core/impl/*.h",
1286        "include/c10/util/*.h",
1287        "include/c10/cuda/*.h",
1288        "include/c10/cuda/impl/*.h",
1289        "include/c10/hip/*.h",
1290        "include/c10/hip/impl/*.h",
1291        "include/c10/xpu/*.h",
1292        "include/c10/xpu/impl/*.h",
1293        "include/torch/*.h",
1294        "include/torch/csrc/*.h",
1295        "include/torch/csrc/api/include/torch/*.h",
1296        "include/torch/csrc/api/include/torch/data/*.h",
1297        "include/torch/csrc/api/include/torch/data/dataloader/*.h",
1298        "include/torch/csrc/api/include/torch/data/datasets/*.h",
1299        "include/torch/csrc/api/include/torch/data/detail/*.h",
1300        "include/torch/csrc/api/include/torch/data/samplers/*.h",
1301        "include/torch/csrc/api/include/torch/data/transforms/*.h",
1302        "include/torch/csrc/api/include/torch/detail/*.h",
1303        "include/torch/csrc/api/include/torch/detail/ordered_dict.h",
1304        "include/torch/csrc/api/include/torch/nn/*.h",
1305        "include/torch/csrc/api/include/torch/nn/functional/*.h",
1306        "include/torch/csrc/api/include/torch/nn/options/*.h",
1307        "include/torch/csrc/api/include/torch/nn/modules/*.h",
1308        "include/torch/csrc/api/include/torch/nn/modules/container/*.h",
1309        "include/torch/csrc/api/include/torch/nn/parallel/*.h",
1310        "include/torch/csrc/api/include/torch/nn/utils/*.h",
1311        "include/torch/csrc/api/include/torch/optim/*.h",
1312        "include/torch/csrc/api/include/torch/optim/schedulers/*.h",
1313        "include/torch/csrc/api/include/torch/serialize/*.h",
1314        "include/torch/csrc/autograd/*.h",
1315        "include/torch/csrc/autograd/functions/*.h",
1316        "include/torch/csrc/autograd/generated/*.h",
1317        "include/torch/csrc/autograd/utils/*.h",
1318        "include/torch/csrc/cuda/*.h",
1319        "include/torch/csrc/distributed/c10d/*.h",
1320        "include/torch/csrc/distributed/c10d/*.hpp",
1321        "include/torch/csrc/distributed/rpc/*.h",
1322        "include/torch/csrc/distributed/autograd/context/*.h",
1323        "include/torch/csrc/distributed/autograd/functions/*.h",
1324        "include/torch/csrc/distributed/autograd/rpc_messages/*.h",
1325        "include/torch/csrc/dynamo/*.h",
1326        "include/torch/csrc/inductor/*.h",
1327        "include/torch/csrc/inductor/aoti_runner/*.h",
1328        "include/torch/csrc/inductor/aoti_runtime/*.h",
1329        "include/torch/csrc/inductor/aoti_torch/*.h",
1330        "include/torch/csrc/inductor/aoti_torch/c/*.h",
1331        "include/torch/csrc/inductor/aoti_torch/generated/*.h",
1332        "include/torch/csrc/jit/*.h",
1333        "include/torch/csrc/jit/backends/*.h",
1334        "include/torch/csrc/jit/generated/*.h",
1335        "include/torch/csrc/jit/passes/*.h",
1336        "include/torch/csrc/jit/passes/quantization/*.h",
1337        "include/torch/csrc/jit/passes/utils/*.h",
1338        "include/torch/csrc/jit/runtime/*.h",
1339        "include/torch/csrc/jit/ir/*.h",
1340        "include/torch/csrc/jit/frontend/*.h",
1341        "include/torch/csrc/jit/api/*.h",
1342        "include/torch/csrc/jit/serialization/*.h",
1343        "include/torch/csrc/jit/python/*.h",
1344        "include/torch/csrc/jit/mobile/*.h",
1345        "include/torch/csrc/jit/testing/*.h",
1346        "include/torch/csrc/jit/tensorexpr/*.h",
1347        "include/torch/csrc/jit/tensorexpr/operators/*.h",
1348        "include/torch/csrc/jit/codegen/cuda/*.h",
1349        "include/torch/csrc/onnx/*.h",
1350        "include/torch/csrc/profiler/*.h",
1351        "include/torch/csrc/profiler/orchestration/*.h",
1352        "include/torch/csrc/profiler/standalone/*.h",
1353        "include/torch/csrc/profiler/stubs/*.h",
1354        "include/torch/csrc/profiler/unwind/*.h",
1355        "include/torch/csrc/profiler/python/*.h",
1356        "include/torch/csrc/utils/*.h",
1357        "include/torch/csrc/tensor/*.h",
1358        "include/torch/csrc/lazy/backend/*.h",
1359        "include/torch/csrc/lazy/core/*.h",
1360        "include/torch/csrc/lazy/core/internal_ops/*.h",
1361        "include/torch/csrc/lazy/core/ops/*.h",
1362        "include/torch/csrc/lazy/python/python_util.h",
1363        "include/torch/csrc/lazy/ts_backend/*.h",
1364        "include/torch/csrc/xpu/*.h",
1365        "include/pybind11/*.h",
1366        "include/pybind11/detail/*.h",
1367        "include/pybind11/eigen/*.h",
1368        "include/TH/*.h*",
1369        "include/TH/generic/*.h*",
1370        "include/THC/*.cuh",
1371        "include/THC/*.h*",
1372        "include/THC/generic/*.h",
1373        "include/THH/*.cuh",
1374        "include/THH/*.h*",
1375        "include/THH/generic/*.h",
1376        "include/sleef.h",
1377        "_inductor/codegen/*.h",
1378        "_inductor/codegen/aoti_runtime/*.cpp",
1379        "_export/serde/*.yaml",
1380        "share/cmake/ATen/*.cmake",
1381        "share/cmake/Caffe2/*.cmake",
1382        "share/cmake/Caffe2/public/*.cmake",
1383        "share/cmake/Caffe2/Modules_CUDA_fix/*.cmake",
1384        "share/cmake/Caffe2/Modules_CUDA_fix/upstream/*.cmake",
1385        "share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/*.cmake",
1386        "share/cmake/Gloo/*.cmake",
1387        "share/cmake/Tensorpipe/*.cmake",
1388        "share/cmake/Torch/*.cmake",
1389        "utils/benchmark/utils/*.cpp",
1390        "utils/benchmark/utils/valgrind_wrapper/*.cpp",
1391        "utils/benchmark/utils/valgrind_wrapper/*.h",
1392        "utils/model_dump/skeleton.html",
1393        "utils/model_dump/code.js",
1394        "utils/model_dump/*.mjs",
1395    ]
1396
1397    if not BUILD_LIBTORCH_WHL:
1398        torch_package_data.extend(
1399            [
1400                "lib/libtorch_python.so",
1401                "lib/libtorch_python.dylib",
1402                "lib/libtorch_python.dll",
1403            ]
1404        )
1405    if not BUILD_PYTHON_ONLY:
1406        torch_package_data.extend(
1407            [
1408                "lib/*.so*",
1409                "lib/*.dylib*",
1410                "lib/*.dll",
1411                "lib/*.lib",
1412            ]
1413        )
1414    if get_cmake_cache_vars()["BUILD_CAFFE2"]:
1415        torch_package_data.extend(
1416            [
1417                "include/caffe2/**/*.h",
1418                "include/caffe2/utils/*.h",
1419                "include/caffe2/utils/**/*.h",
1420            ]
1421        )
1422    if get_cmake_cache_vars()["USE_TENSORPIPE"]:
1423        torch_package_data.extend(
1424            [
1425                "include/tensorpipe/*.h",
1426                "include/tensorpipe/channel/*.h",
1427                "include/tensorpipe/channel/basic/*.h",
1428                "include/tensorpipe/channel/cma/*.h",
1429                "include/tensorpipe/channel/mpt/*.h",
1430                "include/tensorpipe/channel/xth/*.h",
1431                "include/tensorpipe/common/*.h",
1432                "include/tensorpipe/core/*.h",
1433                "include/tensorpipe/transport/*.h",
1434                "include/tensorpipe/transport/ibv/*.h",
1435                "include/tensorpipe/transport/shm/*.h",
1436                "include/tensorpipe/transport/uv/*.h",
1437            ]
1438        )
1439    if get_cmake_cache_vars()["USE_KINETO"]:
1440        torch_package_data.extend(
1441            [
1442                "include/kineto/*.h",
1443            ]
1444        )
1445    torchgen_package_data = [
1446        # Recursive glob doesn't work in setup.py,
1447        # https://github.com/pypa/setuptools/issues/1806
1448        # To make this robust we should replace it with some code that
1449        # returns a list of everything under packaged/
1450        "packaged/ATen/*",
1451        "packaged/ATen/native/*",
1452        "packaged/ATen/templates/*",
1453        "packaged/autograd/*",
1454        "packaged/autograd/templates/*",
1455    ]
1456    package_data = {
1457        "torch": torch_package_data,
1458    }
1459
1460    if not BUILD_LIBTORCH_WHL:
1461        package_data["torchgen"] = torchgen_package_data
1462        package_data["caffe2"] = [
1463            "python/serialized_test/data/operator_test/*.zip",
1464        ]
1465    else:
1466        # no extensions in BUILD_LIBTORCH_WHL mode
1467        extensions = []
1468
1469    setup(
1470        name=package_name,
1471        version=version,
1472        description=(
1473            "Tensors and Dynamic neural networks in "
1474            "Python with strong GPU acceleration"
1475        ),
1476        long_description=long_description,
1477        long_description_content_type="text/markdown",
1478        ext_modules=extensions,
1479        cmdclass=cmdclass,
1480        packages=packages,
1481        entry_points=entry_points,
1482        install_requires=install_requires,
1483        extras_require=extras_require,
1484        package_data=package_data,
1485        url="https://pytorch.org/",
1486        download_url="https://github.com/pytorch/pytorch/tags",
1487        author="PyTorch Team",
1488        author_email="[email protected]",
1489        python_requires=f">={python_min_version_str}",
1490        # PyPI package information.
1491        classifiers=[
1492            "Development Status :: 5 - Production/Stable",
1493            "Intended Audience :: Developers",
1494            "Intended Audience :: Education",
1495            "Intended Audience :: Science/Research",
1496            "License :: OSI Approved :: BSD License",
1497            "Topic :: Scientific/Engineering",
1498            "Topic :: Scientific/Engineering :: Mathematics",
1499            "Topic :: Scientific/Engineering :: Artificial Intelligence",
1500            "Topic :: Software Development",
1501            "Topic :: Software Development :: Libraries",
1502            "Topic :: Software Development :: Libraries :: Python Modules",
1503            "Programming Language :: C++",
1504            "Programming Language :: Python :: 3",
1505        ]
1506        + [
1507            f"Programming Language :: Python :: 3.{i}"
1508            for i in range(python_min_version[1], version_range_max)
1509        ],
1510        license="BSD-3-Clause",
1511        keywords="pytorch, machine learning",
1512    )
1513    if EMIT_BUILD_WARNING:
1514        print_box(build_update_message)
1515
1516
1517if __name__ == "__main__":
1518    main()
1519