Repository: cupy/cupy

Issue:
Support multi-threaded Cythonize (`nthreads` option)
### Description

In general, `cythonize` is reasonably fast as-is, but it can be slow when cross-compiling aarch64 binary on x86_64 host. I think this option can improve the build speed by utilizing multiple cores on host.

### Additional Information

_No response_


Discussion:

Resolving patch:
diff --git a/install/cupy_builder/_command.py b/install/cupy_builder/_command.py
--- a/install/cupy_builder/_command.py
+++ b/install/cupy_builder/_command.py
@@ -1,13 +1,14 @@
+import json
 import os
 import sys
-from typing import List, Tuple
+from typing import Any, Dict, List, Tuple
 
 import setuptools
 import setuptools.command.build_ext
 
 import cupy_builder
+import cupy_builder.install_build as build
 from cupy_builder._context import Context
-from cupy_builder.cupy_setup_build import cythonize
 from cupy_builder._compiler import DeviceCompilerUnix, DeviceCompilerWin32
 
 
@@ -73,6 +74,54 @@ class custom_build_ext(setuptools.command.build_ext.build_ext):
 
     """Custom `build_ext` command to include CUDA C source files."""
 
+    def _cythonize(self, nthreads: int) -> None:
+        # Defer importing Cython as it may be installed via setup_requires if
+        # the user does not have Cython installed.
+        import Cython.Build
+
+        ctx = cupy_builder.get_context()
+        compiler_directives = {
+            'linetrace': ctx.linetrace,
+            'profile': ctx.profile,
+            # Embed signatures for Sphinx documentation.
+            'embedsignature': True,
+        }
+
+        # Compile-time constants to be used in Cython code
+        compile_time_env: Dict[str, Any] = {}
+
+        # Enable CUDA Python.
+        # TODO: add `cuda` to `setup_requires` only when this flag is set
+        use_cuda_python = ctx.use_cuda_python
+        compile_time_env['CUPY_USE_CUDA_PYTHON'] = use_cuda_python
+        if use_cuda_python:
+            print('Using CUDA Python')
+
+        compile_time_env['CUPY_CUFFT_STATIC'] = False
+        compile_time_env['CUPY_CYTHON_VERSION'] = Cython.__version__
+        if ctx.use_stub:  # on RTD
+            compile_time_env['CUPY_CUDA_VERSION'] = 0
+            compile_time_env['CUPY_HIP_VERSION'] = 0
+        elif ctx.use_hip:  # on ROCm/HIP
+            compile_time_env['CUPY_CUDA_VERSION'] = 0
+            compile_time_env['CUPY_HIP_VERSION'] = build.get_hip_version()
+        else:  # on CUDA
+            compile_time_env['CUPY_CUDA_VERSION'] = (
+                ctx.features['cuda'].get_version())
+            compile_time_env['CUPY_HIP_VERSION'] = 0
+
+        print('Compile-time constants: ' +
+              json.dumps(compile_time_env, indent=4))
+
+        if sys.platform == 'win32':
+            # Disable multiprocessing on Windows (spawn)
+            nthreads = 0
+
+        Cython.Build.cythonize(
+            self.extensions, verbose=True, nthreads=nthreads, language_level=3,
+            compiler_directives=compiler_directives, annotate=ctx.annotate,
+            compile_time_env=compile_time_env)
+
     def build_extensions(self) -> None:
         num_jobs = int(os.environ.get('CUPY_NUM_BUILD_JOBS', '4'))
         if num_jobs > 1:
@@ -89,7 +138,7 @@ def build_extensions(self) -> None:
 
         # Compile "*.pyx" files into "*.cpp" files.
         print('Cythonizing...')
-        cythonize(self.extensions, cupy_builder.get_context())
+        self._cythonize(num_jobs)
 
         # Change an extension in each source filenames from "*.pyx" to "*.cpp".
         # c.f. `Cython.Distutils.old_build_ext`
diff --git a/install/cupy_builder/cupy_setup_build.py b/install/cupy_builder/cupy_setup_build.py
--- a/install/cupy_builder/cupy_setup_build.py
+++ b/install/cupy_builder/cupy_setup_build.py
@@ -7,10 +7,8 @@
 import shutil
 import sys
 
-import pkg_resources
 import setuptools
 
-import cupy_builder
 import cupy_builder.install_build as build
 from cupy_builder._context import Context
 from cupy_builder.install_build import PLATFORM_LINUX
@@ -434,54 +432,6 @@ def prepare_wheel_libs(ctx: Context):
     return [os.path.relpath(x[1], 'cupy') for x in files_to_copy]
 
 
-def cythonize(extensions, ctx: Context):
-    # Delay importing Cython as it may be installed via setup_requires if
-    # the user does not have Cython installed.
-    import Cython
-    import Cython.Build
-    cython_version = pkg_resources.parse_version(Cython.__version__)
-
-    directives = {
-        'linetrace': ctx.linetrace,
-        'profile': ctx.profile,
-        # Embed signatures for Sphinx documentation.
-        'embedsignature': True,
-    }
-
-    cythonize_options = {
-        'annotate': ctx.annotate
-    }
-
-    # Compile-time constants to be used in Cython code
-    compile_time_env = cythonize_options.get('compile_time_env')
-    if compile_time_env is None:
-        compile_time_env = {}
-        cythonize_options['compile_time_env'] = compile_time_env
-
-    # Enable CUDA Python.
-    # TODO: add `cuda` to `setup_requires` only when this flag is set
-    use_cuda_python = cupy_builder.get_context().use_cuda_python
-    compile_time_env['CUPY_USE_CUDA_PYTHON'] = use_cuda_python
-    if use_cuda_python:
-        print('Using CUDA Python')
-
-    compile_time_env['CUPY_CUFFT_STATIC'] = False
-    compile_time_env['CUPY_CYTHON_VERSION'] = str(cython_version)
-    if ctx.use_stub:  # on RTD
-        compile_time_env['CUPY_CUDA_VERSION'] = 0
-        compile_time_env['CUPY_HIP_VERSION'] = 0
-    elif ctx.use_hip:  # on ROCm/HIP
-        compile_time_env['CUPY_CUDA_VERSION'] = 0
-        compile_time_env['CUPY_HIP_VERSION'] = build.get_hip_version()
-    else:  # on CUDA
-        compile_time_env['CUPY_CUDA_VERSION'] = ctx.features['cuda'].get_version()  # NOQA
-        compile_time_env['CUPY_HIP_VERSION'] = 0
-
-    return Cython.Build.cythonize(
-        extensions, verbose=True, language_level=3,
-        compiler_directives=directives, **cythonize_options)
-
-
 def get_ext_modules(use_cython: bool, ctx: Context):
     # We need to call get_config_vars to initialize _config_vars in distutils
     # see #1849


Test patch: