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