Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Typing][B-53] Add type annotations for python/paddle/utils/cpp_extension/cpp_extension.py #65818

Merged
merged 1 commit into from
Jul 9, 2024
Merged

[Typing][B-53] Add type annotations for python/paddle/utils/cpp_extension/cpp_extension.py #65818

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 44 additions & 33 deletions python/paddle/utils/cpp_extension/cpp_extension.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@

# isort: skip_file

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Sequence

import os
import copy
import re
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
"""
Expand All @@ -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:

Expand All @@ -388,18 +399,18 @@ 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`
# and `load` interface rely on this attribute.
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()

Expand Down Expand Up @@ -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 = '.'
Expand All @@ -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
"""
Expand All @@ -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.
"""
Expand All @@ -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.
"""
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
"""
Expand All @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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 <https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html>`_
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 <https://gcc.gnu.org/onlinedocs/gcc/Link-Options.html>`_ 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.
Expand Down