diff --git a/python/paddle/utils/cpp_extension/cpp_extension.py b/python/paddle/utils/cpp_extension/cpp_extension.py index 827bb5ccf4b31..a9015803d8d26 100644 --- a/python/paddle/utils/cpp_extension/cpp_extension.py +++ b/python/paddle/utils/cpp_extension/cpp_extension.py @@ -14,6 +14,10 @@ # isort: skip_file +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Sequence + import os import copy import re @@ -61,6 +65,9 @@ from ...base import core +if TYPE_CHECKING: + 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 @@ -76,7 +83,7 @@ CUDA_HOME = ROCM_HOME -def setup(**attr): +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 @@ -234,7 +241,9 @@ def setup(**attr): setuptools.setup(**attr) -def CppExtension(sources, *args, **kwargs): +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 @@ -284,7 +293,9 @@ def CppExtension(sources, *args, **kwargs): return setuptools.Extension(name, sources, *args, **kwargs) -def CUDAExtension(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 @@ -359,7 +370,7 @@ class BuildExtension(build_ext): """ @classmethod - def with_options(cls, **options): + def with_options(cls, **options: Any) -> type[BuildExtension]: """ Returns a BuildExtension subclass containing use-defined options. """ @@ -371,7 +382,7 @@ def __init__(self, *args, **kwargs): return cls_with_options - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: """ Attributes is initialized with following order: @@ -388,10 +399,10 @@ def __init__(self, *args, **kwargs): # whether containing cuda source file in Extensions self.contain_cuda_file = False - def initialize_options(self): + def initialize_options(self) -> None: super().initialize_options() - def finalize_options(self): + 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` @@ -399,7 +410,7 @@ def finalize_options(self): if self.output_dir is not None: self.build_lib = self.output_dir - def build_extensions(self): + def build_extensions(self) -> None: if OS_NAME.startswith("darwin"): self._valid_clang_compiler() @@ -640,7 +651,7 @@ def wrapper(source_filenames, strip_dir=0, output_dir=''): so_path = self.get_ext_fullpath(self.extensions[0]._full_name) _reset_so_rpath(so_path) - def get_ext_filename(self, fullname): + 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 = '.' @@ -658,7 +669,7 @@ def get_ext_filename(self, fullname): ext_name = split_str.join(name_items) return ext_name - def _valid_clang_compiler(self): + def _valid_clang_compiler(self) -> None: """ Make sure to use Clang as compiler on Mac platform """ @@ -672,7 +683,7 @@ def _valid_clang_compiler(self): linker_so=linker_infos, ) - def _check_abi(self): + def _check_abi(self) -> None: """ Check ABI Compatibility. """ @@ -697,7 +708,7 @@ def _check_abi(self): ) raise UserWarning(msg) - def _record_op_info(self): + def _record_op_info(self) -> None: """ Record custom op information. """ @@ -729,11 +740,11 @@ class EasyInstallCommand(easy_install): library file after extracting egg-info into site-packages. """ - def __init__(self, *args, **kwargs): + 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, **kwargs): + 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 @@ -764,7 +775,7 @@ class BuildCommand(build): """ @classmethod - def with_options(cls, **options): + def with_options(cls, **options: Any) -> type[BuildCommand]: """ Returns a BuildCommand subclass containing use-defined options. """ @@ -776,13 +787,13 @@ def __init__(self, *args, **kwargs): return cls_with_options - def __init__(self, *args, **kwargs): + 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): + 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. @@ -793,16 +804,16 @@ def initialize_options(self): def load( - name, - sources, - extra_cxx_cflags=None, - extra_cuda_cflags=None, - extra_ldflags=None, - extra_include_paths=None, - extra_library_paths=None, - build_directory=None, - verbose=False, -): + 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 @@ -863,22 +874,22 @@ def load( 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], optional): Specify additional flags used to compile CPP files. By default + 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], optional): Specify additional flags used to compile CUDA files. By default + 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], optional): Specify additional flags used to link shared library. See + 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], optional): Specify additional include path used to search header files. By default + 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], optional): Specify additional library path used to search library files. By default + 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, optional): Specify root directory path to put shared library file. If set 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.