Skip to content

Commit

Permalink
Add parallel compilation support
Browse files Browse the repository at this point in the history
  • Loading branch information
LinZhihao-723 committed Mar 16, 2024
1 parent f55d9dd commit 5883587
Showing 1 changed file with 31 additions and 1 deletion.
32 changes: 31 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import distutils.ccompiler
import multiprocessing.pool
import os
import platform
import sys
from setuptools import setup, Extension
from typing import List, Optional
from typing import List, Optional, Tuple

ir_native: Extension = Extension(
name="clp_ffi_py.ir.native",
Expand Down Expand Up @@ -41,13 +43,41 @@
define_macros=[("SOURCE_PATH_SIZE", str(len(os.path.abspath("./src/clp/components/core"))))],
)

def _parallel_compile(
self: distutils.ccompiler.CCompiler,
sources: List[str],
output_dir: Optional[str] = None,
macros: Optional[List[Tuple[str, Optional[str]]]] = None,
include_dirs: Optional[List[str]] = None,
debug: int = 0,
extra_preargs: Optional[List[str]] = None,
extra_postargs: Optional[List[str]] = None,
depends: Optional[List[str]] = None,
) -> List[str]:
macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
output_dir, macros, include_dirs, sources, depends, extra_postargs
)
cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)

def _compile_single_file(obj: str) -> None:
try:
src, ext = build[obj]
except KeyError:
return
self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)

num_cores: int = multiprocessing.cpu_count()
multiprocessing.pool.ThreadPool(num_cores).map(_compile_single_file, objects)
return objects

if "__main__" == __name__:
try:
if "Darwin" == platform.system():
target: Optional[str] = os.environ.get("MACOSX_DEPLOYMENT_TARGET")
if None is target or float(target) < 10.15:
os.environ["MACOSX_DEPLOYMENT_TARGET"] = "10.15"

distutils.ccompiler.CCompiler.compile = _parallel_compile
project_name: str = "clp_ffi_py"
description: str = "CLP FFI Python Interface"
extension_modules: List[Extension] = [ir_native]
Expand Down

0 comments on commit 5883587

Please sign in to comment.