How to build a c++ extension in parallel using setuptools?
I followed the tutorial here : https://pytorch.org/tutorials/advanced/cpp_extension.html.
I have a quite similar setup.py (see below) with several c++ source files for one python module.
I build the extension using the command: pip install -e ..
But the build is sequential and every source files are rebuilt, even if only one file changes… which is very time consuming.
How to perform incremental parallel build, in the spirit of make -j16?
I am on Ubuntu 22.04 using PyTorch 2.0.1.
I checked online but could not find any solution yet.
I only saw that it is possible to build several c++ extensions in parallel, but not one unique extension.
I saw a parallel option in the class build_ext of setuptools, but I am not sure how to set it when using the BuildExtension of PyTorch.
Sorry in advance if this not a PyTorch issue but a setuptools one
Here is my setup.py from a minimal example:
import torch
from setuptools import find_packages, setup
from torch.__config__ import parallel_info
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
from pathlib import Path
def get_extensions():
all_sources = [
'csrc/src.cpp',
'csrc/a.cpp',
'csrc/b.cpp',
'csrc/c.cpp']
define_macros = [('WITH_PYTHON', None)]
undef_macros = []
include_dirs = []
extra_compile_args = {'cxx': ['-O3']}
extra_link_args = ['-s']
info = parallel_info()
if ('backend: OpenMP' in info and 'OpenMP not found' not in info):
extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
extra_compile_args['cxx'] += ['-fopenmp']
else:
print('Compiling without OpenMP...')
extension = CppExtension(
'torch_test._C',
all_sources,
include_dirs=include_dirs,
define_macros=define_macros,
undef_macros=undef_macros,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
)
return [extension]
setup(
name='torch_test',
version='none',
description=('none'),
author='none',
author_email='none',
url='none',
download_url='none',
keywords=[],
python_requires='>=3.8',
install_requires=[],
extras_require={'test': []},
ext_modules=get_extensions(),
cmdclass={
'build_ext':
BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)
},
packages=find_packages(),
include_package_data=True,
)