File size: 5,473 Bytes
ef296aa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
#!/usr/bin/env python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import glob
import os
import runpy
import sys
import warnings
from typing import List, Optional
import torch
from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDA_HOME, CUDAExtension
def get_existing_ccbin(nvcc_args: List[str]) -> Optional[str]:
"""
Given a list of nvcc arguments, return the compiler if specified.
Note from CUDA doc: Single value options and list options must have
arguments, which must follow the name of the option itself by either
one of more spaces or an equals character.
"""
last_arg = None
for arg in reversed(nvcc_args):
if arg == "-ccbin":
return last_arg
if arg.startswith("-ccbin="):
return arg[7:]
last_arg = arg
return None
def get_extensions():
no_extension = os.getenv("PYTORCH3D_NO_EXTENSION", "0") == "1"
if no_extension:
msg = "SKIPPING EXTENSION BUILD. PYTORCH3D WILL NOT WORK!"
print(msg, file=sys.stderr)
warnings.warn(msg)
return []
this_dir = os.path.dirname(os.path.abspath(__file__))
extensions_dir = os.path.join(this_dir, "csrc")
sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp"), recursive=True)
source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu"), recursive=True)
extension = CppExtension
extra_compile_args = {"cxx": ["-std=c++14"]}
define_macros = []
include_dirs = [extensions_dir]
force_cuda = os.getenv("FORCE_CUDA", "0") == "1"
force_no_cuda = os.getenv("PYTORCH3D_FORCE_NO_CUDA", "0") == "1"
if (
not force_no_cuda and torch.cuda.is_available() and CUDA_HOME is not None
) or force_cuda:
extension = CUDAExtension
sources += source_cuda
define_macros += [("WITH_CUDA", None)]
# Thrust is only used for its tuple objects.
# With CUDA 11.0 we can't use the cudatoolkit's version of cub.
# We take the risk that CUB and Thrust are incompatible, because
# we aren't using parts of Thrust which actually use CUB.
define_macros += [("THRUST_IGNORE_CUB_VERSION_CHECK", None)]
cub_home = os.environ.get("CUB_HOME", None)
nvcc_args = [
"-DCUDA_HAS_FP16=1",
"-D__CUDA_NO_HALF_OPERATORS__",
"-D__CUDA_NO_HALF_CONVERSIONS__",
"-D__CUDA_NO_HALF2_OPERATORS__",
]
if os.name != "nt":
nvcc_args.append("-std=c++14")
if cub_home is None:
prefix = os.environ.get("CONDA_PREFIX", None)
if prefix is not None and os.path.isdir(prefix + "/include/cub"):
cub_home = prefix + "/include"
if cub_home is None:
warnings.warn(
"The environment variable `CUB_HOME` was not found. "
"NVIDIA CUB is required for compilation and can be downloaded "
"from `https://github.com/NVIDIA/cub/releases`. You can unpack "
"it to a location of your choice and set the environment variable "
"`CUB_HOME` to the folder containing the `CMakeListst.txt` file."
)
else:
include_dirs.append(os.path.realpath(cub_home).replace("\\ ", " "))
nvcc_flags_env = os.getenv("NVCC_FLAGS", "")
if nvcc_flags_env != "":
nvcc_args.extend(nvcc_flags_env.split(" "))
# This is needed for pytorch 1.6 and earlier. See e.g.
# https://github.com/facebookresearch/pytorch3d/issues/436
# It is harmless after https://github.com/pytorch/pytorch/pull/47404 .
# But it can be problematic in torch 1.7.0 and 1.7.1
if torch.__version__[:4] != "1.7.":
CC = os.environ.get("CC", None)
if CC is not None:
existing_CC = get_existing_ccbin(nvcc_args)
if existing_CC is None:
CC_arg = "-ccbin={}".format(CC)
nvcc_args.append(CC_arg)
elif existing_CC != CC:
msg = f"Inconsistent ccbins: {CC} and {existing_CC}"
raise ValueError(msg)
extra_compile_args["nvcc"] = nvcc_args
ext_modules = [
extension(
"_torch3d._C",
sources,
include_dirs=include_dirs,
define_macros=define_macros,
extra_compile_args=extra_compile_args,
)
]
return ext_modules
'''
Usage:
python setup.py build_ext --inplace # build extensions locally, do not install (only can be used from the parent directory)
python setup.py install # build extensions and install (copy) to PATH.
pip install . # ditto but better (e.g., dependency & metadata handling)
python setup.py develop # build extensions and install (symbolic) to PATH.
pip install -e . # ditto but better (e.g., dependency & metadata handling)
'''
setup(
name='torch3d', # package name, import this to use python API
ext_modules=get_extensions(),
packages=find_packages(),
cmdclass={
'build_ext': BuildExtension,
}
) |