frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 13:44:15 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
7ed2bc7c67dce9ec0edc642bd0df5170b89ff30b
pytorch/setup.py
hongxyan 66a76516bf [ROCm] Disabling Kernel Asserts for ROCm by default - fix and clean up and refactoring (#114660) 
Related to #103973  #110532 #108404 #94891

**Context:**
As commented in 6ae0554d11/cmake/Dependencies.cmake (L1198)
Kernel asserts are enabled by default for CUDA and disabled for ROCm.
However it is somewhat broken, and Kernel assert was still enabled for ROCm.

Disabling kernel assert is also needed for users who do not have PCIe atomics support. These community users have verified that disabling the kernel assert in PyTorch/ROCm platform fixed their pytorch workflow, like torch.sum script, stable-diffusion. (see the related issues)

**Changes:**

This pull request serves the following purposes:
* Refactor and clean up the logic,  make it simpler for ROCm to enable and disable Kernel Asserts
* Fix the bug that Kernel Asserts for ROCm was not disabled by default.

Specifically,
- Renamed `TORCH_DISABLE_GPU_ASSERTS` to `C10_USE_ROCM_KERNEL_ASSERT` for the following reasons:
(1) This variable only applies to ROCm.
(2) The new name is more align with #define CUDA_KERNEL_ASSERT function.
(3) With USE_ in front of the name, we can easily control it with environment variable to turn on and off this feature during build (e.g. `USE_ROCM_KERNEL_ASSERT=1 python setup.py develop` will enable kernel assert for ROCm build).
- Get rid of the `ROCM_FORCE_ENABLE_GPU_ASSERTS' to simplify the logic and make it easier to understand and maintain
- Added `#cmakedefine` to carry over the CMake variable to C++

**Tests:**
(1) build with default mode and verify that USE_ROCM_KERNEL_ASSERT  is OFF(0), and kernel assert is disabled:

```
python setup.py develop
```
Verify CMakeCache.txt has correct value.
```
/xxxx/pytorch/build$ grep USE_ROCM_KERNEL_ASSERT CMakeCache.txt
USE_ROCM_KERNEL_ASSERT:BOOL=0
```
Tested the following code in ROCm build and CUDA build, and expected the return code differently.

```
subprocess.call([sys.executable, '-c', "import torch;torch._assert_async(torch.tensor(0,device='cuda'));torch.cuda.synchronize()"])
```
This piece of code is adapted from below unit test to get around the limitation that this unit test now was skipped for ROCm. (We will check to enable this unit test in the future)

```
python test/test_cuda_expandable_segments.py -k test_fixed_cuda_assert_async
```

Ran the following script, expecting r ==0 since the CUDA_KERNEL_ASSERT is defined as nothing:
```
>> import sys
>>> import subprocess
>>> r=subprocess.call([sys.executable, '-c', "import torch;torch._assert_async(torch.tensor(0,device='cuda'));torch.cuda.synchronize()"])
>>> r
0
```

(2) Enable the kernel assert by building with USE_ROCM_KERNEL_ASSERT=1, or USE_ROCM_KERNEL_ASSERT=ON
```
USE_ROCM_KERNEL_ASSERT=1 python setup.py develop
```

Verify `USE_ROCM_KERNEL_ASSERT` is `1`
```
/xxxx/pytorch/build$ grep USE_ROCM_KERNEL_ASSERT CMakeCache.txt
USE_ROCM_KERNEL_ASSERT:BOOL=1
```

Run the assert test, and expected return code not equal to 0.

```
>> import sys
>>> import subprocess
>>> r=subprocess.call([sys.executable, '-c', "import torch;torch._assert_async(torch.tensor(0,device='cuda'));torch.cuda.synchronize()"])
>>>/xxxx/pytorch/aten/src/ATen/native/hip/TensorCompare.hip:108: _assert_async_cuda_kernel: Device-side assertion `input[0] != 0' failed.
:0:rocdevice.cpp            :2690: 2435301199202 us: [pid:206019 tid:0x7f6cf0a77700] Callback: Queue 0x7f64e8400000 aborting with error : HSA_STATUS_ERROR_EXCEPTION: An HSAIL operation resulted in a hardware exception. code: 0x1016

>>> r
-6
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/114660
Approved by: https://github.com/jeffdaily, https://github.com/malfet, https://github.com/jithunnair-amd
2023-12-13 15:44:53 +00:00

1376 lines
47 KiB
Python
Raw Blame History

# Welcome to the PyTorch setup.py.
#
# Environment variables you are probably interested in:
#
# DEBUG
# build with -O0 and -g (debug symbols)
#
# REL_WITH_DEB_INFO
# build with optimizations and -g (debug symbols)
#
# USE_CUSTOM_DEBINFO="path/to/file1.cpp;path/to/file2.cpp"
# build with debug info only for specified files
#
# MAX_JOBS
# maximum number of compile jobs we should use to compile your code
#
# USE_CUDA=0
# disables CUDA build
#
# CFLAGS
# flags to apply to both C and C++ files to be compiled (a quirk of setup.py
# which we have faithfully adhered to in our build system is that CFLAGS
# also applies to C++ files (unless CXXFLAGS is set), in contrast to the
# default behavior of autogoo and cmake build systems.)
#
# CC
# the C/C++ compiler to use
#
# Environment variables for feature toggles:
#
# DEBUG_CUDA=1
# if used in conjunction with DEBUG or REL_WITH_DEB_INFO, will also
# build CUDA kernels with -lineinfo --source-in-ptx. Note that
# on CUDA 12 this may cause nvcc to OOM, so this is disabled by default.
# USE_CUDNN=0
# disables the cuDNN build
#
# USE_CUSPARSELT=0
# disables the cuSPARSELt build
#
# USE_FBGEMM=0
# disables the FBGEMM build
#
# USE_KINETO=0
# disables usage of libkineto library for profiling
#
# USE_NUMPY=0
# disables the NumPy build
#
# BUILD_TEST=0
# disables the test build
#
# USE_MKLDNN=0
# disables use of MKLDNN
#
# USE_MKLDNN_ACL
# enables use of Compute Library backend for MKLDNN on Arm;
# USE_MKLDNN must be explicitly enabled.
#
# MKLDNN_CPU_RUNTIME
# MKL-DNN threading mode: TBB or OMP (default)
#
# USE_STATIC_MKL
# Prefer to link with MKL statically - Unix only
# USE_ITT=0
# disable use of Intel(R) VTune Profiler's ITT functionality
#
# USE_NNPACK=0
# disables NNPACK build
#
# USE_QNNPACK=0
# disables QNNPACK build (quantized 8-bit operators)
#
# USE_DISTRIBUTED=0
# disables distributed (c10d, gloo, mpi, etc.) build
#
# USE_TENSORPIPE=0
# disables distributed Tensorpipe backend build
#
# USE_GLOO=0
# disables distributed gloo backend build
#
# USE_MPI=0
# disables distributed MPI backend build
#
# USE_SYSTEM_NCCL=0
# disables use of system-wide nccl (we will use our submoduled
# copy in third_party/nccl)
#
# BUILD_CAFFE2_OPS=0
# disable Caffe2 operators build
#
# BUILD_CAFFE2=0
# disable Caffe2 build
#
# USE_IBVERBS
# toggle features related to distributed support
#
# USE_OPENCV
# enables use of OpenCV for additional operators
#
# USE_OPENMP=0
# disables use of OpenMP for parallelization
#
# USE_FFMPEG
# enables use of ffmpeg for additional operators
#
# USE_FLASH_ATTENTION=0
# disables building flash attention for scaled dot product attention
#
# USE_MEM_EFF_ATTENTION=0
# disables building memory efficient attention for scaled dot product attention
#
# USE_LEVELDB
# enables use of LevelDB for storage
#
# USE_LMDB
# enables use of LMDB for storage
#
# BUILD_BINARY
# enables the additional binaries/ build
#
# ATEN_AVX512_256=TRUE
# ATen AVX2 kernels can use 32 ymm registers, instead of the default 16.
# This option can be used if AVX512 doesn't perform well on a machine.
# The FBGEMM library also uses AVX512_256 kernels on Xeon D processors,
# but it also has some (optimized) assembly code.
#
# PYTORCH_BUILD_VERSION
# PYTORCH_BUILD_NUMBER
# specify the version of PyTorch, rather than the hard-coded version
# in this file; used when we're building binaries for distribution
#
# TORCH_CUDA_ARCH_LIST
# specify which CUDA architectures to build for.
# ie `TORCH_CUDA_ARCH_LIST="6.0;7.0"`
# These are not CUDA versions, instead, they specify what
# classes of NVIDIA hardware we should generate PTX for.
#
# PYTORCH_ROCM_ARCH
# specify which AMD GPU targets to build for.
# ie `PYTORCH_ROCM_ARCH="gfx900;gfx906"`
#
# ONNX_NAMESPACE
# specify a namespace for ONNX built here rather than the hard-coded
# one in this file; needed to build with other frameworks that share ONNX.
#
# BLAS
# BLAS to be used by Caffe2. Can be MKL, Eigen, ATLAS, FlexiBLAS, or OpenBLAS. If set
# then the build will fail if the requested BLAS is not found, otherwise
# the BLAS will be chosen based on what is found on your system.
#
# MKL_THREADING
# MKL threading mode: SEQ, TBB or OMP (default)
#
# USE_REDIS
# Whether to use Redis for distributed workflows (Linux only)
#
# USE_ZSTD
# Enables use of ZSTD, if the libraries are found
#
# USE_ROCM_KERNEL_ASSERT=1
# Enable kernel assert in ROCm platform
#
# Environment variables we respect (these environment variables are
# conventional and are often understood/set by other software.)
#
# CUDA_HOME (Linux/OS X)
# CUDA_PATH (Windows)
# specify where CUDA is installed; usually /usr/local/cuda or
# /usr/local/cuda-x.y
# CUDAHOSTCXX
# specify a different compiler than the system one to use as the CUDA
# host compiler for nvcc.
#
# CUDA_NVCC_EXECUTABLE
# Specify a NVCC to use. This is used in our CI to point to a cached nvcc
#
# CUDNN_LIB_DIR
# CUDNN_INCLUDE_DIR
# CUDNN_LIBRARY
# specify where cuDNN is installed
#
# MIOPEN_LIB_DIR
# MIOPEN_INCLUDE_DIR
# MIOPEN_LIBRARY
# specify where MIOpen is installed
#
# NCCL_ROOT
# NCCL_LIB_DIR
# NCCL_INCLUDE_DIR
# specify where nccl is installed
#
# NVTOOLSEXT_PATH (Windows only)
# specify where nvtoolsext is installed
#
# ACL_ROOT_DIR
# specify where Compute Library is installed
#
# LIBRARY_PATH
# LD_LIBRARY_PATH
# we will search for libraries in these paths
#
# ATEN_THREADING
# ATen parallel backend to use for intra- and inter-op parallelism
# possible values:
# OMP - use OpenMP for intra-op and native backend for inter-op tasks
# NATIVE - use native thread pool for both intra- and inter-op tasks
# TBB - using TBB for intra- and native thread pool for inter-op parallelism
#
# USE_TBB
# enable TBB support
#
# USE_SYSTEM_TBB
# Use system-provided Intel TBB.
#
# USE_SYSTEM_LIBS (work in progress)
# Use system-provided libraries to satisfy the build dependencies.
# When turned on, the following cmake variables will be toggled as well:
# USE_SYSTEM_CPUINFO=ON USE_SYSTEM_SLEEF=ON BUILD_CUSTOM_PROTOBUF=OFF
#
# USE_MIMALLOC
# Static link mimalloc into C10, and use mimalloc in alloc_cpu & alloc_free.
# By default, It is only enabled on Windows.
import sys
if sys.platform == "win32" and sys.maxsize.bit_length() == 31:
print(
"32-bit Windows Python runtime is not supported. Please switch to 64-bit Python."
)
sys.exit(-1)
import platform
python_min_version = (3, 8, 0)
python_min_version_str = ".".join(map(str, python_min_version))
if sys.version_info < python_min_version:
print(
f"You are using Python {platform.python_version()}. Python >={python_min_version_str} is required."
)
sys.exit(-1)
import filecmp
import glob
import importlib
import json
import os
import shutil
import subprocess
import sysconfig
import time
from collections import defaultdict
import setuptools.command.build_ext
import setuptools.command.install
import setuptools.command.sdist
from setuptools import Extension, find_packages, setup
from setuptools.dist import Distribution
from tools.build_pytorch_libs import build_caffe2
from tools.generate_torch_version import get_torch_version
from tools.setup_helpers.cmake import CMake
from tools.setup_helpers.env import build_type, IS_DARWIN, IS_LINUX, IS_WINDOWS
################################################################################
# Parameters parsed from environment
################################################################################
VERBOSE_SCRIPT = True
RUN_BUILD_DEPS = True
# see if the user passed a quiet flag to setup.py arguments and respect
# that in our parts of the build
EMIT_BUILD_WARNING = False
RERUN_CMAKE = False
CMAKE_ONLY = False
filtered_args = []
for i, arg in enumerate(sys.argv):
if arg == "--cmake":
RERUN_CMAKE = True
continue
if arg == "--cmake-only":
# Stop once cmake terminates. Leave users a chance to adjust build
# options.
CMAKE_ONLY = True
continue
if arg == "rebuild" or arg == "build":
arg = "build" # rebuild is gone, make it build
EMIT_BUILD_WARNING = True
if arg == "--":
filtered_args += sys.argv[i:]
break
if arg == "-q" or arg == "--quiet":
VERBOSE_SCRIPT = False
if arg in ["clean", "egg_info", "sdist"]:
RUN_BUILD_DEPS = False
filtered_args.append(arg)
sys.argv = filtered_args
if VERBOSE_SCRIPT:
def report(*args):
print(*args)
else:
def report(*args):
pass
# Make distutils respect --quiet too
setuptools.distutils.log.warn = report
# Constant known variables used throughout this file
cwd = os.path.dirname(os.path.abspath(__file__))
lib_path = os.path.join(cwd, "torch", "lib")
third_party_path = os.path.join(cwd, "third_party")
caffe2_build_dir = os.path.join(cwd, "build")
# CMAKE: full path to python library
if IS_WINDOWS:
cmake_python_library = "{}/libs/python{}.lib".format(
sysconfig.get_config_var("prefix"), sysconfig.get_config_var("VERSION")
)
# Fix virtualenv builds
if not os.path.exists(cmake_python_library):
cmake_python_library = "{}/libs/python{}.lib".format(
sys.base_prefix, sysconfig.get_config_var("VERSION")
)
else:
cmake_python_library = "{}/{}".format(
sysconfig.get_config_var("LIBDIR"), sysconfig.get_config_var("INSTSONAME")
)
cmake_python_include_dir = sysconfig.get_path("include")
################################################################################
# Version, create_version_file, and package_name
################################################################################
package_name = os.getenv("TORCH_PACKAGE_NAME", "torch")
package_type = os.getenv("PACKAGE_TYPE", "wheel")
version = get_torch_version()
report(f"Building wheel {package_name}-{version}")
cmake = CMake()
def get_submodule_folders():
git_modules_path = os.path.join(cwd, ".gitmodules")
default_modules_path = [
os.path.join(third_party_path, name)
for name in [
"gloo",
"cpuinfo",
"tbb",
"onnx",
"foxi",
"QNNPACK",
"fbgemm",
"cutlass",
]
]
if not os.path.exists(git_modules_path):
return default_modules_path
with open(git_modules_path) as f:
return [
os.path.join(cwd, line.split("=", 1)[1].strip())
for line in f.readlines()
if line.strip().startswith("path")
]
def check_submodules():
def check_for_files(folder, files):
if not any(os.path.exists(os.path.join(folder, f)) for f in files):
report("Could not find any of {} in {}".format(", ".join(files), folder))
report("Did you run 'git submodule update --init --recursive'?")
sys.exit(1)
def not_exists_or_empty(folder):
return not os.path.exists(folder) or (
os.path.isdir(folder) and len(os.listdir(folder)) == 0
)
if bool(os.getenv("USE_SYSTEM_LIBS", False)):
return
folders = get_submodule_folders()
# If none of the submodule folders exists, try to initialize them
if all(not_exists_or_empty(folder) for folder in folders):
try:
print(" --- Trying to initialize submodules")
start = time.time()
subprocess.check_call(
["git", "submodule", "update", "--init", "--recursive"], cwd=cwd
)
end = time.time()
print(f" --- Submodule initialization took {end - start:.2f} sec")
except Exception:
print(" --- Submodule initalization failed")
print("Please run:\n\tgit submodule update --init --recursive")
sys.exit(1)
for folder in folders:
check_for_files(
folder,
[
"CMakeLists.txt",
"Makefile",
"setup.py",
"LICENSE",
"LICENSE.md",
"LICENSE.txt",
],
)
check_for_files(
os.path.join(third_party_path, "fbgemm", "third_party", "asmjit"),
["CMakeLists.txt"],
)
check_for_files(
os.path.join(third_party_path, "onnx", "third_party", "benchmark"),
["CMakeLists.txt"],
)
# Windows has very bad support for symbolic links.
# Instead of using symlinks, we're going to copy files over
def mirror_files_into_torchgen():
# (new_path, orig_path)
# Directories are OK and are recursively mirrored.
paths = [
(
"torchgen/packaged/ATen/native/native_functions.yaml",
"aten/src/ATen/native/native_functions.yaml",
),
("torchgen/packaged/ATen/native/tags.yaml", "aten/src/ATen/native/tags.yaml"),
("torchgen/packaged/ATen/templates", "aten/src/ATen/templates"),
("torchgen/packaged/autograd", "tools/autograd"),
("torchgen/packaged/autograd/templates", "tools/autograd/templates"),
]
for new_path, orig_path in paths:
# Create the dirs involved in new_path if they don't exist
if not os.path.exists(new_path):
os.makedirs(os.path.dirname(new_path), exist_ok=True)
# Copy the files from the orig location to the new location
if os.path.isfile(orig_path):
shutil.copyfile(orig_path, new_path)
continue
if os.path.isdir(orig_path):
if os.path.exists(new_path):
# copytree fails if the tree exists already, so remove it.
shutil.rmtree(new_path)
shutil.copytree(orig_path, new_path)
continue
raise RuntimeError("Check the file paths in `mirror_files_into_torchgen()`")
# all the work we need to do _before_ setup runs
def build_deps():
report("-- Building version " + version)
check_submodules()
check_pydep("yaml", "pyyaml")
build_caffe2(
version=version,
cmake_python_library=cmake_python_library,
build_python=True,
rerun_cmake=RERUN_CMAKE,
cmake_only=CMAKE_ONLY,
cmake=cmake,
)
if CMAKE_ONLY:
report(
'Finished running cmake. Run "ccmake build" or '
'"cmake-gui build" to adjust build options and '
'"python setup.py install" to build.'
)
sys.exit()
# Use copies instead of symbolic files.
# Windows has very poor support for them.
sym_files = [
"tools/shared/_utils_internal.py",
"torch/utils/benchmark/utils/valgrind_wrapper/callgrind.h",
"torch/utils/benchmark/utils/valgrind_wrapper/valgrind.h",
]
orig_files = [
"torch/_utils_internal.py",
"third_party/valgrind-headers/callgrind.h",
"third_party/valgrind-headers/valgrind.h",
]
for sym_file, orig_file in zip(sym_files, orig_files):
same = False
if os.path.exists(sym_file):
if filecmp.cmp(sym_file, orig_file):
same = True
else:
os.remove(sym_file)
if not same:
shutil.copyfile(orig_file, sym_file)
################################################################################
# Building dependent libraries
################################################################################
missing_pydep = """
Missing build dependency: Unable to `import {importname}`.
Please install it via `conda install {module}` or `pip install {module}`
""".strip()
def check_pydep(importname, module):
try:
importlib.import_module(importname)
except ImportError as e:
raise RuntimeError(
missing_pydep.format(importname=importname, module=module)
) from e
class build_ext(setuptools.command.build_ext.build_ext):
# Copy libiomp5.dylib inside the wheel package on OS X
def _embed_libiomp(self):
lib_dir = os.path.join(self.build_lib, "torch", "lib")
libtorch_cpu_path = os.path.join(lib_dir, "libtorch_cpu.dylib")
if not os.path.exists(libtorch_cpu_path):
return
# Parse libtorch_cpu load commands
otool_cmds = (
subprocess.check_output(["otool", "-l", libtorch_cpu_path])
.decode("utf-8")
.split("\n")
)
rpaths, libs = [], []
for idx, line in enumerate(otool_cmds):
if line.strip() == "cmd LC_LOAD_DYLIB":
lib_name = otool_cmds[idx + 2].strip()
assert lib_name.startswith("name ")
libs.append(lib_name.split(" ", 1)[1].rsplit("(", 1)[0][:-1])
if line.strip() == "cmd LC_RPATH":
rpath = otool_cmds[idx + 2].strip()
assert rpath.startswith("path ")
rpaths.append(rpath.split(" ", 1)[1].rsplit("(", 1)[0][:-1])
omp_lib_name = "libiomp5.dylib"
if os.path.join("@rpath", omp_lib_name) not in libs:
return
# Copy libiomp5 from rpath locations
for rpath in rpaths:
source_lib = os.path.join(rpath, omp_lib_name)
if not os.path.exists(source_lib):
continue
target_lib = os.path.join(self.build_lib, "torch", "lib", omp_lib_name)
self.copy_file(source_lib, target_lib)
break
def run(self):
# Report build options. This is run after the build completes so # `CMakeCache.txt` exists and we can get an
# accurate report on what is used and what is not.
cmake_cache_vars = defaultdict(lambda: False, cmake.get_cmake_cache_variables())
if cmake_cache_vars["USE_NUMPY"]:
report("-- Building with NumPy bindings")
else:
report("-- NumPy not found")
if cmake_cache_vars["USE_CUDNN"]:
report(
"-- Detected cuDNN at "
+ cmake_cache_vars["CUDNN_LIBRARY"]
+ ", "
+ cmake_cache_vars["CUDNN_INCLUDE_DIR"]
)
else:
report("-- Not using cuDNN")
if cmake_cache_vars["USE_CUDA"]:
report("-- Detected CUDA at " + cmake_cache_vars["CUDA_TOOLKIT_ROOT_DIR"])
else:
report("-- Not using CUDA")
if cmake_cache_vars["USE_MKLDNN"]:
report("-- Using MKLDNN")
if cmake_cache_vars["USE_MKLDNN_ACL"]:
report("-- Using Compute Library for the Arm architecture with MKLDNN")
else:
report(
"-- Not using Compute Library for the Arm architecture with MKLDNN"
)
if cmake_cache_vars["USE_MKLDNN_CBLAS"]:
report("-- Using CBLAS in MKLDNN")
else:
report("-- Not using CBLAS in MKLDNN")
else:
report("-- Not using MKLDNN")
if cmake_cache_vars["USE_NCCL"] and cmake_cache_vars["USE_SYSTEM_NCCL"]:
report(
"-- Using system provided NCCL library at {}, {}".format(
cmake_cache_vars["NCCL_LIBRARIES"],
cmake_cache_vars["NCCL_INCLUDE_DIRS"],
)
)
elif cmake_cache_vars["USE_NCCL"]:
report("-- Building NCCL library")
else:
report("-- Not using NCCL")
if cmake_cache_vars["USE_DISTRIBUTED"]:
if IS_WINDOWS:
report("-- Building without distributed package")
else:
report("-- Building with distributed package: ")
report(
" -- USE_TENSORPIPE={}".format(cmake_cache_vars["USE_TENSORPIPE"])
)
report(" -- USE_GLOO={}".format(cmake_cache_vars["USE_GLOO"]))
report(" -- USE_MPI={}".format(cmake_cache_vars["USE_OPENMPI"]))
else:
report("-- Building without distributed package")
if cmake_cache_vars["STATIC_DISPATCH_BACKEND"]:
report(
"-- Using static dispatch with backend {}".format(
cmake_cache_vars["STATIC_DISPATCH_BACKEND"]
)
)
if cmake_cache_vars["USE_LIGHTWEIGHT_DISPATCH"]:
report("-- Using lightweight dispatch")
if cmake_cache_vars["BUILD_EXECUTORCH"]:
report("-- Building Executorch")
if cmake_cache_vars["USE_ITT"]:
report("-- Using ITT")
else:
report("-- Not using ITT")
# Do not use clang to compile extensions if `-fstack-clash-protection` is defined
# in system CFLAGS
c_flags = str(os.getenv("CFLAGS", ""))
if (
IS_LINUX
and "-fstack-clash-protection" in c_flags
and "clang" in os.environ.get("CC", "")
):
os.environ["CC"] = str(os.environ["CC"])
# It's an old-style class in Python 2.7...
setuptools.command.build_ext.build_ext.run(self)
if IS_DARWIN and package_type != "conda":
self._embed_libiomp()
# Copy the essential export library to compile C++ extensions.
if IS_WINDOWS:
build_temp = self.build_temp
ext_filename = self.get_ext_filename("_C")
lib_filename = ".".join(ext_filename.split(".")[:-1]) + ".lib"
export_lib = os.path.join(
build_temp, "torch", "csrc", lib_filename
).replace("\\", "/")
build_lib = self.build_lib
target_lib = os.path.join(build_lib, "torch", "lib", "_C.lib").replace(
"\\", "/"
)
# Create "torch/lib" directory if not exists.
# (It is not created yet in "develop" mode.)
target_dir = os.path.dirname(target_lib)
if not os.path.exists(target_dir):
os.makedirs(target_dir)
self.copy_file(export_lib, target_lib)
def build_extensions(self):
self.create_compile_commands()
# The caffe2 extensions are created in
# tmp_install/lib/pythonM.m/site-packages/caffe2/python/
# and need to be copied to build/lib.linux.... , which will be a
# platform dependent build folder created by the "build" command of
# setuptools. Only the contents of this folder are installed in the
# "install" command by default.
# We only make this copy for Caffe2's pybind extensions
caffe2_pybind_exts = [
"caffe2.python.caffe2_pybind11_state",
"caffe2.python.caffe2_pybind11_state_gpu",
"caffe2.python.caffe2_pybind11_state_hip",
]
i = 0
while i < len(self.extensions):
ext = self.extensions[i]
if ext.name not in caffe2_pybind_exts:
i += 1
continue
fullname = self.get_ext_fullname(ext.name)
filename = self.get_ext_filename(fullname)
report(f"\nCopying extension {ext.name}")
relative_site_packages = (
sysconfig.get_path("purelib")
.replace(sysconfig.get_path("data"), "")
.lstrip(os.path.sep)
)
src = os.path.join("torch", relative_site_packages, filename)
if not os.path.exists(src):
report(f"{src} does not exist")
del self.extensions[i]
else:
dst = os.path.join(os.path.realpath(self.build_lib), filename)
report(f"Copying {ext.name} from {src} to {dst}")
dst_dir = os.path.dirname(dst)
if not os.path.exists(dst_dir):
os.makedirs(dst_dir)
self.copy_file(src, dst)
i += 1
# Copy functorch extension
for i, ext in enumerate(self.extensions):
if ext.name != "functorch._C":
continue
fullname = self.get_ext_fullname(ext.name)
filename = self.get_ext_filename(fullname)
fileext = os.path.splitext(filename)[1]
src = os.path.join(os.path.dirname(filename), "functorch" + fileext)
dst = os.path.join(os.path.realpath(self.build_lib), filename)
if os.path.exists(src):
report(f"Copying {ext.name} from {src} to {dst}")
dst_dir = os.path.dirname(dst)
if not os.path.exists(dst_dir):
os.makedirs(dst_dir)
self.copy_file(src, dst)
setuptools.command.build_ext.build_ext.build_extensions(self)
def get_outputs(self):
outputs = setuptools.command.build_ext.build_ext.get_outputs(self)
outputs.append(os.path.join(self.build_lib, "caffe2"))
report(f"setup.py::get_outputs returning {outputs}")
return outputs
def create_compile_commands(self):
def load(filename):
with open(filename) as f:
return json.load(f)
ninja_files = glob.glob("build/*compile_commands.json")
cmake_files = glob.glob("torch/lib/build/*/compile_commands.json")
all_commands = [entry for f in ninja_files + cmake_files for entry in load(f)]
# cquery does not like c++ compiles that start with gcc.
# It forgets to include the c++ header directories.
# We can work around this by replacing the gcc calls that python
# setup.py generates with g++ calls instead
for command in all_commands:
if command["command"].startswith("gcc "):
command["command"] = "g++ " + command["command"][4:]
new_contents = json.dumps(all_commands, indent=2)
contents = ""
if os.path.exists("compile_commands.json"):
with open("compile_commands.json") as f:
contents = f.read()
if contents != new_contents:
with open("compile_commands.json", "w") as f:
f.write(new_contents)
class concat_license_files:
"""Merge LICENSE and LICENSES_BUNDLED.txt as a context manager
LICENSE is the main PyTorch license, LICENSES_BUNDLED.txt is auto-generated
from all the licenses found in ./third_party/. We concatenate them so there
is a single license file in the sdist and wheels with all of the necessary
licensing info.
"""
def __init__(self, include_files=False):
self.f1 = "LICENSE"
self.f2 = "third_party/LICENSES_BUNDLED.txt"
self.include_files = include_files
def __enter__(self):
"""Concatenate files"""
old_path = sys.path
sys.path.append(third_party_path)
try:
from build_bundled import create_bundled
finally:
sys.path = old_path
with open(self.f1) as f1:
self.bsd_text = f1.read()
with open(self.f1, "a") as f1:
f1.write("\n\n")
create_bundled(
os.path.relpath(third_party_path), f1, include_files=self.include_files
)
def __exit__(self, exception_type, exception_value, traceback):
"""Restore content of f1"""
with open(self.f1, "w") as f:
f.write(self.bsd_text)
try:
from wheel.bdist_wheel import bdist_wheel
except ImportError:
# This is useful when wheel is not installed and bdist_wheel is not
# specified on the command line. If it _is_ specified, parsing the command
# line will fail before wheel_concatenate is needed
wheel_concatenate = None
else:
# Need to create the proper LICENSE.txt for the wheel
class wheel_concatenate(bdist_wheel):
"""check submodules on sdist to prevent incomplete tarballs"""
def run(self):
with concat_license_files(include_files=True):
super().run()
class install(setuptools.command.install.install):
def run(self):
super().run()
class clean(setuptools.Command):
user_options = []
def initialize_options(self):
pass
def finalize_options(self):
pass
def run(self):
import glob
import re
with open(".gitignore") as f:
ignores = f.read()
pat = re.compile(r"^#( BEGIN NOT-CLEAN-FILES )?")
for wildcard in filter(None, ignores.split("\n")):
match = pat.match(wildcard)
if match:
if match.group(1):
# Marker is found and stop reading .gitignore.
break
# Ignore lines which begin with '#'.
else:
# Don't remove absolute paths from the system
wildcard = wildcard.lstrip("./")
for filename in glob.glob(wildcard):
try:
os.remove(filename)
except OSError:
shutil.rmtree(filename, ignore_errors=True)
class sdist(setuptools.command.sdist.sdist):
def run(self):
with concat_license_files():
super().run()
def get_cmake_cache_vars():
try:
return defaultdict(lambda: False, cmake.get_cmake_cache_variables())
except FileNotFoundError:
# CMakeCache.txt does not exist. Probably running "python setup.py clean" over a clean directory.
return defaultdict(lambda: False)
def configure_extension_build():
r"""Configures extension build options according to system environment and user's choice.
Returns:
The input to parameters ext_modules, cmdclass, packages, and entry_points as required in setuptools.setup.
"""
cmake_cache_vars = get_cmake_cache_vars()
################################################################################
# Configure compile flags
################################################################################
library_dirs = []
extra_install_requires = []
if IS_WINDOWS:
# /NODEFAULTLIB makes sure we only link to DLL runtime
# and matches the flags set for protobuf and ONNX
extra_link_args = ["/NODEFAULTLIB:LIBCMT.LIB"]
# /MD links against DLL runtime
# and matches the flags set for protobuf and ONNX
# /EHsc is about standard C++ exception handling
extra_compile_args = ["/MD", "/FS", "/EHsc"]
else:
extra_link_args = []
extra_compile_args = [
"-Wall",
"-Wextra",
"-Wno-strict-overflow",
"-Wno-unused-parameter",
"-Wno-missing-field-initializers",
"-Wno-unknown-pragmas",
# Python 2.6 requires -fno-strict-aliasing, see
# http://legacy.python.org/dev/peps/pep-3123/
# We also depend on it in our code (even Python 3).
"-fno-strict-aliasing",
]
library_dirs.append(lib_path)
main_compile_args = []
main_libraries = ["torch_python"]
main_link_args = []
main_sources = ["torch/csrc/stub.c"]
if cmake_cache_vars["USE_CUDA"]:
library_dirs.append(os.path.dirname(cmake_cache_vars["CUDA_CUDA_LIB"]))
if build_type.is_debug():
if IS_WINDOWS:
extra_compile_args.append("/Z7")
extra_link_args.append("/DEBUG:FULL")
else:
extra_compile_args += ["-O0", "-g"]
extra_link_args += ["-O0", "-g"]
if build_type.is_rel_with_deb_info():
if IS_WINDOWS:
extra_compile_args.append("/Z7")
extra_link_args.append("/DEBUG:FULL")
else:
extra_compile_args += ["-g"]
extra_link_args += ["-g"]
# pypi cuda package that requires installation of cuda runtime, cudnn and cublas
# should be included in all wheels uploaded to pypi
pytorch_extra_install_requirements = os.getenv(
"PYTORCH_EXTRA_INSTALL_REQUIREMENTS", ""
)
if pytorch_extra_install_requirements:
report(
f"pytorch_extra_install_requirements: {pytorch_extra_install_requirements}"
)
extra_install_requires += pytorch_extra_install_requirements.split("|")
# Cross-compile for M1
if IS_DARWIN:
macos_target_arch = os.getenv("CMAKE_OSX_ARCHITECTURES", "")
if macos_target_arch in ["arm64", "x86_64"]:
macos_sysroot_path = os.getenv("CMAKE_OSX_SYSROOT")
if macos_sysroot_path is None:
macos_sysroot_path = (
subprocess.check_output(
["xcrun", "--show-sdk-path", "--sdk", "macosx"]
)
.decode("utf-8")
.strip()
)
extra_compile_args += [
"-arch",
macos_target_arch,
"-isysroot",
macos_sysroot_path,
]
extra_link_args += ["-arch", macos_target_arch]
def make_relative_rpath_args(path):
if IS_DARWIN:
return ["-Wl,-rpath,@loader_path/" + path]
elif IS_WINDOWS:
return []
else:
return ["-Wl,-rpath,$ORIGIN/" + path]
################################################################################
# Declare extensions and package
################################################################################
extensions = []
excludes = ["tools", "tools.*"]
if not cmake_cache_vars["BUILD_CAFFE2"]:
excludes.extend(["caffe2", "caffe2.*"])
if not cmake_cache_vars["BUILD_FUNCTORCH"]:
excludes.extend(["functorch", "functorch.*"])
packages = find_packages(exclude=excludes)
C = Extension(
"torch._C",
libraries=main_libraries,
sources=main_sources,
language="c",
extra_compile_args=main_compile_args + extra_compile_args,
include_dirs=[],
library_dirs=library_dirs,
extra_link_args=extra_link_args
+ main_link_args
+ make_relative_rpath_args("lib"),
)
extensions.append(C)
# These extensions are built by cmake and copied manually in build_extensions()
# inside the build_ext implementation
if cmake_cache_vars["BUILD_CAFFE2"]:
extensions.append(
Extension(name="caffe2.python.caffe2_pybind11_state", sources=[]),
)
if cmake_cache_vars["USE_CUDA"]:
extensions.append(
Extension(name="caffe2.python.caffe2_pybind11_state_gpu", sources=[]),
)
if cmake_cache_vars["USE_ROCM"]:
extensions.append(
Extension(name="caffe2.python.caffe2_pybind11_state_hip", sources=[]),
)
if cmake_cache_vars["BUILD_FUNCTORCH"]:
extensions.append(
Extension(name="functorch._C", sources=[]),
)
cmdclass = {
"bdist_wheel": wheel_concatenate,
"build_ext": build_ext,
"clean": clean,
"install": install,
"sdist": sdist,
}
entry_points = {
"console_scripts": [
"convert-caffe2-to-onnx = caffe2.python.onnx.bin.conversion:caffe2_to_onnx",
"convert-onnx-to-caffe2 = caffe2.python.onnx.bin.conversion:onnx_to_caffe2",
"torchrun = torch.distributed.run:main",
]
}
return extensions, cmdclass, packages, entry_points, extra_install_requires
# post run, warnings, printed at the end to make them more visible
build_update_message = """
It is no longer necessary to use the 'build' or 'rebuild' targets
To install:
$ python setup.py install
To develop locally:
$ python setup.py develop
To force cmake to re-generate native build files (off by default):
$ python setup.py develop --cmake
"""
def print_box(msg):
lines = msg.split("\n")
size = max(len(l) + 1 for l in lines)
print("-" * (size + 2))
for l in lines:
print("|{}{}|".format(l, " " * (size - len(l))))
print("-" * (size + 2))
def main():
# the list of runtime dependencies required by this built package
install_requires = [
"filelock",
"typing-extensions>=4.8.0",
"sympy",
"networkx",
"jinja2",
"fsspec",
]
# Parse the command line and check the arguments before we proceed with
# building deps and setup. We need to set values so `--help` works.
dist = Distribution()
dist.script_name = os.path.basename(sys.argv[0])
dist.script_args = sys.argv[1:]
try:
dist.parse_command_line()
except setuptools.distutils.errors.DistutilsArgError as e:
print(e)
sys.exit(1)
mirror_files_into_torchgen()
if RUN_BUILD_DEPS:
build_deps()
(
extensions,
cmdclass,
packages,
entry_points,
extra_install_requires,
) = configure_extension_build()
install_requires += extra_install_requires
extras_require = {
"optree": ["optree>=0.9.1"],
"opt-einsum": ["opt-einsum>=3.3"],
}
# Read in README.md for our long_description
with open(os.path.join(cwd, "README.md"), encoding="utf-8") as f:
long_description = f.read()
version_range_max = max(sys.version_info[1], 10) + 1
torch_package_data = [
"py.typed",
"bin/*",
"test/*",
"*.pyi",
"_C/*.pyi",
"cuda/*.pyi",
"fx/*.pyi",
"optim/*.pyi",
"autograd/*.pyi",
"nn/*.pyi",
"nn/modules/*.pyi",
"nn/parallel/*.pyi",
"utils/data/*.pyi",
"utils/data/datapipes/*.pyi",
"lib/*.so*",
"lib/*.dylib*",
"lib/*.dll",
"lib/*.lib",
"lib/*.pdb",
"lib/torch_shm_manager",
"lib/*.h",
"include/*.h",
"include/ATen/*.h",
"include/ATen/cpu/*.h",
"include/ATen/cpu/vec/vec256/*.h",
"include/ATen/cpu/vec/vec256/vsx/*.h",
"include/ATen/cpu/vec/vec256/zarch/*.h",
"include/ATen/cpu/vec/vec512/*.h",
"include/ATen/cpu/vec/*.h",
"include/ATen/core/*.h",
"include/ATen/cuda/*.cuh",
"include/ATen/cuda/*.h",
"include/ATen/cuda/detail/*.cuh",
"include/ATen/cuda/detail/*.h",
"include/ATen/cudnn/*.h",
"include/ATen/functorch/*.h",
"include/ATen/ops/*.h",
"include/ATen/hip/*.cuh",
"include/ATen/hip/*.h",
"include/ATen/hip/detail/*.cuh",
"include/ATen/hip/detail/*.h",
"include/ATen/hip/impl/*.h",
"include/ATen/mps/*.h",
"include/ATen/miopen/*.h",
"include/ATen/detail/*.h",
"include/ATen/native/*.h",
"include/ATen/native/cpu/*.h",
"include/ATen/native/cuda/*.h",
"include/ATen/native/cuda/*.cuh",
"include/ATen/native/hip/*.h",
"include/ATen/native/hip/*.cuh",
"include/ATen/native/mps/*.h",
"include/ATen/native/quantized/*.h",
"include/ATen/native/quantized/cpu/*.h",
"include/ATen/native/utils/*.h",
"include/ATen/quantized/*.h",
"include/caffe2/serialize/*.h",
"include/c10/*.h",
"include/c10/macros/*.h",
"include/c10/core/*.h",
"include/ATen/core/boxing/*.h",
"include/ATen/core/boxing/impl/*.h",
"include/ATen/core/dispatch/*.h",
"include/ATen/core/op_registration/*.h",
"include/c10/core/impl/*.h",
"include/c10/core/impl/cow/*.h",
"include/c10/util/*.h",
"include/c10/cuda/*.h",
"include/c10/cuda/impl/*.h",
"include/c10/hip/*.h",
"include/c10/hip/impl/*.h",
"include/torch/*.h",
"include/torch/csrc/*.h",
"include/torch/csrc/api/include/torch/*.h",
"include/torch/csrc/api/include/torch/data/*.h",
"include/torch/csrc/api/include/torch/data/dataloader/*.h",
"include/torch/csrc/api/include/torch/data/datasets/*.h",
"include/torch/csrc/api/include/torch/data/detail/*.h",
"include/torch/csrc/api/include/torch/data/samplers/*.h",
"include/torch/csrc/api/include/torch/data/transforms/*.h",
"include/torch/csrc/api/include/torch/detail/*.h",
"include/torch/csrc/api/include/torch/detail/ordered_dict.h",
"include/torch/csrc/api/include/torch/nn/*.h",
"include/torch/csrc/api/include/torch/nn/functional/*.h",
"include/torch/csrc/api/include/torch/nn/options/*.h",
"include/torch/csrc/api/include/torch/nn/modules/*.h",
"include/torch/csrc/api/include/torch/nn/modules/container/*.h",
"include/torch/csrc/api/include/torch/nn/parallel/*.h",
"include/torch/csrc/api/include/torch/nn/utils/*.h",
"include/torch/csrc/api/include/torch/optim/*.h",
"include/torch/csrc/api/include/torch/optim/schedulers/*.h",
"include/torch/csrc/api/include/torch/serialize/*.h",
"include/torch/csrc/autograd/*.h",
"include/torch/csrc/autograd/functions/*.h",
"include/torch/csrc/autograd/generated/*.h",
"include/torch/csrc/autograd/utils/*.h",
"include/torch/csrc/cuda/*.h",
"include/torch/csrc/distributed/c10d/*.h",
"include/torch/csrc/distributed/c10d/*.hpp",
"include/torch/csrc/distributed/rpc/*.h",
"include/torch/csrc/distributed/autograd/context/*.h",
"include/torch/csrc/distributed/autograd/functions/*.h",
"include/torch/csrc/distributed/autograd/rpc_messages/*.h",
"include/torch/csrc/dynamo/*.h",
"include/torch/csrc/inductor/*.h",
"include/torch/csrc/inductor/aoti_runtime/*.h",
"include/torch/csrc/inductor/aoti_torch/*.h",
"include/torch/csrc/inductor/aoti_torch/c/*.h",
"include/torch/csrc/jit/*.h",
"include/torch/csrc/jit/backends/*.h",
"include/torch/csrc/jit/generated/*.h",
"include/torch/csrc/jit/passes/*.h",
"include/torch/csrc/jit/passes/quantization/*.h",
"include/torch/csrc/jit/passes/utils/*.h",
"include/torch/csrc/jit/runtime/*.h",
"include/torch/csrc/jit/ir/*.h",
"include/torch/csrc/jit/frontend/*.h",
"include/torch/csrc/jit/api/*.h",
"include/torch/csrc/jit/serialization/*.h",
"include/torch/csrc/jit/python/*.h",
"include/torch/csrc/jit/mobile/*.h",
"include/torch/csrc/jit/testing/*.h",
"include/torch/csrc/jit/tensorexpr/*.h",
"include/torch/csrc/jit/tensorexpr/operators/*.h",
"include/torch/csrc/jit/codegen/cuda/*.h",
"include/torch/csrc/onnx/*.h",
"include/torch/csrc/profiler/*.h",
"include/torch/csrc/profiler/orchestration/*.h",
"include/torch/csrc/profiler/stubs/*.h",
"include/torch/csrc/profiler/unwind/*.h",
"include/torch/csrc/utils/*.h",
"include/torch/csrc/tensor/*.h",
"include/torch/csrc/lazy/backend/*.h",
"include/torch/csrc/lazy/core/*.h",
"include/torch/csrc/lazy/core/internal_ops/*.h",
"include/torch/csrc/lazy/core/ops/*.h",
"include/torch/csrc/lazy/python/python_util.h",
"include/torch/csrc/lazy/ts_backend/*.h",
"include/pybind11/*.h",
"include/pybind11/detail/*.h",
"include/pybind11/eigen/*.h",
"include/TH/*.h*",
"include/TH/generic/*.h*",
"include/THC/*.cuh",
"include/THC/*.h*",
"include/THC/generic/*.h",
"include/THH/*.cuh",
"include/THH/*.h*",
"include/THH/generic/*.h",
"include/sleef.h",
"_inductor/codegen/*.h",
"_inductor/codegen/aoti_runtime/*.cpp",
"share/cmake/ATen/*.cmake",
"share/cmake/Caffe2/*.cmake",
"share/cmake/Caffe2/public/*.cmake",
"share/cmake/Caffe2/Modules_CUDA_fix/*.cmake",
"share/cmake/Caffe2/Modules_CUDA_fix/upstream/*.cmake",
"share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/*.cmake",
"share/cmake/Gloo/*.cmake",
"share/cmake/Tensorpipe/*.cmake",
"share/cmake/Torch/*.cmake",
"utils/benchmark/utils/*.cpp",
"utils/benchmark/utils/valgrind_wrapper/*.cpp",
"utils/benchmark/utils/valgrind_wrapper/*.h",
"utils/model_dump/skeleton.html",
"utils/model_dump/code.js",
"utils/model_dump/*.mjs",
]
if get_cmake_cache_vars()["BUILD_CAFFE2"]:
torch_package_data.extend(
[
"include/caffe2/**/*.h",
"include/caffe2/utils/*.h",
"include/caffe2/utils/**/*.h",
]
)
if get_cmake_cache_vars()["USE_TENSORPIPE"]:
torch_package_data.extend(
[
"include/tensorpipe/*.h",
"include/tensorpipe/channel/*.h",
"include/tensorpipe/channel/basic/*.h",
"include/tensorpipe/channel/cma/*.h",
"include/tensorpipe/channel/mpt/*.h",
"include/tensorpipe/channel/xth/*.h",
"include/tensorpipe/common/*.h",
"include/tensorpipe/core/*.h",
"include/tensorpipe/transport/*.h",
"include/tensorpipe/transport/ibv/*.h",
"include/tensorpipe/transport/shm/*.h",
"include/tensorpipe/transport/uv/*.h",
]
)
torchgen_package_data = [
# Recursive glob doesn't work in setup.py,
# https://github.com/pypa/setuptools/issues/1806
# To make this robust we should replace it with some code that
# returns a list of everything under packaged/
"packaged/ATen/*",
"packaged/ATen/native/*",
"packaged/ATen/templates/*",
"packaged/autograd/*",
"packaged/autograd/templates/*",
]
setup(
name=package_name,
version=version,
description=(
"Tensors and Dynamic neural networks in "
"Python with strong GPU acceleration"
),
long_description=long_description,
long_description_content_type="text/markdown",
ext_modules=extensions,
cmdclass=cmdclass,
packages=packages,
entry_points=entry_points,
install_requires=install_requires,
extras_require=extras_require,
package_data={
"torch": torch_package_data,
"torchgen": torchgen_package_data,
"caffe2": [
"python/serialized_test/data/operator_test/*.zip",
],
},
url="https://pytorch.org/",
download_url="https://github.com/pytorch/pytorch/tags",
author="PyTorch Team",
author_email="packages@pytorch.org",
python_requires=f">={python_min_version_str}",
# PyPI package information.
classifiers=[
"Development Status :: 5 - Production/Stable",
"Intended Audience :: Developers",
"Intended Audience :: Education",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: BSD License",
"Topic :: Scientific/Engineering",
"Topic :: Scientific/Engineering :: Mathematics",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Software Development",
"Topic :: Software Development :: Libraries",
"Topic :: Software Development :: Libraries :: Python Modules",
"Programming Language :: C++",
"Programming Language :: Python :: 3",
]
+ [
f"Programming Language :: Python :: 3.{i}"
for i in range(python_min_version[1], version_range_max)
],
license="BSD-3",
keywords="pytorch, machine learning",
)
if EMIT_BUILD_WARNING:
print_box(build_update_message)
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink