pytorch

1
# Welcome to the PyTorch setup.py.
2
#
3
# Environment variables you are probably interested in:
4
#
5
#   DEBUG
6
#     build with -O0 and -g (debug symbols)
7
#
8
#   REL_WITH_DEB_INFO
9
#     build with optimizations and -g (debug symbols)
10
#
11
#   USE_CUSTOM_DEBINFO="path/to/file1.cpp;path/to/file2.cpp"
12
#     build with debug info only for specified files
13
#
14
#   MAX_JOBS
15
#     maximum number of compile jobs we should use to compile your code
16
#
17
#   USE_CUDA=0
18
#     disables CUDA build
19
#
20
#   CFLAGS
21
#     flags to apply to both C and C++ files to be compiled (a quirk of setup.py
22
#     which we have faithfully adhered to in our build system is that CFLAGS
23
#     also applies to C++ files (unless CXXFLAGS is set), in contrast to the
24
#     default behavior of autogoo and cmake build systems.)
25
#
26
#   CC
27
#     the C/C++ compiler to use
28
#
29
# Environment variables for feature toggles:
30
#
31
#   DEBUG_CUDA=1
32
#     if used in conjunction with DEBUG or REL_WITH_DEB_INFO, will also
33
#     build CUDA kernels with -lineinfo --source-in-ptx.  Note that
34
#     on CUDA 12 this may cause nvcc to OOM, so this is disabled by default.
35

36
#   USE_CUDNN=0
37
#     disables the cuDNN build
38
#
39
#   USE_CUSPARSELT=0
40
#     disables the cuSPARSELt build
41
#
42
#   USE_FBGEMM=0
43
#     disables the FBGEMM build
44
#
45
#   USE_KINETO=0
46
#     disables usage of libkineto library for profiling
47
#
48
#   USE_NUMPY=0
49
#     disables the NumPy build
50
#
51
#   BUILD_TEST=0
52
#     disables the test build
53
#
54
#   USE_MKLDNN=0
55
#     disables use of MKLDNN
56
#
57
#   USE_MKLDNN_ACL
58
#     enables use of Compute Library backend for MKLDNN on Arm;
59
#     USE_MKLDNN must be explicitly enabled.
60
#
61
#   MKLDNN_CPU_RUNTIME
62
#     MKL-DNN threading mode: TBB or OMP (default)
63
#
64
#   USE_STATIC_MKL
65
#     Prefer to link with MKL statically - Unix only
66
#   USE_ITT=0
67
#     disable use of Intel(R) VTune Profiler's ITT functionality
68
#
69
#   USE_NNPACK=0
70
#     disables NNPACK build
71
#
72
#   USE_QNNPACK=0
73
#     disables QNNPACK build (quantized 8-bit operators)
74
#
75
#   USE_DISTRIBUTED=0
76
#     disables distributed (c10d, gloo, mpi, etc.) build
77
#
78
#   USE_TENSORPIPE=0
79
#     disables distributed Tensorpipe backend build
80
#
81
#   USE_GLOO=0
82
#     disables distributed gloo backend build
83
#
84
#   USE_MPI=0
85
#     disables distributed MPI backend build
86
#
87
#   USE_SYSTEM_NCCL=0
88
#     disables use of system-wide nccl (we will use our submoduled
89
#     copy in third_party/nccl)
90
#
91
#   BUILD_CAFFE2_OPS=0
92
#     disable Caffe2 operators build
93
#
94
#   BUILD_CAFFE2=0
95
#     disable Caffe2 build
96
#
97
#   USE_IBVERBS
98
#     toggle features related to distributed support
99
#
100
#   USE_OPENCV
101
#     enables use of OpenCV for additional operators
102
#
103
#   USE_OPENMP=0
104
#     disables use of OpenMP for parallelization
105
#
106
#   USE_FFMPEG
107
#     enables use of ffmpeg for additional operators
108
#
109
#   USE_FLASH_ATTENTION=0
110
#     disables building flash attention for scaled dot product attention
111
#
112
#   USE_MEM_EFF_ATTENTION=0
113
#    disables building memory efficient attention for scaled dot product attention
114
#
115
#   USE_LEVELDB
116
#     enables use of LevelDB for storage
117
#
118
#   USE_LMDB
119
#     enables use of LMDB for storage
120
#
121
#   BUILD_BINARY
122
#     enables the additional binaries/ build
123
#
124
#   ATEN_AVX512_256=TRUE
125
#     ATen AVX2 kernels can use 32 ymm registers, instead of the default 16.
126
#     This option can be used if AVX512 doesn't perform well on a machine.
127
#     The FBGEMM library also uses AVX512_256 kernels on Xeon D processors,
128
#     but it also has some (optimized) assembly code.
129
#
130
#   PYTORCH_BUILD_VERSION
131
#   PYTORCH_BUILD_NUMBER
132
#     specify the version of PyTorch, rather than the hard-coded version
133
#     in this file; used when we're building binaries for distribution
134
#
135
#   TORCH_CUDA_ARCH_LIST
136
#     specify which CUDA architectures to build for.
137
#     ie `TORCH_CUDA_ARCH_LIST="6.0;7.0"`
138
#     These are not CUDA versions, instead, they specify what
139
#     classes of NVIDIA hardware we should generate PTX for.
140
#
141
#   PYTORCH_ROCM_ARCH
142
#     specify which AMD GPU targets to build for.
143
#     ie `PYTORCH_ROCM_ARCH="gfx900;gfx906"`
144
#
145
#   ONNX_NAMESPACE
146
#     specify a namespace for ONNX built here rather than the hard-coded
147
#     one in this file; needed to build with other frameworks that share ONNX.
148
#
149
#   BLAS
150
#     BLAS to be used by Caffe2. Can be MKL, Eigen, ATLAS, FlexiBLAS, or OpenBLAS. If set
151
#     then the build will fail if the requested BLAS is not found, otherwise
152
#     the BLAS will be chosen based on what is found on your system.
153
#
154
#   MKL_THREADING
155
#     MKL threading mode: SEQ, TBB or OMP (default)
156
#
157
#   USE_REDIS
158
#     Whether to use Redis for distributed workflows (Linux only)
159
#
160
#   USE_ZSTD
161
#     Enables use of ZSTD, if the libraries are found
162
#
163
#   USE_ROCM_KERNEL_ASSERT=1
164
#     Enable kernel assert in ROCm platform
165
#
166
# Environment variables we respect (these environment variables are
167
# conventional and are often understood/set by other software.)
168
#
169
#   CUDA_HOME (Linux/OS X)
170
#   CUDA_PATH (Windows)
171
#     specify where CUDA is installed; usually /usr/local/cuda or
172
#     /usr/local/cuda-x.y
173
#   CUDAHOSTCXX
174
#     specify a different compiler than the system one to use as the CUDA
175
#     host compiler for nvcc.
176
#
177
#   CUDA_NVCC_EXECUTABLE
178
#     Specify a NVCC to use. This is used in our CI to point to a cached nvcc
179
#
180
#   CUDNN_LIB_DIR
181
#   CUDNN_INCLUDE_DIR
182
#   CUDNN_LIBRARY
183
#     specify where cuDNN is installed
184
#
185
#   MIOPEN_LIB_DIR
186
#   MIOPEN_INCLUDE_DIR
187
#   MIOPEN_LIBRARY
188
#     specify where MIOpen is installed
189
#
190
#   NCCL_ROOT
191
#   NCCL_LIB_DIR
192
#   NCCL_INCLUDE_DIR
193
#     specify where nccl is installed
194
#
195
#   NVTOOLSEXT_PATH (Windows only)
196
#     specify where nvtoolsext is installed
197
#
198
#   ACL_ROOT_DIR
199
#     specify where Compute Library is installed
200
#
201
#   LIBRARY_PATH
202
#   LD_LIBRARY_PATH
203
#     we will search for libraries in these paths
204
#
205
#   ATEN_THREADING
206
#     ATen parallel backend to use for intra- and inter-op parallelism
207
#     possible values:
208
#       OMP - use OpenMP for intra-op and native backend for inter-op tasks
209
#       NATIVE - use native thread pool for both intra- and inter-op tasks
210
#       TBB - using TBB for intra- and native thread pool for inter-op parallelism
211
#
212
#   USE_TBB
213
#      enable TBB support
214
#
215
#   USE_SYSTEM_TBB
216
#      Use system-provided Intel TBB.
217
#
218
#   USE_SYSTEM_LIBS (work in progress)
219
#      Use system-provided libraries to satisfy the build dependencies.
220
#      When turned on, the following cmake variables will be toggled as well:
221
#        USE_SYSTEM_CPUINFO=ON USE_SYSTEM_SLEEF=ON BUILD_CUSTOM_PROTOBUF=OFF
222
#
223
#   USE_MIMALLOC
224
#      Static link mimalloc into C10, and use mimalloc in alloc_cpu & alloc_free.
225
#      By default, It is only enabled on Windows.
226

227
import sys
228

229
if sys.platform == "win32" and sys.maxsize.bit_length() == 31:
230
    print(
231
        "32-bit Windows Python runtime is not supported. Please switch to 64-bit Python."
232
    )
233
    sys.exit(-1)
234

235
import platform
236

237
python_min_version = (3, 8, 0)
238
python_min_version_str = ".".join(map(str, python_min_version))
239
if sys.version_info < python_min_version:
240
    print(
241
        f"You are using Python {platform.python_version()}. Python >={python_min_version_str} is required."
242
    )
243
    sys.exit(-1)
244

245
import filecmp
246
import glob
247
import importlib
248
import json
249
import os
250
import shutil
251
import subprocess
252
import sysconfig
253
import time
254
from collections import defaultdict
255

256
import setuptools.command.build_ext
257
import setuptools.command.install
258
import setuptools.command.sdist
259
from setuptools import Extension, find_packages, setup
260
from setuptools.dist import Distribution
261

262
from tools.build_pytorch_libs import build_caffe2
263
from tools.generate_torch_version import get_torch_version
264
from tools.setup_helpers.cmake import CMake
265
from tools.setup_helpers.env import build_type, IS_DARWIN, IS_LINUX, IS_WINDOWS
266

267
################################################################################
268
# Parameters parsed from environment
269
################################################################################
270

271
VERBOSE_SCRIPT = True
272
RUN_BUILD_DEPS = True
273
# see if the user passed a quiet flag to setup.py arguments and respect
274
# that in our parts of the build
275
EMIT_BUILD_WARNING = False
276
RERUN_CMAKE = False
277
CMAKE_ONLY = False
278
filtered_args = []
279
for i, arg in enumerate(sys.argv):
280
    if arg == "--cmake":
281
        RERUN_CMAKE = True
282
        continue
283
    if arg == "--cmake-only":
284
        # Stop once cmake terminates. Leave users a chance to adjust build
285
        # options.
286
        CMAKE_ONLY = True
287
        continue
288
    if arg == "rebuild" or arg == "build":
289
        arg = "build"  # rebuild is gone, make it build
290
        EMIT_BUILD_WARNING = True
291
    if arg == "--":
292
        filtered_args += sys.argv[i:]
293
        break
294
    if arg == "-q" or arg == "--quiet":
295
        VERBOSE_SCRIPT = False
296
    if arg in ["clean", "egg_info", "sdist"]:
297
        RUN_BUILD_DEPS = False
298
    filtered_args.append(arg)
299
sys.argv = filtered_args
300

301
if VERBOSE_SCRIPT:
302

303
    def report(*args):
304
        print(*args)
305

306
else:
307

308
    def report(*args):
309
        pass
310

311
    # Make distutils respect --quiet too
312
    setuptools.distutils.log.warn = report
313

314
# Constant known variables used throughout this file
315
cwd = os.path.dirname(os.path.abspath(__file__))
316
lib_path = os.path.join(cwd, "torch", "lib")
317
third_party_path = os.path.join(cwd, "third_party")
318
caffe2_build_dir = os.path.join(cwd, "build")
319

320
# CMAKE: full path to python library
321
if IS_WINDOWS:
322
    cmake_python_library = "{}/libs/python{}.lib".format(
323
        sysconfig.get_config_var("prefix"), sysconfig.get_config_var("VERSION")
324
    )
325
    # Fix virtualenv builds
326
    if not os.path.exists(cmake_python_library):
327
        cmake_python_library = "{}/libs/python{}.lib".format(
328
            sys.base_prefix, sysconfig.get_config_var("VERSION")
329
        )
330
else:
331
    cmake_python_library = "{}/{}".format(
332
        sysconfig.get_config_var("LIBDIR"), sysconfig.get_config_var("INSTSONAME")
333
    )
334
cmake_python_include_dir = sysconfig.get_path("include")
335

336

337
################################################################################
338
# Version, create_version_file, and package_name
339
################################################################################
340
package_name = os.getenv("TORCH_PACKAGE_NAME", "torch")
341
package_type = os.getenv("PACKAGE_TYPE", "wheel")
342
version = get_torch_version()
343
report(f"Building wheel {package_name}-{version}")
344

345
cmake = CMake()
346

347

348
def get_submodule_folders():
349
    git_modules_path = os.path.join(cwd, ".gitmodules")
350
    default_modules_path = [
351
        os.path.join(third_party_path, name)
352
        for name in [
353
            "gloo",
354
            "cpuinfo",
355
            "tbb",
356
            "onnx",
357
            "foxi",
358
            "QNNPACK",
359
            "fbgemm",
360
            "cutlass",
361
        ]
362
    ]
363
    if not os.path.exists(git_modules_path):
364
        return default_modules_path
365
    with open(git_modules_path) as f:
366
        return [
367
            os.path.join(cwd, line.split("=", 1)[1].strip())
368
            for line in f
369
            if line.strip().startswith("path")
370
        ]
371

372

373
def check_submodules():
374
    def check_for_files(folder, files):
375
        if not any(os.path.exists(os.path.join(folder, f)) for f in files):
376
            report("Could not find any of {} in {}".format(", ".join(files), folder))
377
            report("Did you run 'git submodule update --init --recursive'?")
378
            sys.exit(1)
379

380
    def not_exists_or_empty(folder):
381
        return not os.path.exists(folder) or (
382
            os.path.isdir(folder) and len(os.listdir(folder)) == 0
383
        )
384

385
    if bool(os.getenv("USE_SYSTEM_LIBS", False)):
386
        return
387
    folders = get_submodule_folders()
388
    # If none of the submodule folders exists, try to initialize them
389
    if all(not_exists_or_empty(folder) for folder in folders):
390
        try:
391
            print(" --- Trying to initialize submodules")
392
            start = time.time()
393
            subprocess.check_call(
394
                ["git", "submodule", "update", "--init", "--recursive"], cwd=cwd
395
            )
396
            end = time.time()
397
            print(f" --- Submodule initialization took {end - start:.2f} sec")
398
        except Exception:
399
            print(" --- Submodule initalization failed")
400
            print("Please run:\n\tgit submodule update --init --recursive")
401
            sys.exit(1)
402
    for folder in folders:
403
        check_for_files(
404
            folder,
405
            [
406
                "CMakeLists.txt",
407
                "Makefile",
408
                "setup.py",
409
                "LICENSE",
410
                "LICENSE.md",
411
                "LICENSE.txt",
412
            ],
413
        )
414
    check_for_files(
415
        os.path.join(third_party_path, "fbgemm", "third_party", "asmjit"),
416
        ["CMakeLists.txt"],
417
    )
418
    check_for_files(
419
        os.path.join(third_party_path, "onnx", "third_party", "benchmark"),
420
        ["CMakeLists.txt"],
421
    )
422

423

424
# Windows has very bad support for symbolic links.
425
# Instead of using symlinks, we're going to copy files over
426
def mirror_files_into_torchgen():
427
    # (new_path, orig_path)
428
    # Directories are OK and are recursively mirrored.
429
    paths = [
430
        (
431
            "torchgen/packaged/ATen/native/native_functions.yaml",
432
            "aten/src/ATen/native/native_functions.yaml",
433
        ),
434
        ("torchgen/packaged/ATen/native/tags.yaml", "aten/src/ATen/native/tags.yaml"),
435
        ("torchgen/packaged/ATen/templates", "aten/src/ATen/templates"),
436
        ("torchgen/packaged/autograd", "tools/autograd"),
437
        ("torchgen/packaged/autograd/templates", "tools/autograd/templates"),
438
    ]
439
    for new_path, orig_path in paths:
440
        # Create the dirs involved in new_path if they don't exist
441
        if not os.path.exists(new_path):
442
            os.makedirs(os.path.dirname(new_path), exist_ok=True)
443

444
        # Copy the files from the orig location to the new location
445
        if os.path.isfile(orig_path):
446
            shutil.copyfile(orig_path, new_path)
447
            continue
448
        if os.path.isdir(orig_path):
449
            if os.path.exists(new_path):
450
                # copytree fails if the tree exists already, so remove it.
451
                shutil.rmtree(new_path)
452
            shutil.copytree(orig_path, new_path)
453
            continue
454
        raise RuntimeError("Check the file paths in `mirror_files_into_torchgen()`")
455

456

457
# all the work we need to do _before_ setup runs
458
def build_deps():
459
    report("-- Building version " + version)
460

461
    check_submodules()
462
    check_pydep("yaml", "pyyaml")
463

464
    build_caffe2(
465
        version=version,
466
        cmake_python_library=cmake_python_library,
467
        build_python=True,
468
        rerun_cmake=RERUN_CMAKE,
469
        cmake_only=CMAKE_ONLY,
470
        cmake=cmake,
471
    )
472

473
    if CMAKE_ONLY:
474
        report(
475
            'Finished running cmake. Run "ccmake build" or '
476
            '"cmake-gui build" to adjust build options and '
477
            '"python setup.py install" to build.'
478
        )
479
        sys.exit()
480

481
    # Use copies instead of symbolic files.
482
    # Windows has very poor support for them.
483
    sym_files = [
484
        "tools/shared/_utils_internal.py",
485
        "torch/utils/benchmark/utils/valgrind_wrapper/callgrind.h",
486
        "torch/utils/benchmark/utils/valgrind_wrapper/valgrind.h",
487
    ]
488
    orig_files = [
489
        "torch/_utils_internal.py",
490
        "third_party/valgrind-headers/callgrind.h",
491
        "third_party/valgrind-headers/valgrind.h",
492
    ]
493
    for sym_file, orig_file in zip(sym_files, orig_files):
494
        same = False
495
        if os.path.exists(sym_file):
496
            if filecmp.cmp(sym_file, orig_file):
497
                same = True
498
            else:
499
                os.remove(sym_file)
500
        if not same:
501
            shutil.copyfile(orig_file, sym_file)
502

503

504
################################################################################
505
# Building dependent libraries
506
################################################################################
507

508
missing_pydep = """
509
Missing build dependency: Unable to `import {importname}`.
510
Please install it via `conda install {module}` or `pip install {module}`
511
""".strip()
512

513

514
def check_pydep(importname, module):
515
    try:
516
        importlib.import_module(importname)
517
    except ImportError as e:
518
        raise RuntimeError(
519
            missing_pydep.format(importname=importname, module=module)
520
        ) from e
521

522

523
class build_ext(setuptools.command.build_ext.build_ext):
524
    def _embed_libomp(self):
525
        # Copy libiomp5.dylib/libomp.dylib inside the wheel package on MacOS
526
        lib_dir = os.path.join(self.build_lib, "torch", "lib")
527
        libtorch_cpu_path = os.path.join(lib_dir, "libtorch_cpu.dylib")
528
        if not os.path.exists(libtorch_cpu_path):
529
            return
530
        # Parse libtorch_cpu load commands
531
        otool_cmds = (
532
            subprocess.check_output(["otool", "-l", libtorch_cpu_path])
533
            .decode("utf-8")
534
            .split("\n")
535
        )
536
        rpaths, libs = [], []
537
        for idx, line in enumerate(otool_cmds):
538
            if line.strip() == "cmd LC_LOAD_DYLIB":
539
                lib_name = otool_cmds[idx + 2].strip()
540
                assert lib_name.startswith("name ")
541
                libs.append(lib_name.split(" ", 1)[1].rsplit("(", 1)[0][:-1])
542

543
            if line.strip() == "cmd LC_RPATH":
544
                rpath = otool_cmds[idx + 2].strip()
545
                assert rpath.startswith("path ")
546
                rpaths.append(rpath.split(" ", 1)[1].rsplit("(", 1)[0][:-1])
547

548
        omp_lib_name = (
549
            "libomp.dylib" if os.uname().machine == "arm64" else "libiomp5.dylib"
550
        )
551
        if os.path.join("@rpath", omp_lib_name) not in libs:
552
            return
553

554
        # Copy libomp/libiomp5 from rpath locations
555
        for rpath in rpaths:
556
            source_lib = os.path.join(rpath, omp_lib_name)
557
            if not os.path.exists(source_lib):
558
                continue
559
            target_lib = os.path.join(self.build_lib, "torch", "lib", omp_lib_name)
560
            self.copy_file(source_lib, target_lib)
561
            break
562

563
        # Copy omp.h from OpenMP_C_FLAGS and copy it into include folder
564
        omp_cflags = get_cmake_cache_vars()["OpenMP_C_FLAGS"]
565
        if not omp_cflags:
566
            return
567
        for include_dir in [f[2:] for f in omp_cflags.split(" ") if f.startswith("-I")]:
568
            omp_h = os.path.join(include_dir, "omp.h")
569
            if not os.path.exists(omp_h):
570
                continue
571
            target_omp_h = os.path.join(self.build_lib, "torch", "include", "omp.h")
572
            self.copy_file(omp_h, target_omp_h)
573
            break
574

575
    def run(self):
576
        # Report build options. This is run after the build completes so # `CMakeCache.txt` exists and we can get an
577
        # accurate report on what is used and what is not.
578
        cmake_cache_vars = defaultdict(lambda: False, cmake.get_cmake_cache_variables())
579
        if cmake_cache_vars["USE_NUMPY"]:
580
            report("-- Building with NumPy bindings")
581
        else:
582
            report("-- NumPy not found")
583
        if cmake_cache_vars["USE_CUDNN"]:
584
            report(
585
                "-- Detected cuDNN at "
586
                + cmake_cache_vars["CUDNN_LIBRARY"]
587
                + ", "
588
                + cmake_cache_vars["CUDNN_INCLUDE_DIR"]
589
            )
590
        else:
591
            report("-- Not using cuDNN")
592
        if cmake_cache_vars["USE_CUDA"]:
593
            report("-- Detected CUDA at " + cmake_cache_vars["CUDA_TOOLKIT_ROOT_DIR"])
594
        else:
595
            report("-- Not using CUDA")
596
        if cmake_cache_vars["USE_XPU"]:
597
            report("-- Detected XPU runtime at " + cmake_cache_vars["SYCL_LIBRARY_DIR"])
598
        else:
599
            report("-- Not using XPU")
600
        if cmake_cache_vars["USE_MKLDNN"]:
601
            report("-- Using MKLDNN")
602
            if cmake_cache_vars["USE_MKLDNN_ACL"]:
603
                report("-- Using Compute Library for the Arm architecture with MKLDNN")
604
            else:
605
                report(
606
                    "-- Not using Compute Library for the Arm architecture with MKLDNN"
607
                )
608
            if cmake_cache_vars["USE_MKLDNN_CBLAS"]:
609
                report("-- Using CBLAS in MKLDNN")
610
            else:
611
                report("-- Not using CBLAS in MKLDNN")
612
        else:
613
            report("-- Not using MKLDNN")
614
        if cmake_cache_vars["USE_NCCL"] and cmake_cache_vars["USE_SYSTEM_NCCL"]:
615
            report(
616
                "-- Using system provided NCCL library at {}, {}".format(
617
                    cmake_cache_vars["NCCL_LIBRARIES"],
618
                    cmake_cache_vars["NCCL_INCLUDE_DIRS"],
619
                )
620
            )
621
        elif cmake_cache_vars["USE_NCCL"]:
622
            report("-- Building NCCL library")
623
        else:
624
            report("-- Not using NCCL")
625
        if cmake_cache_vars["USE_DISTRIBUTED"]:
626
            if IS_WINDOWS:
627
                report("-- Building without distributed package")
628
            else:
629
                report("-- Building with distributed package: ")
630
                report(
631
                    "  -- USE_TENSORPIPE={}".format(cmake_cache_vars["USE_TENSORPIPE"])
632
                )
633
                report("  -- USE_GLOO={}".format(cmake_cache_vars["USE_GLOO"]))
634
                report("  -- USE_MPI={}".format(cmake_cache_vars["USE_OPENMPI"]))
635
        else:
636
            report("-- Building without distributed package")
637
        if cmake_cache_vars["STATIC_DISPATCH_BACKEND"]:
638
            report(
639
                "-- Using static dispatch with backend {}".format(
640
                    cmake_cache_vars["STATIC_DISPATCH_BACKEND"]
641
                )
642
            )
643
        if cmake_cache_vars["USE_LIGHTWEIGHT_DISPATCH"]:
644
            report("-- Using lightweight dispatch")
645
        if cmake_cache_vars["BUILD_EXECUTORCH"]:
646
            report("-- Building Executorch")
647

648
        if cmake_cache_vars["USE_ITT"]:
649
            report("-- Using ITT")
650
        else:
651
            report("-- Not using ITT")
652

653
        # Do not use clang to compile extensions if `-fstack-clash-protection` is defined
654
        # in system CFLAGS
655
        c_flags = str(os.getenv("CFLAGS", ""))
656
        if (
657
            IS_LINUX
658
            and "-fstack-clash-protection" in c_flags
659
            and "clang" in os.environ.get("CC", "")
660
        ):
661
            os.environ["CC"] = str(os.environ["CC"])
662

663
        # It's an old-style class in Python 2.7...
664
        setuptools.command.build_ext.build_ext.run(self)
665

666
        if IS_DARWIN and package_type != "conda":
667
            self._embed_libomp()
668

669
        # Copy the essential export library to compile C++ extensions.
670
        if IS_WINDOWS:
671
            build_temp = self.build_temp
672

673
            ext_filename = self.get_ext_filename("_C")
674
            lib_filename = ".".join(ext_filename.split(".")[:-1]) + ".lib"
675

676
            export_lib = os.path.join(
677
                build_temp, "torch", "csrc", lib_filename
678
            ).replace("\\", "/")
679

680
            build_lib = self.build_lib
681

682
            target_lib = os.path.join(build_lib, "torch", "lib", "_C.lib").replace(
683
                "\\", "/"
684
            )
685

686
            # Create "torch/lib" directory if not exists.
687
            # (It is not created yet in "develop" mode.)
688
            target_dir = os.path.dirname(target_lib)
689
            if not os.path.exists(target_dir):
690
                os.makedirs(target_dir)
691

692
            self.copy_file(export_lib, target_lib)
693

694
    def build_extensions(self):
695
        self.create_compile_commands()
696
        # The caffe2 extensions are created in
697
        # tmp_install/lib/pythonM.m/site-packages/caffe2/python/
698
        # and need to be copied to build/lib.linux.... , which will be a
699
        # platform dependent build folder created by the "build" command of
700
        # setuptools. Only the contents of this folder are installed in the
701
        # "install" command by default.
702
        # We only make this copy for Caffe2's pybind extensions
703
        caffe2_pybind_exts = [
704
            "caffe2.python.caffe2_pybind11_state",
705
            "caffe2.python.caffe2_pybind11_state_gpu",
706
            "caffe2.python.caffe2_pybind11_state_hip",
707
        ]
708
        i = 0
709
        while i < len(self.extensions):
710
            ext = self.extensions[i]
711
            if ext.name not in caffe2_pybind_exts:
712
                i += 1
713
                continue
714
            fullname = self.get_ext_fullname(ext.name)
715
            filename = self.get_ext_filename(fullname)
716
            report(f"\nCopying extension {ext.name}")
717

718
            relative_site_packages = (
719
                sysconfig.get_path("purelib")
720
                .replace(sysconfig.get_path("data"), "")
721
                .lstrip(os.path.sep)
722
            )
723
            src = os.path.join("torch", relative_site_packages, filename)
724
            if not os.path.exists(src):
725
                report(f"{src} does not exist")
726
                del self.extensions[i]
727
            else:
728
                dst = os.path.join(os.path.realpath(self.build_lib), filename)
729
                report(f"Copying {ext.name} from {src} to {dst}")
730
                dst_dir = os.path.dirname(dst)
731
                if not os.path.exists(dst_dir):
732
                    os.makedirs(dst_dir)
733
                self.copy_file(src, dst)
734
                i += 1
735

736
        # Copy functorch extension
737
        for i, ext in enumerate(self.extensions):
738
            if ext.name != "functorch._C":
739
                continue
740
            fullname = self.get_ext_fullname(ext.name)
741
            filename = self.get_ext_filename(fullname)
742
            fileext = os.path.splitext(filename)[1]
743
            src = os.path.join(os.path.dirname(filename), "functorch" + fileext)
744
            dst = os.path.join(os.path.realpath(self.build_lib), filename)
745
            if os.path.exists(src):
746
                report(f"Copying {ext.name} from {src} to {dst}")
747
                dst_dir = os.path.dirname(dst)
748
                if not os.path.exists(dst_dir):
749
                    os.makedirs(dst_dir)
750
                self.copy_file(src, dst)
751

752
        setuptools.command.build_ext.build_ext.build_extensions(self)
753

754
    def get_outputs(self):
755
        outputs = setuptools.command.build_ext.build_ext.get_outputs(self)
756
        outputs.append(os.path.join(self.build_lib, "caffe2"))
757
        report(f"setup.py::get_outputs returning {outputs}")
758
        return outputs
759

760
    def create_compile_commands(self):
761
        def load(filename):
762
            with open(filename) as f:
763
                return json.load(f)
764

765
        ninja_files = glob.glob("build/*compile_commands.json")
766
        cmake_files = glob.glob("torch/lib/build/*/compile_commands.json")
767
        all_commands = [entry for f in ninja_files + cmake_files for entry in load(f)]
768

769
        # cquery does not like c++ compiles that start with gcc.
770
        # It forgets to include the c++ header directories.
771
        # We can work around this by replacing the gcc calls that python
772
        # setup.py generates with g++ calls instead
773
        for command in all_commands:
774
            if command["command"].startswith("gcc "):
775
                command["command"] = "g++ " + command["command"][4:]
776

777
        new_contents = json.dumps(all_commands, indent=2)
778
        contents = ""
779
        if os.path.exists("compile_commands.json"):
780
            with open("compile_commands.json") as f:
781
                contents = f.read()
782
        if contents != new_contents:
783
            with open("compile_commands.json", "w") as f:
784
                f.write(new_contents)
785

786

787
class concat_license_files:
788
    """Merge LICENSE and LICENSES_BUNDLED.txt as a context manager
789

790
    LICENSE is the main PyTorch license, LICENSES_BUNDLED.txt is auto-generated
791
    from all the licenses found in ./third_party/. We concatenate them so there
792
    is a single license file in the sdist and wheels with all of the necessary
793
    licensing info.
794
    """
795

796
    def __init__(self, include_files=False):
797
        self.f1 = "LICENSE"
798
        self.f2 = "third_party/LICENSES_BUNDLED.txt"
799
        self.include_files = include_files
800

801
    def __enter__(self):
802
        """Concatenate files"""
803

804
        old_path = sys.path
805
        sys.path.append(third_party_path)
806
        try:
807
            from build_bundled import create_bundled
808
        finally:
809
            sys.path = old_path
810

811
        with open(self.f1) as f1:
812
            self.bsd_text = f1.read()
813

814
        with open(self.f1, "a") as f1:
815
            f1.write("\n\n")
816
            create_bundled(
817
                os.path.relpath(third_party_path), f1, include_files=self.include_files
818
            )
819

820
    def __exit__(self, exception_type, exception_value, traceback):
821
        """Restore content of f1"""
822
        with open(self.f1, "w") as f:
823
            f.write(self.bsd_text)
824

825

826
try:
827
    from wheel.bdist_wheel import bdist_wheel
828
except ImportError:
829
    # This is useful when wheel is not installed and bdist_wheel is not
830
    # specified on the command line. If it _is_ specified, parsing the command
831
    # line will fail before wheel_concatenate is needed
832
    wheel_concatenate = None
833
else:
834
    # Need to create the proper LICENSE.txt for the wheel
835
    class wheel_concatenate(bdist_wheel):
836
        """check submodules on sdist to prevent incomplete tarballs"""
837

838
        def run(self):
839
            with concat_license_files(include_files=True):
840
                super().run()
841

842

843
class install(setuptools.command.install.install):
844
    def run(self):
845
        super().run()
846

847

848
class clean(setuptools.Command):
849
    user_options = []
850

851
    def initialize_options(self):
852
        pass
853

854
    def finalize_options(self):
855
        pass
856

857
    def run(self):
858
        import glob
859
        import re
860

861
        with open(".gitignore") as f:
862
            ignores = f.read()
863
            pat = re.compile(r"^#( BEGIN NOT-CLEAN-FILES )?")
864
            for wildcard in filter(None, ignores.split("\n")):
865
                match = pat.match(wildcard)
866
                if match:
867
                    if match.group(1):
868
                        # Marker is found and stop reading .gitignore.
869
                        break
870
                    # Ignore lines which begin with '#'.
871
                else:
872
                    # Don't remove absolute paths from the system
873
                    wildcard = wildcard.lstrip("./")
874

875
                    for filename in glob.glob(wildcard):
876
                        try:
877
                            os.remove(filename)
878
                        except OSError:
879
                            shutil.rmtree(filename, ignore_errors=True)
880

881

882
class sdist(setuptools.command.sdist.sdist):
883
    def run(self):
884
        with concat_license_files():
885
            super().run()
886

887

888
def get_cmake_cache_vars():
889
    try:
890
        return defaultdict(lambda: False, cmake.get_cmake_cache_variables())
891
    except FileNotFoundError:
892
        # CMakeCache.txt does not exist. Probably running "python setup.py clean" over a clean directory.
893
        return defaultdict(lambda: False)
894

895

896
def configure_extension_build():
897
    r"""Configures extension build options according to system environment and user's choice.
898

899
    Returns:
900
      The input to parameters ext_modules, cmdclass, packages, and entry_points as required in setuptools.setup.
901
    """
902

903
    cmake_cache_vars = get_cmake_cache_vars()
904

905
    ################################################################################
906
    # Configure compile flags
907
    ################################################################################
908

909
    library_dirs = []
910
    extra_install_requires = []
911

912
    if IS_WINDOWS:
913
        # /NODEFAULTLIB makes sure we only link to DLL runtime
914
        # and matches the flags set for protobuf and ONNX
915
        extra_link_args = ["/NODEFAULTLIB:LIBCMT.LIB"]
916
        # /MD links against DLL runtime
917
        # and matches the flags set for protobuf and ONNX
918
        # /EHsc is about standard C++ exception handling
919
        extra_compile_args = ["/MD", "/FS", "/EHsc"]
920
    else:
921
        extra_link_args = []
922
        extra_compile_args = [
923
            "-Wall",
924
            "-Wextra",
925
            "-Wno-strict-overflow",
926
            "-Wno-unused-parameter",
927
            "-Wno-missing-field-initializers",
928
            "-Wno-unknown-pragmas",
929
            # Python 2.6 requires -fno-strict-aliasing, see
930
            # http://legacy.python.org/dev/peps/pep-3123/
931
            # We also depend on it in our code (even Python 3).
932
            "-fno-strict-aliasing",
933
        ]
934

935
    library_dirs.append(lib_path)
936

937
    main_compile_args = []
938
    main_libraries = ["torch_python"]
939
    main_link_args = []
940
    main_sources = ["torch/csrc/stub.c"]
941

942
    if cmake_cache_vars["USE_CUDA"]:
943
        library_dirs.append(os.path.dirname(cmake_cache_vars["CUDA_CUDA_LIB"]))
944

945
    if build_type.is_debug():
946
        if IS_WINDOWS:
947
            extra_compile_args.append("/Z7")
948
            extra_link_args.append("/DEBUG:FULL")
949
        else:
950
            extra_compile_args += ["-O0", "-g"]
951
            extra_link_args += ["-O0", "-g"]
952

953
    if build_type.is_rel_with_deb_info():
954
        if IS_WINDOWS:
955
            extra_compile_args.append("/Z7")
956
            extra_link_args.append("/DEBUG:FULL")
957
        else:
958
            extra_compile_args += ["-g"]
959
            extra_link_args += ["-g"]
960

961
    # pypi cuda package that requires installation of cuda runtime, cudnn and cublas
962
    # should be included in all wheels uploaded to pypi
963
    pytorch_extra_install_requirements = os.getenv(
964
        "PYTORCH_EXTRA_INSTALL_REQUIREMENTS", ""
965
    )
966
    if pytorch_extra_install_requirements:
967
        report(
968
            f"pytorch_extra_install_requirements: {pytorch_extra_install_requirements}"
969
        )
970
        extra_install_requires += pytorch_extra_install_requirements.split("|")
971

972
    # Cross-compile for M1
973
    if IS_DARWIN:
974
        macos_target_arch = os.getenv("CMAKE_OSX_ARCHITECTURES", "")
975
        if macos_target_arch in ["arm64", "x86_64"]:
976
            macos_sysroot_path = os.getenv("CMAKE_OSX_SYSROOT")
977
            if macos_sysroot_path is None:
978
                macos_sysroot_path = (
979
                    subprocess.check_output(
980
                        ["xcrun", "--show-sdk-path", "--sdk", "macosx"]
981
                    )
982
                    .decode("utf-8")
983
                    .strip()
984
                )
985
            extra_compile_args += [
986
                "-arch",
987
                macos_target_arch,
988
                "-isysroot",
989
                macos_sysroot_path,
990
            ]
991
            extra_link_args += ["-arch", macos_target_arch]
992

993
    def make_relative_rpath_args(path):
994
        if IS_DARWIN:
995
            return ["-Wl,-rpath,@loader_path/" + path]
996
        elif IS_WINDOWS:
997
            return []
998
        else:
999
            return ["-Wl,-rpath,$ORIGIN/" + path]
1000

1001
    ################################################################################
1002
    # Declare extensions and package
1003
    ################################################################################
1004

1005
    extensions = []
1006
    excludes = ["tools", "tools.*"]
1007
    if not cmake_cache_vars["BUILD_CAFFE2"]:
1008
        excludes.extend(["caffe2", "caffe2.*"])
1009
    if not cmake_cache_vars["BUILD_FUNCTORCH"]:
1010
        excludes.extend(["functorch", "functorch.*"])
1011
    packages = find_packages(exclude=excludes)
1012
    C = Extension(
1013
        "torch._C",
1014
        libraries=main_libraries,
1015
        sources=main_sources,
1016
        language="c",
1017
        extra_compile_args=main_compile_args + extra_compile_args,
1018
        include_dirs=[],
1019
        library_dirs=library_dirs,
1020
        extra_link_args=extra_link_args
1021
        + main_link_args
1022
        + make_relative_rpath_args("lib"),
1023
    )
1024
    extensions.append(C)
1025

1026
    # These extensions are built by cmake and copied manually in build_extensions()
1027
    # inside the build_ext implementation
1028
    if cmake_cache_vars["BUILD_CAFFE2"]:
1029
        extensions.append(
1030
            Extension(name="caffe2.python.caffe2_pybind11_state", sources=[]),
1031
        )
1032
        if cmake_cache_vars["USE_CUDA"]:
1033
            extensions.append(
1034
                Extension(name="caffe2.python.caffe2_pybind11_state_gpu", sources=[]),
1035
            )
1036
        if cmake_cache_vars["USE_ROCM"]:
1037
            extensions.append(
1038
                Extension(name="caffe2.python.caffe2_pybind11_state_hip", sources=[]),
1039
            )
1040
    if cmake_cache_vars["BUILD_FUNCTORCH"]:
1041
        extensions.append(
1042
            Extension(name="functorch._C", sources=[]),
1043
        )
1044

1045
    cmdclass = {
1046
        "bdist_wheel": wheel_concatenate,
1047
        "build_ext": build_ext,
1048
        "clean": clean,
1049
        "install": install,
1050
        "sdist": sdist,
1051
    }
1052

1053
    entry_points = {
1054
        "console_scripts": [
1055
            "convert-caffe2-to-onnx = caffe2.python.onnx.bin.conversion:caffe2_to_onnx",
1056
            "convert-onnx-to-caffe2 = caffe2.python.onnx.bin.conversion:onnx_to_caffe2",
1057
            "torchrun = torch.distributed.run:main",
1058
        ],
1059
        "torchrun.logs_specs": [
1060
            "default = torch.distributed.elastic.multiprocessing:DefaultLogsSpecs",
1061
        ],
1062
    }
1063

1064
    return extensions, cmdclass, packages, entry_points, extra_install_requires
1065

1066

1067
# post run, warnings, printed at the end to make them more visible
1068
build_update_message = """
1069
    It is no longer necessary to use the 'build' or 'rebuild' targets
1070

1071
    To install:
1072
      $ python setup.py install
1073
    To develop locally:
1074
      $ python setup.py develop
1075
    To force cmake to re-generate native build files (off by default):
1076
      $ python setup.py develop --cmake
1077
"""
1078

1079

1080
def print_box(msg):
1081
    lines = msg.split("\n")
1082
    size = max(len(l) + 1 for l in lines)
1083
    print("-" * (size + 2))
1084
    for l in lines:
1085
        print("|{}{}|".format(l, " " * (size - len(l))))
1086
    print("-" * (size + 2))
1087

1088

1089
def main():
1090
    # the list of runtime dependencies required by this built package
1091
    install_requires = [
1092
        "filelock",
1093
        "typing-extensions>=4.8.0",
1094
        "sympy",
1095
        "networkx",
1096
        "jinja2",
1097
        "fsspec",
1098
    ]
1099
    if IS_WINDOWS:
1100
        install_requires.append("mkl>=2021.1.1,<=2021.4.0")
1101

1102
    # Parse the command line and check the arguments before we proceed with
1103
    # building deps and setup. We need to set values so `--help` works.
1104
    dist = Distribution()
1105
    dist.script_name = os.path.basename(sys.argv[0])
1106
    dist.script_args = sys.argv[1:]
1107
    try:
1108
        dist.parse_command_line()
1109
    except setuptools.distutils.errors.DistutilsArgError as e:
1110
        print(e)
1111
        sys.exit(1)
1112

1113
    mirror_files_into_torchgen()
1114
    if RUN_BUILD_DEPS:
1115
        build_deps()
1116

1117
    (
1118
        extensions,
1119
        cmdclass,
1120
        packages,
1121
        entry_points,
1122
        extra_install_requires,
1123
    ) = configure_extension_build()
1124

1125
    install_requires += extra_install_requires
1126

1127
    extras_require = {
1128
        "optree": ["optree>=0.9.1"],
1129
        "opt-einsum": ["opt-einsum>=3.3"],
1130
    }
1131

1132
    # Read in README.md for our long_description
1133
    with open(os.path.join(cwd, "README.md"), encoding="utf-8") as f:
1134
        long_description = f.read()
1135

1136
    version_range_max = max(sys.version_info[1], 12) + 1
1137
    torch_package_data = [
1138
        "py.typed",
1139
        "bin/*",
1140
        "test/*",
1141
        "*.pyi",
1142
        "_C/*.pyi",
1143
        "cuda/*.pyi",
1144
        "fx/*.pyi",
1145
        "optim/*.pyi",
1146
        "autograd/*.pyi",
1147
        "nn/*.pyi",
1148
        "nn/modules/*.pyi",
1149
        "nn/parallel/*.pyi",
1150
        "utils/data/*.pyi",
1151
        "utils/data/datapipes/*.pyi",
1152
        "lib/*.so*",
1153
        "lib/*.dylib*",
1154
        "lib/*.dll",
1155
        "lib/*.lib",
1156
        "lib/*.pdb",
1157
        "lib/torch_shm_manager",
1158
        "lib/*.h",
1159
        "include/*.h",
1160
        "include/ATen/*.h",
1161
        "include/ATen/cpu/*.h",
1162
        "include/ATen/cpu/vec/vec256/*.h",
1163
        "include/ATen/cpu/vec/vec256/vsx/*.h",
1164
        "include/ATen/cpu/vec/vec256/zarch/*.h",
1165
        "include/ATen/cpu/vec/vec512/*.h",
1166
        "include/ATen/cpu/vec/*.h",
1167
        "include/ATen/core/*.h",
1168
        "include/ATen/cuda/*.cuh",
1169
        "include/ATen/cuda/*.h",
1170
        "include/ATen/cuda/detail/*.cuh",
1171
        "include/ATen/cuda/detail/*.h",
1172
        "include/ATen/cuda/tunable/*.h",
1173
        "include/ATen/cudnn/*.h",
1174
        "include/ATen/functorch/*.h",
1175
        "include/ATen/ops/*.h",
1176
        "include/ATen/hip/*.cuh",
1177
        "include/ATen/hip/*.h",
1178
        "include/ATen/hip/detail/*.cuh",
1179
        "include/ATen/hip/detail/*.h",
1180
        "include/ATen/hip/impl/*.h",
1181
        "include/ATen/hip/tunable/*.h",
1182
        "include/ATen/mps/*.h",
1183
        "include/ATen/miopen/*.h",
1184
        "include/ATen/detail/*.h",
1185
        "include/ATen/native/*.h",
1186
        "include/ATen/native/cpu/*.h",
1187
        "include/ATen/native/cuda/*.h",
1188
        "include/ATen/native/cuda/*.cuh",
1189
        "include/ATen/native/hip/*.h",
1190
        "include/ATen/native/hip/*.cuh",
1191
        "include/ATen/native/mps/*.h",
1192
        "include/ATen/native/nested/*.h",
1193
        "include/ATen/native/quantized/*.h",
1194
        "include/ATen/native/quantized/cpu/*.h",
1195
        "include/ATen/native/sparse/*.h",
1196
        "include/ATen/native/utils/*.h",
1197
        "include/ATen/quantized/*.h",
1198
        "include/ATen/xpu/*.h",
1199
        "include/ATen/xpu/detail/*.h",
1200
        "include/caffe2/serialize/*.h",
1201
        "include/c10/*.h",
1202
        "include/c10/macros/*.h",
1203
        "include/c10/core/*.h",
1204
        "include/ATen/core/boxing/*.h",
1205
        "include/ATen/core/boxing/impl/*.h",
1206
        "include/ATen/core/dispatch/*.h",
1207
        "include/ATen/core/op_registration/*.h",
1208
        "include/c10/core/impl/*.h",
1209
        "include/c10/util/*.h",
1210
        "include/c10/cuda/*.h",
1211
        "include/c10/cuda/impl/*.h",
1212
        "include/c10/hip/*.h",
1213
        "include/c10/hip/impl/*.h",
1214
        "include/c10/xpu/*.h",
1215
        "include/c10/xpu/impl/*.h",
1216
        "include/torch/*.h",
1217
        "include/torch/csrc/*.h",
1218
        "include/torch/csrc/api/include/torch/*.h",
1219
        "include/torch/csrc/api/include/torch/data/*.h",
1220
        "include/torch/csrc/api/include/torch/data/dataloader/*.h",
1221
        "include/torch/csrc/api/include/torch/data/datasets/*.h",
1222
        "include/torch/csrc/api/include/torch/data/detail/*.h",
1223
        "include/torch/csrc/api/include/torch/data/samplers/*.h",
1224
        "include/torch/csrc/api/include/torch/data/transforms/*.h",
1225
        "include/torch/csrc/api/include/torch/detail/*.h",
1226
        "include/torch/csrc/api/include/torch/detail/ordered_dict.h",
1227
        "include/torch/csrc/api/include/torch/nn/*.h",
1228
        "include/torch/csrc/api/include/torch/nn/functional/*.h",
1229
        "include/torch/csrc/api/include/torch/nn/options/*.h",
1230
        "include/torch/csrc/api/include/torch/nn/modules/*.h",
1231
        "include/torch/csrc/api/include/torch/nn/modules/container/*.h",
1232
        "include/torch/csrc/api/include/torch/nn/parallel/*.h",
1233
        "include/torch/csrc/api/include/torch/nn/utils/*.h",
1234
        "include/torch/csrc/api/include/torch/optim/*.h",
1235
        "include/torch/csrc/api/include/torch/optim/schedulers/*.h",
1236
        "include/torch/csrc/api/include/torch/serialize/*.h",
1237
        "include/torch/csrc/autograd/*.h",
1238
        "include/torch/csrc/autograd/functions/*.h",
1239
        "include/torch/csrc/autograd/generated/*.h",
1240
        "include/torch/csrc/autograd/utils/*.h",
1241
        "include/torch/csrc/cuda/*.h",
1242
        "include/torch/csrc/distributed/c10d/*.h",
1243
        "include/torch/csrc/distributed/c10d/*.hpp",
1244
        "include/torch/csrc/distributed/rpc/*.h",
1245
        "include/torch/csrc/distributed/autograd/context/*.h",
1246
        "include/torch/csrc/distributed/autograd/functions/*.h",
1247
        "include/torch/csrc/distributed/autograd/rpc_messages/*.h",
1248
        "include/torch/csrc/dynamo/*.h",
1249
        "include/torch/csrc/inductor/*.h",
1250
        "include/torch/csrc/inductor/aoti_runner/*.h",
1251
        "include/torch/csrc/inductor/aoti_runtime/*.h",
1252
        "include/torch/csrc/inductor/aoti_torch/*.h",
1253
        "include/torch/csrc/inductor/aoti_torch/c/*.h",
1254
        "include/torch/csrc/jit/*.h",
1255
        "include/torch/csrc/jit/backends/*.h",
1256
        "include/torch/csrc/jit/generated/*.h",
1257
        "include/torch/csrc/jit/passes/*.h",
1258
        "include/torch/csrc/jit/passes/quantization/*.h",
1259
        "include/torch/csrc/jit/passes/utils/*.h",
1260
        "include/torch/csrc/jit/runtime/*.h",
1261
        "include/torch/csrc/jit/ir/*.h",
1262
        "include/torch/csrc/jit/frontend/*.h",
1263
        "include/torch/csrc/jit/api/*.h",
1264
        "include/torch/csrc/jit/serialization/*.h",
1265
        "include/torch/csrc/jit/python/*.h",
1266
        "include/torch/csrc/jit/mobile/*.h",
1267
        "include/torch/csrc/jit/testing/*.h",
1268
        "include/torch/csrc/jit/tensorexpr/*.h",
1269
        "include/torch/csrc/jit/tensorexpr/operators/*.h",
1270
        "include/torch/csrc/jit/codegen/cuda/*.h",
1271
        "include/torch/csrc/onnx/*.h",
1272
        "include/torch/csrc/profiler/*.h",
1273
        "include/torch/csrc/profiler/orchestration/*.h",
1274
        "include/torch/csrc/profiler/stubs/*.h",
1275
        "include/torch/csrc/profiler/unwind/*.h",
1276
        "include/torch/csrc/utils/*.h",
1277
        "include/torch/csrc/tensor/*.h",
1278
        "include/torch/csrc/lazy/backend/*.h",
1279
        "include/torch/csrc/lazy/core/*.h",
1280
        "include/torch/csrc/lazy/core/internal_ops/*.h",
1281
        "include/torch/csrc/lazy/core/ops/*.h",
1282
        "include/torch/csrc/lazy/python/python_util.h",
1283
        "include/torch/csrc/lazy/ts_backend/*.h",
1284
        "include/torch/csrc/xpu/*.h",
1285
        "include/pybind11/*.h",
1286
        "include/pybind11/detail/*.h",
1287
        "include/pybind11/eigen/*.h",
1288
        "include/TH/*.h*",
1289
        "include/TH/generic/*.h*",
1290
        "include/THC/*.cuh",
1291
        "include/THC/*.h*",
1292
        "include/THC/generic/*.h",
1293
        "include/THH/*.cuh",
1294
        "include/THH/*.h*",
1295
        "include/THH/generic/*.h",
1296
        "include/sleef.h",
1297
        "_inductor/codegen/*.h",
1298
        "_inductor/codegen/aoti_runtime/*.cpp",
1299
        "_export/serde/*.yaml",
1300
        "share/cmake/ATen/*.cmake",
1301
        "share/cmake/Caffe2/*.cmake",
1302
        "share/cmake/Caffe2/public/*.cmake",
1303
        "share/cmake/Caffe2/Modules_CUDA_fix/*.cmake",
1304
        "share/cmake/Caffe2/Modules_CUDA_fix/upstream/*.cmake",
1305
        "share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/*.cmake",
1306
        "share/cmake/Gloo/*.cmake",
1307
        "share/cmake/Tensorpipe/*.cmake",
1308
        "share/cmake/Torch/*.cmake",
1309
        "utils/benchmark/utils/*.cpp",
1310
        "utils/benchmark/utils/valgrind_wrapper/*.cpp",
1311
        "utils/benchmark/utils/valgrind_wrapper/*.h",
1312
        "utils/model_dump/skeleton.html",
1313
        "utils/model_dump/code.js",
1314
        "utils/model_dump/*.mjs",
1315
    ]
1316

1317
    if get_cmake_cache_vars()["BUILD_CAFFE2"]:
1318
        torch_package_data.extend(
1319
            [
1320
                "include/caffe2/**/*.h",
1321
                "include/caffe2/utils/*.h",
1322
                "include/caffe2/utils/**/*.h",
1323
            ]
1324
        )
1325
    if get_cmake_cache_vars()["USE_TENSORPIPE"]:
1326
        torch_package_data.extend(
1327
            [
1328
                "include/tensorpipe/*.h",
1329
                "include/tensorpipe/channel/*.h",
1330
                "include/tensorpipe/channel/basic/*.h",
1331
                "include/tensorpipe/channel/cma/*.h",
1332
                "include/tensorpipe/channel/mpt/*.h",
1333
                "include/tensorpipe/channel/xth/*.h",
1334
                "include/tensorpipe/common/*.h",
1335
                "include/tensorpipe/core/*.h",
1336
                "include/tensorpipe/transport/*.h",
1337
                "include/tensorpipe/transport/ibv/*.h",
1338
                "include/tensorpipe/transport/shm/*.h",
1339
                "include/tensorpipe/transport/uv/*.h",
1340
            ]
1341
        )
1342
    torchgen_package_data = [
1343
        # Recursive glob doesn't work in setup.py,
1344
        # https://github.com/pypa/setuptools/issues/1806
1345
        # To make this robust we should replace it with some code that
1346
        # returns a list of everything under packaged/
1347
        "packaged/ATen/*",
1348
        "packaged/ATen/native/*",
1349
        "packaged/ATen/templates/*",
1350
        "packaged/autograd/*",
1351
        "packaged/autograd/templates/*",
1352
    ]
1353
    setup(
1354
        name=package_name,
1355
        version=version,
1356
        description=(
1357
            "Tensors and Dynamic neural networks in "
1358
            "Python with strong GPU acceleration"
1359
        ),
1360
        long_description=long_description,
1361
        long_description_content_type="text/markdown",
1362
        ext_modules=extensions,
1363
        cmdclass=cmdclass,
1364
        packages=packages,
1365
        entry_points=entry_points,
1366
        install_requires=install_requires,
1367
        extras_require=extras_require,
1368
        package_data={
1369
            "torch": torch_package_data,
1370
            "torchgen": torchgen_package_data,
1371
            "caffe2": [
1372
                "python/serialized_test/data/operator_test/*.zip",
1373
            ],
1374
        },
1375
        url="https://pytorch.org/",
1376
        download_url="https://github.com/pytorch/pytorch/tags",
1377
        author="PyTorch Team",
1378
        author_email="packages@pytorch.org",
1379
        python_requires=f">={python_min_version_str}",
1380
        # PyPI package information.
1381
        classifiers=[
1382
            "Development Status :: 5 - Production/Stable",
1383
            "Intended Audience :: Developers",
1384
            "Intended Audience :: Education",
1385
            "Intended Audience :: Science/Research",
1386
            "License :: OSI Approved :: BSD License",
1387
            "Topic :: Scientific/Engineering",
1388
            "Topic :: Scientific/Engineering :: Mathematics",
1389
            "Topic :: Scientific/Engineering :: Artificial Intelligence",
1390
            "Topic :: Software Development",
1391
            "Topic :: Software Development :: Libraries",
1392
            "Topic :: Software Development :: Libraries :: Python Modules",
1393
            "Programming Language :: C++",
1394
            "Programming Language :: Python :: 3",
1395
        ]
1396
        + [
1397
            f"Programming Language :: Python :: 3.{i}"
1398
            for i in range(python_min_version[1], version_range_max)
1399
        ],
1400
        license="BSD-3",
1401
        keywords="pytorch, machine learning",
1402
    )
1403
    if EMIT_BUILD_WARNING:
1404
        print_box(build_update_message)
1405

1406

1407
if __name__ == "__main__":
1408
    main()
1409