# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # isort: skip_file from __future__ import annotations from typing import TYPE_CHECKING, Any import os import copy import concurrent import functools import re import setuptools import sys import paddle import site from distutils.errors import DistutilsExecError, LinkError from setuptools.command.easy_install import easy_install from setuptools.command.build_ext import build_ext from distutils.command.build import build from setuptools.command.install import install from .extension_utils import ( add_compile_flag, find_cuda_home, find_ccache_home, find_rocm_home, normalize_extension_kwargs, define_paddle_extension_name, ) from .extension_utils import ( is_cuda_file, prepare_unix_cudaflags, prepare_win_cudaflags, ) from .extension_utils import ( _import_module_from_library, _write_setup_file, _jit_compile, ) from .extension_utils import ( check_abi_compatibility, log_v, CustomOpInfo, parse_op_name_from, ) from .extension_utils import _reset_so_rpath, clean_object_if_change_cflags from .extension_utils import ( get_build_directory, add_std_without_repeat, custom_write_stub, ) from .extension_utils import ( IS_WINDOWS, OS_NAME, MSVC_COMPILE_FLAGS, ) from .extension_utils import CLANG_COMPILE_FLAGS, CLANG_LINK_FLAGS from ...base import core from concurrent.futures import ThreadPoolExecutor if TYPE_CHECKING: from collections.abc import Sequence from types import ModuleType # Note(zhouwei): On windows, it will export function 'PyInit_[name]' by default, # The solution is: 1.User add function PyInit_[name] 2. set not to export # refer to https://stackoverflow.com/questions/34689210/error-exporting-symbol-when-building-python-c-extension-in-windows if IS_WINDOWS: from distutils.command.build_ext import build_ext as _du_build_ext from unittest.mock import Mock _du_build_ext.get_export_symbols = Mock(return_value=None) CUDA_HOME = find_cuda_home() if core.is_compiled_with_rocm(): ROCM_HOME = find_rocm_home() CUDA_HOME = ROCM_HOME @functools.cache def _get_ccache_home(): return find_ccache_home() def setup(**attr: Any) -> None: """ The interface is used to config the process of compiling customized operators, mainly includes how to compile shared library, automatically generate python API and install it into site-package. It supports using customized operators directly with ``import`` statement. It encapsulates the python built-in ``setuptools.setup`` function and keeps arguments and usage same as the native interface. Meanwhile, it hides Paddle inner framework concepts, such as necessary compiling flags, included paths of head files, and linking flags. It also will automatically search and valid local environment and versions of ``cc(Linux)`` , ``cl.exe(Windows)`` and ``nvcc`` , then compiles customized operators supporting CPU or GPU device according to the specified Extension type. Moreover, `ABI compatibility `_ will be checked to ensure that compiler version from ``cc(Linux)`` , ``cl.exe(Windows)`` on local machine is compatible with pre-installed Paddle whl in python site-packages. For Linux, GCC version will be checked . For example if Paddle with CUDA 10.1 is built with GCC 8.2, then the version of user's local machine should satisfy GCC >= 8.2. For Windows, Visual Studio version will be checked, and it should be greater than or equal to that of PaddlePaddle (Visual Studio 2017). If the above conditions are not met, the corresponding warning will be printed, and a fatal error may occur because of ABI compatibility. Note: 1. Currently we support Linux, MacOS and Windows platform. 2. On Linux platform, we recommend to use GCC 8.2 as soft linking candidate of ``/usr/bin/cc`` . Then, Use ``which cc`` to ensure location of ``cc`` and using ``cc --version`` to ensure linking GCC version. 3. On Windows platform, we recommend to install `` Visual Studio`` (>=2017). Compared with Just-In-Time ``load`` interface, it only compiles once by executing ``python setup.py install`` . Then customized operators API will be available everywhere after importing it. A simple example of ``setup.py`` as followed: .. code-block:: text # setup.py # Case 1: Compiling customized operators supporting CPU and GPU devices from paddle.utils.cpp_extension import CUDAExtension, setup setup( name='custom_op', # name of package used by "import" ext_modules=CUDAExtension( sources=['relu_op.cc', 'relu_op.cu', 'tanh_op.cc', 'tanh_op.cu'] # Support for compilation of multiple OPs ) ) # Case 2: Compiling customized operators supporting only CPU device from paddle.utils.cpp_extension import CppExtension, setup setup( name='custom_op', # name of package used by "import" ext_modules=CppExtension( sources=['relu_op.cc', 'tanh_op.cc'] # Support for compilation of multiple OPs ) ) Applying compilation and installation by executing ``python setup.py install`` under source files directory. Then we can use the layer api as followed: .. code-block:: text import paddle from custom_op import relu, tanh x = paddle.randn([4, 10], dtype='float32') relu_out = relu(x) tanh_out = tanh(x) Args: name(str): Specify the name of shared library file and installed python package. ext_modules(Extension): Specify the Extension instance including customized operator source files, compiling flags et.al. If only compile operator supporting CPU device, please use ``CppExtension`` ; If compile operator supporting CPU and GPU devices, please use ``CUDAExtension`` . include_dirs(list[str], optional): Specify the extra include directories to search head files. The interface will automatically add ``site-package/paddle/include`` . Please add the corresponding directory path if including third-party head files. Default is None. extra_compile_args(list[str] | dict, optional): Specify the extra compiling flags such as ``-O3`` . If set ``list[str]`` , all these flags will be applied for ``cc`` and ``nvcc`` compiler. It supports specify flags only applied ``cc`` or ``nvcc`` compiler using dict type with ``{'cxx': [...], 'nvcc': [...]}`` . Default is None. **attr(dict, optional): Specify other arguments same as ``setuptools.setup`` . Returns: None """ cmdclass = attr.get('cmdclass', {}) assert isinstance(cmdclass, dict) # if not specific cmdclass in setup, add it automatically. if 'build_ext' not in cmdclass: cmdclass['build_ext'] = BuildExtension.with_options( no_python_abi_suffix=True ) attr['cmdclass'] = cmdclass error_msg = """ Required to specific `name` argument in paddle.utils.cpp_extension.setup. It's used as `import XXX` when you want install and import your custom operators.\n For Example: # setup.py file from paddle.utils.cpp_extension import CUDAExtension, setup setup(name='custom_module', ext_modules=CUDAExtension( sources=['relu_op.cc', 'relu_op.cu']) # After running `python setup.py install` from custom_module import relu """ # name argument is required if 'name' not in attr: raise ValueError(error_msg) assert not attr['name'].endswith('module'), ( "Please don't use 'module' as suffix in `name` argument, " ) "it will be stripped in setuptools.bdist_egg and cause import error." ext_modules = attr.get('ext_modules', []) if not isinstance(ext_modules, list): ext_modules = [ext_modules] assert len(ext_modules) == 1, ( f"Required only one Extension, but received {len(ext_modules)}. If you want to compile multi operators, you can include all necessary source files in one Extension." ) # replace Extension.name with attr['name] to keep consistent with Package name. for ext_module in ext_modules: ext_module.name = attr['name'] attr['ext_modules'] = ext_modules # Add rename .so hook in easy_install assert 'easy_install' not in cmdclass cmdclass['easy_install'] = EasyInstallCommand # Compatible with wheel installation via `pip install .` # Note: This is rarely used with modern pip, which uses bdist_wheel instead assert 'install' not in cmdclass cmdclass['install'] = InstallCommand # Note(Aurelius84): Add rename build_base directory hook in build command. # To avoid using same build directory that will lead to remove the directory # by mistake while parallelling execute setup.py, for example on CI. assert 'build' not in cmdclass build_base = os.path.join('build', attr['name']) cmdclass['build'] = BuildCommand.with_options(build_base=build_base) # Always set zip_safe=False to make compatible in PY2 and PY3 # See http://peak.telecommunity.com/DevCenter/setuptools#setting-the-zip-safe-flag attr['zip_safe'] = False setuptools.setup(**attr) def CppExtension( sources: Sequence[str], *args: Any, **kwargs: Any ) -> setuptools.Extension: """ The interface is used to config source files of customized operators and complies Op Kernel only supporting CPU device. Please use ``CUDAExtension`` if you want to compile Op Kernel that supports both CPU and GPU devices. It further encapsulates python built-in ``setuptools.Extension`` .The arguments and usage are same as the native interface, except for no need to explicitly specify ``name`` . **A simple example:** .. code-block:: text # setup.py # Compiling customized operators supporting only CPU device from paddle.utils.cpp_extension import CppExtension, setup setup( name='custom_op', ext_modules=CppExtension(sources=['relu_op.cc']) ) Note: It is mainly used in ``setup`` and the name of built shared library keeps same as ``name`` argument specified in ``setup`` interface. Args: sources(list[str]): Specify the C++/CUDA source files of customized operators. *args(list[options], optional): Specify other arguments same as ``setuptools.Extension`` . **kwargs(dict[option], optional): Specify other arguments same as ``setuptools.Extension`` . Returns: setuptools.Extension: An instance of ``setuptools.Extension`` """ kwargs = normalize_extension_kwargs(kwargs, use_cuda=False) # Note(Aurelius84): While using `setup` and `jit`, the Extension `name` will # be replaced as `setup.name` to keep consistent with package. Because we allow # users can not specific name in Extension. # See `paddle.utils.cpp_extension.setup` for details. name = kwargs.pop('name', None) if name is None: name = _generate_extension_name(sources) return setuptools.Extension(name, sources, *args, **kwargs) def CUDAExtension( sources: Sequence[str], *args: Any, **kwargs: Any ) -> setuptools.Extension: """ The interface is used to config source files of customized operators and complies Op Kernel supporting both CPU and GPU devices. Please use ``CppExtension`` if you want to compile Op Kernel that supports only CPU device. It further encapsulates python built-in ``setuptools.Extension`` .The arguments and usage are same as the native interface, except for no need to explicitly specify ``name`` . **A simple example:** .. code-block:: text # setup.py # Compiling customized operators supporting CPU and GPU devices from paddle.utils.cpp_extension import CUDAExtension, setup setup( name='custom_op', ext_modules=CUDAExtension( sources=['relu_op.cc', 'relu_op.cu'] ) ) Note: It is mainly used in ``setup`` and the name of built shared library keeps same as ``name`` argument specified in ``setup`` interface. Args: sources(list[str]): Specify the C++/CUDA source files of customized operators. *args(list[options], optional): Specify other arguments same as ``setuptools.Extension`` . **kwargs(dict[option], optional): Specify other arguments same as ``setuptools.Extension`` . Returns: setuptools.Extension: An instance of setuptools.Extension. """ kwargs = normalize_extension_kwargs(kwargs, use_cuda=True) # Note(Aurelius84): While using `setup` and `jit`, the Extension `name` will # be replaced as `setup.name` to keep consistent with package. Because we allow # users can not specific name in Extension. # See `paddle.utils.cpp_extension.setup` for details. name = kwargs.pop('name', None) if name is None: name = _generate_extension_name(sources) return setuptools.Extension(name, sources, *args, **kwargs) def _generate_extension_name(sources): """ Generate extension name by source files. """ assert len(sources) > 0, "source files is empty" file_prefix = [] for source in sources: source = os.path.basename(source) filename, _ = os.path.splitext(source) # Use list to generate same order. if filename not in file_prefix: file_prefix.append(filename) return '_'.join(file_prefix) class BuildExtension(build_ext): """ Inherited from setuptools.command.build_ext to customize how to apply compilation process with share library. """ @classmethod def with_options(cls, **options: Any) -> type[BuildExtension]: """ Returns a BuildExtension subclass containing use-defined options. """ class cls_with_options(cls): def __init__(self, *args, **kwargs): kwargs.update(options) cls.__init__(self, *args, **kwargs) return cls_with_options def __init__(self, *args: Any, **kwargs: Any) -> None: """ Attributes is initialized with following order: 1. super().__init__() 2. initialize_options(self) 3. the reset of current __init__() 4. finalize_options(self) So, it is recommended to set attribute value in `finalize_options`. """ super().__init__(*args, **kwargs) self.no_python_abi_suffix = kwargs.get("no_python_abi_suffix", True) self.output_dir = kwargs.get("output_dir", None) # whether containing cuda source file in Extensions self.contain_cuda_file = False # Initialize ccache_home to avoid race condition in multi-thread compilation _get_ccache_home() def initialize_options(self) -> None: super().initialize_options() def finalize_options(self) -> None: super().finalize_options() # NOTE(Aurelius84): Set location of compiled shared library. # Carefully to modify this because `setup.py build/install` # and `load` interface rely on this attribute. if self.output_dir is not None: self.build_lib = self.output_dir def build_extensions(self) -> None: if OS_NAME.startswith("darwin"): self._valid_clang_compiler() self._check_abi() current_extension_builder = self # Check nvcc_dlink ext = self.extensions[0] if ( isinstance(ext.extra_compile_args, dict) and 'nvcc_dlink' in ext.extra_compile_args ): cuda_dlink_post_cflags = prepare_unix_cudaflags( copy.deepcopy(ext.extra_compile_args['nvcc_dlink']) ) else: cuda_dlink_post_cflags = None # Note(Aurelius84): If already compiling source before, we should check whether # cflags have changed and delete the built shared library to re-compile the source # even though source file content keep unchanged. so_name = self.get_ext_fullpath(self.extensions[0].name) clean_object_if_change_cflags( os.path.abspath(so_name), self.extensions[0] ) # Consider .cu, .cu.cc as valid source extensions. self.compiler.src_extensions += ['.cu', '.cu.cc'] original_compile = None original_link = None # Save the original _compile method for later. if self.compiler.compiler_type == 'msvc': self.compiler._cpp_extensions += ['.cu', '.cuh'] original_compile = self.compiler.compile original_spawn = self.compiler.spawn else: original_compile = self.compiler.__class__.compile for extension in self.extensions: define_paddle_extension_name(extension) def unix_custom_compile_single_file( self, obj, src, ext, cc_args, extra_postargs, pp_opts ): """ Monkey patch mechanism to replace inner compiler to custom compile process on Unix platform. """ # use abspath to ensure no warning and don't remove deepcopy because modify params # with dict type is dangerous. src = os.path.abspath(src) cflags = copy.deepcopy(extra_postargs) try: original_compiler = self.compiler_so # nvcc or hipcc compile CUDA source if is_cuda_file(src): if core.is_compiled_with_rocm(): assert ROCM_HOME is not None, ( "Not found ROCM runtime, \ please use `export ROCM_PATH= XXX` to specify it." ) ccache_home = _get_ccache_home() if ccache_home is not None: hipcc_cmd = os.path.join(ROCM_HOME, 'bin', 'hipcc') hipcc_cmd = f'{ccache_home} {hipcc_cmd}' else: hipcc_cmd = os.path.join(ROCM_HOME, 'bin', 'hipcc') self.set_executable('compiler_so', hipcc_cmd) # {'nvcc': {}, 'cxx: {}} if isinstance(cflags, dict): cflags = cflags['hipcc'] elif core.is_compiled_with_custom_device("iluvatar_gpu"): ixcc_cmd = os.path.join( os.getenv("COREX_HOME", "/usr/local/corex/"), 'bin', 'clang++', ) if not os.path.isfile(ixcc_cmd): raise ValueError( "Corex compiler is unavailable, please set `COREX_HOME` to specify it." ) self.set_executable('compiler_so', ixcc_cmd) # {'nvcc': {}, 'cxx: {}} if isinstance(cflags, dict): cflags = cflags['nvcc'] else: assert CUDA_HOME is not None, ( "Not found CUDA runtime, \ please use `export CUDA_HOME= XXX` to specify it." ) ccache_home = _get_ccache_home() if ccache_home is not None: nvcc_cmd = os.path.join(CUDA_HOME, 'bin', 'nvcc') nvcc_cmd = f'{ccache_home} {nvcc_cmd}' else: nvcc_cmd = os.path.join(CUDA_HOME, 'bin', 'nvcc') self.set_executable('compiler_so', nvcc_cmd) # {'nvcc': {}, 'cxx: {}} if isinstance(cflags, dict): cflags = cflags['nvcc'] cflags = prepare_unix_cudaflags(cflags) # cxx compile Cpp source else: ccache_home = _get_ccache_home() if ccache_home is not None: self.set_executable( 'compiler_so', [ccache_home, *self.compiler_so] ) if isinstance(cflags, dict): cflags = cflags['cxx'] # Note(qili93): HIP require some additional flags for CMAKE_C_FLAGS if core.is_compiled_with_rocm(): cflags.append('-D__HIP_PLATFORM_HCC__') # cflags.append('-D__HIP_NO_HALF_CONVERSIONS__=1') cflags.append( '-DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_HIP' ) # NOTE(Aurelius84): Since Paddle 2.0, we require gcc version > 5.x, # so we add this flag to ensure the symbol names from user compiled # shared library have same ABI suffix with libpaddle.so. # See https://stackoverflow.com/questions/34571583/understanding-gcc-5s-glibcxx-use-cxx11-abi-or-the-new-abi add_compile_flag(cflags, ['-D_GLIBCXX_USE_CXX11_ABI=1']) # Append this macro only when jointly compiling .cc with .cu if ( not is_cuda_file(src) and current_extension_builder.contain_cuda_file ): if core.is_compiled_with_rocm(): cflags.append('-DPADDLE_WITH_HIP') else: cflags.append('-DPADDLE_WITH_CUDA') add_std_without_repeat( cflags, self.compiler_type, use_std17=True ) self._compile(obj, src, ext, cc_args, cflags, pp_opts) except Exception as e: print(f'{src} compile failed, {e}') finally: # restore original_compiler self.set_executable('compiler_so', original_compiler) def unix_custom_link_shared_object( self, objects: list[str] | tuple[str, ...], output_filename: str, output_dir: str | None = None, libraries: list[str] | tuple[str, ...] | None = None, library_dirs: list[str] | tuple[str, ...] | None = None, runtime_library_dirs: list[str] | tuple[str, ...] | None = None, export_symbols: Any | None = None, debug: bool = False, extra_preargs: list[str] | None = None, extra_postargs: list[str] | None = None, build_temp: str | os.PathLike[str] | None = None, target_lang: str | None = None, ): # Get extension dlink_dir = os.path.dirname(objects[0]) dlink_object = os.path.join(dlink_dir, 'dlink.o') # Construct command # nvcc -o if CUDA_HOME is None: raise RuntimeError("CUDA_HOME is not found, please set it.") nvcc_cmd = os.path.join(CUDA_HOME, 'bin', 'nvcc') cmd = [] ccache_home = _get_ccache_home() if ccache_home: cmd.append(ccache_home) cmd.append(nvcc_cmd) cmd.extend(objects) cmd.extend(['-o', dlink_object]) cmd.extend(cuda_dlink_post_cflags) # Execute try: self.spawn(cmd) except DistutilsExecError as msg: raise LinkError(msg) # Add dlink object to objects objects = [*list(objects), dlink_object] return original_link( self, objects, output_filename, output_dir, libraries, library_dirs, runtime_library_dirs, export_symbols, debug, extra_preargs, extra_postargs, build_temp, target_lang, ) def unix_custom_single_compiler( self, sources, output_dir=None, macros=None, include_dirs=None, debug=False, extra_preargs=None, extra_postargs=None, depends=None, ): # A concrete compiler class can either override this method # entirely or implement _compile(). 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) # Create a thread pool requested_workers = _get_num_workers(verbose=bool(self.verbose)) worker_number = _compute_worker_number( requested_workers, os.cpu_count(), len(objects) ) print( f"Using {worker_number} workers for compilation. HINT: export MAX_JOBS=n to set the number of workers" ) with ThreadPoolExecutor(max_workers=worker_number) as executor: # Submit all compilation tasks to the thread pool. futures = { executor.submit( unix_custom_compile_single_file, copy.copy(self), obj, build[obj][0], build[obj][1], cc_args, extra_postargs, pp_opts, ): obj for obj in objects } for future in concurrent.futures.as_completed(futures): obj = futures[future] try: future.result() except Exception as exc: print(f'{obj!r} generated an exception: {exc}') else: print(f'{obj} is compiled') # Return *all* object filenames, not just the ones we just built. return objects def win_custom_single_compiler( sources, output_dir=None, macros=None, include_dirs=None, debug=0, extra_preargs=None, extra_postargs=None, depends=None, ): self.cflags = copy.deepcopy(extra_postargs) extra_postargs = None def win_custom_spawn(cmd): # Using regex to modify compile options compile_options = self.compiler.compile_options for i in range(len(cmd)): if re.search('/MD', cmd[i]) is not None: cmd[i] = '/MT' if re.search('/W[1-4]', cmd[i]) is not None: cmd[i] = '/W0' # Using regex to match src, obj and include files src_regex = re.compile('/T(p|c)(.*)') src_list = [ m.group(2) for m in (src_regex.match(elem) for elem in cmd) if m ] obj_regex = re.compile('/Fo(.*)') obj_list = [ m.group(1) for m in (obj_regex.match(elem) for elem in cmd) if m ] include_regex = re.compile(r'((\-|\/)I.*)') include_list = [ m.group(1) for m in (include_regex.match(elem) for elem in cmd) if m ] assert len(src_list) == 1 and len(obj_list) == 1 src = src_list[0] obj = obj_list[0] if is_cuda_file(src): assert CUDA_HOME is not None, ( "Not found CUDA runtime, \ please use `export CUDA_HOME= XXX` to specify it." ) nvcc_cmd = os.path.join(CUDA_HOME, 'bin', 'nvcc') if isinstance(self.cflags, dict): cflags = self.cflags['nvcc'] elif isinstance(self.cflags, list): cflags = self.cflags else: cflags = [] cflags = [*prepare_win_cudaflags(cflags), '--use-local-env'] for flag in MSVC_COMPILE_FLAGS: cflags = ['-Xcompiler', flag, *cflags] cmd = [ nvcc_cmd, '-c', src, '-o', obj, *include_list, *cflags, ] elif isinstance(self.cflags, dict): cflags = MSVC_COMPILE_FLAGS + self.cflags['cxx'] cmd += cflags elif isinstance(self.cflags, list): cflags = MSVC_COMPILE_FLAGS + self.cflags cmd += cflags # Append this macro only when jointly compiling .cc with .cu if not is_cuda_file(src) and self.contain_cuda_file: cmd.append('-DPADDLE_WITH_CUDA') return original_spawn(cmd) try: self.compiler.spawn = win_custom_spawn return original_compile( sources, output_dir, macros, include_dirs, debug, extra_preargs, extra_postargs, depends, ) finally: self.compiler.spawn = original_spawn def object_filenames_with_cuda(original_func, build_directory): """ Decorated the function to add customized naming mechanism. Originally, both .cc/.cu will have .o object output that will bring file override problem. Use .cu.o as CUDA object suffix. """ def wrapper(source_filenames, strip_dir=0, output_dir=''): try: objects = original_func( source_filenames, strip_dir, output_dir ) for i, source in enumerate(source_filenames): # modify xx.o -> xx.cu.o/xx.cu.obj if is_cuda_file(source): old_obj = objects[i] if self.compiler.compiler_type == 'msvc': objects[i] = old_obj[:-3] + 'cu.obj' else: objects[i] = old_obj[:-1] + 'cu.o' # if user set build_directory, output objects there. if build_directory is not None: objects = [ os.path.join(build_directory, obj) for obj in objects ] # ensure to use abspath objects = [os.path.abspath(obj) for obj in objects] finally: self.compiler.object_filenames = original_func return objects return wrapper # customized compile process if self.compiler.compiler_type == 'msvc': self.compiler.compile = win_custom_single_compiler else: self.compiler.__class__.compile = unix_custom_single_compiler # Ensure object files are generated under build_temp, not build_lib, # to avoid accidental inclusion into wheel contents. self.compiler.object_filenames = object_filenames_with_cuda( self.compiler.object_filenames, self.build_temp ) self._record_op_info() try: if cuda_dlink_post_cflags and self.compiler.compiler_type != 'msvc': original_link = self.compiler.__class__.link_shared_object self.compiler.__class__.link_shared_object = ( unix_custom_link_shared_object ) print("Compiling user custom op, it will cost a few seconds.....") build_ext.build_extensions(self) finally: if self.compiler.compiler_type == 'msvc': self.compiler.compile = original_compile else: self.compiler.__class__.compile = original_compile if original_link: self.compiler.__class__.link_shared_object = original_link # Reset runtime library path on MacOS platform so_path = self.get_ext_fullpath(self.extensions[0]._full_name) _reset_so_rpath(so_path) def get_ext_filename(self, fullname: str) -> str: # for example: customized_extension.cpython-37m-x86_64-linux-gnu.so ext_name = super().get_ext_filename(fullname) split_str = '.' name_items = ext_name.split(split_str) if self.no_python_abi_suffix: assert len(name_items) > 2, ( f"Expected len(name_items) > 2, but received {len(name_items)}" ) name_items.pop(-2) ext_name = split_str.join(name_items) # customized_extension.dylib if OS_NAME.startswith('darwin'): name_items[-1] = 'dylib' ext_name = split_str.join(name_items) return ext_name def _valid_clang_compiler(self) -> None: """ Make sure to use Clang as compiler on Mac platform """ compiler_infos = ['clang', *CLANG_COMPILE_FLAGS] linker_infos = ['clang', *CLANG_LINK_FLAGS] self.compiler.set_executables( compiler=compiler_infos, compiler_so=compiler_infos, compiler_cxx=['clang'], linker_exe=['clang'], linker_so=linker_infos, ) def _check_abi(self) -> None: """ Check ABI Compatibility. """ if hasattr(self.compiler, 'compiler_cxx'): compiler = self.compiler.compiler_cxx[0] elif IS_WINDOWS: compiler = os.environ.get('CXX', 'cl') else: compiler = os.environ.get('CXX', 'c++') check_abi_compatibility(compiler) # Warn user if VC env is activated but `DISTUTILS_USE_SDK` is not set. if ( IS_WINDOWS and 'VSCMD_ARG_TGT_ARCH' in os.environ and 'DISTUTILS_USE_SDK' not in os.environ ): msg = ( 'It seems that the VC environment is activated but DISTUTILS_USE_SDK is not set.' 'This may lead to multiple activations of the VC env.' 'Please run `set DISTUTILS_USE_SDK=1` and try again.' ) raise UserWarning(msg) def _record_op_info(self) -> None: """ Record custom op information. """ # parse shared library abs path outputs = self.get_outputs() assert len(outputs) == 1 # multi operators built into same one .so file so_path = os.path.abspath(outputs[0]) so_name = os.path.basename(so_path) for i, extension in enumerate(self.extensions): sources = [os.path.abspath(s) for s in extension.sources] if not self.contain_cuda_file: self.contain_cuda_file = any(is_cuda_file(s) for s in sources) op_names = parse_op_name_from(sources) for op_name in op_names: CustomOpInfo.instance().add( op_name, so_name=so_name, so_path=so_path ) def _clean_intermediate_files(self): for ext in self.extensions: build_dir = os.path.dirname(self.get_ext_fullpath(ext.name)) for root, _, files in os.walk(build_dir): for file in files: if file.endswith(".cu.o") or file.endswith('.o'): os.remove(os.path.join(root, file)) print(f"Removed: {os.path.join(root, file)}") def _generate_python_api_file(self) -> None: """ Generate the top-level python api file (package stub) alongside the built shared library in build_lib. This replaces the legacy bdist_egg write_stub mechanism that is no longer triggered in setuptools >= 80. """ try: if not self.extensions: return # We only support a single extension per setup() ext = self.extensions[0] # Use get_ext_fullpath to handle both standard and inplace builds correctly so_path = os.path.abspath(self.get_ext_fullpath(ext.name)) so_name = os.path.basename(so_path) build_dir = os.path.dirname(so_path) # Get the extension name from the extension module, not the distribution name # This ensures we use the correct package name from setup.py ext_name = ext.name # Extract the last part of the extension name for the Python file # For example, from "custom_setup_ops.my_ops.custom_relu" we get "custom_relu" lib_name = ext_name.split('.')[-1] if '.' in ext_name else ext_name pyfile = os.path.join(build_dir, f"{lib_name}.py") # Write stub; it will reference the _pd_ renamed resource at import time custom_write_stub(so_name, pyfile) except Exception as e: raise RuntimeError( f"Failed to generate python api file: {e}" ) from e def _rename_inplace_shared_library(self) -> None: """ Rename the shared library to *_pd_.so if it is an inplace build. This is necessary for editable installs to work correctly with the python stub. """ # We only support a single extension per setup() if not self.extensions: return ext = self.extensions[0] fullpath = self.get_ext_fullpath(ext.name) filename = os.path.basename(fullpath) dirname = os.path.dirname(fullpath) name, ext_suffix = os.path.splitext(filename) will_rename = False if OS_NAME.startswith('linux') and ext_suffix == '.so': will_rename = True elif OS_NAME.startswith('darwin') and ( ext_suffix == '.dylib' or ext_suffix == '.so' ): will_rename = True elif IS_WINDOWS and ext_suffix == '.pyd': will_rename = True if will_rename: new_name = f"{name}_pd_{ext_suffix}" new_path = os.path.join(dirname, new_name) if os.path.exists(fullpath): if os.path.exists(new_path): os.remove(new_path) os.rename(fullpath, new_path) print( f"Renaming {fullpath} to {new_path} for editable install compatibility" ) def run(self): super().run() # Compatible with wheel installation via `pip install .` self._generate_python_api_file() if self.inplace: self._rename_inplace_shared_library() self._clean_intermediate_files() class EasyInstallCommand(easy_install): """ Extend easy_install Command to control the behavior of naming shared library file. NOTE(Aurelius84): This is a hook subclass inherited Command used to rename shared library file after extracting egg-info into site-packages. """ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) # NOTE(Aurelius84): Add args and kwargs to make compatible with PY2/PY3 def run(self, *args: Any, **kwargs: Any) -> None: super().run(*args, **kwargs) # NOTE: To avoid failing import .so file instead of # python file because they have same name, we rename # .so shared library to another name. for egg_file in self.outputs: filename, ext = os.path.splitext(egg_file) will_rename = False if OS_NAME.startswith('linux') and ext == '.so': will_rename = True elif OS_NAME.startswith('darwin') and ext == '.dylib': will_rename = True elif IS_WINDOWS and ext == '.pyd': will_rename = True if will_rename: new_so_path = filename + "_pd_" + ext if not os.path.exists(new_so_path): os.rename(rf'{egg_file}', rf'{new_so_path}') assert os.path.exists(new_so_path) class BuildCommand(build): """ Extend build Command to control the behavior of specifying `build_base` root directory. NOTE(Aurelius84): This is a hook subclass inherited Command used to specify customized build_base directory. """ @classmethod def with_options(cls, **options: Any) -> type[BuildCommand]: """ Returns a BuildCommand subclass containing use-defined options. """ class cls_with_options(cls): def __init__(self, *args, **kwargs): kwargs.update(options) cls.__init__(self, *args, **kwargs) return cls_with_options def __init__(self, *args: Any, **kwargs: Any) -> None: # Note: shall put before super() self._specified_build_base = kwargs.get('build_base', None) super().__init__(*args, **kwargs) def initialize_options(self) -> None: """ build_base is root directory for all sub-command, such as build_lib, build_temp. See `distutils.command.build` for details. """ super().initialize_options() if self._specified_build_base is not None: self.build_base = self._specified_build_base class InstallCommand(install): """ Extend install Command to: 1) choose an install dir that is actually importable (on sys.path) 2) ensure a single top-level entry for the package in site/dist-packages so legacy tests that expect a sole artifact (egg/package) keep working 3) rename the compiled library to *_pd_.so to avoid shadowing the python stub """ def _get_extension_name(self) -> str: """ Get the extension name from the extension module, not the distribution name. This ensures we use the correct package name from setup.py. Note: This assumes there is only one extension module (len(ext_modules) == 1). Returns: str: The extension name """ return self.distribution.ext_modules[0].name def finalize_options(self) -> None: super().finalize_options() install_dir = ( getattr(self, 'install_lib', None) or getattr(self, 'install_purelib', None) or getattr(self, 'install_platlib', None) ) if not install_dir or not os.path.isdir(install_dir): return # Get the extension name ext_name = self._get_extension_name() # Extract the first part of the extension name for the shared library # For example, from "custom_setup_ops.my_ops.custom_relu" we get "custom_setup_ops" pkg_name = ext_name.split('.')[0] if '.' in ext_name else ext_name # Check if dist-info exists has_dist_info = any( name.endswith('.dist-info') and name.startswith(pkg_name) for name in os.listdir(install_dir) ) # If dist-info exists, we are installing a wheel, so we are done if has_dist_info: return # Build candidate site dirs: global + user + entries already on sys.path candidates = [] candidates.extend(site.getsitepackages()) usp = site.getusersitepackages() if usp: candidates.append(usp) for sp in sys.path: if isinstance(sp, str) and sp.endswith( ('site-packages', 'dist-packages') ): candidates.append(sp) # De-dup while preserving order seen = set() ordered = [] for c in candidates: if c and c not in seen: seen.add(c) ordered.append(c) # Prefer a candidate that is actually on sys.path target = None for c in ordered: if c in sys.path and os.path.isdir(c): target = c break # Fallback: pick the first existing candidate if target is None: for c in ordered: if os.path.isdir(c): target = c break if target: option_dict = self.distribution.get_option_dict('install') if 'install_lib' not in option_dict: self.install_lib = target if 'install_purelib' not in option_dict: self.install_purelib = target if 'install_platlib' not in option_dict: self.install_platlib = target def run(self, *args: Any, **kwargs: Any) -> None: super().run(*args, **kwargs) install_dir = ( getattr(self, 'install_lib', None) or getattr(self, 'install_purelib', None) or getattr(self, 'install_platlib', None) ) if not install_dir or not os.path.isdir(install_dir): return # Get the extension name ext_name = self._get_extension_name() # Extract the first part of the extension name for the shared library # For example, from "custom_setup_ops.my_ops.custom_relu" we get "custom_setup_ops" pkg_name = ext_name.split('.')[0] if '.' in ext_name else ext_name # Check if dist-info exists has_egg_info = any( name.endswith('.egg-info') and name.startswith(pkg_name) for name in os.listdir(install_dir) ) # If egg-info exists, we are installing a source distribution, we need to # reorganize the files if has_egg_info: # First rename the shared library if present at top-level self._rename_shared_library() # Then canonicalize layout to a single top-level entry for this package self._single_entry_layout() def _rename_shared_library(self) -> None: install_dir = ( getattr(self, 'install_lib', None) or getattr(self, 'install_purelib', None) or getattr(self, 'install_platlib', None) ) if not install_dir or not os.path.isdir(install_dir): return # Get the extension name ext_name = self._get_extension_name() # Extract the last part of the extension name for the shared library # For example, from "custom_setup_ops.my_ops.custom_relu" we get "custom_relu" names = ext_name.split('.') if '.' in ext_name else [ext_name] lib_name = names[-1] suffix = ( '.pyd' if IS_WINDOWS else ('.dylib' if OS_NAME.startswith('darwin') else '.so') ) # Build the directory path for the shared library # For single-level: names[:-1] is empty, so dir_path = install_dir # For multi-level: names[:-1] contains the package path dir_path = os.path.join(install_dir, *names[:-1]) old = os.path.join(dir_path, f"{lib_name}{suffix}") new = os.path.join(dir_path, f"{lib_name}_pd_{suffix}") if os.path.exists(old): if os.path.exists(new): os.remove(new) os.rename(old, new) def _single_entry_layout(self) -> None: """ Ensure only one top-level item in install_dir contains the package name by: - moving {pkg}.py -> {pkg}/__init__.py - moving {pkg}_pd_.so -> {pkg}/{pkg}_pd_.so - removing any {pkg}-*.egg-info left by setuptools install (only if dist-info exists) This keeps legacy tests that scan os.listdir(site_dir) happy. """ install_dir = ( getattr(self, 'install_lib', None) or getattr(self, 'install_purelib', None) or getattr(self, 'install_platlib', None) ) if not install_dir or not os.path.isdir(install_dir): return # Get the extension name ext_name = self._get_extension_name() # Extract the package path from the extension name # For example, from "custom_setup_ops.my_ops.custom_relu" we get "custom_setup_ops/my_ops" pkg_path_parts = ( ext_name.split('.')[:-1] if '.' in ext_name else [ext_name] ) pkg_path = os.path.join(*pkg_path_parts) # Extract the last part of the extension name for the Python file and shared library # For example, from "custom_setup_ops.my_ops.custom_relu" we get "custom_relu" lib_name = ext_name.split('.')[-1] if '.' in ext_name else ext_name # Prepare paths pkg_dir = os.path.join(install_dir, pkg_path) py_src = os.path.join(install_dir, f"{lib_name}.py") # Find compiled lib (renamed or not) suf_so = ( '.pyd' if IS_WINDOWS else ('.dylib' if OS_NAME.startswith('darwin') else '.so') ) so_candidates = [ os.path.join(install_dir, f"{lib_name}_pd_{suf_so}"), os.path.join(install_dir, f"{lib_name}{suf_so}"), ] so_src = next((p for p in so_candidates if os.path.exists(p)), None) # Create package dir if not os.path.isdir(pkg_dir): os.makedirs(pkg_dir, exist_ok=True) # Move python stub to package/__init__.py if exists if os.path.exists(py_src): py_dst = os.path.join(pkg_dir, "__init__.py") if os.path.exists(py_dst): os.remove(py_dst) os.replace(py_src, py_dst) # Move shared lib into the package dir if exists if so_src and os.path.exists(so_src): so_dst = os.path.join(pkg_dir, os.path.basename(so_src)) if os.path.exists(so_dst): os.remove(so_dst) os.replace(so_src, so_dst) def load( name: str, sources: Sequence[str], extra_cxx_cflags: Sequence[str] | None = None, extra_cuda_cflags: Sequence[str] | None = None, extra_ldflags: Sequence[str] | None = None, extra_include_paths: Sequence[str] | None = None, extra_library_paths: Sequence[str] | None = None, build_directory: str | None = None, verbose: bool = False, ) -> ModuleType: """ An Interface to automatically compile C++/CUDA source files Just-In-Time and return callable python function as other Paddle layers API. It will append user defined custom operators in background while building models. It will perform compiling, linking, Python API generation and module loading processes under a individual subprocess. It does not require CMake or Ninja environment. On Linux platform, it requires GCC compiler whose version is greater than 5.4 and it should be soft linked to ``/usr/bin/cc`` . On Windows platform, it requires Visual Studio whose version is greater than 2017. On MacOS, clang++ is requited. In addition, if compiling Operators supporting GPU device, please make sure ``nvcc`` compiler is installed in local environment. Moreover, `ABI compatibility `_ will be checked to ensure that compiler version from ``cc(Linux)`` , ``cl.exe(Windows)`` on local machine is compatible with pre-installed Paddle whl in python site-packages. For Linux, GCC version will be checked . For example if Paddle with CUDA 10.1 is built with GCC 8.2, then the version of user's local machine should satisfy GCC >= 8.2. For Windows, Visual Studio version will be checked, and it should be greater than or equal to that of PaddlePaddle (Visual Studio 2017). If the above conditions are not met, the corresponding warning will be printed, and a fatal error may occur because of ABI compatibility. Compared with ``setup`` interface, it doesn't need extra ``setup.py`` and execute ``python setup.py install`` command. The interface contains all compiling and installing process underground. Note: 1. Currently we support Linux, MacOS and Windows platform. 2. On Linux platform, we recommend to use GCC 8.2 as soft linking candidate of ``/usr/bin/cc`` . Then, Use ``which cc`` to ensure location of ``cc`` and using ``cc --version`` to ensure linking GCC version. 3. On Windows platform, we recommend to install `` Visual Studio`` (>=2017). **A simple example:** .. code-block:: text import paddle from paddle.utils.cpp_extension import load custom_op_module = load( name="op_shared_library_name", # name of shared library sources=['relu_op.cc', 'relu_op.cu'], # source files of customized op extra_cxx_cflags=['-g', '-w'], # optional, specify extra flags to compile .cc/.cpp file extra_cuda_cflags=['-O2'], # optional, specify extra flags to compile .cu file verbose=True # optional, specify to output log information ) x = paddle.randn([4, 10], dtype='float32') out = custom_op_module.relu(x) Args: name(str): Specify the name of generated shared library file name, not including ``.so`` and ``.dll`` suffix. sources(list[str]): Specify source files name of customized operators. Supporting ``.cc`` , ``.cpp`` for CPP file and ``.cu`` for CUDA file. extra_cxx_cflags(list[str]|None, optional): Specify additional flags used to compile CPP files. By default all basic and framework related flags have been included. extra_cuda_cflags(list[str]|None, optional): Specify additional flags used to compile CUDA files. By default all basic and framework related flags have been included. See `Cuda Compiler Driver NVCC `_ for details. Default is None. extra_ldflags(list[str]|None, optional): Specify additional flags used to link shared library. See `GCC Link Options `_ for details. Default is None. extra_include_paths(list[str]|None, optional): Specify additional include path used to search header files. By default all basic headers are included implicitly from ``site-package/paddle/include`` . Default is None. extra_library_paths(list[str]|None, optional): Specify additional library path used to search library files. By default all basic libraries are included implicitly from ``site-packages/paddle/libs`` . Default is None. build_directory(str|None, optional): Specify root directory path to put shared library file. If set None, it will use ``PADDLE_EXTENSION_DIR`` from os.environ. Use ``paddle.utils.cpp_extension.get_build_directory()`` to see the location. Default is None. verbose(bool, optional): whether to verbose compiled log information. Default is False. Returns: Module: A callable python module contains all CustomOp Layer APIs. """ if build_directory is None: build_directory = get_build_directory(verbose) # ensure to use abs path build_directory = os.path.abspath(build_directory) log_v(f"build_directory: {build_directory}", verbose) file_path = os.path.join(build_directory, f"{name}_setup.py") sources = [os.path.abspath(source) for source in sources] if extra_cxx_cflags is None: extra_cxx_cflags = [] if extra_cuda_cflags is None: extra_cuda_cflags = [] assert isinstance(extra_cxx_cflags, list), ( f"Required type(extra_cxx_cflags) == list[str], but received {extra_cxx_cflags}" ) assert isinstance(extra_cuda_cflags, list), ( f"Required type(extra_cuda_cflags) == list[str], but received {extra_cuda_cflags}" ) log_v( "additional extra_cxx_cflags: [{}], extra_cuda_cflags: [{}]".format( ' '.join(extra_cxx_cflags), ' '.join(extra_cuda_cflags) ), verbose, ) # write setup.py file and compile it build_base_dir = os.path.join(build_directory, name) _write_setup_file( name, sources, file_path, build_base_dir, extra_include_paths, extra_library_paths, extra_cxx_cflags, extra_cuda_cflags, extra_ldflags, verbose, ) _jit_compile(file_path, verbose) # import as callable python api custom_op_api = _import_module_from_library(name, build_base_dir, verbose) return custom_op_api def _get_pybind11_abi_build_flags(): abi_cflags = [] for pname in ["COMPILER_TYPE", "STDLIB", "BUILD_ABI"]: pval = getattr(paddle._C, f"_PYBIND11_{pname}") if pval is not None and not IS_WINDOWS: abi_cflags.append(f'-DPYBIND11_{pname}=\\"{pval}\\"') return abi_cflags def _get_num_workers(verbose: bool) -> int | None: max_jobs = os.environ.get('MAX_JOBS') if max_jobs is not None and max_jobs.isdigit(): if verbose: print( f'Using envvar MAX_JOBS ({max_jobs}) as the number of workers...', file=sys.stderr, ) return int(max_jobs) if verbose: print( 'Allowing ninja to set a default number of workers... ' '(overridable by setting the environment variable MAX_JOBS=N)', file=sys.stderr, ) return None def _compute_worker_number( requested_workers: int | None, cpu_count: int | None, num_objects: int ) -> int: cpu_count = cpu_count or 1 if requested_workers is None: worker_number = min(cpu_count, num_objects) else: worker_number = max(1, min(requested_workers, cpu_count, num_objects)) return worker_number