|
|
|
|
|
|
| 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},
|
| )
|
|
|
