| | |
| | |
| |
|
| | import glob |
| | import os |
| | from setuptools import find_packages, setup |
| | import torch |
| | from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension |
| |
|
| |
|
| | def get_extensions(): |
| | this_dir = os.path.dirname(os.path.abspath(__file__)) |
| | extensions_dir = os.path.join(this_dir, "tensormask", "layers", "csrc") |
| |
|
| | main_source = os.path.join(extensions_dir, "vision.cpp") |
| | sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp")) |
| | source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu")) + glob.glob( |
| | os.path.join(extensions_dir, "*.cu") |
| | ) |
| |
|
| | sources = [main_source] + sources |
| |
|
| | extension = CppExtension |
| |
|
| | extra_compile_args = {"cxx": []} |
| | define_macros = [] |
| |
|
| | if (torch.cuda.is_available() and CUDA_HOME is not None) or os.getenv("FORCE_CUDA", "0") == "1": |
| | extension = CUDAExtension |
| | sources += source_cuda |
| | define_macros += [("WITH_CUDA", None)] |
| | extra_compile_args["nvcc"] = [ |
| | "-DCUDA_HAS_FP16=1", |
| | "-D__CUDA_NO_HALF_OPERATORS__", |
| | "-D__CUDA_NO_HALF_CONVERSIONS__", |
| | "-D__CUDA_NO_HALF2_OPERATORS__", |
| | ] |
| |
|
| | |
| | CC = os.environ.get("CC", None) |
| | if CC is not None: |
| | extra_compile_args["nvcc"].append("-ccbin={}".format(CC)) |
| |
|
| | sources = [os.path.join(extensions_dir, s) for s in sources] |
| |
|
| | include_dirs = [extensions_dir] |
| |
|
| | ext_modules = [ |
| | extension( |
| | "tensormask._C", |
| | sources, |
| | include_dirs=include_dirs, |
| | define_macros=define_macros, |
| | extra_compile_args=extra_compile_args, |
| | ) |
| | ] |
| |
|
| | return ext_modules |
| |
|
| |
|
| | setup( |
| | name="tensormask", |
| | version="0.1", |
| | author="FAIR", |
| | packages=find_packages(exclude=("configs", "tests")), |
| | python_requires=">=3.6", |
| | ext_modules=get_extensions(), |
| | cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}, |
| | ) |
| |
|
