| # To install, run | |
| # python setup.py build_ext -i | |
| # Ref: https://github.com/pytorch/pytorch/blob/11a40410e755b1fe74efe9eaa635e7ba5712846b/test/cpp_extensions/setup.py#L62 | |
| from setuptools import setup | |
| from torch.utils.cpp_extension import BuildExtension, CUDAExtension | |
| import os | |
| # USE_NINJA = os.getenv('USE_NINJA') == '1' | |
| os.environ["CC"] = "gcc-7" | |
| os.environ["CXX"] = "gcc-7" | |
| USE_NINJA = os.getenv('USE_NINJA') == '1' | |
| setup( | |
| name='standard_rasterize_cuda', | |
| ext_modules=[ | |
| CUDAExtension('standard_rasterize_cuda', [ | |
| 'standard_rasterize_cuda.cpp', | |
| 'standard_rasterize_cuda_kernel.cu', | |
| ]) | |
| ], | |
| cmdclass={'build_ext': BuildExtension.with_options(use_ninja=USE_NINJA)} | |
| ) | |