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
|
import distutils.command.clean
import shutil
from pathlib import Path
from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
PACKAGE_NAME = "pytorch_openreg"
version = 1.0
ROOT_DIR = Path(__file__).absolute().parent
CSRS_DIR = ROOT_DIR / "pytorch_openreg/csrc"
class clean(distutils.command.clean.clean):
def run(self):
# Run default behavior first
distutils.command.clean.clean.run(self)
# Remove pytorch_openreg extension
for path in (ROOT_DIR / "pytorch_openreg").glob("**/*.so"):
path.unlink()
# Remove build directory
build_dirs = [
ROOT_DIR / "build",
]
for path in build_dirs:
if path.exists():
shutil.rmtree(str(path), ignore_errors=True)
if __name__ == "__main__":
sources = list(CSRS_DIR.glob("*.cpp"))
# Note that we always compile with debug info
ext_modules = [
CppExtension(
name="pytorch_openreg._C",
sources=sorted(str(s) for s in sources),
include_dirs=[CSRS_DIR],
extra_compile_args={"cxx": ["-g", "-Wall", "-Werror"]},
)
]
setup(
name=PACKAGE_NAME,
version=version,
author="PyTorch Core Team",
description="Example for PyTorch out of tree registration",
packages=find_packages(exclude=("test",)),
package_data={PACKAGE_NAME: ["*.dll", "*.dylib", "*.so"]},
install_requires=[
"torch",
],
ext_modules=ext_modules,
python_requires=">=3.8",
cmdclass={
"build_ext": BuildExtension.with_options(no_python_abi_suffix=True),
"clean": clean,
},
)
|