chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
# Copyright (c) 2021 NVIDIA Corporation. 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.
|
||||
|
||||
from . import marker as marker
|
||||
from .api import (
|
||||
ignore_module,
|
||||
json_to_pdmodel, # noqa: F401
|
||||
load,
|
||||
save,
|
||||
to_static,
|
||||
)
|
||||
from .dy2static.logging_utils import set_code_level, set_verbosity
|
||||
from .dy2static.program_translator import enable_to_static
|
||||
from .marker import (
|
||||
not_to_static,
|
||||
)
|
||||
from .translated_layer import TranslatedLayer
|
||||
|
||||
__all__ = [
|
||||
'save',
|
||||
'load',
|
||||
'to_static',
|
||||
'ignore_module',
|
||||
'TranslatedLayer',
|
||||
'set_code_level',
|
||||
'set_verbosity',
|
||||
'not_to_static',
|
||||
'enable_to_static',
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,38 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from .convert_call_func import convert_call as Call # noqa: F401
|
||||
from .convert_operators import ( # noqa: F401
|
||||
convert_assert as Assert,
|
||||
convert_attr as Attr,
|
||||
convert_ifelse as IfElse,
|
||||
convert_len as Len,
|
||||
convert_load as Ld,
|
||||
convert_logical_and as And,
|
||||
convert_logical_not as Not,
|
||||
convert_logical_or as Or,
|
||||
convert_shape as Shape,
|
||||
convert_super as WrapSuper,
|
||||
convert_var_dtype as AsDtype,
|
||||
convert_while_loop as While,
|
||||
create_bool_as_type,
|
||||
indexable as Indexable,
|
||||
to_static_variable,
|
||||
unpack_by_structure as Unpack,
|
||||
)
|
||||
from .program_translator import convert_to_static # noqa: F401
|
||||
from .transformers import DygraphToStaticAst # noqa: F401
|
||||
from .utils import UndefinedVar, ast_to_source_code, saw # noqa: F401
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,33 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
|
||||
import ast
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
|
||||
def ast_to_source_code(ast_node):
|
||||
"""
|
||||
Transforms ast node into source code.
|
||||
"""
|
||||
if not isinstance(ast_node, (gast.AST, ast.AST)):
|
||||
raise TypeError(
|
||||
f"Type of ast_root should be gast.AST or ast.AST, but received {type(ast_node)}."
|
||||
)
|
||||
if isinstance(ast_node, gast.AST):
|
||||
ast_node = gast.gast_to_ast(ast_node)
|
||||
|
||||
ast.fix_missing_locations(ast_node)
|
||||
return ast.unparse(ast_node)
|
||||
@@ -0,0 +1,373 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import copy
|
||||
import functools
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import pdb # noqa: T100
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy
|
||||
|
||||
from paddle.nn import Layer
|
||||
|
||||
from .convert_operators import (
|
||||
convert_enumerate,
|
||||
convert_len,
|
||||
convert_print,
|
||||
convert_range,
|
||||
convert_zip,
|
||||
)
|
||||
from .logging_utils import TranslatorLogger
|
||||
from .program_translator import (
|
||||
StaticFunction,
|
||||
convert_to_static,
|
||||
unwrap_decorators,
|
||||
)
|
||||
from .utils import (
|
||||
TransformOptions,
|
||||
is_builtin,
|
||||
is_paddle_func,
|
||||
patch_method_guard,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
from types import ModuleType
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
translator_logger = TranslatorLogger()
|
||||
|
||||
|
||||
def builtin_modules():
|
||||
"""
|
||||
Return builtin modules.
|
||||
"""
|
||||
modules = [
|
||||
copy,
|
||||
collections,
|
||||
inspect,
|
||||
logging,
|
||||
numpy,
|
||||
os,
|
||||
pdb,
|
||||
re,
|
||||
]
|
||||
try:
|
||||
import six
|
||||
|
||||
modules.append(six)
|
||||
except ImportError:
|
||||
pass # do nothing
|
||||
|
||||
return modules
|
||||
|
||||
|
||||
BUILTIN_LIKELY_MODULES = builtin_modules()
|
||||
|
||||
|
||||
def add_ignore_module(modules: list[ModuleType]):
|
||||
"""
|
||||
Adds modules that ignore transcription
|
||||
"""
|
||||
global BUILTIN_LIKELY_MODULES
|
||||
for module in modules:
|
||||
if module not in BUILTIN_LIKELY_MODULES:
|
||||
BUILTIN_LIKELY_MODULES.append(module)
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def get_module_functions(module: ModuleType) -> list[Callable[..., Any]]:
|
||||
visited = set()
|
||||
|
||||
def _try_get_members(module) -> list[tuple[str, Any]]:
|
||||
try:
|
||||
return inspect.getmembers(module)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
def _get_module_functions(module):
|
||||
if module in visited:
|
||||
return []
|
||||
visited.add(module)
|
||||
results = []
|
||||
for _member_name, member in _try_get_members(module):
|
||||
if callable(member):
|
||||
results.append(member)
|
||||
if inspect.ismodule(member):
|
||||
results.extend(_get_module_functions(member))
|
||||
return results
|
||||
|
||||
return _get_module_functions(module)
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def get_module_defining_path(module: ModuleType) -> str | None:
|
||||
def _remove_module_init_suffix(file_path: str) -> str:
|
||||
return file_path.removesuffix("__init__.py")
|
||||
|
||||
if not hasattr(module, "__file__") or module.__file__ is None:
|
||||
return None
|
||||
return _remove_module_init_suffix(module.__file__)
|
||||
|
||||
|
||||
def is_unsupported(func):
|
||||
"""
|
||||
Checks whether the func is supported by dygraph to static graph.
|
||||
"""
|
||||
|
||||
builtin_module_paths = [
|
||||
module_path
|
||||
for module in BUILTIN_LIKELY_MODULES
|
||||
if (module_path := get_module_defining_path(module)) is not None
|
||||
]
|
||||
|
||||
# Skip module function by function defining path (For Python functions)
|
||||
if hasattr(func, "__code__") and func.__code__.co_filename:
|
||||
func_path = func.__code__.co_filename
|
||||
if any(
|
||||
func_path.startswith(module_path)
|
||||
for module_path in builtin_module_paths
|
||||
):
|
||||
translator_logger.log(
|
||||
2,
|
||||
"Whitelist: %s is part of built-in module and does not have to be transformed.",
|
||||
func,
|
||||
)
|
||||
return True
|
||||
|
||||
builtin_functions = [
|
||||
func
|
||||
for module in BUILTIN_LIKELY_MODULES
|
||||
for func in get_module_functions(module)
|
||||
]
|
||||
|
||||
# Skip module function by module members (For C/C++ binding functions)
|
||||
for builtin_fn in builtin_functions:
|
||||
if func is builtin_fn:
|
||||
translator_logger.log(
|
||||
2,
|
||||
"Whitelist: %s is part of built-in module and does not have to be transformed.",
|
||||
func,
|
||||
)
|
||||
return True
|
||||
|
||||
if is_paddle_func(func):
|
||||
translator_logger.log(
|
||||
2,
|
||||
"Whitelist: %s is part of Paddle module and does not have to be transformed.",
|
||||
func,
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
class StaticLayerWrapper:
|
||||
def __init__(self, layer):
|
||||
self.layer = layer
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
with patch_method_guard(
|
||||
self.layer, "forward", convert_call(self.layer.forward)
|
||||
):
|
||||
return self.layer(*args, **kwargs)
|
||||
|
||||
|
||||
def convert_call(func):
|
||||
"""
|
||||
Converts a function call which needs to be transformed to static function.
|
||||
|
||||
Args:
|
||||
func (callable): A callable function or method to convert.
|
||||
|
||||
Returns:
|
||||
Callable: A converted function.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
|
||||
>>> import paddle
|
||||
>>> from paddle.jit.dy2static import Call
|
||||
|
||||
>>> paddle.enable_static()
|
||||
>>> def dyfunc(x):
|
||||
... if paddle.mean(x) < 0:
|
||||
... x_v = x - 1
|
||||
... else:
|
||||
... x_v = x + 1
|
||||
... return x_v
|
||||
>>> new_func = Call(dyfunc)
|
||||
>>> x = paddle.tensor.manipulation.fill_constant(shape=[3, 3], value=0, dtype='float64')
|
||||
>>> x_v = new_func(x)
|
||||
|
||||
>>> exe = paddle.static.Executor(paddle.CPUPlace())
|
||||
>>> out = exe.run(fetch_list=[x_v])
|
||||
>>> print(out[0])
|
||||
[[1. 1. 1.]
|
||||
[1. 1. 1.]
|
||||
[1. 1. 1.]]
|
||||
|
||||
"""
|
||||
translator_logger.log(1, "Convert callable object: convert %s.", func)
|
||||
func_self = None
|
||||
converted_call = None
|
||||
|
||||
# Function in convert_call may be decorated by another `@to_static`,
|
||||
# in this case, unwraps it into a raw method or function.
|
||||
_, func = unwrap_decorators(func)
|
||||
|
||||
if not TransformOptions.check_fn_need_transform(
|
||||
func, TransformOptions.ToStaticMode.AST
|
||||
):
|
||||
translator_logger.log(
|
||||
2,
|
||||
"%s is not converted when it is decorated by 'paddle.jit.not_to_static'.",
|
||||
func,
|
||||
)
|
||||
return func
|
||||
|
||||
if is_builtin(func, "len"):
|
||||
return convert_len
|
||||
|
||||
if is_builtin(func, "zip"):
|
||||
return convert_zip
|
||||
|
||||
if is_builtin(func, "range"):
|
||||
return convert_range
|
||||
|
||||
if is_builtin(func, "enumerate"):
|
||||
return convert_enumerate
|
||||
|
||||
if is_builtin(func, "print"):
|
||||
return convert_print
|
||||
|
||||
if is_builtin(func) or is_unsupported(func):
|
||||
return func
|
||||
|
||||
if inspect.isgeneratorfunction(func):
|
||||
# NOTE(xiongkun03): inspect.isfunction() will return True even though func is a generator function.
|
||||
# If we don't deal generator function here, we will regard it as normal function and get errors in some
|
||||
# occasion.
|
||||
number_of_stars = 30
|
||||
translator_logger.warn(
|
||||
"\n\n"
|
||||
+ "*" * number_of_stars
|
||||
+ f"\nYour function:`{func.__name__}` doesn't support to transform to static function because it is a generator function, it will be run as-is."
|
||||
+ "\n"
|
||||
+ "*" * number_of_stars
|
||||
+ "\n\n"
|
||||
)
|
||||
return func
|
||||
|
||||
if inspect.isfunction(func):
|
||||
# TODO(liym27): If func is a lambda function, special conversion is needed.
|
||||
if func.__name__ == '<lambda>':
|
||||
return func
|
||||
try:
|
||||
# Note(Aurelius84): Because `@to_static` returns a class instance instead of
|
||||
# a function. This will modify the value referring to itself in `__globals__`.
|
||||
|
||||
# For example:
|
||||
#
|
||||
# @to_static
|
||||
# def foo(x):
|
||||
# return x
|
||||
#
|
||||
# `foo` will be converted into a wrapper class, suppose as `StaticFunction`.
|
||||
# And `foo.__globals__['foo']` will still return this `StaticFunction` instead of
|
||||
# `foo` function. So `isinstance(fn, StaticFunction)` is added here.
|
||||
_origfunc = inspect.unwrap(func)
|
||||
global_functions = set()
|
||||
for fn in _origfunc.__globals__.values():
|
||||
if inspect.isfunction(fn):
|
||||
global_functions.add(fn)
|
||||
elif isinstance(fn, StaticFunction):
|
||||
_, fn = unwrap_decorators(fn)
|
||||
global_functions.add(fn)
|
||||
elif inspect.isclass(fn):
|
||||
if isinstance(
|
||||
fn.__dict__.get(func.__name__, None), staticmethod
|
||||
):
|
||||
global_functions.add(
|
||||
func
|
||||
) # Add func to ensure that we will convert
|
||||
|
||||
if func in global_functions:
|
||||
converted_call = convert_to_static(func)
|
||||
func_self = getattr(func, '__self__', None)
|
||||
else:
|
||||
# NOTE:
|
||||
# If func is not in __globals__, it does not need to be transformed
|
||||
# because it has been transformed before.
|
||||
translator_logger.warn(
|
||||
"%s doesn't have to be transformed to static function because it has been transformed before, it will be run as-is.",
|
||||
func,
|
||||
)
|
||||
converted_call = func
|
||||
except AttributeError:
|
||||
# NOTE:
|
||||
# If func is not in __globals__, it does not need to be transformed
|
||||
# because it has been transformed before.
|
||||
converted_call = None
|
||||
except OSError:
|
||||
# NOTE:
|
||||
# If func has been decorated, its source code can not be get
|
||||
# so that it can not be transformed to static function.
|
||||
converted_call = None
|
||||
elif inspect.ismethod(func):
|
||||
try:
|
||||
converted_call = convert_to_static(func)
|
||||
func_self = getattr(func, '__self__', None)
|
||||
except OSError:
|
||||
# NOTE: func may have been decorated.
|
||||
converted_call = None
|
||||
|
||||
elif hasattr(func, '__class__') and callable(func.__class__):
|
||||
if hasattr(func, 'forward') and isinstance(func, Layer):
|
||||
return StaticLayerWrapper(func)
|
||||
else:
|
||||
try:
|
||||
call_func = func.__class__.__call__
|
||||
converted_call = convert_to_static(call_func)
|
||||
func_self = func
|
||||
except (OSError, TypeError):
|
||||
# NOTE:
|
||||
# If `func` is a class which is being initialized, for example `convert_call(Foo)()`,
|
||||
# it doesn't need to be transformed
|
||||
func_self = None if func_self else func_self
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Callable {func} can not be transformed at present."
|
||||
)
|
||||
|
||||
if converted_call is None:
|
||||
translator_logger.warn(
|
||||
"%s doesn't have to be transformed to static function, and it will be run as-is.",
|
||||
func,
|
||||
)
|
||||
return func
|
||||
|
||||
if func_self is not None:
|
||||
converted_call = functools.partial(converted_call, func_self)
|
||||
return converted_call
|
||||
@@ -0,0 +1,878 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
|
||||
import paddle
|
||||
from paddle.autograd.py_layer import PyLayerMeta
|
||||
from paddle.base.data_feeder import convert_dtype
|
||||
from paddle.base.dygraph.base import _convert_into_variable, in_to_static_mode
|
||||
from paddle.base.framework import Variable, core, default_main_program
|
||||
from paddle.framework import use_pir_api
|
||||
from paddle.jit.utils import OrderedSet
|
||||
from paddle.pir import Value
|
||||
from paddle.static.amp.fp16_utils import AmpOptions
|
||||
from paddle.utils import is_sequence, map_structure
|
||||
|
||||
from .py_layer import StaticPyLayer
|
||||
from .utils import (
|
||||
RETURN_NO_VALUE_VAR_NAME,
|
||||
Dygraph2StaticException,
|
||||
GetterSetterHelper,
|
||||
UndefinedVar,
|
||||
create_undefined_variable,
|
||||
)
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def to_static_variable(x, dtype=None):
|
||||
'''
|
||||
Translate a Python Tensor to PaddlePaddle static graph Tensor
|
||||
'''
|
||||
if isinstance(x, bool):
|
||||
dtype = 'bool' if dtype is None else dtype
|
||||
return paddle.full(shape=[], dtype=dtype, fill_value=x)
|
||||
if isinstance(x, float):
|
||||
dtype = 'float64' if dtype is None else dtype
|
||||
return paddle.full(shape=[], dtype=dtype, fill_value=x)
|
||||
if isinstance(x, int):
|
||||
dtype = 'int64' if dtype is None else dtype
|
||||
return paddle.full(shape=[], dtype=dtype, fill_value=x)
|
||||
if not use_pir_api() and (isinstance(x, UndefinedVar) or x is None):
|
||||
"""
|
||||
for early return case, we need a variable to represent None, current we use data_layer_not_check.
|
||||
"""
|
||||
return create_undefined_variable()
|
||||
if is_sequence(x):
|
||||
return map_structure(to_static_variable, x)
|
||||
return x
|
||||
|
||||
|
||||
def convert_attr(x, attr):
|
||||
# TODO(cleanup-legacy-ir): In PIR mode, the size attr in
|
||||
# Value and Tensor are unified. So we don't need to transform
|
||||
# the size attr into a method call. The AttributeJstTransformer and
|
||||
# convert_attr can be safely removed.
|
||||
if (
|
||||
isinstance(x, Variable)
|
||||
and not isinstance(x, paddle.Tensor)
|
||||
and attr == "size"
|
||||
):
|
||||
return x.size()
|
||||
else:
|
||||
return getattr(x, attr)
|
||||
|
||||
|
||||
def convert_load(x):
|
||||
# convert dygraph `PyLayer` into StaticPyLayer
|
||||
if isinstance(x, PyLayerMeta):
|
||||
return StaticPyLayer(x)
|
||||
|
||||
if in_to_static_mode():
|
||||
if isinstance(x, paddle.Tensor):
|
||||
"""
|
||||
TODO:(@xiongkun) may run convert_load in dygraph mode, which should be fixed.
|
||||
"""
|
||||
return _convert_into_variable(x)
|
||||
|
||||
# get the new output of the var
|
||||
if isinstance(x, Value):
|
||||
from paddle.jit.dy2static.parameter_recorder import (
|
||||
_global_inplace_map,
|
||||
)
|
||||
|
||||
new_var = _global_inplace_map.get(
|
||||
paddle.static.default_main_program(), x
|
||||
)
|
||||
if new_var is not None:
|
||||
return new_var
|
||||
|
||||
if isinstance(x, Variable):
|
||||
cur_block = default_main_program().current_block()
|
||||
|
||||
from paddle.jit.dy2static.program_translator import (
|
||||
ProgramTranslator,
|
||||
)
|
||||
|
||||
new_var = ProgramTranslator.get_instance()._inplace_map.get(
|
||||
cur_block.program, x.desc.id()
|
||||
)
|
||||
if new_var is not None:
|
||||
return new_var
|
||||
|
||||
if x is paddle.amp.auto_cast and not use_pir_api():
|
||||
return convert_auto_cast
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def indexable(x, code=None):
|
||||
if isinstance(x, (Variable, Value)):
|
||||
return x
|
||||
elif hasattr(x, '__iter__'):
|
||||
return list(x)
|
||||
elif hasattr(x, '__len__') and hasattr(
|
||||
x, '__getitem__'
|
||||
): # used for customed type and non-iterable type.
|
||||
return x
|
||||
else:
|
||||
raise RuntimeError("X can't be convert into indexable.")
|
||||
|
||||
|
||||
def unpack_by_structure(target, structure):
|
||||
"""unified unpack interface for paddle and python."""
|
||||
if isinstance(target, (Variable, Value)):
|
||||
return _unpack_by_structure_paddle(target, structure)
|
||||
else:
|
||||
return _unpack_by_structure_python(target, structure)
|
||||
|
||||
|
||||
def _unpack_by_structure_python(target, structure):
|
||||
"""TODO(xiongkun): analysis the differences between python and paddle unpack."""
|
||||
return _unpack_by_structure_paddle(target, structure)
|
||||
|
||||
|
||||
def _unpack_by_structure_paddle(target, structure):
|
||||
if structure == 1:
|
||||
return target
|
||||
ret = []
|
||||
for idx, ele in enumerate(structure):
|
||||
if ele == 1:
|
||||
ret.append(target[idx])
|
||||
continue
|
||||
if isinstance(ele, list):
|
||||
ret.append(unpack_by_structure(target[idx], ele))
|
||||
continue
|
||||
raise AssertionError("structure element must be 1 or list")
|
||||
return ret
|
||||
|
||||
|
||||
def convert_while_loop(
|
||||
cond, body, getter, setter, return_name_ids=None, push_pop_names=None
|
||||
):
|
||||
"""
|
||||
A function representation of a Python ``while`` statement.
|
||||
|
||||
Args:
|
||||
cond(Callable): A callable object that returns a boolean variable to control whether to execute the loop body. It takes ``loop_vars`` as arguments.
|
||||
body(Callable): A callable object that returns a tuple or list of variables with the same arguments ``loops_vars`` as ``cond`` .
|
||||
get_args(callable): Get all arguments that needed in true_fn and false_fn.
|
||||
set_args(callable): Update arguments that modified in trure_fn and false_fn.
|
||||
return_name_ids(list[string], optional): the returned names.
|
||||
push_pop_names(list[string], optional): the names on which called .append() or .pop().
|
||||
|
||||
Returns:
|
||||
A list or tuple of variables which returned by ``body``.
|
||||
"""
|
||||
|
||||
# NOTE: It may be slower if cond is very expensive, but usually cond is just O(1).
|
||||
# If loop_vars is changed during cond callable, then it causes bug, but current logical_and/logical_not/... doesn't change the loop_vars.
|
||||
pred = cond()
|
||||
if isinstance(pred, (Variable, Value)):
|
||||
_run_paddle_while(
|
||||
cond, body, getter, setter, return_name_ids, push_pop_names
|
||||
)
|
||||
else:
|
||||
_run_py_while(cond, body, getter, setter)
|
||||
|
||||
|
||||
def _convert_tensor_array_if_necessary(setterhelper, push_pop_names):
|
||||
push_pop_vars = setterhelper.get(push_pop_names)
|
||||
if push_pop_vars is None:
|
||||
return
|
||||
|
||||
def maybe_to_tensor_array(v):
|
||||
if isinstance(v, list):
|
||||
dtype = (
|
||||
paddle.base.libpaddle.DataType.UNDEFINED
|
||||
if use_pir_api()
|
||||
else "float32"
|
||||
)
|
||||
return paddle.tensor.create_array(dtype, initialized_list=v)
|
||||
else:
|
||||
return v
|
||||
|
||||
setterhelper.set(
|
||||
push_pop_names, [maybe_to_tensor_array(v) for v in push_pop_vars]
|
||||
)
|
||||
|
||||
|
||||
def _run_paddle_while(
|
||||
cond, body, getter, setter, return_name_ids, push_pop_names
|
||||
):
|
||||
# NOTE: loop_vars of Paddle op `control_flow.while_loop` must be Paddle Tensors.
|
||||
helper = GetterSetterHelper(getter, setter, return_name_ids, push_pop_names)
|
||||
_convert_tensor_array_if_necessary(helper, push_pop_names)
|
||||
|
||||
union_name = (
|
||||
OrderedSet(return_name_ids) if return_name_ids else OrderedSet()
|
||||
) | (OrderedSet(push_pop_names) if push_pop_names else OrderedSet())
|
||||
union_name = list(union_name)
|
||||
|
||||
def new_body_fn(*args):
|
||||
"""wrap the body() and add return value for `while_loop`
|
||||
the args may be differ from getter().
|
||||
"""
|
||||
mutable_loop_vars = args
|
||||
helper.set(union_name, mutable_loop_vars)
|
||||
body()
|
||||
return helper.get(union_name)
|
||||
|
||||
def new_cond_fn(*args):
|
||||
"""cond is a zero-args function, which is not
|
||||
compatible with `while_loop`.
|
||||
"""
|
||||
return cond()
|
||||
|
||||
# UndefinedVar will become data layer not check variable with value=NO_VALUE_MAGIC.
|
||||
loop_vars = [
|
||||
to_static_variable(var) if not isinstance(var, UndefinedVar) else var
|
||||
for var in helper.get(union_name)
|
||||
]
|
||||
helper.set(union_name, loop_vars) # change the non-local var to variable
|
||||
# variable maybe modified to inner var. change it into
|
||||
from paddle.static.nn import while_loop
|
||||
|
||||
loop_vars = while_loop(new_cond_fn, new_body_fn, loop_vars)
|
||||
helper.set(union_name, loop_vars)
|
||||
return loop_vars
|
||||
|
||||
|
||||
def _run_py_while(cond, body, getter, setter):
|
||||
while True:
|
||||
pred = cond()
|
||||
if isinstance(pred, (Variable, Value)):
|
||||
raise Dygraph2StaticException(
|
||||
"python while pred change from bool to variable."
|
||||
)
|
||||
if not pred:
|
||||
break
|
||||
body()
|
||||
|
||||
|
||||
def convert_logical_and(x_func, y_func):
|
||||
"""
|
||||
A function representation of a Python ``and`` statement.
|
||||
|
||||
Args:
|
||||
x_func(callable): x_func() is the left hand operand of ``and`` operator. x_func() is bool or Tensor.
|
||||
y_func(callable): y_func() is the right hand operand of ``and`` operator. y_func() is bool or Tensor.
|
||||
|
||||
Returns:
|
||||
A python bool variable or a bool Tensor.
|
||||
|
||||
NOTE(liym27):
|
||||
1) The operands are executed sequentially according to the running logic of Python. So here the arguments
|
||||
should be callable.
|
||||
2) If the left hand operand is False, the right hand operand should be executed.
|
||||
|
||||
For example:
|
||||
a = x > 1 and y < 1
|
||||
Transformed code:
|
||||
a = paddle.jit.dy2static.convert_logical_and(lambda:x>1, lambda:y<1)
|
||||
|
||||
In `convert_logical_and(lambda:x>1, lambda:y<1)`, `lambda:y<1` must be run after `lambda:x>1`. And
|
||||
if `x>1` is False, `y<1` should NOT be run.
|
||||
"""
|
||||
x_value = x_func()
|
||||
if not isinstance(x_value, (Variable, Value)):
|
||||
return _run_py_logical_and(lambda: x_value, y_func)
|
||||
|
||||
y_value = y_func()
|
||||
if not isinstance(y_value, (Variable, Value)):
|
||||
return _run_py_logical_and(lambda: y_value, lambda: x_value)
|
||||
|
||||
return _run_paddle_logical_and(x_value, y_value)
|
||||
|
||||
|
||||
def _run_paddle_logical_and(x, y):
|
||||
x = cast_bool_if_necessary(x)
|
||||
y = cast_bool_if_necessary(y)
|
||||
return paddle.logical_and(x, y)
|
||||
|
||||
|
||||
def _run_py_logical_and(x_func, y_func):
|
||||
x_value = x_func()
|
||||
assert not isinstance(x_value, (Variable, Value))
|
||||
|
||||
# NOTE(liym27):
|
||||
# 1. Returns y_func() if x_value is False;
|
||||
# 2. If x_value is False, y_func() should not be run.
|
||||
return x_value and y_func()
|
||||
|
||||
|
||||
def convert_logical_or(x_func, y_func):
|
||||
"""
|
||||
A function representation of a Python ``or`` statement.
|
||||
|
||||
Args:
|
||||
x_func(callable): x_func() is the left hand operand of ``or`` operator. x_func() is bool or Tensor.
|
||||
y_func(callable): y_func() is the right hand operand of ``or`` operator. y_func() is bool or Tensor.
|
||||
|
||||
Returns:
|
||||
A python bool variable or a bool Tensor.
|
||||
|
||||
NOTE(liym27):
|
||||
1) The operands are executed sequentially according to the running logic of Python. So here the arguments
|
||||
should be callable.
|
||||
2) If the left hand operand is True, the right hand operand should be executed.
|
||||
|
||||
For example:
|
||||
a = x > 1 or y < 1
|
||||
Transformed code:
|
||||
a = paddle.jit.dy2static.convert_logical_or(lambda:x>1, lambda:y<1)
|
||||
|
||||
In `convert_logical_or(lambda:x>1, lambda:y<1)`, `lambda:y<1` must be run after `lambda:x>1`. And
|
||||
if `x>1` is True, `y<1` should NOT be run.
|
||||
"""
|
||||
x_value = x_func()
|
||||
if not isinstance(x_value, (Variable, Value)):
|
||||
return _run_py_logical_or(lambda: x_value, y_func)
|
||||
|
||||
y_value = y_func()
|
||||
if not isinstance(y_value, (Variable, Value)):
|
||||
return _run_py_logical_or(lambda: y_value, lambda: x_value)
|
||||
|
||||
return _run_paddle_logical_or(x_value, y_value)
|
||||
|
||||
|
||||
def _run_paddle_logical_or(x, y):
|
||||
x = cast_bool_if_necessary(x)
|
||||
y = cast_bool_if_necessary(y)
|
||||
return paddle.logical_or(x, y)
|
||||
|
||||
|
||||
def _run_py_logical_or(x_func, y_func):
|
||||
x_value = x_func()
|
||||
assert not isinstance(x_value, (Variable, Value))
|
||||
|
||||
# NOTE(liym27):
|
||||
# 1. Returns y_func() if x_value is False;
|
||||
# 2. If x_value is True, y_func() should not be run.
|
||||
return x_value or y_func()
|
||||
|
||||
|
||||
def convert_logical_not(x):
|
||||
"""
|
||||
A function representation of a Python ``not`` statement.
|
||||
|
||||
Args:
|
||||
x(bool|Tensor): Operand of ``not`` operator.
|
||||
|
||||
Returns:
|
||||
A python bool variable or a bool Tensor.
|
||||
"""
|
||||
|
||||
if isinstance(x, (Variable, Value)):
|
||||
return _run_paddle_logical_not(x)
|
||||
else:
|
||||
return _run_py_logical_not(x)
|
||||
|
||||
|
||||
def _run_paddle_logical_not(x):
|
||||
x = cast_bool_if_necessary(x)
|
||||
return paddle.logical_not(x)
|
||||
|
||||
|
||||
def _run_py_logical_not(x):
|
||||
return not x
|
||||
|
||||
|
||||
def convert_ifelse(
|
||||
pred,
|
||||
true_fn,
|
||||
false_fn,
|
||||
get_args,
|
||||
set_args,
|
||||
return_name_ids,
|
||||
push_pop_names=None,
|
||||
):
|
||||
"""
|
||||
A function representation of a Python ``if/else`` statement.
|
||||
|
||||
Args:
|
||||
pred(bool|Tensor): A boolean Tensor which determines whether to return the result of ``true_fn`` or ``false_fn`` .
|
||||
true_fn(callable): A callable to be performed if ``pred`` is true.
|
||||
false_fn(callable): A callable to be performed if ``pred`` is false.
|
||||
get_args(callable): Get all arguments that needed in true_fn and false_fn.
|
||||
set_args(callable): Update arguments that modified in trure_fn and false_fn.
|
||||
return_name_ids(list[string], optional): the returned names.
|
||||
push_pop_names(list[string], optional): the names on which called .append() or .pop().
|
||||
|
||||
Returns:
|
||||
``true_fn()`` if the predicate ``pred`` is true else ``false_fn()`` .
|
||||
|
||||
"""
|
||||
if isinstance(pred, (Variable, Value)):
|
||||
out = _run_paddle_cond(
|
||||
pred,
|
||||
true_fn,
|
||||
false_fn,
|
||||
get_args,
|
||||
set_args,
|
||||
return_name_ids,
|
||||
push_pop_names,
|
||||
)
|
||||
else:
|
||||
out = _run_py_ifelse(
|
||||
pred, true_fn, false_fn, get_args, set_args, return_name_ids
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def _run_paddle_cond(
|
||||
pred, true_fn, false_fn, get_args, set_args, return_name_ids, push_pop_names
|
||||
):
|
||||
"""
|
||||
Paddle cond API will evaluate both true_fn and false_fn codes.
|
||||
"""
|
||||
helper = GetterSetterHelper(
|
||||
get_args, set_args, return_name_ids, push_pop_names
|
||||
)
|
||||
_convert_tensor_array_if_necessary(helper, push_pop_names)
|
||||
pred = cast_bool_if_necessary(pred)
|
||||
init_args = helper.get(return_name_ids)
|
||||
from paddle.jit.dy2static.parameter_recorder import _global_inplace_map
|
||||
from paddle.jit.dy2static.program_translator import ProgramTranslator
|
||||
|
||||
if use_pir_api():
|
||||
inplace_map = _global_inplace_map
|
||||
else:
|
||||
inplace_map = ProgramTranslator.get_instance()._inplace_map
|
||||
union_name = None
|
||||
# TODO(@xiongkun) lambda can have push_pop_names, which will cause error.
|
||||
if return_name_ids is None and push_pop_names is None:
|
||||
union_name = None
|
||||
else:
|
||||
union_name = (
|
||||
OrderedSet(return_name_ids) if return_name_ids else OrderedSet()
|
||||
) | (OrderedSet(push_pop_names) if push_pop_names else OrderedSet())
|
||||
union_name = list(union_name)
|
||||
|
||||
def new_true_fn():
|
||||
nonlocal union_name
|
||||
# init args may contain mutable python container like [var, 2], we copy then like in while_loop
|
||||
inplace_map_checkpoint = inplace_map.save_checkpoint()
|
||||
helper.set(
|
||||
return_name_ids,
|
||||
paddle.utils.copy_mutable_vars(init_args),
|
||||
)
|
||||
ret = true_fn()
|
||||
# IfExpr will return a non-None return value, so we just return ret.
|
||||
# We assume normal return has no return value.
|
||||
if ret is None:
|
||||
ret = helper.get(union_name)
|
||||
inplace_map.restore_checkpoint(inplace_map_checkpoint)
|
||||
return ret
|
||||
|
||||
def new_false_fn():
|
||||
nonlocal union_name
|
||||
# init args may contain mutable python container like [var, 2], we copy then like in while_loop
|
||||
inplace_map_checkpoint = inplace_map.save_checkpoint()
|
||||
helper.set(
|
||||
return_name_ids,
|
||||
paddle.utils.copy_mutable_vars(init_args),
|
||||
)
|
||||
ret = false_fn()
|
||||
if ret is None:
|
||||
ret = helper.get(union_name)
|
||||
inplace_map.restore_checkpoint(inplace_map_checkpoint)
|
||||
return ret
|
||||
|
||||
try:
|
||||
cond_outs = paddle.static.nn.cond(
|
||||
pred, new_true_fn, new_false_fn, None, union_name
|
||||
)
|
||||
except Exception as e:
|
||||
if re.search(
|
||||
"Unsupported return type of true_fn and false_fn in cond", str(e)
|
||||
):
|
||||
raise Dygraph2StaticException(
|
||||
f"Your if/else have different return type. TODO: add link to modify. {e}"
|
||||
)
|
||||
if re.search("Incompatible return values of", str(e)):
|
||||
raise Dygraph2StaticException(
|
||||
f"Your if/else have different number of return value. TODO: add link to modify. {e}"
|
||||
)
|
||||
raise e
|
||||
get_args = lambda: helper.get(union_name)
|
||||
set_args = lambda vs: helper.set(union_name, vs)
|
||||
return _recover_args_state(cond_outs, get_args, set_args, union_name)
|
||||
|
||||
|
||||
def _run_py_ifelse(
|
||||
pred, true_fn, false_fn, get_args, set_args, return_name_ids
|
||||
):
|
||||
"""
|
||||
Evaluate python original branch function if-else.
|
||||
"""
|
||||
py_outs = true_fn() if pred else false_fn()
|
||||
return py_outs
|
||||
|
||||
|
||||
def _remove_no_value_return_var(out):
|
||||
if isinstance(out, tuple) and len(out) > 0:
|
||||
processed_out = out
|
||||
align_ret = out[0]
|
||||
if isinstance(align_ret, tuple):
|
||||
for index, item in enumerate(align_ret):
|
||||
if isinstance(item, (Variable, Value)) and (
|
||||
RETURN_NO_VALUE_VAR_NAME in item.name
|
||||
):
|
||||
# return None
|
||||
if index == 0:
|
||||
processed_out = (None, *out[1:])
|
||||
elif index == 1:
|
||||
processed_out = align_ret[:1] + out[1:]
|
||||
else:
|
||||
processed_out = (align_ret[:index], *out[1:])
|
||||
break
|
||||
|
||||
for index, item in enumerate(processed_out):
|
||||
if isinstance(item, (Variable, Value)) and (
|
||||
RETURN_NO_VALUE_VAR_NAME in item.name
|
||||
):
|
||||
processed_out = processed_out[:index]
|
||||
|
||||
if not processed_out:
|
||||
return None
|
||||
elif len(processed_out) == 1:
|
||||
return processed_out[0]
|
||||
else:
|
||||
return processed_out
|
||||
|
||||
else:
|
||||
return out
|
||||
|
||||
|
||||
def _check_no_undefined_var(outs, names, branch_name):
|
||||
if names is None:
|
||||
return
|
||||
if not isinstance(outs, (list, tuple)):
|
||||
outs = [outs]
|
||||
for var, name in zip(list(outs), names):
|
||||
if isinstance(var, UndefinedVar):
|
||||
raise ValueError(
|
||||
f"Required '{name}' must be initialized both in if-else branch, but found it not initialized in '{branch_name}'."
|
||||
)
|
||||
|
||||
|
||||
def _recover_args_state(outs, get_args, set_args, return_name_ids):
|
||||
"""
|
||||
Currently we support variant length of early return statement by padding
|
||||
_no_return_value.
|
||||
|
||||
# TODO(dev): We shall consider to evaluate whether should support this for Python if-else?
|
||||
"""
|
||||
# IfExpr's return_name_ids maybe None
|
||||
if return_name_ids is None:
|
||||
return outs
|
||||
|
||||
init_args = get_args()
|
||||
# recover args state
|
||||
num_outs = len(return_name_ids)
|
||||
num_args = len(init_args)
|
||||
assert num_outs <= num_args
|
||||
|
||||
if num_args == 1:
|
||||
final_outs = (
|
||||
(outs,) if not isinstance(outs, (list, tuple)) else tuple(outs)
|
||||
)
|
||||
else:
|
||||
outs = (outs,) if num_outs == 1 else tuple(outs)
|
||||
final_outs = outs + init_args[num_outs:]
|
||||
|
||||
set_args(final_outs)
|
||||
return final_outs
|
||||
|
||||
|
||||
def convert_len(var):
|
||||
"""
|
||||
Returns variable(length) from shape ops based on var.type
|
||||
|
||||
Note: In addition to some ast transformations, some block-related
|
||||
operations are added in `len` transformation, such as appending
|
||||
`shape_op` in var.block.
|
||||
"""
|
||||
if isinstance(var, Variable):
|
||||
assert var.ndim > 0, "len() of a 0-D tensor is wrong"
|
||||
if var.type in [
|
||||
core.VarDesc.VarType.DENSE_TENSOR,
|
||||
core.VarDesc.VarType.SELECTED_ROWS,
|
||||
]:
|
||||
# Note: Length of var may be known ahead of time in dygraph,
|
||||
# but it probably represents batch size which can be variant.
|
||||
# so we return a variable dynamically inferred from var.shape.
|
||||
if (
|
||||
var.shape[0] > 0
|
||||
and var.type == core.VarDesc.VarType.DENSE_TENSOR
|
||||
):
|
||||
return var.shape[0]
|
||||
return paddle.shape(var)[0]
|
||||
elif var.type == core.VarDesc.VarType.DENSE_TENSOR_ARRAY:
|
||||
return paddle.tensor.array_length(var)
|
||||
else:
|
||||
raise TypeError(
|
||||
f'len(var) only supports DenseTensor/DenseTensorArray/SelectedRows, but received {type(var)}.'
|
||||
)
|
||||
elif isinstance(var, Value):
|
||||
if var.is_dense_tensor_type() or var.is_selected_row_type():
|
||||
assert var.ndim > 0, "len() of a 0-D tensor is wrong"
|
||||
# Note: Length of var may be known ahead of time in dygraph,
|
||||
# but it probably represents batch size which can be variant.
|
||||
# so we return a variable dynamically inferred from var.shape.
|
||||
if var.shape[0] > 0 and var.is_dense_tensor_type():
|
||||
return var.shape[0]
|
||||
return paddle.shape(var)[0]
|
||||
elif var.is_dense_tensor_array_type():
|
||||
return paddle.tensor.array_length(var)
|
||||
else:
|
||||
raise TypeError(
|
||||
'len(var) only supports DenseTensor/DenseTensorArray/SelectedRows, '
|
||||
+ f'but received {type(var)}.'
|
||||
)
|
||||
else:
|
||||
if isinstance(var, VariableTuple):
|
||||
return var.__len__()
|
||||
return len(var)
|
||||
|
||||
|
||||
def convert_zip(*args):
|
||||
for i, arg in enumerate(args):
|
||||
if isinstance(arg, (Variable, Value)) and arg.shape[0] == -1:
|
||||
raise RuntimeError(
|
||||
"Not support zip(tensor, ...) when tensor.shape[0] == -1, "
|
||||
f"but found args[{i}].shape[0] == -1 in 'zip'"
|
||||
)
|
||||
return zip(*args)
|
||||
|
||||
|
||||
def convert_super(super_fn):
|
||||
if super_fn is super:
|
||||
return super_fn
|
||||
return lambda cls, instance: super_fn()
|
||||
|
||||
|
||||
# TODO(xiongkun): delete when list<variable> is ready.
|
||||
class VariableTuple:
|
||||
"""
|
||||
this class will cause enumerate can't be wrapped by other iterator change function.
|
||||
this will be fixed when list<Variable> is produced.
|
||||
VariableTuple can only deal with variables which is fixed.
|
||||
"""
|
||||
|
||||
def __init__(self, var, start=0):
|
||||
self.var = var
|
||||
self.len = convert_len(var)
|
||||
if isinstance(self.len, (Variable, Value)):
|
||||
self.rag = paddle.arange(start, start + self.len, 1, "int64")
|
||||
else:
|
||||
self.rag = range(start, start + self.len)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.rag[idx], self.var[idx]
|
||||
|
||||
def __len__(self):
|
||||
return self.len
|
||||
|
||||
|
||||
def convert_enumerate(*args):
|
||||
has_variable = any(isinstance(x, (Variable, Value)) for x in args)
|
||||
if has_variable:
|
||||
return VariableTuple(*args)
|
||||
return enumerate(*args)
|
||||
|
||||
|
||||
def convert_range(*args):
|
||||
has_variable = any(isinstance(x, (Variable, Value)) for x in args)
|
||||
# NOTE(SigureMo): Add an `Assign` OP after the Tensor input to mark it as a variable, which can
|
||||
# avoid confusing it with the scalar case in `arange` API.
|
||||
# For example:
|
||||
# ```python
|
||||
# l = []
|
||||
# for i in range(n):
|
||||
# l.append(i)
|
||||
# ```
|
||||
# - If `n` is a scalar (e.g., `n=10`), we expect to create an `ArangeOp` with a fixed output shape [10].
|
||||
# - If `n` is a Tensor (e.g., `n=full([], 10, "int32")`), we expect to create an `ArangeOp` with a dynamic
|
||||
# output shape [-1]. To ensure the python level and graph level all recognize this is data-dependent control
|
||||
# flow.
|
||||
# However, we can't distinguish the scalar case and the Tensor case when creating the `ArangeOp`. Because
|
||||
# the scalar case also be convert as a `Full` OP output.
|
||||
# So we add an `Assign` OP after the Tensor input to **mark** it as a variable, which can avoid confusing
|
||||
# it with the scalar case.
|
||||
is_full_op_output = lambda x: (
|
||||
isinstance(x, Value)
|
||||
and x.get_defining_op()
|
||||
and x.get_defining_op().name() == "pd_op.full"
|
||||
)
|
||||
args = [
|
||||
paddle.assign(arg) if is_full_op_output(arg) else arg for arg in args
|
||||
]
|
||||
if has_variable:
|
||||
if len(args) == 1:
|
||||
return paddle.arange(0, args[0], 1, "int64")
|
||||
if len(args) == 2:
|
||||
return paddle.arange(args[0], args[1], 1, "int64")
|
||||
if len(args) == 3:
|
||||
return paddle.arange(args[0], args[1], args[2], "int64")
|
||||
return range(*args)
|
||||
|
||||
|
||||
def convert_shape(x):
|
||||
"""
|
||||
A function representation of the shape of variable.
|
||||
"""
|
||||
|
||||
def has_negative(list_shape):
|
||||
return any(x < 0 for x in list_shape)
|
||||
|
||||
# When `x` is Variable:
|
||||
# (1) if x.shape contains -1, such as [2, -1, 64], returns [2, var, 64],
|
||||
# where var = paddle.shape(x)[1]
|
||||
|
||||
# (2) if x.shape does not contains -1, return list(x.shape) directly
|
||||
|
||||
if isinstance(x, (Variable, Value)):
|
||||
values = list(x.shape)
|
||||
if has_negative(values):
|
||||
shape_tensor = paddle.shape(x)
|
||||
for i, v in enumerate(values):
|
||||
if v is None or v < 0:
|
||||
values[i] = shape_tensor[i]
|
||||
return values
|
||||
else:
|
||||
return x.shape
|
||||
|
||||
|
||||
def cast_bool_if_necessary(var):
|
||||
assert isinstance(var, (Variable, Value))
|
||||
if convert_dtype(var.dtype) not in ['bool']:
|
||||
var = paddle.cast(var, dtype="bool")
|
||||
return var
|
||||
|
||||
|
||||
def convert_var_dtype(var, dtype):
|
||||
if isinstance(var, (Variable, Value)):
|
||||
src_dtype = convert_dtype(var.dtype)
|
||||
assert src_dtype in [
|
||||
'bool',
|
||||
'float16',
|
||||
'float32',
|
||||
'float64',
|
||||
'int32',
|
||||
'int64',
|
||||
'uint8',
|
||||
], (
|
||||
f"The dtype of var {var.name} is {src_dtype}, which is not supported in the cast op."
|
||||
)
|
||||
assert dtype in [
|
||||
'bool',
|
||||
'int',
|
||||
'float',
|
||||
'complex',
|
||||
], (
|
||||
f"The casted target dtype is {dtype}, which is not supported in type casting."
|
||||
)
|
||||
cast_map = {
|
||||
'bool': 'bool',
|
||||
'int': 'int32',
|
||||
'float': 'float32',
|
||||
'complex': 'complex64',
|
||||
}
|
||||
return paddle.cast(var, dtype=cast_map[dtype])
|
||||
else:
|
||||
assert dtype in [
|
||||
'bool',
|
||||
'int',
|
||||
'float',
|
||||
'complex',
|
||||
], (
|
||||
f"The casted target dtype is {dtype}, which is not supported in type casting."
|
||||
)
|
||||
return eval(dtype)(var)
|
||||
|
||||
|
||||
def convert_assert(cond, message=""):
|
||||
"""
|
||||
A function representation of a Python ``assert`` statement.
|
||||
"""
|
||||
if isinstance(cond, (Variable, Value)):
|
||||
cond = paddle.cast(cond, "bool")
|
||||
# NOTE: message is not used because Paddle Assert has no corresponding parameter to use.
|
||||
from paddle.static.nn.control_flow import Assert
|
||||
|
||||
return Assert(cond)
|
||||
else:
|
||||
assert cond, message
|
||||
|
||||
|
||||
def convert_print(*objects, sep=' ', end='\n', file=None, flush=False):
|
||||
"""
|
||||
A function representing Python ``print`` function. It will print all arguments
|
||||
at compile time and only print the Tensor values at runtime.
|
||||
"""
|
||||
for obj in objects:
|
||||
if isinstance(obj, (Variable, Value)):
|
||||
paddle.static.Print(obj)
|
||||
print(*objects, sep=sep, end=end, file=file, flush=flush)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def convert_auto_cast(
|
||||
enable=True,
|
||||
custom_white_list=None,
|
||||
custom_black_list=None,
|
||||
level='O1',
|
||||
dtype='float16',
|
||||
use_promote=True,
|
||||
):
|
||||
from .program_translator import ProgramTranslator
|
||||
|
||||
warnings.warn(
|
||||
"paddle.amp.auto_cast is an experimental features in auto parallel."
|
||||
+ "This will take no effect in normal dy2static."
|
||||
)
|
||||
|
||||
amp_records = ProgramTranslator.get_instance()._amp_records
|
||||
main_program = paddle.static.default_main_program()
|
||||
current_block_idx = main_program.current_block_idx
|
||||
current_block = main_program.current_block()
|
||||
start_op_idx = len(current_block.ops)
|
||||
amp_options = AmpOptions(
|
||||
enable, custom_white_list, custom_black_list, level, dtype, use_promote
|
||||
)
|
||||
yield
|
||||
end_op_idx = len(current_block.ops)
|
||||
if current_block_idx not in amp_records:
|
||||
amp_records[current_block_idx] = []
|
||||
amp_records[current_block_idx].append(
|
||||
(amp_options, start_op_idx, end_op_idx)
|
||||
)
|
||||
|
||||
|
||||
def create_bool_as_type(x, value=True):
|
||||
'''
|
||||
Create a bool variable, which type is the same as x.
|
||||
'''
|
||||
if isinstance(x, (Variable, Value)):
|
||||
return paddle.full(shape=[], fill_value=value, dtype="bool")
|
||||
else:
|
||||
return value
|
||||
@@ -0,0 +1,454 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
import linecache
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .origin_info import Location, OriginInfo, global_origin_info_map
|
||||
from .utils import (
|
||||
RE_PYMODULE,
|
||||
is_api_in_module_helper,
|
||||
)
|
||||
|
||||
__all__ = []
|
||||
|
||||
ERROR_DATA = "Error data about original source code information and traceback."
|
||||
|
||||
# A flag to set whether to open the dygraph2static error reporting module
|
||||
SIMPLIFY_ERROR_ENV_NAME = "TRANSLATOR_SIMPLIFY_NEW_ERROR"
|
||||
DEFAULT_SIMPLIFY_NEW_ERROR = 1
|
||||
|
||||
# A flag to set whether to display the simplified error stack
|
||||
DISABLE_ERROR_ENV_NAME = "TRANSLATOR_DISABLE_NEW_ERROR"
|
||||
DEFAULT_DISABLE_NEW_ERROR = 0
|
||||
|
||||
SOURCE_CODE_RANGE = 5
|
||||
BLANK_COUNT_BEFORE_FILE_STR = 4
|
||||
|
||||
|
||||
def attach_error_data(error, in_runtime=False):
|
||||
"""
|
||||
Attaches error data about original source code information and traceback to an error.
|
||||
|
||||
Args:
|
||||
error(Exception): An native error.
|
||||
in_runtime(bool): `error` is raised in runtime if in_runtime is True, otherwise in compile time
|
||||
Returns:
|
||||
An error attached data about original source code information and traceback.
|
||||
"""
|
||||
|
||||
e_type, e_value, e_traceback = sys.exc_info()
|
||||
tb = traceback.extract_tb(e_traceback)[1:]
|
||||
|
||||
error_data = ErrorData(e_type, e_value, tb, global_origin_info_map)
|
||||
error_data.in_runtime = in_runtime
|
||||
|
||||
setattr(error, ERROR_DATA, error_data)
|
||||
|
||||
return error
|
||||
|
||||
|
||||
class TraceBackFrame(OriginInfo):
|
||||
"""
|
||||
Traceback frame information.
|
||||
"""
|
||||
|
||||
def __init__(self, location, function_name, source_code):
|
||||
self.location = location
|
||||
self.function_name = function_name
|
||||
self.source_code = source_code
|
||||
self.error_line = ''
|
||||
|
||||
def formatted_message(self):
|
||||
# self.source_code may be empty in some functions.
|
||||
# For example, decorator generated function
|
||||
return (
|
||||
' ' * BLANK_COUNT_BEFORE_FILE_STR
|
||||
+ 'File "{}", line {}, in {}\n\t{}'.format(
|
||||
self.location.filepath,
|
||||
self.location.lineno,
|
||||
self.function_name,
|
||||
(
|
||||
self.source_code.lstrip()
|
||||
if isinstance(self.source_code, str)
|
||||
else self.source_code
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TraceBackFrameRange(OriginInfo):
|
||||
"""
|
||||
Traceback frame information.
|
||||
"""
|
||||
|
||||
def __init__(self, location, function_name):
|
||||
self.location = location
|
||||
self.function_name = function_name
|
||||
self.source_code = []
|
||||
self.error_line = ''
|
||||
blank_count = []
|
||||
begin_lineno = max(1, self.location.lineno - int(SOURCE_CODE_RANGE / 2))
|
||||
|
||||
for i in range(begin_lineno, begin_lineno + SOURCE_CODE_RANGE):
|
||||
line = linecache.getline(self.location.filepath, i).rstrip('\n')
|
||||
line_lstrip = line.lstrip()
|
||||
self.source_code.append(line_lstrip)
|
||||
if not line_lstrip: # empty line from source code
|
||||
blank_count.append(-1)
|
||||
else:
|
||||
blank_count.append(len(line) - len(line_lstrip))
|
||||
|
||||
if i == self.location.lineno:
|
||||
self.error_line = self.source_code[-1]
|
||||
hint_msg = '~' * len(self.source_code[-1]) + ' <--- HERE'
|
||||
self.source_code.append(hint_msg)
|
||||
blank_count.append(blank_count[-1])
|
||||
# Note(gouzil): Under jupyter, files are read multiple times,
|
||||
# and we can't actively clean the read cache, which can cause subsequent reads to fail.
|
||||
# It is not possible to modify the contents of the file in the meantime,
|
||||
# so there is no need to clear the cache
|
||||
|
||||
# remove top and bottom empty line in source code
|
||||
while len(self.source_code) > 0 and not self.source_code[0]:
|
||||
self.source_code.pop(0)
|
||||
blank_count.pop(0)
|
||||
while len(self.source_code) > 0 and not self.source_code[-1]:
|
||||
self.source_code.pop(-1)
|
||||
blank_count.pop(-1)
|
||||
|
||||
min_black_count = min([i for i in blank_count if i >= 0])
|
||||
for i in range(len(self.source_code)):
|
||||
# if source_code[i] is empty line between two code line, dont add blank
|
||||
if self.source_code[i]:
|
||||
self.source_code[i] = (
|
||||
' '
|
||||
* (
|
||||
blank_count[i]
|
||||
- min_black_count
|
||||
+ BLANK_COUNT_BEFORE_FILE_STR * 2
|
||||
)
|
||||
+ self.source_code[i]
|
||||
)
|
||||
|
||||
def formatted_message(self):
|
||||
msg = (
|
||||
' ' * BLANK_COUNT_BEFORE_FILE_STR
|
||||
+ f'File "{self.location.filepath}", line {self.location.lineno}, in {self.function_name}\n'
|
||||
)
|
||||
# add empty line after range code
|
||||
return msg + '\n'.join(self.source_code)
|
||||
|
||||
|
||||
class SuggestionDict:
|
||||
def __init__(self):
|
||||
# {(keywords): (suggestions)}
|
||||
self.suggestion_dict = {
|
||||
('is not initialized.', 'Hint:', 'IsInitialized'): (
|
||||
"Please ensure all your sublayers are inherited from nn.Layer.",
|
||||
"Please ensure there is no tensor created explicitly depended on external data, "
|
||||
+ "we suggest to register it as buffer tensor. "
|
||||
+ "See https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/jit/principle_cn.html#buffers for details",
|
||||
)
|
||||
}
|
||||
|
||||
def keys(self):
|
||||
return self.suggestion_dict.keys()
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.suggestion_dict[key]
|
||||
|
||||
|
||||
class Dy2StKeyError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class ErrorData:
|
||||
"""
|
||||
Error data attached to an exception which is raised in un-transformed code.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, error_type, error_value, origin_traceback, origin_info_map
|
||||
):
|
||||
self.error_type = error_type
|
||||
self.error_value = error_value
|
||||
self.origin_traceback = origin_traceback
|
||||
self.origin_info_map = origin_info_map
|
||||
self.in_runtime = False
|
||||
self.suggestion_dict = SuggestionDict()
|
||||
|
||||
def create_exception(self):
|
||||
message = self.create_message()
|
||||
if self.error_type is KeyError:
|
||||
new_exception = Dy2StKeyError(message)
|
||||
else:
|
||||
new_exception = self.error_type(message)
|
||||
setattr(new_exception, ERROR_DATA, self)
|
||||
return new_exception
|
||||
|
||||
def numpy_api_check(self, format_exception, error_line):
|
||||
if self.error_type is not TypeError:
|
||||
return format_exception
|
||||
|
||||
tb = self.origin_traceback
|
||||
func_str = None
|
||||
for frame in tb:
|
||||
searched_name = re.search(
|
||||
rf'({RE_PYMODULE})*{frame.name}',
|
||||
error_line,
|
||||
)
|
||||
if searched_name:
|
||||
func_str = searched_name.group(0)
|
||||
break
|
||||
try:
|
||||
globals = {'np': np}
|
||||
fn = eval(func_str, globals)
|
||||
module_result = is_api_in_module_helper(fn, "numpy")
|
||||
is_numpy_api_err = module_result or (
|
||||
func_str.startswith("numpy.") or func_str.startswith("np.")
|
||||
)
|
||||
except Exception:
|
||||
is_numpy_api_err = False
|
||||
|
||||
if is_numpy_api_err and func_str:
|
||||
return [
|
||||
f"TypeError: Code '{error_line}' called numpy API {func_str}, please use Paddle API to replace it.",
|
||||
" values will be changed to variables by dy2static, numpy api can not handle variables",
|
||||
]
|
||||
else:
|
||||
return format_exception
|
||||
|
||||
def create_message(self):
|
||||
"""
|
||||
Creates a custom error message which includes trace stack with source code information of dygraph from user.
|
||||
"""
|
||||
message_lines = []
|
||||
|
||||
# Step1: Adds header message to prompt users that the following is the original information.
|
||||
header_message = "In transformed code:"
|
||||
message_lines.append(header_message)
|
||||
message_lines.append("")
|
||||
error_line = None
|
||||
|
||||
# Simplify error value to improve readability if error is raised in runtime
|
||||
if self.in_runtime:
|
||||
try:
|
||||
if int(
|
||||
os.getenv(
|
||||
SIMPLIFY_ERROR_ENV_NAME, DEFAULT_SIMPLIFY_NEW_ERROR
|
||||
)
|
||||
):
|
||||
self._simplify_error_value()
|
||||
except:
|
||||
pass
|
||||
else:
|
||||
message_lines.append(str(self.error_value))
|
||||
return '\n'.join(message_lines)
|
||||
|
||||
# Step2: Optimizes stack information with source code information of dygraph from user.
|
||||
user_code_traceback_index = []
|
||||
for i, (filepath, lineno, funcname, code) in enumerate(
|
||||
self.origin_traceback
|
||||
):
|
||||
dygraph_func_info = self.origin_info_map.get(
|
||||
(filepath, lineno), None
|
||||
)
|
||||
if dygraph_func_info:
|
||||
user_code_traceback_index.append(i)
|
||||
|
||||
# Add user code traceback
|
||||
for i in user_code_traceback_index:
|
||||
filepath, lineno, funcname, code = self.origin_traceback[i]
|
||||
dygraph_func_info = self.origin_info_map.get(
|
||||
(filepath, lineno), None
|
||||
)
|
||||
if i == user_code_traceback_index[-1]:
|
||||
traceback_frame = TraceBackFrameRange(
|
||||
dygraph_func_info.location, dygraph_func_info.function_name
|
||||
)
|
||||
else:
|
||||
traceback_frame = TraceBackFrame(
|
||||
dygraph_func_info.location,
|
||||
dygraph_func_info.function_name,
|
||||
dygraph_func_info.source_code,
|
||||
)
|
||||
|
||||
message_lines.append(traceback_frame.formatted_message())
|
||||
error_line = traceback_frame.error_line
|
||||
message_lines.append("")
|
||||
|
||||
# Add paddle traceback after user code traceback
|
||||
paddle_traceback_start_index = (
|
||||
user_code_traceback_index[-1] + 1
|
||||
if user_code_traceback_index
|
||||
else 0
|
||||
)
|
||||
for filepath, lineno, funcname, code in self.origin_traceback[
|
||||
paddle_traceback_start_index:
|
||||
]:
|
||||
traceback_frame = TraceBackFrame(
|
||||
Location(filepath, lineno), funcname, code
|
||||
)
|
||||
message_lines.append(traceback_frame.formatted_message())
|
||||
message_lines.append("")
|
||||
|
||||
# Step3: Adds error message like "TypeError: dtype must be int32, but received float32".
|
||||
# NOTE: `format_exception` is a list, its length is 1 in most cases, but sometimes its length
|
||||
# is gather than 1, for example, the error_type is IndentationError.
|
||||
format_exception = traceback.format_exception_only(
|
||||
self.error_type, self.error_value
|
||||
)
|
||||
if error_line is not None:
|
||||
format_exception = self.numpy_api_check(
|
||||
format_exception, error_line
|
||||
)
|
||||
|
||||
error_message = [
|
||||
" " * BLANK_COUNT_BEFORE_FILE_STR + line
|
||||
for line in format_exception
|
||||
]
|
||||
message_lines.extend(error_message)
|
||||
|
||||
return '\n'.join(message_lines)
|
||||
|
||||
def _create_revise_suggestion(self, bottom_error_message):
|
||||
revise_suggestions = [
|
||||
'',
|
||||
' ' * BLANK_COUNT_BEFORE_FILE_STR + 'Revise suggestion: ',
|
||||
]
|
||||
for keywords in self.suggestion_dict.keys():
|
||||
contain_keywords = [
|
||||
True for i in keywords if i in ''.join(bottom_error_message)
|
||||
]
|
||||
if len(contain_keywords) == len(
|
||||
keywords
|
||||
): # all keywords should be in bottom_error_message
|
||||
for suggestion in self.suggestion_dict[keywords]:
|
||||
suggestion_msg = (
|
||||
' ' * BLANK_COUNT_BEFORE_FILE_STR * 2
|
||||
+ f'{len(revise_suggestions) - 1}. {suggestion}'
|
||||
)
|
||||
revise_suggestions.append(suggestion_msg)
|
||||
return revise_suggestions if len(revise_suggestions) > 2 else []
|
||||
|
||||
def _simplify_error_value(self):
|
||||
"""
|
||||
Simplifies error value to improve readability if error is raised in runtime.
|
||||
|
||||
NOTE(liym27): The op callstack information about transformed static code has been replaced with original dygraph code.
|
||||
|
||||
TODO(liym27):
|
||||
1. Need a more robust way because the code of start_trace may change.
|
||||
2. Set the switch to determine whether to simplify error_value
|
||||
"""
|
||||
|
||||
assert self.in_runtime is True
|
||||
|
||||
error_value_lines = str(self.error_value).split("\n")
|
||||
error_value_lines_strip = [mes.lstrip(" ") for mes in error_value_lines]
|
||||
|
||||
start_trace = "outputs = static_func(*inputs)"
|
||||
start_idx = error_value_lines_strip.index(start_trace)
|
||||
|
||||
error_value_lines = error_value_lines[start_idx + 1 :]
|
||||
error_value_lines_strip = error_value_lines_strip[start_idx + 1 :]
|
||||
|
||||
# use empty line to locate the bottom_error_message
|
||||
empty_line_idx = error_value_lines_strip.index('')
|
||||
bottom_error_message = error_value_lines[empty_line_idx + 1 :]
|
||||
revise_suggestion = self._create_revise_suggestion(bottom_error_message)
|
||||
|
||||
error_traceback = []
|
||||
user_code_traceback_index = []
|
||||
pattern = 'File "(?P<filepath>.+)", line (?P<lineno>.+), in (?P<function_name>.+)'
|
||||
|
||||
# Distinguish user code and framework code using static_info_map
|
||||
static_info_map = {}
|
||||
for k, v in self.origin_info_map.items():
|
||||
origin_filepath = v.location.filepath
|
||||
origin_lineno = v.location.lineno
|
||||
static_info_map[(origin_filepath, origin_lineno)] = k
|
||||
|
||||
for i in range(0, len(error_value_lines_strip), 2):
|
||||
if error_value_lines_strip[i].startswith("File "):
|
||||
re_result = re.search(pattern, error_value_lines_strip[i])
|
||||
tmp_filepath, lineno_str, function_name = re_result.groups()
|
||||
code = (
|
||||
error_value_lines_strip[i + 1]
|
||||
if i + 1 < len(error_value_lines_strip)
|
||||
else ''
|
||||
)
|
||||
|
||||
if static_info_map.get((tmp_filepath, int(lineno_str))):
|
||||
user_code_traceback_index.append(len(error_traceback))
|
||||
|
||||
error_traceback.append(
|
||||
(tmp_filepath, int(lineno_str), function_name, code)
|
||||
)
|
||||
|
||||
error_frame = []
|
||||
# Add user code traceback
|
||||
for i in user_code_traceback_index:
|
||||
filepath, lineno, funcname, code = error_traceback[i]
|
||||
if i == user_code_traceback_index[-1]:
|
||||
traceback_frame = TraceBackFrameRange(
|
||||
Location(filepath, lineno), funcname
|
||||
)
|
||||
else:
|
||||
traceback_frame = TraceBackFrame(
|
||||
Location(filepath, lineno), funcname, code
|
||||
)
|
||||
error_frame.append(traceback_frame.formatted_message())
|
||||
error_frame.append("")
|
||||
|
||||
# Add paddle traceback after user code traceback
|
||||
paddle_traceback_start_index = (
|
||||
user_code_traceback_index[-1] + 1
|
||||
if user_code_traceback_index
|
||||
else 0
|
||||
)
|
||||
for filepath, lineno, funcname, code in error_traceback[
|
||||
paddle_traceback_start_index:
|
||||
]:
|
||||
traceback_frame = TraceBackFrame(
|
||||
Location(filepath, lineno), funcname, code
|
||||
)
|
||||
error_frame.append(traceback_frame.formatted_message())
|
||||
error_frame.append("")
|
||||
|
||||
error_frame.extend(bottom_error_message)
|
||||
error_frame.extend(revise_suggestion)
|
||||
error_value_str = '\n'.join(error_frame)
|
||||
self.error_value = self.error_type(error_value_str)
|
||||
|
||||
def raise_new_exception(self):
|
||||
# Raises the origin error if disable dygraph2static error module,
|
||||
if int(os.getenv(DISABLE_ERROR_ENV_NAME, DEFAULT_DISABLE_NEW_ERROR)):
|
||||
raise self.error_value
|
||||
|
||||
new_exception = self.create_exception()
|
||||
# NOTE(liym27):
|
||||
# Why `raise new_exception from None`?
|
||||
#
|
||||
# In Python 3, by default, an new exception is raised with trace information of the caught exception.
|
||||
# This only raises new_exception and hides unwanted implementation details from tracebacks of the
|
||||
# caught exception.
|
||||
|
||||
raise new_exception from None
|
||||
@@ -0,0 +1,630 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
import collections
|
||||
import inspect
|
||||
import weakref
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
import paddle.pir.core as ir_static
|
||||
from paddle.base import core
|
||||
from paddle.base.data_feeder import convert_dtype
|
||||
from paddle.base.dygraph.base import switch_to_static_graph
|
||||
from paddle.jit.pir_translated_layer import PirTranslatedLayer
|
||||
from paddle.jit.translated_layer import TranslatedLayer
|
||||
from paddle.nn.layer import layers
|
||||
|
||||
from . import logging_utils
|
||||
from .utils import (
|
||||
func_to_source_code,
|
||||
parse_arg_and_kwargs,
|
||||
parse_varargs_name,
|
||||
type_name,
|
||||
)
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class FunctionSpec:
|
||||
"""
|
||||
Wrapper class for a function for class method.
|
||||
"""
|
||||
|
||||
def __init__(self, function, input_spec=None):
|
||||
if inspect.ismethod(function):
|
||||
self._dygraph_function = function.__func__
|
||||
self._class_instance = weakref.ref(function.__self__)
|
||||
else:
|
||||
self._dygraph_function = function
|
||||
self._class_instance = None
|
||||
if input_spec is None:
|
||||
self._input_spec = None
|
||||
self._flat_input_spec = None
|
||||
else:
|
||||
self._input_spec = self._verify_input_spec(input_spec)
|
||||
self._flat_input_spec = paddle.utils.flatten(self._input_spec)
|
||||
|
||||
# parse full argument names list.
|
||||
self._arg_names, self._default_kwargs = parse_arg_and_kwargs(function)
|
||||
# parse *args
|
||||
self.varargs_name = parse_varargs_name(function)
|
||||
if self.varargs_name is not None and isinstance(
|
||||
getattr(function, '__self__', None),
|
||||
(TranslatedLayer, PirTranslatedLayer),
|
||||
):
|
||||
self._arg_names += function.__self__._input_args_names
|
||||
|
||||
def unified_args_and_kwargs(self, args, kwargs):
|
||||
"""
|
||||
Moves kwargs with default value into arguments list to keep `args` contain the same length
|
||||
value as function definition.
|
||||
|
||||
For example:
|
||||
|
||||
Given function definition: `def foo(x, a=1, b=2)`,
|
||||
when calling it by `foo(23)`, the args is `[23]`, kwargs is `{a=1, b=2}`.
|
||||
In this function, it will return args with `[23, 1, 2]`, kwargs with `{}`
|
||||
|
||||
Args:
|
||||
args(tuple): tuple of input arguments value of decorated function.
|
||||
kwargs(dict): dict of input keyword arguments value of decorated function.
|
||||
|
||||
Return:
|
||||
New arguments tuple containing default kwargs value.
|
||||
"""
|
||||
if len(self._arg_names) < len(args):
|
||||
error_msg = f"The decorated function `{self.dygraph_function.__name__}` requires {len(self._arg_names)} arguments: {self._arg_names}, but received {len(args)} with {args}."
|
||||
if args and inspect.isclass(args[0]):
|
||||
error_msg += "\n\tMaybe the function has more than one decorator, we don't support this for now."
|
||||
raise NotImplementedError(error_msg)
|
||||
else:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
args = list(args)
|
||||
|
||||
for i in range(len(args), len(self._arg_names)):
|
||||
arg_name = self._arg_names[i]
|
||||
if arg_name in kwargs:
|
||||
args.append(kwargs[arg_name])
|
||||
del kwargs[arg_name]
|
||||
else:
|
||||
if arg_name not in self._default_kwargs:
|
||||
raise ValueError(
|
||||
f"`{self.dygraph_function.__name__}()` requires `{arg_name}` arguments, but not found in input `args`: {args} and `kwargs`: {kwargs}."
|
||||
)
|
||||
args.append(self._default_kwargs[arg_name])
|
||||
|
||||
return tuple(args), kwargs
|
||||
|
||||
def args_to_input_spec(self, args, kwargs):
|
||||
"""
|
||||
Converts input arguments into InputSpec.
|
||||
|
||||
1. If specific input_spec, use them to construct feed layers.
|
||||
2. If input_spec is None, consider all Tensor and Numpy.ndarray as feed layers
|
||||
|
||||
Args:
|
||||
args(tuple): tuple of input arguments value of function containing default kwargs value.
|
||||
kwargs(dict): kwargs arguments received by **kwargs.
|
||||
|
||||
Return:
|
||||
Same nest structure with args and kwargs by replacing value with InputSpec.
|
||||
"""
|
||||
|
||||
args_with_spec = []
|
||||
kwargs_with_spec = []
|
||||
if self._input_spec is not None:
|
||||
# Note: Because the value type and length of `kwargs` is uncertain.
|
||||
# So we don't support to deal this case while specifying `input_spec` currently.
|
||||
if kwargs:
|
||||
raise ValueError(
|
||||
f"{self.dygraph_function.__name__} got unexpected keyword arguments: {kwargs}. Cannot trace the function when `input_spec` is specified."
|
||||
)
|
||||
|
||||
# Note: The length of `input_spec` can be greater than `args`,
|
||||
# because `args` may contains non-tensor value merged form `kwargs`
|
||||
# after `unified_args_and_kwargs`.
|
||||
if len(args) < len(self._input_spec):
|
||||
raise ValueError(
|
||||
f"Requires len(arguments) >= len(input_spec), but received len(args):{len(args)} < len(InputSpec): {len(self._input_spec)}"
|
||||
)
|
||||
|
||||
# replace argument with corresponding InputSpec.
|
||||
args_with_spec = convert_to_input_spec(args, self._input_spec)
|
||||
else:
|
||||
args_with_spec = _replace_to_input_spec_with_new_name(
|
||||
args, self._arg_names
|
||||
)
|
||||
kwarg_names = ["kwargs." + key for key in kwargs.keys()]
|
||||
kwargs_list_with_spec = _replace_to_input_spec_with_new_name(
|
||||
list(kwargs.values()), kwarg_names
|
||||
)
|
||||
kwargs_with_spec = {
|
||||
key: kwargs_list_with_spec[idx]
|
||||
for idx, key in enumerate(kwargs)
|
||||
}
|
||||
|
||||
# If without specifying name in input_spec, add default name
|
||||
# according to argument name from decorated function.
|
||||
args_with_spec = replace_spec_empty_name(
|
||||
self._arg_names, args_with_spec
|
||||
)
|
||||
|
||||
return args_with_spec, kwargs_with_spec
|
||||
|
||||
@switch_to_static_graph
|
||||
def pir_to_static_inputs_with_spec(self, input_with_spec, main_program):
|
||||
"""
|
||||
Constructs feed layer by inputs with InputSpec information for main program.
|
||||
|
||||
Args:
|
||||
input_with_spec(tuple): input arguments by replacing argument with InputSpec.
|
||||
main_program(Program): main program for inserting feed layer.
|
||||
"""
|
||||
from paddle.distributed.auto_parallel.placement_type import (
|
||||
to_placements,
|
||||
)
|
||||
|
||||
flat_input_spec = paddle.utils.flatten(input_with_spec)
|
||||
|
||||
# NOTE(zhangbo): Why do we need function_args and program_inputs: The primary function of this module is to construct the corresponding DataOp based on the inputSpec. The output of DataOp serves as the input to the Program and will also become the arguments for subsequent Static functions to construct the static graph. When the input is a DistributedInputSpec, the shard_tensor operation will be performed on the output of DataOp to obtain the corresponding distributed Value. The input to the Program will still be the output of DataOp, but in this case, the arguments for the Static functions will be the output of shard_tensor.
|
||||
function_args = []
|
||||
program_inputs = []
|
||||
with ir_static.program_guard(main_program):
|
||||
for i, var_spec in enumerate(flat_input_spec):
|
||||
if isinstance(var_spec, paddle.static.InputSpec):
|
||||
stop_gradient = getattr(var_spec, 'stop_gradient', False)
|
||||
feed_value = paddle.static.input.data(
|
||||
name=var_spec.name or f"feed_{i}",
|
||||
shape=var_spec.shape,
|
||||
dtype=convert_dtype(var_spec.dtype),
|
||||
)
|
||||
feed_value.stop_gradient = stop_gradient
|
||||
|
||||
# warp dist tensor
|
||||
from paddle.distributed.auto_parallel.static.dist_input_spec import (
|
||||
DistributedInputSpec,
|
||||
)
|
||||
|
||||
if isinstance(var_spec, DistributedInputSpec):
|
||||
placements = to_placements(
|
||||
var_spec.dims_mapping, var_spec
|
||||
)
|
||||
dist_feed_value = paddle._pir_ops.shard_tensor(
|
||||
feed_value, var_spec.mesh, placements
|
||||
)
|
||||
function_args.append(dist_feed_value)
|
||||
else:
|
||||
function_args.append(feed_value)
|
||||
else:
|
||||
feed_value = var_spec
|
||||
function_args.append(feed_value)
|
||||
|
||||
program_inputs.append(feed_value)
|
||||
|
||||
return paddle.utils.pack_sequence_as(
|
||||
input_with_spec, function_args
|
||||
), paddle.utils.pack_sequence_as(input_with_spec, program_inputs)
|
||||
|
||||
@switch_to_static_graph
|
||||
def to_static_inputs_with_spec(self, input_with_spec, main_program):
|
||||
"""
|
||||
Constructs feed layer by inputs with InputSpec information for main program.
|
||||
|
||||
Args:
|
||||
input_with_spec(tuple): input arguments by replacing argument with InputSpec.
|
||||
main_program(Program): main program for inserting feed layer.
|
||||
"""
|
||||
flat_input_spec = paddle.utils.flatten(input_with_spec)
|
||||
|
||||
inputs = []
|
||||
block = main_program.global_block()
|
||||
for i, var_spec in enumerate(flat_input_spec):
|
||||
if isinstance(var_spec, paddle.static.InputSpec):
|
||||
stop_gradient = getattr(var_spec, 'stop_gradient', False)
|
||||
feed_layer = block.create_var(
|
||||
# TODO(Aurelius84): consider a more elegant way to name this
|
||||
name=var_spec.name or f"feed_{i}",
|
||||
shape=var_spec.shape,
|
||||
dtype=var_spec.dtype,
|
||||
is_data=True,
|
||||
need_check_feed=False,
|
||||
stop_gradient=stop_gradient,
|
||||
)
|
||||
# warp dist tensor
|
||||
from paddle.distributed.auto_parallel.static.dist_input_spec import (
|
||||
DistributedInputSpec,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.dist_tensor import (
|
||||
DistributedTensor,
|
||||
)
|
||||
|
||||
if isinstance(var_spec, DistributedInputSpec):
|
||||
from paddle.distributed.auto_parallel.static.dist_context import (
|
||||
get_default_distributed_context,
|
||||
)
|
||||
|
||||
default_dist_ctx = get_default_distributed_context()
|
||||
dist_tensor = DistributedTensor(feed_layer)
|
||||
dist_tensor.dist_attr.process_mesh = var_spec.mesh
|
||||
dist_tensor.dist_attr.dims_mapping = var_spec.dims_mapping
|
||||
dist_tensor.dist_attr.mark_annotated("process_mesh")
|
||||
dist_tensor.dist_attr.mark_annotated("dims_mapping")
|
||||
default_dist_ctx.add_dist_tensor_for_program(dist_tensor)
|
||||
else:
|
||||
feed_layer = var_spec
|
||||
|
||||
inputs.append(feed_layer)
|
||||
|
||||
return paddle.utils.pack_sequence_as(input_with_spec, inputs)
|
||||
|
||||
def _verify_input_spec(self, input_spec):
|
||||
"""
|
||||
Verifies the `input_spec` and its element type is valid.
|
||||
"""
|
||||
if not isinstance(input_spec, (tuple, list)):
|
||||
raise TypeError(
|
||||
f"The type(input_spec) should be one of (tuple, list), but received {type_name(input_spec)}."
|
||||
)
|
||||
|
||||
return tuple(input_spec)
|
||||
|
||||
def __repr__(self):
|
||||
return "function: {}({}), input_spec: {}".format(
|
||||
self.dygraph_function.__name__,
|
||||
','.join(self._arg_names),
|
||||
self._input_spec,
|
||||
)
|
||||
|
||||
@property
|
||||
def class_instance(self):
|
||||
if self._class_instance is None:
|
||||
return None
|
||||
if self._class_instance() is None:
|
||||
raise RuntimeError(
|
||||
"The instance of class has been deleted, please re-create the instance."
|
||||
)
|
||||
return self._class_instance()
|
||||
|
||||
@property
|
||||
def dygraph_function(self):
|
||||
if self.class_instance is not None:
|
||||
return self._dygraph_function.__get__(self.class_instance)
|
||||
else:
|
||||
return self._dygraph_function
|
||||
|
||||
@property
|
||||
def args_name(self):
|
||||
return self._arg_names
|
||||
|
||||
@property
|
||||
def input_spec(self):
|
||||
return self._input_spec
|
||||
|
||||
@property
|
||||
def flat_input_spec(self):
|
||||
return self._flat_input_spec
|
||||
|
||||
@property
|
||||
def code(self):
|
||||
return func_to_source_code(self.dygraph_function)
|
||||
|
||||
|
||||
def get_parameters(layer_instance, include_sublayer=True):
|
||||
"""
|
||||
Returns parameters of decorated layers. If set `include_sublayer` True,
|
||||
the parameters created in sub layers will be added.
|
||||
"""
|
||||
params = collections.OrderedDict()
|
||||
if layer_instance is not None:
|
||||
if isinstance(layer_instance, layers.Layer):
|
||||
if include_sublayer:
|
||||
params = layer_instance.parameters()
|
||||
names = [p.name for p in params]
|
||||
params = collections.OrderedDict(zip(names, params))
|
||||
else:
|
||||
params = layer_instance._parameters
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Type of `layer_instance` should be nn.Layer, but received {type_name(layer_instance)}"
|
||||
)
|
||||
|
||||
return params
|
||||
|
||||
|
||||
def get_buffers(layer_instance, include_sublayer=True):
|
||||
"""
|
||||
Returns Variable buffers of decorated layers. If set `include_sublayer` True,
|
||||
the Variable buffers created in sub layers will be added.
|
||||
"""
|
||||
buffers = collections.OrderedDict()
|
||||
if layer_instance is not None:
|
||||
if isinstance(layer_instance, layers.Layer):
|
||||
if include_sublayer:
|
||||
buffers = layer_instance.buffers()
|
||||
names = [buffer.name for buffer in buffers]
|
||||
buffers = collections.OrderedDict(zip(names, buffers))
|
||||
else:
|
||||
buffers = layer_instance._buffers
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Type of `layer_instance` should be nn.Layer, but received {type_name(layer_instance)}"
|
||||
)
|
||||
return buffers
|
||||
|
||||
|
||||
def _replace_value_with_input_spec(args):
|
||||
from paddle.distributed.auto_parallel.placement_type import (
|
||||
to_placements,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.dist_input_spec import (
|
||||
DistributedInputSpec,
|
||||
)
|
||||
|
||||
args_with_spec = []
|
||||
for idx, input_var in enumerate(paddle.utils.flatten(args)):
|
||||
if isinstance(input_var, np.ndarray):
|
||||
input_var = paddle.static.InputSpec.from_numpy(input_var)
|
||||
input_var.stop_gradient = True
|
||||
elif isinstance(input_var, core.eager.Tensor):
|
||||
stop_gradient = input_var.stop_gradient
|
||||
if input_var.is_dist():
|
||||
input_var = DistributedInputSpec.from_dtensor(input_var)
|
||||
else:
|
||||
input_var = paddle.static.InputSpec.from_tensor(input_var)
|
||||
input_var.stop_gradient = stop_gradient
|
||||
elif isinstance(
|
||||
input_var, (paddle.base.framework.Variable, paddle.pir.Value)
|
||||
):
|
||||
stop_gradient = input_var.stop_gradient
|
||||
if input_var.is_dist():
|
||||
mesh = input_var.dist_attr().process_mesh
|
||||
placements = to_placements(
|
||||
input_var.dist_attr().dims_mapping, mesh
|
||||
)
|
||||
input_var = DistributedInputSpec(
|
||||
input_var.shape,
|
||||
dtype=input_var.dtype,
|
||||
name=input_var.name,
|
||||
mesh=mesh,
|
||||
placements=placements,
|
||||
local_shape=input_var._local_shape,
|
||||
)
|
||||
else:
|
||||
input_var = paddle.static.InputSpec(
|
||||
input_var.shape, input_var.dtype, input_var.name
|
||||
)
|
||||
input_var.stop_gradient = stop_gradient
|
||||
|
||||
args_with_spec.append(input_var)
|
||||
args_with_spec = paddle.utils.pack_sequence_as(args, args_with_spec)
|
||||
return args_with_spec
|
||||
|
||||
|
||||
def _replace_to_input_spec_with_new_name(args, arg_names):
|
||||
assert len(args) == len(arg_names)
|
||||
order_digit = len(str(len(arg_names) - 1))
|
||||
args_with_spec = []
|
||||
for order, (arg, name_prefix) in enumerate(zip(args, arg_names)):
|
||||
index = 0
|
||||
for idx, origin_input in enumerate(paddle.utils.flatten(arg)):
|
||||
if isinstance(origin_input, np.ndarray):
|
||||
input_var = paddle.static.InputSpec.from_numpy(origin_input)
|
||||
input_var.stop_gradient = True
|
||||
elif isinstance(origin_input, core.eager.Tensor):
|
||||
stop_gradient = origin_input.stop_gradient
|
||||
input_var = paddle.static.InputSpec.from_tensor(origin_input)
|
||||
input_var.stop_gradient = stop_gradient
|
||||
elif isinstance(origin_input, paddle.base.framework.Variable):
|
||||
stop_gradient = origin_input.stop_gradient
|
||||
input_var = paddle.static.InputSpec(
|
||||
origin_input.shape, origin_input.dtype, origin_input.name
|
||||
)
|
||||
input_var.stop_gradient = stop_gradient
|
||||
else:
|
||||
input_var = origin_input
|
||||
|
||||
if isinstance(
|
||||
origin_input,
|
||||
(
|
||||
np.ndarray,
|
||||
core.eager.Tensor,
|
||||
paddle.base.framework.Variable,
|
||||
),
|
||||
):
|
||||
input_var.name = f"_jst.{str(order).zfill(order_digit)}.{name_prefix}.{index}"
|
||||
index += 1
|
||||
args_with_spec.append(input_var)
|
||||
args_with_spec = paddle.utils.pack_sequence_as(args, args_with_spec)
|
||||
return args_with_spec
|
||||
|
||||
|
||||
def convert_to_input_spec(inputs, input_spec):
|
||||
"""
|
||||
Replaces tensor in structured `inputs` by InputSpec in `input_spec`.
|
||||
|
||||
Args:
|
||||
inputs(list|dict): nested structure list or dict.
|
||||
input_spec(list|dict): same nested structure list or dict as inputs.
|
||||
|
||||
|
||||
Return:
|
||||
Same structure with inputs by replacing the element with specified InputSpec.
|
||||
"""
|
||||
|
||||
def check_type_and_len(input, spec, check_length=False):
|
||||
if type(input) is not type(spec):
|
||||
raise TypeError(
|
||||
f'type(input) should be {type(spec)}, but received {type(input)}.'
|
||||
)
|
||||
if check_length and len(input) < len(spec):
|
||||
raise ValueError(
|
||||
f'Requires len(inputs) >= len(input_spec), but received len(inputs):{len(inputs)} < len(input_spec):{len(input_spec)}'
|
||||
)
|
||||
|
||||
if isinstance(input_spec, (tuple, list)):
|
||||
input_with_spec = []
|
||||
check_type_and_len(inputs, input_spec, True)
|
||||
|
||||
for i, spec in enumerate(input_spec):
|
||||
out_spec = convert_to_input_spec(inputs[i], spec)
|
||||
input_with_spec.append(out_spec)
|
||||
|
||||
# Note: If the rest inputs contain tensor or numpy.ndarray
|
||||
# without specific InputSpec, raise warning.
|
||||
if len(inputs) > len(input_spec):
|
||||
for rest_input in inputs[len(input_spec) :]:
|
||||
if isinstance(rest_input, (core.eager.Tensor, np.ndarray)):
|
||||
logging_utils.warn(
|
||||
f"The inputs contain `{type_name(rest_input)}` without specifying InputSpec, its shape and dtype will be treated immutable. "
|
||||
"Please specific InputSpec information in `@to_static` if you expect them as mutable inputs."
|
||||
)
|
||||
input_with_spec.extend(inputs[len(input_spec) :])
|
||||
|
||||
return input_with_spec
|
||||
elif isinstance(input_spec, dict):
|
||||
input_with_spec = {}
|
||||
check_type_and_len(inputs, input_spec, True)
|
||||
for name, input in inputs.items():
|
||||
if name in input_spec:
|
||||
input_with_spec[name] = convert_to_input_spec(
|
||||
input, input_spec[name]
|
||||
)
|
||||
else:
|
||||
input_with_spec[name] = input
|
||||
return input_with_spec
|
||||
elif isinstance(input_spec, paddle.static.InputSpec):
|
||||
"""we compare input_spec with real_input_spec constructed from arguments."""
|
||||
real_spec = _replace_value_with_input_spec([inputs])[0]
|
||||
if not isinstance(real_spec, paddle.static.InputSpec):
|
||||
raise RuntimeError(
|
||||
f"Give input spec into a non-tensorable arguments `{inputs}`."
|
||||
)
|
||||
real_spec.name = input_spec.name
|
||||
if spec_greater(input_spec, real_spec):
|
||||
# change shape but keep the others (stop_gradient / dtype) .
|
||||
real_spec.shape = input_spec.shape
|
||||
else:
|
||||
logging_utils.warn(
|
||||
f"input spec is not compatible with real inputs. input_spec: {input_spec} , real_spec: {real_spec} "
|
||||
)
|
||||
return real_spec
|
||||
else:
|
||||
# NOTE(Aurelius84): Support non-Tensor type as input spec info
|
||||
return input_spec
|
||||
|
||||
|
||||
def replace_spec_empty_name(args_name, input_with_spec):
|
||||
"""
|
||||
Adds default name according to argument name from decorated function
|
||||
if without specifying InputSpec.name
|
||||
|
||||
The naming rule are as followed:
|
||||
1. If InputSpec.name is not None, do nothing.
|
||||
2. If each argument `x` corresponds to an InputSpec, using the argument name like `x`
|
||||
3. If the arguments `inputs` corresponds to a list(InputSpec), using name like `inputs_0`, `inputs_1`
|
||||
4. If the arguments `input_dic` corresponds to a dict(InputSpec), using key as name.
|
||||
|
||||
For example:
|
||||
|
||||
# case 1: foo(x, y)
|
||||
foo = to_static(foo, input_spec=[InputSpec([None, 10]), InputSpec([None])])
|
||||
print([in_var.name for in_var in foo.inputs]) # [x, y]
|
||||
|
||||
# case 2: foo(inputs) where inputs is a list
|
||||
foo = to_static(foo, input_spec=[[InputSpec([None, 10]), InputSpec([None])]])
|
||||
print([in_var.name for in_var in foo.inputs]) # [inputs_0, inputs_1]
|
||||
|
||||
# case 3: foo(inputs) where inputs is a dict
|
||||
foo = to_static(foo, input_spec=[{'x': InputSpec([None, 10]), 'y': InputSpec([None])}])
|
||||
print([in_var.name for in_var in foo.inputs]) # [x, y]
|
||||
"""
|
||||
input_with_spec = list(input_with_spec)
|
||||
candidate_arg_names = args_name[: len(input_with_spec)]
|
||||
|
||||
for i, arg_name in enumerate(candidate_arg_names):
|
||||
input_spec = input_with_spec[i]
|
||||
input_with_spec[i] = _replace_spec_name(arg_name, input_spec)
|
||||
|
||||
return input_with_spec
|
||||
|
||||
|
||||
def _replace_spec_name(name, input_spec):
|
||||
"""
|
||||
Replaces InputSpec.name with given `name` while not specifying it.
|
||||
"""
|
||||
if isinstance(input_spec, paddle.static.InputSpec):
|
||||
if input_spec.name is None:
|
||||
input_spec.name = name
|
||||
return input_spec
|
||||
elif isinstance(input_spec, (list, tuple)):
|
||||
processed_specs = []
|
||||
for i, spec in enumerate(input_spec):
|
||||
new_name = f"{name}_{i}"
|
||||
processed_specs.append(_replace_spec_name(new_name, spec))
|
||||
return processed_specs
|
||||
elif isinstance(input_spec, dict):
|
||||
processed_specs = {}
|
||||
for key, spec in input_spec.items():
|
||||
processed_specs[key] = _replace_spec_name(key, spec)
|
||||
return processed_specs
|
||||
else:
|
||||
return input_spec
|
||||
|
||||
|
||||
def _hash_spec_names(args_specs, kwargs_specs):
|
||||
"""
|
||||
Generator hash spec with args/kwargs InputSpec names.
|
||||
Consider the following InputSpecs with same shape/dtype except for name:
|
||||
1. [InputSpec([3,3], 'float32', 'x'), InputSpec([3,3], 'float32', 'x')]
|
||||
2. [InputSpec([3,3], 'float32', 'x'), InputSpec([3,3], 'float32', 'y')]
|
||||
Under @to_static, we should generate two different program not just one, because
|
||||
the former has one input ('x'), but the latter has two input ('x', 'y').
|
||||
"""
|
||||
spec_names = [
|
||||
spec.name
|
||||
for spec in paddle.utils.flatten(args_specs)
|
||||
if isinstance(spec, paddle.static.InputSpec)
|
||||
]
|
||||
spec_names += [
|
||||
spec.name
|
||||
for spec in paddle.utils.flatten(kwargs_specs)
|
||||
if isinstance(spec, paddle.static.InputSpec)
|
||||
]
|
||||
i, name_ids = 0, {}
|
||||
|
||||
def to_idx(name):
|
||||
nonlocal i
|
||||
if name not in name_ids:
|
||||
name_ids[name] = i
|
||||
i += 1
|
||||
return name_ids[name]
|
||||
|
||||
value = [to_idx(name) for name in spec_names]
|
||||
|
||||
return tuple(value)
|
||||
|
||||
|
||||
def spec_greater(first, other):
|
||||
def _shape_greater(first_shape, second_shape):
|
||||
if len(first_shape) != len(second_shape):
|
||||
return False
|
||||
for first_n, second_n in zip(first_shape, second_shape):
|
||||
if first_n != -1 and first_n != second_n:
|
||||
return False
|
||||
return True
|
||||
|
||||
return _shape_greater(first.shape, other.shape)
|
||||
@@ -0,0 +1,279 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
import os
|
||||
import threading
|
||||
|
||||
from paddle.base import log_helper
|
||||
|
||||
from .ast_utils import ast_to_source_code
|
||||
|
||||
__all__ = []
|
||||
|
||||
VERBOSITY_ENV_NAME = 'TRANSLATOR_VERBOSITY'
|
||||
CODE_LEVEL_ENV_NAME = 'TRANSLATOR_CODE_LEVEL'
|
||||
DEFAULT_VERBOSITY = -1
|
||||
DEFAULT_CODE_LEVEL = -1
|
||||
|
||||
LOG_AllTransformer = 100
|
||||
|
||||
|
||||
def synchronized(func):
|
||||
def wrapper(*args, **kwargs):
|
||||
with threading.Lock():
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class TranslatorLogger:
|
||||
"""
|
||||
class for Logging and debugging during the transformation from dygraph to static graph.
|
||||
The object of this class is a singleton.
|
||||
"""
|
||||
|
||||
@synchronized
|
||||
def __new__(cls, *args, **kwargs):
|
||||
if not hasattr(cls, '_instance'):
|
||||
cls._instance = object.__new__(cls, *args, **kwargs)
|
||||
cls._instance._initialized = False
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
self._initialized = True
|
||||
self.logger_name = "Dynamic-to-Static"
|
||||
self._logger = log_helper.get_logger(
|
||||
self.logger_name,
|
||||
1,
|
||||
fmt='%(asctime)s %(name)s %(levelname)s: %(message)s',
|
||||
)
|
||||
self._verbosity_level = None
|
||||
self._transformed_code_level = None
|
||||
self._need_to_echo_log_to_stdout = None
|
||||
self._need_to_echo_code_to_stdout = None
|
||||
|
||||
@property
|
||||
def logger(self):
|
||||
return self._logger
|
||||
|
||||
@property
|
||||
def verbosity_level(self):
|
||||
if self._verbosity_level is not None:
|
||||
return self._verbosity_level
|
||||
else:
|
||||
return int(os.getenv(VERBOSITY_ENV_NAME, DEFAULT_VERBOSITY))
|
||||
|
||||
@verbosity_level.setter
|
||||
def verbosity_level(self, level):
|
||||
self.check_level(level)
|
||||
self._verbosity_level = level
|
||||
|
||||
@property
|
||||
def transformed_code_level(self):
|
||||
if self._transformed_code_level is not None:
|
||||
return self._transformed_code_level
|
||||
else:
|
||||
return int(os.getenv(CODE_LEVEL_ENV_NAME, DEFAULT_CODE_LEVEL))
|
||||
|
||||
@transformed_code_level.setter
|
||||
def transformed_code_level(self, level):
|
||||
self.check_level(level)
|
||||
self._transformed_code_level = level
|
||||
|
||||
@property
|
||||
def need_to_echo_log_to_stdout(self):
|
||||
if self._need_to_echo_log_to_stdout is not None:
|
||||
return self._need_to_echo_log_to_stdout
|
||||
return False
|
||||
|
||||
@need_to_echo_log_to_stdout.setter
|
||||
def need_to_echo_log_to_stdout(self, log_to_stdout):
|
||||
assert isinstance(log_to_stdout, (bool, type(None)))
|
||||
self._need_to_echo_log_to_stdout = log_to_stdout
|
||||
|
||||
@property
|
||||
def need_to_echo_code_to_stdout(self):
|
||||
if self._need_to_echo_code_to_stdout is not None:
|
||||
return self._need_to_echo_code_to_stdout
|
||||
return False
|
||||
|
||||
@need_to_echo_code_to_stdout.setter
|
||||
def need_to_echo_code_to_stdout(self, code_to_stdout):
|
||||
assert isinstance(code_to_stdout, (bool, type(None)))
|
||||
self._need_to_echo_code_to_stdout = code_to_stdout
|
||||
|
||||
def check_level(self, level):
|
||||
if isinstance(level, (int, type(None))):
|
||||
rv = level
|
||||
else:
|
||||
raise TypeError(f"Level is not an integer: {level}")
|
||||
return rv
|
||||
|
||||
def has_code_level(self, level):
|
||||
level = self.check_level(level)
|
||||
return level == self.transformed_code_level
|
||||
|
||||
def has_verbosity(self, level):
|
||||
"""
|
||||
Checks whether the verbosity level set by the user is greater than or equal to the log level.
|
||||
Args:
|
||||
level(int): The level of log.
|
||||
Returns:
|
||||
True if the verbosity level set by the user is greater than or equal to the log level, otherwise False.
|
||||
"""
|
||||
level = self.check_level(level)
|
||||
return self.verbosity_level >= level
|
||||
|
||||
def error(self, msg, *args, **kwargs):
|
||||
self.logger.error(msg, *args, **kwargs)
|
||||
if self.need_to_echo_log_to_stdout:
|
||||
self._output_to_stdout('ERROR: ' + msg, *args)
|
||||
|
||||
def warn(self, msg, *args, **kwargs):
|
||||
if self.verbosity_level != -1:
|
||||
self.logger.warning(msg, *args, **kwargs)
|
||||
if self.need_to_echo_log_to_stdout:
|
||||
self._output_to_stdout('WARNING: ' + msg, *args)
|
||||
|
||||
def log(self, level, msg, *args, **kwargs):
|
||||
if self.has_verbosity(level):
|
||||
msg_with_level = f'(Level {level}) {msg}'
|
||||
self.logger.info(msg_with_level, *args, **kwargs)
|
||||
if self.need_to_echo_log_to_stdout:
|
||||
self._output_to_stdout('INFO: ' + msg_with_level, *args)
|
||||
|
||||
def log_transformed_code(
|
||||
self, level, ast_node, transformer_name, *args, **kwargs
|
||||
):
|
||||
if self.has_code_level(level):
|
||||
source_code = ast_to_source_code(ast_node)
|
||||
if level == LOG_AllTransformer:
|
||||
header_msg = f"After the last level ast transformer: '{transformer_name}', the transformed code:\n"
|
||||
else:
|
||||
header_msg = f"After the level {level} ast transformer: '{transformer_name}', the transformed code:\n"
|
||||
|
||||
msg = header_msg + source_code
|
||||
self.logger.info(msg, *args, **kwargs)
|
||||
|
||||
if self.need_to_echo_code_to_stdout:
|
||||
self._output_to_stdout('INFO: ' + msg, *args)
|
||||
|
||||
def _output_to_stdout(self, msg, *args):
|
||||
msg = self.logger_name + ' ' + msg
|
||||
print(msg % args)
|
||||
|
||||
|
||||
_TRANSLATOR_LOGGER = TranslatorLogger()
|
||||
|
||||
|
||||
def set_verbosity(level: int = 0, also_to_stdout: bool = False) -> None:
|
||||
"""
|
||||
Sets the verbosity level of log for dygraph to static graph. Logs can be output to stdout by setting `also_to_stdout`.
|
||||
|
||||
There are two means to set the logging verbosity:
|
||||
|
||||
1. Call function `set_verbosity`
|
||||
|
||||
2. Set environment variable `TRANSLATOR_VERBOSITY`
|
||||
|
||||
|
||||
**Note**:
|
||||
`set_verbosity` has a higher priority than the environment variable.
|
||||
|
||||
Args:
|
||||
level(int): The verbosity level. The larger value indicates more verbosity.
|
||||
The default value is 0, which means no logging.
|
||||
also_to_stdout(bool): Whether to also output log messages to `sys.stdout`.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import os
|
||||
>>> import paddle
|
||||
|
||||
>>> paddle.jit.set_verbosity(1)
|
||||
>>> # The verbosity level is now 1
|
||||
|
||||
>>> os.environ['TRANSLATOR_VERBOSITY'] = '3'
|
||||
>>> # The verbosity level is now 3, but it has no effect because it has a lower priority than `set_verbosity`
|
||||
"""
|
||||
_TRANSLATOR_LOGGER.verbosity_level = level
|
||||
_TRANSLATOR_LOGGER.need_to_echo_log_to_stdout = also_to_stdout
|
||||
|
||||
|
||||
def get_verbosity() -> int:
|
||||
return _TRANSLATOR_LOGGER.verbosity_level
|
||||
|
||||
|
||||
def set_code_level(
|
||||
level: int = LOG_AllTransformer, also_to_stdout: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Sets the level to print code from specific level Ast Transformer. Code can be output to stdout by setting `also_to_stdout`.
|
||||
|
||||
There are two means to set the code level:
|
||||
|
||||
1. Call function `set_code_level`
|
||||
|
||||
2. Set environment variable `TRANSLATOR_CODE_LEVEL`
|
||||
|
||||
|
||||
**Note**:
|
||||
`set_code_level` has a higher priority than the environment variable.
|
||||
|
||||
Args:
|
||||
level(int): The level to print code. Default is 100, which means to print the code after all AST Transformers.
|
||||
also_to_stdout(bool): Whether to also output code to `sys.stdout`.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import os
|
||||
>>> import paddle
|
||||
|
||||
>>> paddle.jit.set_code_level(2)
|
||||
>>> # It will print the transformed code at level 2, which means to print the code after second transformer,
|
||||
>>> # as the date of August 28, 2020, it is CastTransformer.
|
||||
|
||||
>>> os.environ['TRANSLATOR_CODE_LEVEL'] = '3'
|
||||
>>> # The code level is now 3, but it has no effect because it has a lower priority than `set_code_level`
|
||||
|
||||
"""
|
||||
_TRANSLATOR_LOGGER.transformed_code_level = level
|
||||
_TRANSLATOR_LOGGER.need_to_echo_code_to_stdout = also_to_stdout
|
||||
|
||||
|
||||
def get_code_level():
|
||||
return _TRANSLATOR_LOGGER.transformed_code_level
|
||||
|
||||
|
||||
def error(msg, *args, **kwargs):
|
||||
_TRANSLATOR_LOGGER.error(msg, *args, **kwargs)
|
||||
|
||||
|
||||
def warn(msg, *args, **kwargs):
|
||||
_TRANSLATOR_LOGGER.warn(msg, *args, **kwargs)
|
||||
|
||||
|
||||
def log(level, msg, *args, **kwargs):
|
||||
_TRANSLATOR_LOGGER.log(level, msg, *args, **kwargs)
|
||||
|
||||
|
||||
def log_transformed_code(level, ast_node, transformer_name, *args, **kwargs):
|
||||
_TRANSLATOR_LOGGER.log_transformed_code(
|
||||
level, ast_node, transformer_name, *args, **kwargs
|
||||
)
|
||||
@@ -0,0 +1,327 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
import inspect
|
||||
import textwrap
|
||||
from collections.abc import Sequence
|
||||
|
||||
from paddle.base import core
|
||||
from paddle.framework import use_pir_api
|
||||
from paddle.utils import gast
|
||||
|
||||
from .utils import ORIGIN_INFO
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class Location:
|
||||
"""
|
||||
Location information of source code.
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"filepath",
|
||||
"lineno",
|
||||
"col_offset",
|
||||
)
|
||||
|
||||
def __init__(self, filepath, lineno, col_offset=None):
|
||||
self.filepath = filepath
|
||||
self.lineno = lineno
|
||||
self.col_offset = col_offset
|
||||
|
||||
def __str__(self):
|
||||
return f"location: {self.filepath}:{self.lineno}:{self.col_offset}"
|
||||
|
||||
@property
|
||||
def line_location(self):
|
||||
return (self.filepath, self.lineno)
|
||||
|
||||
|
||||
class OriginInfo:
|
||||
"""
|
||||
Original information of source code.
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"location",
|
||||
"function_name",
|
||||
"source_code",
|
||||
)
|
||||
|
||||
def __init__(self, location, function_name, source_code):
|
||||
self.location = location
|
||||
self.function_name = function_name
|
||||
self.source_code = source_code
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.location} \nsource_code: {self.source_code} in function {self.function_name}\n "
|
||||
|
||||
def formatted_message(self):
|
||||
flag_for_origin_info = "(* user code *)"
|
||||
return f' File "{self.location.filepath}", line {self.location.lineno}, in {self.function_name} {flag_for_origin_info}\n\t{self.source_code.lstrip()}'
|
||||
|
||||
def as_frame(self):
|
||||
return (
|
||||
self.location.filepath,
|
||||
self.location.lineno,
|
||||
self.function_name,
|
||||
self.source_code.lstrip(),
|
||||
)
|
||||
|
||||
|
||||
class OriginInfoAttacher(gast.NodeTransformer):
|
||||
"""
|
||||
Attach original source information to AST node according corresponding function.
|
||||
"""
|
||||
|
||||
def __init__(self, root, func):
|
||||
self.root = root
|
||||
self.func = inspect.unwrap(func)
|
||||
self.filepath = inspect.getsourcefile(self.func)
|
||||
self.source_code = inspect.getsource(self.func)
|
||||
self.current_func = []
|
||||
|
||||
def transform(self):
|
||||
source_lines, begin_lineno = inspect.getsourcelines(self.func)
|
||||
begin_line = source_lines[0]
|
||||
self.col_offset = len(begin_line) - len(begin_line.lstrip())
|
||||
self.source_lines = [line.strip("\n") for line in source_lines]
|
||||
self.lineno_offset = begin_lineno - 1
|
||||
self.visit(self.root)
|
||||
|
||||
def visit(self, node):
|
||||
if isinstance(node, gast.FunctionDef):
|
||||
self.current_func.append(node)
|
||||
if getattr(node, "lineno", None) is not None:
|
||||
self._attach_origin_info(node)
|
||||
self.generic_visit(node)
|
||||
|
||||
if isinstance(node, gast.FunctionDef):
|
||||
self.current_func.pop()
|
||||
return node
|
||||
|
||||
def _attach_origin_info(self, node):
|
||||
assert isinstance(node, gast.AST)
|
||||
assert hasattr(node, "lineno")
|
||||
|
||||
lineno = self._abs_lineno(node)
|
||||
col_offset = self._abs_col_offset(node)
|
||||
loc = Location(self.filepath, lineno, col_offset)
|
||||
func_name = self.current_func[-1].name
|
||||
code_line = self.source_lines[node.lineno - 1]
|
||||
|
||||
origin_info = OriginInfo(loc, func_name, code_line)
|
||||
setattr(node, ORIGIN_INFO, origin_info)
|
||||
|
||||
def _abs_lineno(self, node):
|
||||
return self.lineno_offset + node.lineno
|
||||
|
||||
def _abs_col_offset(self, node):
|
||||
return self.col_offset + node.col_offset
|
||||
|
||||
|
||||
global_origin_info_map = {}
|
||||
|
||||
|
||||
def create_and_update_origin_info_map(
|
||||
transformed_node, static_func, is_global=True
|
||||
):
|
||||
"""
|
||||
Creates a original information map between transformed static function and original dygraph function.
|
||||
|
||||
Args:
|
||||
transformed_node(gast.AST): The AST node of transformed dygraph function with attached source information of original dygraph function.
|
||||
static_func(Callable): The static function transformed by dygraph function corresponding to transformed_node.
|
||||
|
||||
Returns:
|
||||
The original information map.
|
||||
"""
|
||||
|
||||
origin_info_map = {}
|
||||
static_source = textwrap.dedent(inspect.getsource(static_func))
|
||||
static_node = gast.parse(static_source)
|
||||
static_node = attach_origin_info(static_node, static_func)
|
||||
|
||||
for t_node, s_node in ast_walk(transformed_node, static_node):
|
||||
assert type(t_node) == type(s_node), (
|
||||
f"The node types should be the same, but received type(t_node) is {type(t_node)}, and type(s_node) is {type(s_node)}."
|
||||
)
|
||||
dygraph_info = getattr(t_node, ORIGIN_INFO, None)
|
||||
static_info = getattr(s_node, ORIGIN_INFO, None)
|
||||
|
||||
if dygraph_info is None or static_info is None:
|
||||
continue
|
||||
static_loc = static_info.location.line_location
|
||||
exist_origin_info = origin_info_map.get(static_loc)
|
||||
|
||||
if exist_origin_info is not None:
|
||||
if (
|
||||
exist_origin_info.location.lineno
|
||||
>= dygraph_info.location.lineno
|
||||
):
|
||||
continue
|
||||
if (
|
||||
exist_origin_info.location.col_offset
|
||||
<= dygraph_info.location.col_offset
|
||||
):
|
||||
continue
|
||||
|
||||
origin_info_map[static_loc] = dygraph_info
|
||||
|
||||
global_origin_info_map.update(origin_info_map)
|
||||
if is_global:
|
||||
return global_origin_info_map
|
||||
|
||||
return origin_info_map
|
||||
|
||||
|
||||
def attach_origin_info(ast_node, func):
|
||||
"""
|
||||
Attach original source information to AST node according corresponding function.
|
||||
|
||||
Args:
|
||||
ast_node(gast.AST): The AST node to attach original source information.
|
||||
func(Callable): The corresponding function of ast_node. Parse the original information from this function.
|
||||
|
||||
Returns:
|
||||
An AST node attached original source information.
|
||||
"""
|
||||
resolver = OriginInfoAttacher(ast_node, func)
|
||||
resolver.transform()
|
||||
return ast_node
|
||||
|
||||
|
||||
def ast_walk(transformed_node, static_node):
|
||||
"""
|
||||
Recursively yield all descendant nodes in the trees starting at transformed_node and static_node (including itself) in parallel.
|
||||
|
||||
NOTE(liym27):
|
||||
Function ast.walk is not used because it yield all descendant nodes in no specified order.
|
||||
"""
|
||||
|
||||
def _as_list(x):
|
||||
if x is None:
|
||||
return []
|
||||
return list(x) if isinstance(x, Sequence) else [x]
|
||||
|
||||
transformed_node_list = _as_list(transformed_node)
|
||||
static_node_list = _as_list(static_node)
|
||||
|
||||
while transformed_node_list:
|
||||
assert len(transformed_node_list) == len(static_node_list)
|
||||
t_node = transformed_node_list.pop()
|
||||
s_node = static_node_list.pop()
|
||||
if type(t_node) != type(s_node):
|
||||
# NOTE(liym27):
|
||||
# Node types should be strictly required, but there is no strict distinction between gast.Load and gast.Param
|
||||
# in the ast transformation process.
|
||||
if isinstance(t_node, (gast.Load, gast.Param)) or isinstance(
|
||||
s_node, (gast.Load, gast.Param)
|
||||
):
|
||||
continue
|
||||
|
||||
assert type(t_node) == type(s_node), (
|
||||
f"The node types should be the same, but received type(t_node) is {type(t_node)}, and type(s_node) is {type(s_node)}."
|
||||
)
|
||||
|
||||
yield t_node, s_node
|
||||
|
||||
for field in t_node._fields:
|
||||
t_node_child = getattr(t_node, field)
|
||||
s_node_child = getattr(s_node, field)
|
||||
|
||||
if isinstance(t_node_child, gast.AST):
|
||||
transformed_node_list.append(t_node_child)
|
||||
static_node_list.append(s_node_child)
|
||||
elif isinstance(t_node_child, (list, tuple)):
|
||||
assert len(t_node_child) == len(s_node_child)
|
||||
for d_item, s_item in zip(t_node_child, s_node_child):
|
||||
if isinstance(d_item, gast.AST):
|
||||
transformed_node_list.append(d_item)
|
||||
static_node_list.append(s_item)
|
||||
|
||||
|
||||
def update_op_callstack_with_origin_info(program):
|
||||
"""
|
||||
Replaces op callstack information about transformed static code with original dygraph code.
|
||||
"""
|
||||
|
||||
def get_new_op_callstack(callstack):
|
||||
"""
|
||||
An example of callstack:
|
||||
|
||||
File "path1/to/file.py", line 10, in func_1
|
||||
y = paddle.tensor.fill_constant(x, shape=[1], dtype="int32")
|
||||
File "path2/to/file.py", line 740, in fill_constant
|
||||
stop_gradient=True)
|
||||
File "path3/to/file.py", line 43, in append_op
|
||||
return self.main_program.current_block().append_op(*args, **kwargs)
|
||||
File "path4/to/file.py", line 2811, in append_op
|
||||
attrs=kwargs.get("attrs", None))
|
||||
File "path5/to/file.py", line 1919, in __init__
|
||||
for frame in traceback.extract_stack():
|
||||
"""
|
||||
|
||||
assert len(callstack) % 2 == 0
|
||||
for i in range(0, len(callstack), 2):
|
||||
file_line = callstack[i].lstrip(" ").split(",")
|
||||
|
||||
filepath = file_line[0][6:-1]
|
||||
lineno = int(file_line[1][6:])
|
||||
funcname = file_line[2][4:]
|
||||
code = callstack[i + 1].lstrip(" ")
|
||||
|
||||
loc = Location(filepath, lineno)
|
||||
dygraph_func_info = global_origin_info_map.get(loc.line_location)
|
||||
if dygraph_func_info:
|
||||
filepath, lineno, funcname, code = dygraph_func_info.as_frame()
|
||||
|
||||
callstack[i] = f' File "{filepath}", line {lineno}, in {funcname}'
|
||||
callstack[i + 1] = f' {code}'
|
||||
|
||||
return callstack
|
||||
|
||||
def get_all_pir_block_ops(block):
|
||||
ops = []
|
||||
for op in block.ops:
|
||||
ops.append(op)
|
||||
for sub_block in op.blocks():
|
||||
ops += get_all_pir_block_ops(sub_block)
|
||||
return ops
|
||||
|
||||
op_maker = core.op_proto_and_checker_maker
|
||||
callstack_var_name = op_maker.kOpCreationCallstackAttrName()
|
||||
|
||||
if use_pir_api():
|
||||
global_block = program.global_block()
|
||||
ops = get_all_pir_block_ops(global_block)
|
||||
for op in ops:
|
||||
if op.has_attr(callstack_var_name):
|
||||
op.callstack = get_new_op_callstack(op.callstack)
|
||||
else:
|
||||
for block in program.blocks:
|
||||
for i, op in enumerate(block.ops):
|
||||
if op.has_attr(callstack_var_name):
|
||||
callstack = op.attr(callstack_var_name)
|
||||
|
||||
callstack = get_new_op_callstack(callstack)
|
||||
|
||||
try:
|
||||
# (@xiongkun) In 2-order derivative for paddle science, there may exists `pow_grad`
|
||||
# which has op_proto == nullptr and causes _set_attr failed. so we add a try...except.
|
||||
op._set_attr(callstack_var_name, callstack)
|
||||
except:
|
||||
pass
|
||||
return program
|
||||
@@ -0,0 +1,141 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
import paddle
|
||||
from paddle.autograd.backward_utils import ValueDict
|
||||
from paddle.framework import core
|
||||
|
||||
from ..dy2static.program_translator import _program_hash, synchronized
|
||||
|
||||
|
||||
@contextmanager
|
||||
def append_op_in_top_block():
|
||||
current_insertion_point = paddle.pir.get_current_insertion_point()
|
||||
top_block = paddle.static.default_main_program().global_block()
|
||||
paddle.pir.set_insertion_point_to_block_end(top_block)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
paddle.pir.set_insertion_point(current_insertion_point)
|
||||
|
||||
|
||||
class ParametersRecorder:
|
||||
def __init__(self):
|
||||
self.params_dict = {}
|
||||
self.tensor2value = {}
|
||||
|
||||
@synchronized
|
||||
def get(self, program, tensor):
|
||||
from paddle.pir.core import create_parameter, vartype_to_datatype
|
||||
|
||||
"""use the default_program as key, append tensor the parameter list."""
|
||||
key = _program_hash(program)
|
||||
if key not in self.params_dict:
|
||||
self.params_dict[key] = set()
|
||||
self.tensor2value[key] = {}
|
||||
|
||||
params = self.params_dict[key]
|
||||
mappings = self.tensor2value[key]
|
||||
if id(tensor) not in mappings:
|
||||
non_used_initializer = paddle.nn.initializer.Constant(0.0)
|
||||
dtype = tensor.dtype
|
||||
if isinstance(dtype, core.VarDesc.VarType):
|
||||
dtype = vartype_to_datatype[dtype]
|
||||
with append_op_in_top_block():
|
||||
value = create_parameter(
|
||||
dtype=dtype,
|
||||
shape=tensor.shape,
|
||||
type=tensor.type,
|
||||
name=tensor.name,
|
||||
initializer=non_used_initializer,
|
||||
trainable=(not tensor.stop_gradient),
|
||||
placements=tensor.placements,
|
||||
process_mesh=tensor.process_mesh,
|
||||
)
|
||||
|
||||
if isinstance(tensor, paddle.Tensor):
|
||||
params.add(tensor)
|
||||
mappings[id(tensor)] = value
|
||||
|
||||
return mappings[id(tensor)]
|
||||
|
||||
def pop(self, program):
|
||||
hash_id = _program_hash(program)
|
||||
params = self.params_dict.get(hash_id)
|
||||
if params is None:
|
||||
return [], []
|
||||
params = list(params)
|
||||
params.sort(key=lambda x: x.name)
|
||||
params_values = [self.tensor2value[hash_id][id(x)] for x in params]
|
||||
del self.params_dict[hash_id]
|
||||
del self.tensor2value[hash_id]
|
||||
return params, params_values
|
||||
|
||||
|
||||
class InplaceMap:
|
||||
def __init__(self):
|
||||
self.params_dict = {}
|
||||
|
||||
@synchronized
|
||||
def add(self, program, origin_value, new_value):
|
||||
key = _program_hash(program)
|
||||
if key not in self.params_dict:
|
||||
self.params_dict[key] = ValueDict()
|
||||
inplace_dict = self.params_dict[key]
|
||||
inplace_dict[origin_value] = new_value
|
||||
|
||||
def get(self, program, value):
|
||||
inplace_dict = self.params_dict.get(_program_hash(program))
|
||||
if inplace_dict is None:
|
||||
return None
|
||||
if value not in inplace_dict:
|
||||
return None
|
||||
root_var = inplace_dict[value]
|
||||
saved = []
|
||||
while root_var in inplace_dict:
|
||||
saved.append(root_var)
|
||||
root_var = inplace_dict[root_var]
|
||||
for var in saved:
|
||||
inplace_dict[var] = root_var
|
||||
return root_var
|
||||
|
||||
def pop(self, program):
|
||||
key = _program_hash(program)
|
||||
if key not in self.params_dict:
|
||||
return
|
||||
del self.params_dict[key]
|
||||
|
||||
def restore_checkpoint(self, checkpoint):
|
||||
# InplaceMap is a nested effect.
|
||||
# when enter a block, we should save a checkpoint
|
||||
# when exit a block, we should restore a checkpoint
|
||||
# for example:
|
||||
# if cond > 0:
|
||||
# x [:] = 0
|
||||
# return x
|
||||
# x[:] only effect current cond block, we should restore in false block.
|
||||
self.params_dict = checkpoint
|
||||
|
||||
def save_checkpoint(self):
|
||||
checkpoint = {}
|
||||
for program_id, params in self.params_dict.items():
|
||||
new_params = dict(params.items())
|
||||
checkpoint[program_id] = new_params
|
||||
return checkpoint
|
||||
|
||||
|
||||
_global_parameter_recorder = ParametersRecorder()
|
||||
_global_inplace_map = InplaceMap()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,162 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
import textwrap
|
||||
|
||||
from paddle import pir
|
||||
from paddle.base.framework import Variable, in_pir_mode
|
||||
from paddle.base.libpaddle.pir import build_pipe_for_pylayer
|
||||
from paddle.common_ops_import import LayerHelper
|
||||
from paddle.static.nn import static_pylayer
|
||||
from paddle.utils import flatten, pack_sequence_as
|
||||
|
||||
from .program_translator import convert_to_static, unwrap_decorators
|
||||
|
||||
|
||||
class StaticPyLayerContext:
|
||||
def __init__(self):
|
||||
self.saved_vars = []
|
||||
self.saved_vars_structure = None
|
||||
|
||||
if in_pir_mode():
|
||||
self.tuple_push_op_name = "cf.tuple_push"
|
||||
self.tuple_pop_op_name = "cf.tuple_pop"
|
||||
|
||||
def __setattr__(self, attr: str, value: object):
|
||||
attr_allow_list = ["saved_vars", "saved_vars_structure"]
|
||||
if (
|
||||
in_pir_mode()
|
||||
and attr not in attr_allow_list
|
||||
and isinstance(value, pir.Value)
|
||||
):
|
||||
raise AttributeError(
|
||||
textwrap.dedent(
|
||||
f"""\
|
||||
`ctx.{attr} = tensor` is not allowed in static mode, please use `ctx.save_for_backward(tensor)` instead.
|
||||
|
||||
For example:
|
||||
|
||||
>>> class ExamplePyLayer(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... # ctx.x = x # This is not allowed in static mode, Replace it with `ctx.save_for_backward(x)`
|
||||
... ctx.save_for_backward(x)
|
||||
... x1 = paddle.tanh(x)
|
||||
... return x1
|
||||
|
||||
... @staticmethod
|
||||
... def backward(ctx, grad):
|
||||
... # x = ctx.x # Same as above, replace it with `x, = ctx.saved_tensor()`
|
||||
... x, = ctx.saved_tensor()
|
||||
... x_grad = grad * (1 - paddle.square(x))
|
||||
... return x_grad
|
||||
"""
|
||||
)
|
||||
)
|
||||
super().__setattr__(attr, value)
|
||||
|
||||
def save_for_backward(self, *tensors):
|
||||
if in_pir_mode():
|
||||
self.saved_vars_structure = tensors
|
||||
flatten_tensors = flatten(tensors)
|
||||
tensor_elements = list(
|
||||
filter(lambda x: isinstance(x, pir.Value), flatten_tensors)
|
||||
)
|
||||
current_insert_point = pir.get_current_insertion_point()
|
||||
current_block = current_insert_point.block()
|
||||
build_pipe_for_pylayer(current_block, tensor_elements)
|
||||
else:
|
||||
for tensor in tensors:
|
||||
assert isinstance(tensor, Variable)
|
||||
self.saved_vars.append(tensor)
|
||||
|
||||
def saved_tensor(self):
|
||||
if in_pir_mode():
|
||||
current_insert_point = pir.get_current_insertion_point()
|
||||
current_block = current_insert_point.block()
|
||||
out_list = []
|
||||
for op in current_block.ops:
|
||||
if op.name() == self.tuple_pop_op_name:
|
||||
out_list = op.as_tuple_pop_op().pop_all_values()
|
||||
if self.saved_vars_structure is not None:
|
||||
flattened_structure = flatten(self.saved_vars_structure)
|
||||
value_cursor = 0
|
||||
for i, tensor in enumerate(flattened_structure):
|
||||
if isinstance(tensor, pir.Value):
|
||||
flattened_structure[i] = out_list[value_cursor]
|
||||
value_cursor += 1
|
||||
out_list = pack_sequence_as(
|
||||
self.saved_vars_structure, flattened_structure
|
||||
)
|
||||
else:
|
||||
helper = LayerHelper("StaticPyLayerContext")
|
||||
out_list = []
|
||||
for saved_var in self.saved_vars:
|
||||
out = helper.create_variable(
|
||||
name=saved_var.name,
|
||||
dtype=saved_var.dtype,
|
||||
shape=saved_var.shape,
|
||||
type=saved_var.type,
|
||||
)
|
||||
out_list.append(out)
|
||||
|
||||
return out_list
|
||||
|
||||
# TODO(MarioLulab): support not_inplace
|
||||
def mark_not_inplace(self, *args):
|
||||
raise NotImplementedError
|
||||
|
||||
# TODO(MarioLulab): support non_differentiable
|
||||
def mark_non_differentiable(self, *args):
|
||||
raise NotImplementedError
|
||||
|
||||
# TODO(MarioLulab): support materialize_grads
|
||||
def set_materialize_grads(self, value: bool):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class StaticPyLayer:
|
||||
def __init__(self, dyfunc_self):
|
||||
self.dyfunc_self = dyfunc_self
|
||||
_, self.orig_forward_fn = unwrap_decorators(dyfunc_self.forward)
|
||||
_, self.orig_backward_fn = unwrap_decorators(dyfunc_self.backward)
|
||||
self.static_pylayer_context = StaticPyLayerContext()
|
||||
|
||||
self.forward_fn_with_ctx = functools.partial(
|
||||
convert_to_static(self.orig_forward_fn), self.static_pylayer_context
|
||||
)
|
||||
self.backward_fn_with_ctx = functools.partial(
|
||||
convert_to_static(self.orig_backward_fn),
|
||||
self.static_pylayer_context,
|
||||
)
|
||||
|
||||
# NOTE: only support position args and Variables Now
|
||||
def apply(self, *args, **kwargs):
|
||||
# rearrange `position-args + keyword-args` into `position-args`
|
||||
dyfunc_sig = inspect.signature(self.dyfunc_self.forward)
|
||||
bound_args = dyfunc_sig.bind(self.dyfunc_self, *args, **kwargs)
|
||||
bound_args.apply_defaults()
|
||||
input_args = [
|
||||
item
|
||||
for i, item in enumerate(bound_args.arguments.values())
|
||||
if i > 0
|
||||
] # index 0 indicate `dyfunc_self` which shouldn't be put into `input_args`
|
||||
|
||||
return static_pylayer(
|
||||
forward_fn=self.forward_fn_with_ctx,
|
||||
inputs=input_args,
|
||||
backward_fn=self.backward_fn_with_ctx,
|
||||
)
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from .transform import DygraphToStaticAst # noqa: F401
|
||||
@@ -0,0 +1,46 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.jit.dy2static.utils import ast_to_source_code
|
||||
from paddle.utils import gast
|
||||
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class AssertTransformer(BaseTransformer):
|
||||
"""
|
||||
A class transforms python assert to convert_assert.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_Assert(self, node):
|
||||
convert_assert_node = (
|
||||
gast.parse(
|
||||
'_jst.Assert({test}, {msg})'.format(
|
||||
test=ast_to_source_code(node.test),
|
||||
msg=ast_to_source_code(node.msg) if node.msg else "",
|
||||
)
|
||||
)
|
||||
.body[0]
|
||||
.value
|
||||
)
|
||||
|
||||
return gast.Expr(value=convert_assert_node)
|
||||
@@ -0,0 +1,566 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.base import unique_name
|
||||
from paddle.jit.dy2static.utils import (
|
||||
ORIGIN_INFO,
|
||||
ast_to_source_code,
|
||||
)
|
||||
from paddle.utils import gast
|
||||
|
||||
from .utils import (
|
||||
FOR_ITER_INDEX_PREFIX,
|
||||
FOR_ITER_ITERATOR_PREFIX,
|
||||
FOR_ITER_TARGET_PREFIX,
|
||||
FOR_ITER_VAR_LEN_PREFIX,
|
||||
FOR_ITER_VAR_NAME_PREFIX,
|
||||
FOR_ITER_ZIP_TO_LIST_PREFIX,
|
||||
create_assign_node,
|
||||
)
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class BaseTransformer(gast.NodeTransformer):
|
||||
def visit(self, node):
|
||||
if not isinstance(node, gast.AST):
|
||||
msg = f'Expected "gast.AST", but got "{type(node)}".'
|
||||
raise ValueError(msg)
|
||||
origin_info = getattr(node, ORIGIN_INFO, None)
|
||||
|
||||
result = super().visit(node)
|
||||
|
||||
iter_result = result
|
||||
if iter_result is not node and iter_result is not None:
|
||||
if not isinstance(iter_result, (list, tuple)):
|
||||
iter_result = (iter_result,)
|
||||
if origin_info is not None:
|
||||
for n in iter_result:
|
||||
setattr(n, ORIGIN_INFO, origin_info)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class NameNodeReplaceTransformer(BaseTransformer):
|
||||
"""
|
||||
This class replaces specified gast.Name node by replace_node.
|
||||
"""
|
||||
|
||||
def __init__(self, root_node, target_name, replace_node):
|
||||
assert isinstance(target_name, str)
|
||||
|
||||
# NOTE(liym27):
|
||||
# Use gast.Name to replace gast.Name, otherwise, errors may occur.
|
||||
#
|
||||
# For examples:
|
||||
# If using a gast.Subscript to replace gast.Name, and the original gast.Name
|
||||
# is in the arguments of FunctionDef, an exception will be raised.
|
||||
#
|
||||
# ```
|
||||
# def func(x[i])) # x[i] can not be a argument
|
||||
# # ...
|
||||
# ```
|
||||
|
||||
assert isinstance(replace_node, gast.Name)
|
||||
self.target_name = target_name
|
||||
self.replace_node = replace_node
|
||||
|
||||
self.visit(root_node)
|
||||
|
||||
def visit_Name(self, node):
|
||||
if node.id == self.target_name:
|
||||
return self.replace_node
|
||||
return node
|
||||
|
||||
def visit_Nonlocal(self, node):
|
||||
names = node.names
|
||||
|
||||
def replace(s):
|
||||
if s == self.target_name:
|
||||
return self.replace_node.id
|
||||
return s
|
||||
|
||||
node.names = list(map(replace, names))
|
||||
return node
|
||||
|
||||
|
||||
class ForLoopTuplePreTransformer(BaseTransformer):
|
||||
"""pre-process of for loop.
|
||||
>>> for A in B:
|
||||
>>> C
|
||||
|
||||
will be changed into :
|
||||
|
||||
>>> # make iterator-only to indexable list.
|
||||
>>> UUID_iterator = _jst.Indexable(B)
|
||||
>>> for UUID_target in UUID_iterator:
|
||||
>>> A = _jst.Unpack(UUID_target, structure)
|
||||
>>> C
|
||||
|
||||
make the later loop_transform have unified type:
|
||||
>>> for target in iter:
|
||||
>>> body
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_For(self, node):
|
||||
self.generic_visit(node)
|
||||
tuple_target = unique_name.generate(FOR_ITER_TARGET_PREFIX)
|
||||
tuple_iterator = unique_name.generate(FOR_ITER_ITERATOR_PREFIX)
|
||||
origin_tuple_node = node.target
|
||||
assign_iterator_node = gast.parse(
|
||||
f"{tuple_iterator} = _jst.Indexable({ast_to_source_code(node.iter).strip()})"
|
||||
).body[0]
|
||||
node.target = gast.Name(
|
||||
id=tuple_target,
|
||||
ctx=gast.Store(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
)
|
||||
node.iter = gast.Name(
|
||||
id=tuple_iterator,
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
)
|
||||
node.body[0:0] = self.tuple_to_stmts(origin_tuple_node, tuple_target)
|
||||
# return a list will insert a list of node replace the original for node.
|
||||
return [assign_iterator_node, node]
|
||||
|
||||
def tuple_node_to_unpack_structure(self, node):
|
||||
"""Create a sequence to represents the structure of nest.
|
||||
For example: `a, (b,c), [d,e,f]` is represented by
|
||||
`[1, [1,1], [1,1,1]]`. the `1` is just a notation.
|
||||
|
||||
Specially, `a` is represented by `1`.
|
||||
"""
|
||||
ret = []
|
||||
if not isinstance(node, (gast.Tuple, gast.List)):
|
||||
return 1
|
||||
for element in node.elts:
|
||||
ret.append(self.tuple_node_to_unpack_structure(element))
|
||||
return ret
|
||||
|
||||
def tuple_to_stmts(self, node, tuple_name):
|
||||
structure_str = str(self.tuple_node_to_unpack_structure(node))
|
||||
node_str = ast_to_source_code(node).strip()
|
||||
assign_node_str = (
|
||||
f"{node_str} = _jst.Unpack({tuple_name}, {structure_str})"
|
||||
)
|
||||
assign_node = gast.parse(assign_node_str).body[0]
|
||||
return [assign_node]
|
||||
|
||||
|
||||
class ForNodeVisitor:
|
||||
"""
|
||||
This class parses python for statement, get transformed 3 statement components of for node
|
||||
three key statements:
|
||||
1). init_stmts: list[node], prepare nodes of for loop, may not only one
|
||||
2). cond_stmt: node, condition node to judge whether continue loop
|
||||
3). body_stmts: list[node], updated loop body, sometimes we should change
|
||||
the original statement in body, not just append new statement
|
||||
|
||||
In this process, the semantics of for does not change.
|
||||
|
||||
Now only can parse 3 type statements (Here var is Tensor(Tensor) or python variable):
|
||||
1). for x in range(var[*]|var.numpy()[*])
|
||||
2). for x in var|var.numpy()
|
||||
3). for i, x enumerate(var|var.numpy())
|
||||
"""
|
||||
|
||||
def __init__(self, for_node):
|
||||
assert isinstance(for_node, gast.For), (
|
||||
"Input node for the initialization of ForNodeVisitor is not gast.For node."
|
||||
)
|
||||
# 1. original for node
|
||||
self.node = for_node
|
||||
|
||||
# 2. gast.For node main parts
|
||||
self.target = for_node.target
|
||||
# NOTE: type may be Node or list[Node]
|
||||
self.iter_args = (
|
||||
for_node.iter if self.is_for_iter() else for_node.iter.args
|
||||
)
|
||||
self.body = for_node.body
|
||||
|
||||
# 3. key shared node or names
|
||||
# - x:
|
||||
# - for x in range(***)
|
||||
# - for x in var|var.numpy()
|
||||
# - for i, x enumerate(var|var.numpy())
|
||||
self.iter_var_name = self._get_iter_var_name()
|
||||
|
||||
# - created index var to slice Variable: __for_loop_var_index_0
|
||||
# - for x in var|var.numpy()
|
||||
# - for i, x enumerate(var|var.numpy())
|
||||
self.iter_idx_name = unique_name.generate(FOR_ITER_INDEX_PREFIX)
|
||||
|
||||
# - created shape var to build loop condition: __for_loop_var_len_0
|
||||
# - for x in var|var.numpy()
|
||||
# - for i, x enumerate(var|var.numpy())
|
||||
# - for x in var
|
||||
self.iter_var_len_name = unique_name.generate(FOR_ITER_VAR_LEN_PREFIX)
|
||||
# - created zip to list var : __for_loop_iter_zip_0
|
||||
self.iter_zip_to_list_name = unique_name.generate(
|
||||
FOR_ITER_ZIP_TO_LIST_PREFIX
|
||||
)
|
||||
|
||||
# - var.numpy()/var
|
||||
# - for x in var|var.numpy()
|
||||
# - for i, x enumerate(var|var.numpy())
|
||||
self.iter_node = self._get_iter_node()
|
||||
|
||||
# - enumerate i:
|
||||
# - for i, x enumerate(var|var.numpy())
|
||||
self.enum_idx_name = self._get_enum_idx_name()
|
||||
|
||||
# - range/enumerate args length
|
||||
self.args_length = None
|
||||
|
||||
def parse(self):
|
||||
self._args_check()
|
||||
if self.is_for_range_iter():
|
||||
return self._parse_for_range_stmts()
|
||||
elif self.is_for_iter():
|
||||
return self._parse_for_stmts()
|
||||
elif self.is_for_enumerate_iter():
|
||||
return self._parse_for_enumerate_stmts()
|
||||
else:
|
||||
return None
|
||||
|
||||
def is_for_range_iter(self):
|
||||
return (
|
||||
isinstance(self.node.iter, gast.Call)
|
||||
and isinstance(self.node.iter.func, gast.Name)
|
||||
and self.node.iter.func.id == "range"
|
||||
)
|
||||
|
||||
def is_for_iter(self):
|
||||
if isinstance(
|
||||
self.node.iter, (gast.Name, gast.Attribute, gast.List, gast.Tuple)
|
||||
):
|
||||
return True
|
||||
elif (
|
||||
isinstance(self.node.iter, gast.Call)
|
||||
and isinstance(self.node.iter.func, gast.Attribute)
|
||||
and self.node.iter.func.attr == 'numpy'
|
||||
):
|
||||
return True
|
||||
elif isinstance(self.node.iter, gast.Subscript):
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
def is_for_enumerate_iter(self):
|
||||
return (
|
||||
isinstance(self.node.iter, gast.Call)
|
||||
and isinstance(self.node.iter.func, gast.Name)
|
||||
and self.node.iter.func.id == "enumerate"
|
||||
)
|
||||
|
||||
def _args_check(self):
|
||||
if self.is_for_range_iter():
|
||||
self.args_length = len(self.iter_args)
|
||||
assert self.args_length >= 1 and self.args_length <= 3, (
|
||||
"range() function takes 1 to 3 arguments"
|
||||
)
|
||||
elif self.is_for_enumerate_iter():
|
||||
self.args_length = len(self.iter_args)
|
||||
assert self.args_length >= 1 and self.args_length <= 2, (
|
||||
"enumerate() function takes 1 to 2 arguments"
|
||||
)
|
||||
else:
|
||||
self.args_length = None
|
||||
|
||||
def _parse_for_range_stmts(self):
|
||||
init_stmts = []
|
||||
init_stmts.append(self._build_index_init_node())
|
||||
|
||||
compare_node = self._build_compare_node()
|
||||
step_node = self._build_step_node()
|
||||
cond_stmt = self._build_cond_stmt(step_node, compare_node)
|
||||
|
||||
body_stmts = self.body
|
||||
body_stmts.append(self._build_index_increase_node(step_node))
|
||||
|
||||
return init_stmts, cond_stmt, body_stmts
|
||||
|
||||
def _parse_for_stmts(self):
|
||||
init_stmts = []
|
||||
init_stmts.extend(self._build_iter_node())
|
||||
init_stmts.append(self._build_index_init_node())
|
||||
init_stmts.append(self._build_var_len_assign_node())
|
||||
|
||||
compare_node = self._build_compare_node()
|
||||
step_node = self._build_step_node()
|
||||
cond_stmt = self._build_cond_stmt(step_node, compare_node)
|
||||
|
||||
body_stmts = self.body
|
||||
|
||||
# NOTE(liym27): Here add a gast.Assign, and the target of it is gast.Name.
|
||||
# In NameNodeReplaceTransformer, using gast.Name to replace gast.Name is safe.
|
||||
target_node, assign_node = self._build_assign_var_slice_node()
|
||||
body_stmts[0:0] = [assign_node]
|
||||
for body_node in body_stmts:
|
||||
NameNodeReplaceTransformer(
|
||||
body_node, self.iter_var_name, target_node
|
||||
)
|
||||
body_stmts.append(self._build_index_increase_node(step_node))
|
||||
|
||||
return init_stmts, cond_stmt, body_stmts
|
||||
|
||||
def _parse_for_enumerate_stmts(self):
|
||||
init_stmts = []
|
||||
init_stmts.extend(self._build_iter_node())
|
||||
init_stmts.append(self._build_index_init_node())
|
||||
init_stmts.append(self._build_var_len_assign_node())
|
||||
init_stmts.append(self._build_enum_init_node())
|
||||
|
||||
compare_node = self._build_compare_node()
|
||||
step_node = self._build_step_node()
|
||||
cond_stmt = self._build_cond_stmt(step_node, compare_node)
|
||||
|
||||
body_stmts = self.body
|
||||
|
||||
target_node, assign_node = self._build_assign_var_slice_node()
|
||||
body_stmts[0:0] = [assign_node]
|
||||
for body_node in body_stmts:
|
||||
NameNodeReplaceTransformer(
|
||||
body_node, self.iter_var_name, target_node
|
||||
)
|
||||
|
||||
body_stmts.append(self._build_index_increase_node(step_node))
|
||||
body_stmts.append(self._build_enum_increase_node())
|
||||
|
||||
return init_stmts, cond_stmt, body_stmts
|
||||
|
||||
def _build_index_init_node(self):
|
||||
if self.is_for_range_iter():
|
||||
if self.args_length == 1:
|
||||
index_init_value_str = '0'
|
||||
else:
|
||||
index_init_value_str = ast_to_source_code(
|
||||
self.iter_args[0]
|
||||
).strip()
|
||||
|
||||
index_init_var_name = self.iter_var_name
|
||||
else:
|
||||
index_init_value_str = '0'
|
||||
index_init_var_name = self.iter_idx_name
|
||||
|
||||
index_init_node_source_str = (
|
||||
f"{index_init_var_name} = {index_init_value_str}"
|
||||
)
|
||||
|
||||
index_init_node = gast.parse(index_init_node_source_str).body[0]
|
||||
|
||||
return index_init_node
|
||||
|
||||
def _build_var_len_assign_node(self):
|
||||
# get the length of iterable variable
|
||||
if (
|
||||
isinstance(self.iter_node, gast.Call)
|
||||
and isinstance(self.iter_node.func, gast.Attribute)
|
||||
and self.iter_node.func.attr == 'numpy'
|
||||
):
|
||||
iter_var_name = ast_to_source_code(
|
||||
self.iter_node.func.value
|
||||
).strip()
|
||||
else:
|
||||
iter_var_name = ast_to_source_code(self.iter_node).strip()
|
||||
|
||||
convert_len_node_source_str = (
|
||||
f'{self.iter_var_len_name} = _jst.Len({iter_var_name})'
|
||||
)
|
||||
|
||||
convert_len_node = gast.parse(convert_len_node_source_str).body[0]
|
||||
|
||||
return convert_len_node
|
||||
|
||||
def _build_iter_node(self):
|
||||
"""
|
||||
Process special cases for iter_node include:
|
||||
- Case 1 (for zip):
|
||||
|
||||
- for i, val in enumerate(zip(x, y)) # original code:
|
||||
|
||||
- __for_loop_iter_zip_0 = list(zip(x, y))
|
||||
- for i, val in enumerate(__for_loop_iter_zip_0)
|
||||
"""
|
||||
new_nodes = []
|
||||
if isinstance(self.iter_node, gast.Call) and isinstance(
|
||||
self.iter_node.func, gast.Name
|
||||
):
|
||||
if self.iter_node.func.id == 'zip':
|
||||
iter_var_name = ast_to_source_code(self.iter_node).strip()
|
||||
zip_to_list_str = (
|
||||
f"{self.iter_zip_to_list_name} = list({iter_var_name})"
|
||||
)
|
||||
zip_to_list_node = gast.parse(zip_to_list_str).body[0]
|
||||
new_nodes.append(zip_to_list_node)
|
||||
|
||||
self.iter_node = gast.Name(
|
||||
id=self.iter_zip_to_list_name,
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
)
|
||||
|
||||
return new_nodes
|
||||
|
||||
def _build_enum_init_node(self):
|
||||
if self.is_for_enumerate_iter() and self.args_length != 1:
|
||||
init_value_str = ast_to_source_code(self.iter_args[1]).strip()
|
||||
else:
|
||||
init_value_str = '0'
|
||||
|
||||
enum_init_node_source_str = f"{self.enum_idx_name} = {init_value_str}"
|
||||
enum_init_node = gast.parse(enum_init_node_source_str).body[0]
|
||||
return enum_init_node
|
||||
|
||||
def _build_compare_node(self):
|
||||
if self.is_for_range_iter():
|
||||
compare_node = (
|
||||
self.iter_args[0]
|
||||
if self.args_length == 1
|
||||
else self.iter_args[1]
|
||||
)
|
||||
else:
|
||||
compare_node = gast.Name(
|
||||
id=self.iter_var_len_name,
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
)
|
||||
return compare_node
|
||||
|
||||
def _build_step_node(self):
|
||||
if self.is_for_range_iter():
|
||||
step_node = (
|
||||
self.iter_args[2]
|
||||
if self.args_length == 3
|
||||
else gast.Constant(value=1, kind=None)
|
||||
)
|
||||
else:
|
||||
step_node = gast.Constant(value=1, kind=None)
|
||||
return step_node
|
||||
|
||||
def _build_cond_stmt(self, step_node, compare_node):
|
||||
if not isinstance(step_node, (gast.Constant, gast.UnaryOp)):
|
||||
raise NotImplementedError(
|
||||
"Dynamic-to-Static only supports the step value is a constant or negative constant in 'for-range' statements, "
|
||||
f"such as '2', '-3'. But received: '{ast_to_source_code(step_node).strip()}'. Please fix code to be compatible with Dynamic-to-Static."
|
||||
)
|
||||
|
||||
if isinstance(step_node, gast.UnaryOp) or step_node.value < 0:
|
||||
# eg:
|
||||
# range(max, min, -2)
|
||||
# ->
|
||||
# i > min
|
||||
return gast.Compare(
|
||||
left=gast.Name(
|
||||
id=(
|
||||
self.iter_var_name
|
||||
if self.is_for_range_iter()
|
||||
else self.iter_idx_name
|
||||
),
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
ops=[gast.Gt()],
|
||||
comparators=[compare_node],
|
||||
)
|
||||
else:
|
||||
# eg:
|
||||
# range(min, max, 2)
|
||||
# ->
|
||||
# i < max
|
||||
return gast.Compare(
|
||||
left=gast.Name(
|
||||
id=(
|
||||
self.iter_var_name
|
||||
if self.is_for_range_iter()
|
||||
else self.iter_idx_name
|
||||
),
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
ops=[gast.Lt()],
|
||||
comparators=[compare_node],
|
||||
)
|
||||
|
||||
def _build_index_increase_node(self, step_node):
|
||||
return gast.AugAssign(
|
||||
target=gast.Name(
|
||||
id=(
|
||||
self.iter_var_name
|
||||
if self.is_for_range_iter()
|
||||
else self.iter_idx_name
|
||||
),
|
||||
ctx=gast.Store(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
op=gast.Add(),
|
||||
value=step_node,
|
||||
)
|
||||
|
||||
def _build_assign_var_slice_node(self):
|
||||
var_slice_str = f"{ast_to_source_code(self.iter_node).strip()}[{self.iter_idx_name}]"
|
||||
var_slice_node = gast.parse(var_slice_str).body[0].value
|
||||
new_iter_var_name = unique_name.generate(FOR_ITER_VAR_NAME_PREFIX)
|
||||
target_node, assign_node = create_assign_node(
|
||||
new_iter_var_name, var_slice_node
|
||||
)
|
||||
return target_node, assign_node
|
||||
|
||||
def _build_enum_increase_node(self):
|
||||
return gast.AugAssign(
|
||||
target=gast.Name(
|
||||
id=self.enum_idx_name,
|
||||
ctx=gast.Store(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
op=gast.Add(),
|
||||
value=gast.Constant(value=1, kind=None),
|
||||
)
|
||||
|
||||
def _get_iter_var_name(self):
|
||||
if self.is_for_range_iter():
|
||||
return self.target.id
|
||||
elif self.is_for_iter():
|
||||
return self.target.id
|
||||
elif self.is_for_enumerate_iter():
|
||||
return self.target.elts[1].id
|
||||
return None
|
||||
|
||||
def _get_iter_node(self):
|
||||
if self.is_for_iter():
|
||||
return self.iter_args
|
||||
elif self.is_for_enumerate_iter():
|
||||
return self.iter_args[0]
|
||||
return None
|
||||
|
||||
def _get_enum_idx_name(self):
|
||||
if self.is_for_enumerate_iter():
|
||||
return self.target.elts[0].id
|
||||
return None
|
||||
@@ -0,0 +1,419 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.base import unique_name
|
||||
from paddle.utils import gast
|
||||
|
||||
from .base import BaseTransformer, ForNodeVisitor
|
||||
from .utils import BaseNodeVisitor, create_bool_node, index_in_list
|
||||
|
||||
__all__ = []
|
||||
|
||||
BREAK_NAME_PREFIX = '__break'
|
||||
CONTINUE_NAME_PREFIX = '__continue'
|
||||
|
||||
|
||||
class ForToWhileTransformer(BaseTransformer):
|
||||
"""
|
||||
Transform python for loop into while loop and add condition node in the
|
||||
loop test
|
||||
"""
|
||||
|
||||
def __init__(self, parent_node, loop_node, condition_node):
|
||||
assert isinstance(loop_node, gast.For), (
|
||||
"loop_node is not gast.For in ForToWhileTransformer"
|
||||
)
|
||||
self.parent_node = parent_node
|
||||
self.loop_node = loop_node
|
||||
self.condition_node = condition_node
|
||||
|
||||
def transform(self):
|
||||
if hasattr(self.parent_node, 'body'):
|
||||
body_list = self.parent_node.body
|
||||
i = index_in_list(body_list, self.loop_node)
|
||||
if i != -1:
|
||||
new_stmts = self.get_for_stmt_nodes(body_list[i])
|
||||
body_list[i : i + 1] = new_stmts
|
||||
i += len(new_stmts)
|
||||
return new_stmts
|
||||
if hasattr(self.parent_node, 'orelse'):
|
||||
body_list = self.parent_node.orelse
|
||||
i = index_in_list(body_list, self.loop_node)
|
||||
if i != -1:
|
||||
new_stmts = self.get_for_stmt_nodes(body_list[i])
|
||||
body_list[i : i + 1] = new_stmts
|
||||
i += len(new_stmts)
|
||||
return new_stmts
|
||||
raise ValueError(
|
||||
"parent_node doesn't contain the loop_node in ForToWhileTransformer"
|
||||
)
|
||||
|
||||
def get_for_stmt_nodes(self, node):
|
||||
assert isinstance(node, gast.For), (
|
||||
"Input node is NOT gast.For in get_for_stmt_nodes"
|
||||
)
|
||||
|
||||
# 1. parse current gast.For node
|
||||
current_for_node_parser = ForNodeVisitor(node)
|
||||
stmts_tuple = current_for_node_parser.parse()
|
||||
if stmts_tuple is None:
|
||||
return [node]
|
||||
init_stmts, cond_stmt, body_stmts = stmts_tuple
|
||||
|
||||
# 2. append break statement
|
||||
new_cond_stmt = gast.BoolOp(
|
||||
op=gast.And(), values=[cond_stmt, self.condition_node]
|
||||
)
|
||||
|
||||
# 3. construct gast.While node
|
||||
while_node = gast.While(
|
||||
test=new_cond_stmt, body=body_stmts, orelse=node.orelse
|
||||
)
|
||||
init_stmts.append(while_node)
|
||||
return init_stmts
|
||||
|
||||
|
||||
class BreakContinueTransformer(BaseNodeVisitor):
|
||||
"""
|
||||
Rewrite 'break' and 'continue' key words in a if-else python way to make
|
||||
it equivalent to original control flow
|
||||
|
||||
The main idea of this class is:
|
||||
|
||||
1. Map the 'break/continue' stmt with an unique boolean variable V.
|
||||
|
||||
2. Find the first ancestor block containing this 'break/continue', a
|
||||
block can be a node containing stmt list. We should remove all stmts
|
||||
after the 'break/continue' and set the V to True here.
|
||||
|
||||
3. Add 'if V' for stmts in ancestor blocks between the first one
|
||||
(exclusive) and the ancestor loop (inclusive)
|
||||
|
||||
4. For 'break' add break into condition of the loop. For 'continue',
|
||||
set continue to False at the beginning of each loop
|
||||
|
||||
TODO: more details should be summarized as design document
|
||||
|
||||
Note: The class is inherited from BaseNodeVisitor instead of NodeTransformer,
|
||||
because ancestor nodes will be modified inplace for `Break/Continue` here.
|
||||
In general, we recommend to inheriting NodeTransformer to modify node!
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
super().__init__()
|
||||
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_Break(self, node):
|
||||
function_def_node_index = _find_ancestor_function_def_index(
|
||||
self.ancestor_nodes
|
||||
)
|
||||
loop_node_index = _find_ancestor_loop_index(node, self.ancestor_nodes)
|
||||
assert loop_node_index != -1, "SyntaxError: 'break' outside loop"
|
||||
loop_node = self.ancestor_nodes[loop_node_index]
|
||||
function_def_node = self.ancestor_nodes[function_def_node_index]
|
||||
|
||||
# 1. Map the 'break/continue' stmt with an unique boolean variable V.
|
||||
variable_name = unique_name.generate(BREAK_NAME_PREFIX)
|
||||
|
||||
# 2. Find the first ancestor block containing this 'break/continue', a
|
||||
# block can be a node containing stmt list. We should remove all stmts
|
||||
# after the 'break/continue' and set the V to True here.
|
||||
first_block_index = self._remove_stmts_after_break_continue(
|
||||
node, variable_name, loop_node_index
|
||||
)
|
||||
|
||||
# 3. Add 'if not V' for stmts in ancestor blocks between the first one
|
||||
# (exclusive) and the ancestor loop (inclusive)
|
||||
self._replace_if_stmt(loop_node_index, first_block_index, variable_name)
|
||||
|
||||
# 4. For 'break' add break into condition of the loop.
|
||||
assign_false_node = create_bool_node(variable_name, False)
|
||||
function_def_node.body.insert(0, assign_false_node)
|
||||
|
||||
cond_var_node = gast.UnaryOp(
|
||||
op=gast.Not(),
|
||||
operand=gast.Name(
|
||||
id=variable_name,
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
)
|
||||
|
||||
if isinstance(loop_node, gast.While):
|
||||
loop_node.test = gast.BoolOp(
|
||||
op=gast.And(), values=[loop_node.test, cond_var_node]
|
||||
)
|
||||
elif isinstance(loop_node, gast.For):
|
||||
parent_node = self.ancestor_nodes[loop_node_index - 1]
|
||||
for_to_while = ForToWhileTransformer(
|
||||
parent_node, loop_node, cond_var_node
|
||||
)
|
||||
for_to_while.transform()
|
||||
|
||||
def visit_Continue(self, node):
|
||||
function_def_node_index = _find_ancestor_function_def_index(
|
||||
self.ancestor_nodes
|
||||
)
|
||||
loop_node_index = _find_ancestor_loop_index(node, self.ancestor_nodes)
|
||||
assert loop_node_index != -1, "SyntaxError: 'continue' outside loop"
|
||||
loop_node = self.ancestor_nodes[loop_node_index]
|
||||
function_def_node = self.ancestor_nodes[function_def_node_index]
|
||||
|
||||
# 1. Map the 'break/continue' stmt with an unique boolean variable V.
|
||||
variable_name = unique_name.generate(CONTINUE_NAME_PREFIX)
|
||||
|
||||
# 2. Find the first ancestor block containing this 'break/continue', a
|
||||
# block can be a node containing stmt list. We should remove all stmts
|
||||
# after the 'break/continue' and set the V to True here.
|
||||
first_block_index = self._remove_stmts_after_break_continue(
|
||||
node, variable_name, loop_node_index
|
||||
)
|
||||
|
||||
# 3. Add 'if not V' for stmts in ancestor blocks between the first one
|
||||
# (exclusive) and the ancestor loop (inclusive)
|
||||
self._replace_if_stmt(loop_node_index, first_block_index, variable_name)
|
||||
|
||||
# 4. For 'continue', set continue to False at the beginning of each loop
|
||||
assign_false_node = create_bool_node(variable_name, False)
|
||||
loop_node.body.insert(0, assign_false_node)
|
||||
# Add a same assign statement to the beginning of function body to avoid
|
||||
# generate the UndefinedVar
|
||||
function_def_node.body.insert(0, assign_false_node)
|
||||
|
||||
def _remove_stmts_after_break_continue(
|
||||
self, break_continue_node, break_continue_name, loop_node_index
|
||||
):
|
||||
for first_block_index in range(
|
||||
len(self.ancestor_nodes) - 1, loop_node_index - 1, -1
|
||||
):
|
||||
first_block = self.ancestor_nodes[first_block_index]
|
||||
if hasattr(
|
||||
first_block, "body"
|
||||
) and self._replace_break_continue_in_stmt_list(
|
||||
first_block.body, break_continue_node, break_continue_name
|
||||
):
|
||||
return first_block_index
|
||||
|
||||
if hasattr(
|
||||
first_block, "orelse"
|
||||
) and self._replace_break_continue_in_stmt_list(
|
||||
first_block.orelse, break_continue_node, break_continue_name
|
||||
):
|
||||
return first_block_index
|
||||
|
||||
return first_block_index
|
||||
|
||||
def _replace_if_stmt(
|
||||
self, loop_node_index, first_block_index, break_continue_name
|
||||
):
|
||||
for i in range(first_block_index - 1, loop_node_index - 1, -1):
|
||||
cur_node = self.ancestor_nodes[i]
|
||||
son_node = self.ancestor_nodes[i + 1]
|
||||
if hasattr(
|
||||
cur_node, 'body'
|
||||
) and self._replace_after_node_to_if_in_stmt_list(
|
||||
cur_node.body, son_node, break_continue_name
|
||||
):
|
||||
continue
|
||||
if hasattr(
|
||||
cur_node, 'orelse'
|
||||
) and self._replace_after_node_to_if_in_stmt_list(
|
||||
cur_node.orelse, son_node, break_continue_name
|
||||
):
|
||||
continue
|
||||
|
||||
def _replace_break_continue_in_stmt_list(
|
||||
self, stmt_list, break_continue_node, break_continue_name
|
||||
):
|
||||
i = index_in_list(stmt_list, break_continue_node)
|
||||
if i == -1:
|
||||
return False
|
||||
assign_true_node = create_bool_node(break_continue_name, True)
|
||||
stmt_list[i:] = [assign_true_node]
|
||||
return True
|
||||
|
||||
def _replace_after_node_to_if_in_stmt_list(
|
||||
self, stmt_list, node, break_continue_name
|
||||
):
|
||||
i = index_in_list(stmt_list, node)
|
||||
if i == -1:
|
||||
return False
|
||||
|
||||
if i == len(stmt_list) - 1:
|
||||
# No need to add, we consider this as added successfully
|
||||
return True
|
||||
|
||||
if_stmt = gast.If(
|
||||
test=gast.UnaryOp(
|
||||
op=gast.Not(),
|
||||
operand=gast.Name(
|
||||
id=break_continue_name,
|
||||
ctx=gast.Store(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
),
|
||||
body=stmt_list[i + 1 :],
|
||||
orelse=[],
|
||||
)
|
||||
stmt_list[i + 1 :] = []
|
||||
stmt_list.append(if_stmt)
|
||||
return True
|
||||
|
||||
def _add_stmt_before_cur_node(self, cur_node_index, stmt_node):
|
||||
cur_node = self.ancestor_nodes[cur_node_index]
|
||||
parent_node = self.ancestor_nodes[cur_node_index - 1]
|
||||
if hasattr(
|
||||
parent_node, "body"
|
||||
) and self._add_stmt_into_list_before_node(
|
||||
parent_node.body, cur_node, stmt_node
|
||||
):
|
||||
return True
|
||||
if hasattr(
|
||||
parent_node, "orelse"
|
||||
) and self._add_stmt_into_list_before_node(
|
||||
parent_node.orelse, cur_node, stmt_node
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _add_stmt_into_list_before_node(self, stmt_list, node, stmt_node):
|
||||
i = index_in_list(stmt_list, node)
|
||||
if i == -1:
|
||||
return False
|
||||
stmt_list.insert(i, stmt_node)
|
||||
return True
|
||||
|
||||
|
||||
def _find_ancestor_loop_index(node, ancestor_nodes):
|
||||
for i in range(len(ancestor_nodes) - 1, -1, -1):
|
||||
if isinstance(ancestor_nodes[i], (gast.For, gast.While)):
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
def _find_ancestor_function_def_index(ancestor_nodes):
|
||||
for i in range(len(ancestor_nodes) - 1, -1, -1):
|
||||
if isinstance(ancestor_nodes[i], gast.FunctionDef):
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
class BreakTransformOptimizer(BaseNodeVisitor):
|
||||
"""
|
||||
In specific pattern, the transformed code could be optimized by joining the
|
||||
If.test with while.test.
|
||||
|
||||
Currently supported pattern is:
|
||||
```
|
||||
while cond1: while cond1 and not cond2:
|
||||
if cond2: ---> do_something()
|
||||
break
|
||||
do_something()
|
||||
```
|
||||
|
||||
See following example:
|
||||
|
||||
>>> def foo(x):
|
||||
... i = paddle.to_tensor(1, dtype='int32')
|
||||
... while i < 10:
|
||||
... if x.mean() > 5:
|
||||
... break
|
||||
... x += i
|
||||
... i += 1
|
||||
... return x
|
||||
|
||||
The generated code after applying optimization will be:
|
||||
```
|
||||
def foo(x):
|
||||
i = paddle.to_tensor(1, dtype='int32')
|
||||
while i < 10 and not x.mean() > 5:
|
||||
x += i
|
||||
i += 1
|
||||
return x
|
||||
```
|
||||
It can avoid wrapping all ops after `break` statement into `cond_op` that
|
||||
usually brings very heavy overhead.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
super().__init__()
|
||||
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_Break(self, node):
|
||||
loop_node_index = _find_ancestor_loop_index(node, self.ancestor_nodes)
|
||||
assert loop_node_index != -1, "SyntaxError: 'break' outside loop"
|
||||
loop_node = self.ancestor_nodes[loop_node_index]
|
||||
|
||||
if self._is_break_cond_pattern(node, loop_node):
|
||||
cond_var_node = self._join_with_while_cond(node, loop_node)
|
||||
|
||||
if isinstance(loop_node, gast.While):
|
||||
loop_node.test = gast.BoolOp(
|
||||
op=gast.And(), values=[loop_node.test, cond_var_node]
|
||||
)
|
||||
elif isinstance(loop_node, gast.For):
|
||||
parent_node = self.ancestor_nodes[loop_node_index - 1]
|
||||
for_to_while = ForToWhileTransformer(
|
||||
parent_node, loop_node, cond_var_node
|
||||
)
|
||||
for_to_while.transform()
|
||||
|
||||
def _is_break_cond_pattern(self, break_node, loop_node):
|
||||
"""
|
||||
Judge whether if match the pattern to join `If.test` with `while.test`
|
||||
"""
|
||||
# while/for -> if -> break
|
||||
if len(self.ancestor_nodes) < 3 or self.ancestor_nodes[-3] != loop_node:
|
||||
return False
|
||||
|
||||
assert self.ancestor_nodes[-1] == break_node
|
||||
parent_if_node = self.ancestor_nodes[-2]
|
||||
|
||||
is_matched = False
|
||||
if isinstance(parent_if_node, gast.If):
|
||||
# gast.If only contains `break`
|
||||
break_first_in_if = (
|
||||
parent_if_node.body[0] == break_node
|
||||
and len(parent_if_node.orelse) == 0
|
||||
)
|
||||
# gast.If is first node of loop_node
|
||||
if_first_in_loop = loop_node.body[0] == parent_if_node
|
||||
|
||||
is_matched = if_first_in_loop and break_first_in_if
|
||||
|
||||
return is_matched
|
||||
|
||||
def _join_with_while_cond(self, break_node, loop_node):
|
||||
"""
|
||||
Join the `If.test` with `While.test` together.
|
||||
"""
|
||||
parent_if_node = self.ancestor_nodes[-2]
|
||||
|
||||
cond_var_node = gast.UnaryOp(op=gast.Not(), operand=parent_if_node.test)
|
||||
|
||||
# remove the gast.If node that contains the gast.Break.
|
||||
assert loop_node.body[0] == parent_if_node
|
||||
loop_node.body.pop(0)
|
||||
|
||||
return cond_var_node
|
||||
@@ -0,0 +1,81 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import ast_to_source_code, is_builtin
|
||||
from .base import BaseTransformer
|
||||
from .utils import is_paddle_api
|
||||
|
||||
PDB_SET = "pdb.set_trace"
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class CallTransformer(BaseTransformer):
|
||||
"""
|
||||
This class transforms function calls into Static Graph Ast.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def _no_need_convert_call(self, node):
|
||||
"""
|
||||
Determines whether a function needs to be transformed by `convert_call`.
|
||||
It doesn't need to be transformed when a function satisfies the following conditions:
|
||||
1. It's a api of paddle
|
||||
2. It's a python builtin function not include `len`, `zip`, `range` and `enumerate`
|
||||
"""
|
||||
assert isinstance(node, gast.Call)
|
||||
if is_paddle_api(node):
|
||||
return True
|
||||
|
||||
func_str = ast_to_source_code(node.func).strip()
|
||||
try:
|
||||
need_convert_builtin_func_list = {
|
||||
'len',
|
||||
'zip',
|
||||
'range',
|
||||
'enumerate',
|
||||
'print',
|
||||
}
|
||||
fn = eval(func_str)
|
||||
is_builtin_fn = is_builtin(fn)
|
||||
need_convert = func_str in need_convert_builtin_func_list
|
||||
return is_builtin_fn and not need_convert
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_Call(self, node):
|
||||
self.generic_visit(node)
|
||||
|
||||
if self._no_need_convert_call(node):
|
||||
return node
|
||||
|
||||
func_str = ast_to_source_code(node.func).strip()
|
||||
|
||||
# NOTE(liym27): Don't convert `pad.set_trace` even if the conversion doesn't work finally, because
|
||||
# it is clearer to see where it is called from.
|
||||
if PDB_SET in func_str:
|
||||
return node
|
||||
|
||||
new_func_str = f"_jst.Call({func_str})"
|
||||
new_func_ast = gast.parse(new_func_str).body[0].value
|
||||
node.func = new_func_ast
|
||||
|
||||
return node
|
||||
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.jit.dy2static.utils import ast_to_source_code
|
||||
from paddle.utils import gast
|
||||
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class CastTransformer(BaseTransformer):
|
||||
"""
|
||||
This class transforms type casting into Static Graph Ast.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
self._castable_type = {'bool', 'int', 'float', 'complex'}
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_Call(self, node):
|
||||
self.generic_visit(node)
|
||||
func_str = ast_to_source_code(node.func).strip()
|
||||
if func_str in self._castable_type and len(node.args) > 0:
|
||||
args_str = ast_to_source_code(node.args[0]).strip()
|
||||
new_func_str = f"_jst.AsDtype({args_str}, '{func_str}')"
|
||||
new_node = gast.parse(new_func_str).body[0].value
|
||||
return new_node
|
||||
|
||||
return node
|
||||
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from .base import BaseTransformer
|
||||
from .utils import FunctionNameLivenessAnalysis, create_undefined_var
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class CreateVariableTransformer(BaseTransformer):
|
||||
""" """
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
FunctionNameLivenessAnalysis(self.root)
|
||||
|
||||
def transform(self):
|
||||
"""
|
||||
Main function to transform AST.
|
||||
"""
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
# attributes = set(filter(lambda x: '.' in x, node.pd_scope.modified_vars()))
|
||||
self.generic_visit(node)
|
||||
bodys = node.body
|
||||
names = sorted(node.pd_scope.created_vars())
|
||||
for name in names:
|
||||
bodys[0:0] = [create_undefined_var(name)]
|
||||
return node
|
||||
@@ -0,0 +1,140 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
import re
|
||||
import warnings
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import RE_PYMODULE, RE_PYNAME, ast_to_source_code
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
IGNORE_NAMES = [
|
||||
'declarative',
|
||||
'to_static',
|
||||
'wraps',
|
||||
'staticmethod',
|
||||
'classmethod',
|
||||
'decorator',
|
||||
'inference',
|
||||
]
|
||||
|
||||
|
||||
class DecoratorTransformer(BaseTransformer):
|
||||
"""
|
||||
Transform decorators.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
"""
|
||||
Main function to transform AST.
|
||||
"""
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
assert isinstance(node, gast.FunctionDef)
|
||||
self.generic_visit(node)
|
||||
|
||||
deco_list = node.decorator_list
|
||||
node.decorator_list = []
|
||||
|
||||
# every decorator will append a node
|
||||
decofun_nodes = []
|
||||
# func to be decoded next time
|
||||
deco_target = '_orig_' + node.name
|
||||
# last decoded func
|
||||
decoded_func = ''
|
||||
|
||||
for deco in reversed(deco_list):
|
||||
# skip IGNORE_NAMES
|
||||
deco_full_name = ast_to_source_code(deco).strip()
|
||||
if isinstance(deco, gast.Call):
|
||||
# match case like :
|
||||
# 1: @_jst.Call(a.b.c.d.deco)()
|
||||
# 2: @q.w.e.r.deco()
|
||||
re_tmp = re.match(
|
||||
rf'({RE_PYMODULE})*({RE_PYNAME}\(){{0,1}}({RE_PYMODULE})*({RE_PYNAME})(\)){{0,1}}\(.*$',
|
||||
deco_full_name,
|
||||
)
|
||||
deco_name = re_tmp.group(4)
|
||||
else:
|
||||
# match case like:
|
||||
# @a.d.g.deco
|
||||
re_tmp = re.match(
|
||||
rf'({RE_PYMODULE})*({RE_PYNAME})$',
|
||||
deco_full_name,
|
||||
)
|
||||
deco_name = re_tmp.group(2)
|
||||
if deco_name in IGNORE_NAMES:
|
||||
continue
|
||||
elif deco_name == 'contextmanager':
|
||||
warnings.warn(
|
||||
"Dy2Static : A context manager decorator is used, this may not work correctly after transform."
|
||||
)
|
||||
|
||||
decoded_func = '_decoedby_' + deco_name
|
||||
|
||||
# get function after decoration
|
||||
if isinstance(deco, gast.Call):
|
||||
if '_jst.Call' in deco_full_name:
|
||||
# in this case , the deco_full_name will be like:
|
||||
# '_jst.Call(deco)(5)'
|
||||
rematch = re.match(
|
||||
r'\_jst\.Call\((.+?)\)\((.*)\)', deco_full_name
|
||||
)
|
||||
re_name = rematch.group(1)
|
||||
re_args = rematch.group(2)
|
||||
re_args_with_func = deco_target + ', ' + re_args
|
||||
decofun_str = f'try:\n\t{decoded_func} = _jst.Call({re_name})({re_args_with_func})\nexcept:\n\t{decoded_func} = _jst.Call({re_name})({re_args})({deco_target})'
|
||||
else:
|
||||
# paddle api will not be transformed to '_jst.Call'
|
||||
rematch = re.match(r'(.+?)\((.*)\)', deco_full_name)
|
||||
re_name = rematch.group(1)
|
||||
re_args = rematch.group(2)
|
||||
re_args_with_func = deco_target + ', ' + re_args
|
||||
decofun_str = f'try:\n\t{decoded_func} = {re_name}({re_args_with_func})\nexcept:\n\t{decoded_func} = {re_name}({re_args})({deco_target})'
|
||||
|
||||
else:
|
||||
decofun_str = f'{decoded_func} = _jst.Call({deco_full_name})({deco_target})'
|
||||
|
||||
decofun_nodes.extend(gast.parse(decofun_str).body)
|
||||
deco_target = decoded_func
|
||||
|
||||
if not decofun_nodes:
|
||||
return node
|
||||
|
||||
orig_func_node = gast.FunctionDef(
|
||||
name='_orig_' + node.name,
|
||||
args=node.args,
|
||||
body=node.body,
|
||||
decorator_list=[],
|
||||
returns=None,
|
||||
type_comment=None,
|
||||
type_params=[],
|
||||
)
|
||||
|
||||
args = [arg.id for arg in node.args.args]
|
||||
arg_str = ','.join(args)
|
||||
callfun_str = f'return {decoded_func}({arg_str})'
|
||||
callfun_node = gast.parse(callfun_str).body[0]
|
||||
|
||||
node.body = [orig_func_node, *decofun_nodes, callfun_node]
|
||||
|
||||
return node
|
||||
@@ -0,0 +1,86 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class EarlyReturnTransformer(BaseTransformer):
|
||||
"""
|
||||
Transform if/else return statement of Dygraph into Static Graph.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
"""
|
||||
Main function to transform AST.
|
||||
"""
|
||||
self.visit(self.root)
|
||||
|
||||
def is_define_return_in_if(self, node):
|
||||
assert isinstance(node, gast.If), (
|
||||
f"Type of input node should be gast.If, but received {type(node)}."
|
||||
)
|
||||
for child in node.body:
|
||||
if isinstance(child, gast.Return):
|
||||
return True
|
||||
return False
|
||||
|
||||
def visit_block_nodes(self, nodes):
|
||||
result_nodes = []
|
||||
destination_nodes = result_nodes
|
||||
for node in nodes:
|
||||
rewritten_node = self.visit(node)
|
||||
|
||||
if isinstance(rewritten_node, (list, tuple)):
|
||||
destination_nodes.extend(rewritten_node)
|
||||
else:
|
||||
destination_nodes.append(rewritten_node)
|
||||
|
||||
# append other nodes to if.orelse even though if.orelse is not empty
|
||||
if isinstance(node, gast.If) and self.is_define_return_in_if(node):
|
||||
destination_nodes = node.orelse
|
||||
# handle stmt like `if/elif/elif`
|
||||
while (
|
||||
len(destination_nodes) > 0
|
||||
and isinstance(destination_nodes[0], gast.If)
|
||||
and self.is_define_return_in_if(destination_nodes[0])
|
||||
):
|
||||
destination_nodes = destination_nodes[0].orelse
|
||||
|
||||
return result_nodes
|
||||
|
||||
def visit_If(self, node):
|
||||
node.body = self.visit_block_nodes(node.body)
|
||||
node.orelse = self.visit_block_nodes(node.orelse)
|
||||
return node
|
||||
|
||||
def visit_While(self, node):
|
||||
node.body = self.visit_block_nodes(node.body)
|
||||
node.orelse = self.visit_block_nodes(node.orelse)
|
||||
return node
|
||||
|
||||
def visit_For(self, node):
|
||||
node.body = self.visit_block_nodes(node.body)
|
||||
node.orelse = self.visit_block_nodes(node.orelse)
|
||||
return node
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
node.body = self.visit_block_nodes(node.body)
|
||||
return node
|
||||
@@ -0,0 +1,447 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
import copy
|
||||
from collections import defaultdict
|
||||
|
||||
from paddle.base import unique_name
|
||||
from paddle.jit.dy2static.utils import (
|
||||
GetterSetterHelper,
|
||||
ast_to_source_code,
|
||||
)
|
||||
|
||||
# gast is a generic AST to represent Python2 and Python3's Abstract Syntax Tree(AST).
|
||||
# It provides a compatibility layer between the AST of various Python versions,
|
||||
# as produced by ast.parse from the standard ast module.
|
||||
# See details in https://github.com/serge-sans-paille/gast/
|
||||
from paddle.utils import gast
|
||||
|
||||
from .base import BaseTransformer
|
||||
from .utils import (
|
||||
FALSE_FUNC_PREFIX,
|
||||
FOR_ITER_INDEX_PREFIX,
|
||||
FOR_ITER_ITERATOR_PREFIX,
|
||||
FOR_ITER_TARGET_PREFIX,
|
||||
FOR_ITER_TUPLE_INDEX_PREFIX,
|
||||
FOR_ITER_TUPLE_PREFIX,
|
||||
FOR_ITER_VAR_LEN_PREFIX,
|
||||
FOR_ITER_VAR_NAME_PREFIX,
|
||||
FOR_ITER_ZIP_TO_LIST_PREFIX,
|
||||
TRUE_FUNC_PREFIX,
|
||||
FunctionNameLivenessAnalysis,
|
||||
create_function_def_node,
|
||||
create_get_args_node,
|
||||
create_name_str,
|
||||
create_nonlocal_stmt_nodes,
|
||||
create_set_args_node,
|
||||
)
|
||||
|
||||
__all__ = []
|
||||
|
||||
GET_ARGS_FUNC_PREFIX = 'get_args'
|
||||
SET_ARGS_FUNC_PREFIX = 'set_args'
|
||||
ARGS_NAME = '__args'
|
||||
|
||||
|
||||
class IfElseTransformer(BaseTransformer):
|
||||
"""
|
||||
Transform if/else statement of Dygraph into Static Graph.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
FunctionNameLivenessAnalysis(
|
||||
self.root
|
||||
) # name analysis of current ast tree.
|
||||
|
||||
def transform(self):
|
||||
"""
|
||||
Main function to transform AST.
|
||||
"""
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_If(self, node):
|
||||
self.generic_visit(node)
|
||||
(
|
||||
true_func_node,
|
||||
false_func_node,
|
||||
get_args_node,
|
||||
set_args_node,
|
||||
return_name_ids,
|
||||
push_pop_ids,
|
||||
) = transform_if_else(node, self.root)
|
||||
|
||||
new_node = create_convert_ifelse_node(
|
||||
return_name_ids,
|
||||
push_pop_ids,
|
||||
node.test,
|
||||
true_func_node,
|
||||
false_func_node,
|
||||
get_args_node,
|
||||
set_args_node,
|
||||
)
|
||||
|
||||
return [
|
||||
get_args_node,
|
||||
set_args_node,
|
||||
true_func_node,
|
||||
false_func_node,
|
||||
new_node,
|
||||
]
|
||||
|
||||
def visit_Call(self, node):
|
||||
# Remove `numpy()` statement, like `Tensor.numpy()[i]` -> `Tensor[i]`
|
||||
if isinstance(node.func, gast.Attribute):
|
||||
attribute = node.func
|
||||
if attribute.attr == 'numpy':
|
||||
node = attribute.value
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def visit_IfExp(self, node):
|
||||
"""
|
||||
Transformation with `true_fn(x) if Tensor > 0 else false_fn(x)`
|
||||
"""
|
||||
self.generic_visit(node)
|
||||
|
||||
new_node = create_convert_ifelse_node(
|
||||
None, None, node.test, node.body, node.orelse, None, None, True
|
||||
)
|
||||
# Note: A blank line will be added separately if transform gast.Expr
|
||||
# into source code. Using gast.Expr.value instead to avoid syntax error
|
||||
# in python.
|
||||
if isinstance(new_node, gast.Expr):
|
||||
new_node = new_node.value
|
||||
|
||||
return new_node
|
||||
|
||||
|
||||
class NameVisitor(gast.NodeVisitor):
|
||||
def __init__(self, after_node=None, end_node=None):
|
||||
# The start node (exclusive) of the visitor
|
||||
self.after_node = after_node
|
||||
# The terminate node of the visitor.
|
||||
self.end_node = end_node
|
||||
# Dict to store the names and ctxs of vars.
|
||||
self.name_ids = defaultdict(list)
|
||||
# List of current visited nodes
|
||||
self.ancestor_nodes = []
|
||||
# True when in range (after_node, end_node).
|
||||
self._in_range = after_node is None
|
||||
self._candidate_ctxs = (gast.Store, gast.Load, gast.Param)
|
||||
self._def_func_names = set()
|
||||
|
||||
def visit(self, node):
|
||||
"""Visit a node."""
|
||||
if self.after_node is not None and node == self.after_node:
|
||||
self._in_range = True
|
||||
return
|
||||
if node == self.end_node:
|
||||
self._in_range = False
|
||||
return
|
||||
|
||||
self.ancestor_nodes.append(node)
|
||||
method = 'visit_' + node.__class__.__name__
|
||||
visitor = getattr(self, method, self.generic_visit)
|
||||
ret = visitor(node)
|
||||
self.ancestor_nodes.pop()
|
||||
|
||||
return ret
|
||||
|
||||
def visit_If(self, node):
|
||||
"""
|
||||
For nested `if/else`, the created vars are not always visible for parent node.
|
||||
In addition, the vars created in `if.body` are not visible for `if.orelse`.
|
||||
|
||||
Case 1:
|
||||
x = 1
|
||||
if m > 1:
|
||||
res = new_tensor
|
||||
res = res + 1 # Error, `res` is not visible here.
|
||||
|
||||
Case 2:
|
||||
if x_tensor > 0:
|
||||
res = new_tensor
|
||||
else:
|
||||
res = res + 1 # Error, `res` is not visible here.
|
||||
|
||||
In above two cases, we should consider to manage the scope of vars to parsing
|
||||
the arguments and returned vars correctly.
|
||||
"""
|
||||
if not self._in_range or not self.end_node:
|
||||
self.generic_visit(node)
|
||||
return
|
||||
else:
|
||||
before_if_name_ids = copy.deepcopy(self.name_ids)
|
||||
body_name_ids = self._visit_child(node.body)
|
||||
# If traversal process stops early in `if.body`, return the currently seen name_ids.
|
||||
if not self._in_range:
|
||||
self._update_name_ids(before_if_name_ids)
|
||||
else:
|
||||
else_name_ids = self._visit_child(node.orelse)
|
||||
# If traversal process stops early in `if.orelse`, return the currently seen name_ids.
|
||||
if not self._in_range:
|
||||
self._update_name_ids(before_if_name_ids)
|
||||
else:
|
||||
# Blocks the vars in `if.body` and only inserts the vars both created in 'if/else' branch
|
||||
# into name_ids.
|
||||
new_name_ids = self._find_new_name_ids(
|
||||
body_name_ids, else_name_ids
|
||||
)
|
||||
for new_name_id in new_name_ids:
|
||||
before_if_name_ids[new_name_id].append(gast.Store())
|
||||
|
||||
self.name_ids = before_if_name_ids
|
||||
|
||||
def visit_Attribute(self, node):
|
||||
if not self._in_range or not self._is_call_func_name_node(node):
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Name(self, node):
|
||||
if not self._in_range:
|
||||
self.generic_visit(node)
|
||||
return
|
||||
blacklist = {'True', 'False', 'None'}
|
||||
if node.id in blacklist:
|
||||
return
|
||||
if node.id in self._def_func_names:
|
||||
return
|
||||
if not self._is_call_func_name_node(node):
|
||||
if isinstance(node.ctx, self._candidate_ctxs):
|
||||
self.name_ids[node.id].append(node.ctx)
|
||||
|
||||
def visit_Assign(self, node):
|
||||
if not self._in_range:
|
||||
self.generic_visit(node)
|
||||
return
|
||||
# Visit `value` firstly.
|
||||
node._fields = ('value', 'targets')
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
# NOTE: We skip to visit names of get_args and set_args, because they contains
|
||||
# nonlocal statement such as 'nonlocal x, self' where 'self' should not be
|
||||
# parsed as returned value in control flow.
|
||||
if (
|
||||
GET_ARGS_FUNC_PREFIX in node.name
|
||||
or SET_ARGS_FUNC_PREFIX in node.name
|
||||
):
|
||||
return
|
||||
|
||||
if not self._in_range:
|
||||
self.generic_visit(node)
|
||||
return
|
||||
self._def_func_names.add(node.name)
|
||||
if not self.end_node:
|
||||
self.generic_visit(node)
|
||||
else:
|
||||
before_name_ids = copy.deepcopy(self.name_ids)
|
||||
self.name_ids = defaultdict(list)
|
||||
self.generic_visit(node)
|
||||
|
||||
if not self._in_range:
|
||||
self._update_name_ids(before_name_ids)
|
||||
else:
|
||||
self.name_ids = before_name_ids
|
||||
|
||||
def _visit_child(self, node):
|
||||
self.name_ids = defaultdict(list)
|
||||
if isinstance(node, list):
|
||||
for item in node:
|
||||
if isinstance(item, gast.AST):
|
||||
self.visit(item)
|
||||
elif isinstance(node, gast.AST):
|
||||
self.visit(node)
|
||||
|
||||
return copy.deepcopy(self.name_ids)
|
||||
|
||||
def _find_new_name_ids(self, body_name_ids, else_name_ids):
|
||||
def is_required_ctx(ctxs, required_ctx):
|
||||
for ctx in ctxs:
|
||||
if isinstance(ctx, required_ctx):
|
||||
return True
|
||||
return False
|
||||
|
||||
candidate_name_ids = set(body_name_ids.keys()) & set(
|
||||
else_name_ids.keys()
|
||||
)
|
||||
store_ctx = gast.Store
|
||||
new_name_ids = set()
|
||||
for name_id in candidate_name_ids:
|
||||
if is_required_ctx(
|
||||
body_name_ids[name_id], store_ctx
|
||||
) and is_required_ctx(else_name_ids[name_id], store_ctx):
|
||||
new_name_ids.add(name_id)
|
||||
|
||||
return new_name_ids
|
||||
|
||||
def _is_call_func_name_node(self, node):
|
||||
white_func_names = {'append', 'extend'}
|
||||
if len(self.ancestor_nodes) > 1:
|
||||
assert self.ancestor_nodes[-1] == node
|
||||
parent_node = self.ancestor_nodes[-2]
|
||||
if isinstance(parent_node, gast.Call) and parent_node.func == node:
|
||||
# e.g: var_list.append(elem), var_list is also a name_id.
|
||||
should_skip = (
|
||||
isinstance(node, gast.Attribute)
|
||||
and node.attr in white_func_names
|
||||
)
|
||||
if not should_skip:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _update_name_ids(self, new_name_ids):
|
||||
for name_id, ctxs in new_name_ids.items():
|
||||
self.name_ids[name_id] = ctxs + self.name_ids[name_id]
|
||||
|
||||
|
||||
def _valid_nonlocal_names(return_name_ids, nonlocal_names):
|
||||
"""
|
||||
All var in return_name_ids should be in nonlocal_names.
|
||||
Moreover, we will always put return_name_ids in front of nonlocal_names.
|
||||
|
||||
For Example:
|
||||
|
||||
return_name_ids: [x, y]
|
||||
nonlocal_names : [a, y, b, x]
|
||||
|
||||
Return:
|
||||
nonlocal_names : [x, y, a, b]
|
||||
"""
|
||||
assert isinstance(return_name_ids, list)
|
||||
for name in return_name_ids:
|
||||
if name not in nonlocal_names:
|
||||
raise ValueError(
|
||||
f"Required returned var '{name}' must be in 'nonlocal' statement '', but not found."
|
||||
)
|
||||
nonlocal_names.remove(name)
|
||||
|
||||
return return_name_ids + nonlocal_names
|
||||
|
||||
|
||||
def transform_if_else(node, root):
|
||||
"""
|
||||
Transform ast.If into control flow statement of Paddle static graph.
|
||||
"""
|
||||
|
||||
# TODO(liym27): Consider variable like `self.a` modified in if/else node.
|
||||
return_name_ids = sorted(node.pd_scope.modified_vars())
|
||||
push_pop_ids = sorted(node.pd_scope.variadic_length_vars())
|
||||
nonlocal_names = list(return_name_ids)
|
||||
nonlocal_names.sort()
|
||||
# NOTE: All var in return_name_ids should be in nonlocal_names.
|
||||
nonlocal_names = _valid_nonlocal_names(return_name_ids, nonlocal_names)
|
||||
|
||||
# TODO(dev): Need a better way to deal this.
|
||||
# LoopTransformer will create some special vars, which is not visible by users. so we can sure it's safe to remove them.
|
||||
filter_names = [
|
||||
ARGS_NAME,
|
||||
FOR_ITER_INDEX_PREFIX,
|
||||
FOR_ITER_TUPLE_PREFIX,
|
||||
FOR_ITER_TARGET_PREFIX,
|
||||
FOR_ITER_ITERATOR_PREFIX,
|
||||
FOR_ITER_TUPLE_INDEX_PREFIX,
|
||||
FOR_ITER_VAR_LEN_PREFIX,
|
||||
FOR_ITER_VAR_NAME_PREFIX,
|
||||
FOR_ITER_ZIP_TO_LIST_PREFIX,
|
||||
]
|
||||
|
||||
def remove_if(x):
|
||||
for name in filter_names:
|
||||
if x.startswith(name):
|
||||
return False
|
||||
return True
|
||||
|
||||
nonlocal_names = list(filter(remove_if, nonlocal_names))
|
||||
return_name_ids = nonlocal_names
|
||||
|
||||
nonlocal_stmt_node = create_nonlocal_stmt_nodes(nonlocal_names)
|
||||
|
||||
empty_arg_node = gast.arguments(
|
||||
args=[],
|
||||
posonlyargs=[],
|
||||
vararg=None,
|
||||
kwonlyargs=[],
|
||||
kw_defaults=None,
|
||||
kwarg=None,
|
||||
defaults=[],
|
||||
)
|
||||
|
||||
true_func_node = create_function_def_node(
|
||||
nonlocal_stmt_node + node.body,
|
||||
name=unique_name.generate(TRUE_FUNC_PREFIX),
|
||||
input_args=empty_arg_node,
|
||||
return_name_ids=[],
|
||||
)
|
||||
false_func_node = create_function_def_node(
|
||||
nonlocal_stmt_node + node.orelse,
|
||||
name=unique_name.generate(FALSE_FUNC_PREFIX),
|
||||
input_args=empty_arg_node,
|
||||
return_name_ids=[],
|
||||
)
|
||||
|
||||
helper = GetterSetterHelper(None, None, nonlocal_names, push_pop_ids)
|
||||
get_args_node = create_get_args_node(helper.union())
|
||||
set_args_node = create_set_args_node(helper.union())
|
||||
|
||||
return (
|
||||
true_func_node,
|
||||
false_func_node,
|
||||
get_args_node,
|
||||
set_args_node,
|
||||
return_name_ids,
|
||||
push_pop_ids,
|
||||
)
|
||||
|
||||
|
||||
def create_convert_ifelse_node(
|
||||
return_name_ids,
|
||||
push_pop_ids,
|
||||
pred,
|
||||
true_func,
|
||||
false_func,
|
||||
get_args_func,
|
||||
set_args_func,
|
||||
is_if_expr=False,
|
||||
):
|
||||
"""
|
||||
Create `paddle.jit.dy2static.convert_ifelse(
|
||||
pred, true_fn, false_fn, get_args, set_args, return_name_ids)`
|
||||
to replace original `python if/else` statement.
|
||||
"""
|
||||
if is_if_expr:
|
||||
true_func_source = f"lambda : {ast_to_source_code(true_func)}"
|
||||
false_func_source = f"lambda : {ast_to_source_code(false_func)}"
|
||||
else:
|
||||
true_func_source = true_func.name
|
||||
false_func_source = false_func.name
|
||||
|
||||
convert_ifelse_layer = gast.parse(
|
||||
'_jst.IfElse('
|
||||
'{pred}, {true_fn}, {false_fn}, {get_args}, {set_args}, {return_name_ids}, push_pop_names={push_pop_ids})'.format(
|
||||
pred=ast_to_source_code(pred),
|
||||
true_fn=true_func_source,
|
||||
false_fn=false_func_source,
|
||||
get_args=(
|
||||
get_args_func.name if not is_if_expr else 'lambda: None'
|
||||
), # TODO: better way to deal with this
|
||||
set_args=(
|
||||
set_args_func.name if not is_if_expr else 'lambda args: None'
|
||||
),
|
||||
return_name_ids=create_name_str(return_name_ids),
|
||||
push_pop_ids=create_name_str(push_pop_ids),
|
||||
)
|
||||
).body[0]
|
||||
|
||||
return convert_ifelse_layer
|
||||
@@ -0,0 +1,102 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import ast_to_source_code
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
cmpop_type_to_str = {
|
||||
gast.Eq: "==",
|
||||
gast.NotEq: "!=",
|
||||
gast.Lt: "<",
|
||||
gast.LtE: "<=",
|
||||
gast.Gt: ">",
|
||||
gast.GtE: ">=",
|
||||
gast.Is: "is",
|
||||
gast.IsNot: "is not",
|
||||
gast.In: "in",
|
||||
gast.NotIn: "not in",
|
||||
}
|
||||
|
||||
|
||||
def cmpop_node_to_str(node):
|
||||
return cmpop_type_to_str[type(node)]
|
||||
|
||||
|
||||
class LogicalTransformer(BaseTransformer):
|
||||
"""
|
||||
Transform python boolean op into Paddle logical op.
|
||||
|
||||
For example:
|
||||
a = x > 1 and y < 1
|
||||
|
||||
Transformed code:
|
||||
a = _jst.And(lambda:x>1, lambda:y<1)
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
return self.visit(self.root)
|
||||
|
||||
def visit_UnaryOp(self, node):
|
||||
self.generic_visit(node)
|
||||
if isinstance(node.op, gast.Not):
|
||||
arg = ast_to_source_code(node.operand)
|
||||
new_node_str = f"_jst.Not({arg})"
|
||||
# NOTE: gast.parse returns Module(body=[expr(value=...)])
|
||||
new_node = gast.parse(new_node_str).body[0].value
|
||||
return new_node
|
||||
return node
|
||||
|
||||
def visit_BoolOp(self, node):
|
||||
self.generic_visit(node)
|
||||
if isinstance(node.op, gast.And):
|
||||
new_node = self._create_bool_op_node(node.values, 'And')
|
||||
elif isinstance(node.op, gast.Or):
|
||||
new_node = self._create_bool_op_node(node.values, 'Or')
|
||||
else:
|
||||
raise TypeError(
|
||||
"Only supports and/or syntax in control flow if statement."
|
||||
)
|
||||
return new_node
|
||||
|
||||
def _create_bool_op_node(self, nodes, api_type):
|
||||
'''
|
||||
NOTE(liym27):
|
||||
The arguments of function convert_logical_XX should be callable so that they can be run
|
||||
according to the actual order. In `convert_logical_and(lambda:x>1, lambda:y<1)`, `lambda:y<1`
|
||||
must be run after `lambda:x>1`, If `x>1` is False, `y<1` should NOT be run.
|
||||
'''
|
||||
assert len(nodes) > 1, (
|
||||
f"The length of BoolOp should be at least 2, but received {len(nodes)}."
|
||||
)
|
||||
if len(nodes) > 2:
|
||||
# Creates logic_and/logic_or node recursively.
|
||||
pre_logic_node = self._create_bool_op_node(nodes[:2], api_type)
|
||||
if len(nodes[2:]) == 1:
|
||||
post_logic_node = nodes[2]
|
||||
else:
|
||||
post_logic_node = self._create_bool_op_node(nodes[2:], api_type)
|
||||
nodes = [pre_logic_node, post_logic_node]
|
||||
|
||||
args = [ast_to_source_code(child) for child in nodes]
|
||||
new_node_str = f"_jst.{api_type}(lambda:{args[0]}, lambda:{args[1]})"
|
||||
# NOTE: gast.parse return Module(body=[expr(...)])
|
||||
new_node = gast.parse(new_node_str).body[0].value
|
||||
return new_node
|
||||
@@ -0,0 +1,713 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
import copy
|
||||
from collections import defaultdict
|
||||
|
||||
from paddle.base import unique_name
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import (
|
||||
GetterSetterHelper,
|
||||
ast_to_source_code,
|
||||
)
|
||||
from .base import (
|
||||
BaseTransformer,
|
||||
ForLoopTuplePreTransformer,
|
||||
ForNodeVisitor,
|
||||
)
|
||||
from .utils import (
|
||||
ARGS_NAME,
|
||||
FOR_BODY_PREFIX,
|
||||
FOR_CONDITION_PREFIX,
|
||||
WHILE_BODY_PREFIX,
|
||||
WHILE_CONDITION_PREFIX,
|
||||
FunctionNameLivenessAnalysis,
|
||||
create_get_args_node,
|
||||
create_name_str,
|
||||
create_nonlocal_stmt_nodes,
|
||||
create_set_args_node,
|
||||
get_attribute_full_name,
|
||||
get_parent_mapping,
|
||||
)
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def create_while_nodes(
|
||||
condition_name,
|
||||
body_name,
|
||||
loop_var_names,
|
||||
push_pop_names,
|
||||
getter_name,
|
||||
setter_name,
|
||||
):
|
||||
"""
|
||||
Returns a list of gast.Node which represents the calling of Paddle
|
||||
controlflow while_loop.
|
||||
|
||||
Usually, the list just contain 1 statement such as:
|
||||
|
||||
[a, b, c] = paddle.jit.dy2static.convert_while_loop(
|
||||
condition_name, body_name, [a, b, c])
|
||||
|
||||
where a, b, c are in loop_var_names.
|
||||
|
||||
However, if loop_var_names contains property such as foo.x, we cannot
|
||||
assign the property as output of convert_while_loop because Python
|
||||
property is a kind of read-only attribute. To handle the case, we replace
|
||||
the attributes which are output of convert_while_loop with generated
|
||||
variables, then if we know the attribute is not read-only at runtime, we
|
||||
assign the attribute. The created statements are like:
|
||||
|
||||
[a, b, __attribute_variable_1] = paddle.jit.dy2static.convert_while_loop(
|
||||
condition_name, body_name, [a, b, foo.x])
|
||||
if not isinstance(getattr(type(foo), x, None), property): foo.x = __attribute_variable_1
|
||||
|
||||
The number of above statements is not only 1, that's why the return type is
|
||||
a list of gast.Node.
|
||||
"""
|
||||
# NOTE(liym27):
|
||||
# It's better to parse the source code into an AST node than to customize an AST node
|
||||
# including child nodes, because it is easy to mistake the ast node type when customizing the node.
|
||||
#
|
||||
# For example: loop_var_names = [a, b, foo.x], the type of `a` or `b` is gast.Name,
|
||||
# but the type of `foo.x` gast.Attribute.
|
||||
# We have to make loop_var_names and assign_loop_var_names with same order
|
||||
# set doesn't have order so we convert it to list
|
||||
loop_var_names = list(loop_var_names)
|
||||
assign_loop_var_names = []
|
||||
for name in loop_var_names:
|
||||
assign_loop_var_names.append(name)
|
||||
|
||||
while_func_name = "_jst.While"
|
||||
while_node_str = f"{while_func_name}({condition_name}, {body_name}, {getter_name}, {setter_name}, return_name_ids={create_name_str(loop_var_names)}, push_pop_names={create_name_str(push_pop_names)})"
|
||||
while_node = gast.parse(while_node_str).body[0]
|
||||
|
||||
ret = [while_node]
|
||||
return ret
|
||||
|
||||
|
||||
class NameVisitor(gast.NodeVisitor):
|
||||
'''
|
||||
Analysis name liveness for loop transformer
|
||||
'''
|
||||
|
||||
def __init__(self, root_node):
|
||||
# Set of gast.Name or gast.Attribute for variables
|
||||
self.current_seen_vars = set()
|
||||
|
||||
# List of gast.While/gast.For nodes
|
||||
self.current_loop = []
|
||||
|
||||
# List of nodes that have scope of variables.
|
||||
self.nodes_with_scope = []
|
||||
self.blacklist_names = {"False", "True", "None"}
|
||||
|
||||
# Mapping from gast.While/gast.For to variable nodes
|
||||
self.before_loop_body_vars = defaultdict(set)
|
||||
# NOTE: Use ordered list as dict value
|
||||
self.in_loop_vars = defaultdict(list)
|
||||
|
||||
# Mapping from gast.While/gast.For to variable nodes which is condition
|
||||
# of loop or being modified during the loop
|
||||
self.write_in_loop = defaultdict(set)
|
||||
self.condition_vars = defaultdict(set)
|
||||
self.in_condition = False
|
||||
|
||||
# Some names are types, we shouldn't record them as loop var names.
|
||||
self.type_vars = set()
|
||||
|
||||
self.to_parent_mapping = get_parent_mapping(root_node)
|
||||
|
||||
self.visit(root_node)
|
||||
|
||||
def get_loop_var_names(self, node):
|
||||
assert isinstance(node, (gast.While, gast.For)), (
|
||||
"Input node is not gast loop node"
|
||||
)
|
||||
loop_var_names = set()
|
||||
create_var_names = set()
|
||||
read_context = {type(gast.Load()), type(gast.AugLoad())}
|
||||
|
||||
in_loop_vars_list = self.in_loop_vars[node]
|
||||
|
||||
# get dict `var_name_to_ctxs`
|
||||
var_name_to_ctxs = defaultdict(list)
|
||||
for var_node in in_loop_vars_list:
|
||||
var_name_to_ctxs[self._var_node_to_name(var_node)].append(
|
||||
var_node.ctx
|
||||
)
|
||||
|
||||
in_loop_vars = set(in_loop_vars_list)
|
||||
in_loop_vars = self._remove_unnecessary_vars(in_loop_vars, node)
|
||||
in_loop_name_strs = self._var_nodes_to_names(in_loop_vars)
|
||||
|
||||
before_loop_body_vars = self.before_loop_body_vars[node]
|
||||
before_loop_body_vars = self._remove_unnecessary_vars(
|
||||
before_loop_body_vars, node
|
||||
)
|
||||
before_loop_name_strs = self._var_nodes_to_names(before_loop_body_vars)
|
||||
|
||||
after_loop_vars = (
|
||||
self.current_seen_vars - before_loop_body_vars - in_loop_vars
|
||||
)
|
||||
after_loop_vars = self._remove_unnecessary_vars(after_loop_vars, node)
|
||||
after_loop_name_strs = self._var_nodes_to_names(
|
||||
after_loop_vars, read_context
|
||||
)
|
||||
condition_vars = self.condition_vars[node]
|
||||
condition_names = self._var_nodes_to_names(condition_vars)
|
||||
|
||||
write_vars = self.write_in_loop[node]
|
||||
write_names = self._var_nodes_to_names(write_vars)
|
||||
|
||||
for name in in_loop_name_strs:
|
||||
if name in before_loop_name_strs:
|
||||
# If a variable is used in loop and created before loop
|
||||
|
||||
# If this var is a basic variable and read-only and not
|
||||
# condition var, it may not be loop_var else it should
|
||||
# be in loop_var as input
|
||||
if (name not in condition_names) and (name not in write_names):
|
||||
continue
|
||||
loop_var_names.add(name)
|
||||
|
||||
elif name in after_loop_name_strs:
|
||||
# If a variable is created in the while loop and read after
|
||||
# loop, it should be in loop_var and we should create it
|
||||
|
||||
# because name in after_loop_name must be initialized in loop
|
||||
# So it is write-only, we don't have to filter read-only basic
|
||||
# vars out
|
||||
loop_var_names.add(name)
|
||||
create_var_names.add(name)
|
||||
else:
|
||||
# If a variable is used and created in loop, but used before created,
|
||||
# it should be in loop_var and we should create it.
|
||||
|
||||
# For example, `var_a` should be in loop_var and we should create it.
|
||||
#
|
||||
# res = 0
|
||||
# for i, x in enumerate(x_array):
|
||||
# if i > 2:
|
||||
# x = func1(var_a)
|
||||
# var_a = func2(x)
|
||||
#
|
||||
|
||||
is_created = False
|
||||
for ctx in var_name_to_ctxs[name]:
|
||||
if isinstance(ctx, gast.Store):
|
||||
is_created = True
|
||||
|
||||
if (
|
||||
isinstance(var_name_to_ctxs[name][0], gast.Load)
|
||||
and is_created
|
||||
):
|
||||
loop_var_names.add(name)
|
||||
create_var_names.add(name)
|
||||
|
||||
return loop_var_names, create_var_names
|
||||
|
||||
def visit_Name(self, node):
|
||||
if self._is_call_func_name_node(node):
|
||||
self.generic_visit(node)
|
||||
return
|
||||
if node.id in self.blacklist_names:
|
||||
self.generic_visit(node)
|
||||
return
|
||||
|
||||
self.current_seen_vars.add(node)
|
||||
write_context = {
|
||||
type(gast.Store()),
|
||||
type(gast.AugStore()),
|
||||
type(gast.Del()),
|
||||
}
|
||||
|
||||
for loop_node in self.current_loop:
|
||||
self.in_loop_vars[loop_node].append(node)
|
||||
if type(node.ctx) in write_context:
|
||||
self.write_in_loop[loop_node].add(node)
|
||||
if self.in_condition:
|
||||
self.condition_vars[loop_node].add(node)
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
self.nodes_with_scope.append(node)
|
||||
self.blacklist_names.add(node.name)
|
||||
|
||||
# The variables in the function are not visible to the outside scope.
|
||||
before_func_seen_vars = copy.copy(self.current_seen_vars)
|
||||
|
||||
self.generic_visit(node)
|
||||
self.nodes_with_scope.pop()
|
||||
# After exiting the scope of the node, variables in this scope
|
||||
# should be removed from self.current_seen_vars.
|
||||
if self.nodes_with_scope:
|
||||
self.current_seen_vars = before_func_seen_vars
|
||||
|
||||
def visit(self, node):
|
||||
method = 'visit_' + node.__class__.__name__
|
||||
visitor = getattr(self, method, self.generic_visit)
|
||||
ret = visitor(node)
|
||||
return ret
|
||||
|
||||
def visit_Attribute(self, node):
|
||||
if self._is_call_func_name_node(node):
|
||||
return
|
||||
attr_full_name = get_attribute_full_name(node)
|
||||
# Class variables are not allowed to appear in the arguments list
|
||||
# of defined function under class methods in Python.
|
||||
"""
|
||||
def class_func(self):
|
||||
def while_loop_body(self.x, y) # `self.x` is illegal.
|
||||
"""
|
||||
# TODO: If do change the variable with `self.var`, need a better
|
||||
# way to deal with this case.
|
||||
if attr_full_name.startswith("self."):
|
||||
return
|
||||
self.current_seen_vars.add(node)
|
||||
|
||||
for loop_node in self.current_loop:
|
||||
self.in_loop_vars[loop_node].append(node)
|
||||
|
||||
# sub-nodes are visited during get_attribute_full_name and we shouldn't
|
||||
# visit again
|
||||
|
||||
def visit_For(self, node):
|
||||
self.current_loop.append(node)
|
||||
self.in_condition = True
|
||||
self.visit(node.target)
|
||||
self.visit(node.iter)
|
||||
self.in_condition = False
|
||||
self.before_loop_body_vars[node] = copy.copy(self.current_seen_vars)
|
||||
self.generic_visit(node)
|
||||
self.current_loop.pop()
|
||||
|
||||
def visit_While(self, node):
|
||||
self.current_loop.append(node)
|
||||
self.in_condition = True
|
||||
self.visit(node.test)
|
||||
self.in_condition = False
|
||||
self.before_loop_body_vars[node] = copy.copy(self.current_seen_vars)
|
||||
self.generic_visit(node)
|
||||
self.current_loop.pop()
|
||||
|
||||
def visit_Call(self, node):
|
||||
# Store type var names such as "isinstance(x, some_type_names)" and
|
||||
# Remove them later
|
||||
if isinstance(node.func, gast.Name) and node.func.id == 'isinstance':
|
||||
type_node = node.args[1]
|
||||
if isinstance(type_node, gast.Tuple):
|
||||
for element in type_node.elts:
|
||||
self.type_vars.add(ast_to_source_code(element).strip())
|
||||
else:
|
||||
self.type_vars.add(ast_to_source_code(type_node).strip())
|
||||
self.generic_visit(node)
|
||||
|
||||
def _var_nodes_to_names(self, node_set, ctx_filter_set=None):
|
||||
ret = set()
|
||||
for node in node_set:
|
||||
if ctx_filter_set is None or type(node.ctx) in ctx_filter_set:
|
||||
ret.add(self._var_node_to_name(node))
|
||||
return ret
|
||||
|
||||
def _var_node_to_name(self, node):
|
||||
if isinstance(node, gast.Name):
|
||||
return node.id
|
||||
elif isinstance(node, gast.Attribute):
|
||||
return get_attribute_full_name(node)
|
||||
|
||||
def _is_call_func_name_node(self, node):
|
||||
parent_node = self._get_parent_node(node)
|
||||
if isinstance(parent_node, gast.Call) and parent_node.func == node:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _is_global_or_nonlocal(self, node):
|
||||
return False
|
||||
|
||||
def _is_ancestor_node(self, ancestor_node, node):
|
||||
parent_node = self._get_parent_node(node)
|
||||
|
||||
while parent_node is not None:
|
||||
if parent_node == ancestor_node:
|
||||
return True
|
||||
parent_node = self._get_parent_node(parent_node)
|
||||
return False
|
||||
|
||||
def _get_parent_node(self, node):
|
||||
return self.to_parent_mapping.get(node)
|
||||
|
||||
def _remove_unnecessary_vars(self, loop_vars, loop_node):
|
||||
"""
|
||||
Remove unnecessary vars from before_loop_vars, after_loop_vars or in_loop_vars about loop_node.
|
||||
1. Remove target vars of gast.For from before_loop_vars or after_loop_vars.
|
||||
2. Remove vars only in gast.comprehension.
|
||||
3. Remove vars that are type names, for example: "isinstance(x, var_type_name)"
|
||||
:param loop_vars: before_loop_vars, after_loop_vars or in_loop_vars of loop_node.
|
||||
:param loop_node: Current loop node.
|
||||
"""
|
||||
|
||||
def filter_name_nodes_from(root_node, target_var_names):
|
||||
"""
|
||||
Filter children with gast.Name type from node.(inclusivly)
|
||||
"""
|
||||
name_nodes = set()
|
||||
if isinstance(root_node, gast.Name):
|
||||
if node.id in target_var_names:
|
||||
name_nodes.add(root_node)
|
||||
for child_node in gast.walk(root_node):
|
||||
if isinstance(child_node, gast.Name):
|
||||
if child_node.id in target_var_names:
|
||||
name_nodes.add(child_node)
|
||||
|
||||
return name_nodes
|
||||
|
||||
vars_of_list_generator = set()
|
||||
target_vars_of_for_node = set()
|
||||
|
||||
for name_node in loop_vars:
|
||||
if not isinstance(name_node, gast.Name):
|
||||
continue
|
||||
|
||||
parent_node = self._get_parent_node(name_node)
|
||||
|
||||
# NOTE: gast.For.target or gast.comprehension.target can be gast.Tuple.
|
||||
# For examples:
|
||||
# 1) `for i, j in enumerate(x)` has two target vars: i and j
|
||||
# 2) `[x for x,y in array]` has two target vars: x and y
|
||||
if isinstance(parent_node, gast.Tuple):
|
||||
parent_node = self._get_parent_node(parent_node)
|
||||
|
||||
# 1. Get vars only in gast.comprehension.
|
||||
# For examples:
|
||||
# 1) [x for x,y in array] -> x, x, y
|
||||
# 2) [f(x) for x in array] -> x
|
||||
# 3) [func(x, y) for x in array] -> x, x
|
||||
if isinstance(parent_node, gast.comprehension):
|
||||
# 1.1 target vars in list/set comprehensions
|
||||
target_node = parent_node.target
|
||||
if isinstance(target_node, gast.Tuple):
|
||||
target_vars = target_node.elts
|
||||
else:
|
||||
target_vars = [target_node]
|
||||
|
||||
vars_of_list_generator = vars_of_list_generator | set(
|
||||
target_vars
|
||||
)
|
||||
|
||||
# 1.2 vars from target vars used in elt_node
|
||||
target_var_names = {var.id for var in target_vars}
|
||||
comp_node = self._get_parent_node(parent_node)
|
||||
elt_nodes = []
|
||||
if isinstance(comp_node, gast.ListComp):
|
||||
elt_nodes.append(comp_node.elt)
|
||||
elif isinstance(comp_node, gast.DictComp):
|
||||
elt_nodes.extend([comp_node.key, comp_node.value])
|
||||
|
||||
for node in elt_nodes:
|
||||
vars_of_list_generator |= filter_name_nodes_from(
|
||||
node, target_var_names
|
||||
)
|
||||
|
||||
# 2. Get target vars or vars from target vars used in for-loop but the for-loop is
|
||||
# 1) not the "loop_node" itself
|
||||
# 2) not the ancestor of the "loop_node"
|
||||
#
|
||||
# For examples:
|
||||
# for k in range(x): # if it's this "loop_node", i or j both should be target vars.
|
||||
# # do something
|
||||
#
|
||||
# for i in range(a): # if it's this "loop_node", k or j should be in target vars but i should not.
|
||||
# for j in range(a): # if it's this "loop_node", k should be in target_vars but i or j should not.
|
||||
# x = i+j
|
||||
elif isinstance(parent_node, gast.For):
|
||||
if parent_node is loop_node:
|
||||
continue
|
||||
if self._is_ancestor_node(parent_node, loop_node):
|
||||
continue
|
||||
# 2.1 target vars in gast.For node.
|
||||
target_node = parent_node.target
|
||||
if isinstance(target_node, gast.Tuple):
|
||||
target_vars = target_node.elts
|
||||
else:
|
||||
target_vars = [target_node]
|
||||
|
||||
target_vars_of_for_node = target_vars_of_for_node | set(
|
||||
target_vars
|
||||
)
|
||||
|
||||
# 2.2 vars from target vars used in for-loop
|
||||
target_vars_name_strs = {var.id for var in target_vars_of_for_node}
|
||||
for var in loop_vars:
|
||||
if not isinstance(var, gast.Name):
|
||||
continue
|
||||
if (
|
||||
var.id in target_vars_name_strs
|
||||
and var not in self.condition_vars[loop_node]
|
||||
):
|
||||
target_vars_of_for_node.add(var)
|
||||
|
||||
removed_vars = target_vars_of_for_node | vars_of_list_generator
|
||||
|
||||
# 3. Remove var type names which are stored in self.type_vars
|
||||
for var in loop_vars:
|
||||
if ast_to_source_code(var).strip() in self.type_vars:
|
||||
removed_vars.add(var)
|
||||
|
||||
return loop_vars - removed_vars
|
||||
|
||||
|
||||
class LoopTransformer(BaseTransformer):
|
||||
"""
|
||||
This class transforms python while/for statement into Static Graph Ast
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
FunctionNameLivenessAnalysis(self.root)
|
||||
|
||||
def transform(self):
|
||||
ForLoopTuplePreTransformer(self.root).transform()
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_While(self, node):
|
||||
self.generic_visit(node)
|
||||
new_stmts = self.get_while_stmt_nodes(node)
|
||||
return new_stmts
|
||||
|
||||
def visit_For(self, node):
|
||||
self.generic_visit(node)
|
||||
new_stmts = self.get_for_stmt_nodes(node)
|
||||
return new_stmts
|
||||
|
||||
def replace_stmt_list(self, body_list):
|
||||
if not isinstance(body_list, list):
|
||||
return
|
||||
|
||||
i = 0
|
||||
while i < len(body_list):
|
||||
if isinstance(body_list[i], gast.While):
|
||||
new_stmts = self.get_while_stmt_nodes(body_list[i])
|
||||
body_list[i : i + 1] = new_stmts
|
||||
i += len(new_stmts)
|
||||
elif isinstance(body_list[i], gast.For):
|
||||
new_stmts = self.get_for_stmt_nodes(body_list[i])
|
||||
body_list[i : i + 1] = new_stmts
|
||||
i += len(new_stmts)
|
||||
else:
|
||||
i += 1
|
||||
|
||||
def get_for_stmt_nodes(self, node):
|
||||
# TODO: consider for - else in python
|
||||
|
||||
# 1. get key statements for different cases
|
||||
# NOTE 1: three key statements:
|
||||
# 1). init_stmts: list[node], prepare nodes of for loop, may not only one
|
||||
# 2). cond_stmt: node, condition node to judge whether continue loop
|
||||
# 3). body_stmts: list[node], updated loop body, sometimes we should change
|
||||
# the original statement in body, not just append new statement
|
||||
#
|
||||
# NOTE 2: The following `for` statements will be transformed to `while` statements:
|
||||
# 1). for x in range(*)
|
||||
# 2). for x in iter_var
|
||||
# 3). for i, x in enumerate(*)
|
||||
|
||||
current_for_node_parser = ForNodeVisitor(node)
|
||||
stmts_tuple = current_for_node_parser.parse()
|
||||
if stmts_tuple is None:
|
||||
return [node]
|
||||
init_stmts, cond_stmt, body_stmts = stmts_tuple
|
||||
# 2. get original loop vars
|
||||
loop_var_names, create_var_names = (
|
||||
node.pd_scope.modified_vars(),
|
||||
node.pd_scope.created_vars(),
|
||||
)
|
||||
push_pop_names = list(node.pd_scope.variadic_length_vars())
|
||||
# TODO: Remove the bunch of code? We have the unique format `for A in B:`
|
||||
# NOTE: in 'for x in var' or 'for i, x in enumerate(var)' cases,
|
||||
# we need append new loop var & remove useless loop var
|
||||
# 1. for x in var -> x is no need
|
||||
# 2. for i, x in enumerate(var) -> x is no need
|
||||
if current_for_node_parser.is_for_iter():
|
||||
iter_var_name = current_for_node_parser.iter_var_name
|
||||
iter_idx_name = current_for_node_parser.iter_idx_name
|
||||
loop_var_names.add(iter_idx_name)
|
||||
if current_for_node_parser.enum_idx_name is not None:
|
||||
loop_var_names.add(current_for_node_parser.enum_idx_name)
|
||||
|
||||
# 3. prepare result statement list
|
||||
new_stmts = []
|
||||
# Python can create variable in loop and use it out of loop, E.g.
|
||||
#
|
||||
# for x in range(10):
|
||||
# y += x
|
||||
# print(x) # x = 10
|
||||
#
|
||||
# We don't need to create static variable for them, because
|
||||
# we do this in CreateUndefinedVarTransformer
|
||||
|
||||
# create non-local statement for body and cond.
|
||||
nonlocal_names = list(loop_var_names | create_var_names)
|
||||
nonlocal_names.sort()
|
||||
# TODO(dev): Need a better way to deal this.
|
||||
if ARGS_NAME in nonlocal_names:
|
||||
nonlocal_names.remove(ARGS_NAME)
|
||||
|
||||
nonlocal_stmt_node = create_nonlocal_stmt_nodes(nonlocal_names)
|
||||
|
||||
# 4. append init statements
|
||||
new_stmts.extend(init_stmts)
|
||||
|
||||
# 5. create & append condition function node
|
||||
condition_func_node = gast.FunctionDef(
|
||||
name=unique_name.generate(FOR_CONDITION_PREFIX),
|
||||
args=gast.arguments(
|
||||
args=[],
|
||||
posonlyargs=[],
|
||||
vararg=None,
|
||||
kwonlyargs=[],
|
||||
kw_defaults=None,
|
||||
kwarg=None,
|
||||
defaults=[],
|
||||
),
|
||||
body=[*nonlocal_stmt_node, gast.Return(value=cond_stmt)],
|
||||
decorator_list=[],
|
||||
returns=None,
|
||||
type_comment=None,
|
||||
type_params=[],
|
||||
)
|
||||
new_stmts.append(condition_func_node)
|
||||
|
||||
# 6. create & append loop body function node
|
||||
# append return values for loop body
|
||||
body_func_node = gast.FunctionDef(
|
||||
name=unique_name.generate(FOR_BODY_PREFIX),
|
||||
args=gast.arguments(
|
||||
args=[],
|
||||
posonlyargs=[],
|
||||
vararg=None,
|
||||
kwonlyargs=[],
|
||||
kw_defaults=None,
|
||||
kwarg=None,
|
||||
defaults=[],
|
||||
),
|
||||
body=nonlocal_stmt_node + body_stmts,
|
||||
decorator_list=[],
|
||||
returns=None,
|
||||
type_comment=None,
|
||||
type_params=[],
|
||||
)
|
||||
new_stmts.append(body_func_node)
|
||||
|
||||
helper = GetterSetterHelper(None, None, nonlocal_names, push_pop_names)
|
||||
get_args_node = create_get_args_node(helper.union())
|
||||
set_args_node = create_set_args_node(helper.union())
|
||||
# 7. create & append while loop node
|
||||
while_loop_nodes = create_while_nodes(
|
||||
condition_func_node.name,
|
||||
body_func_node.name,
|
||||
nonlocal_names,
|
||||
push_pop_names,
|
||||
get_args_node.name,
|
||||
set_args_node.name,
|
||||
)
|
||||
new_stmts.extend([get_args_node, set_args_node])
|
||||
new_stmts.extend(while_loop_nodes)
|
||||
|
||||
return new_stmts
|
||||
|
||||
def get_while_stmt_nodes(self, node):
|
||||
loop_var_names, create_var_names = (
|
||||
node.pd_scope.modified_vars(),
|
||||
node.pd_scope.created_vars(),
|
||||
)
|
||||
push_pop_names = list(node.pd_scope.variadic_length_vars())
|
||||
new_stmts = []
|
||||
|
||||
# create non-local statement for body and cond.
|
||||
nonlocal_names = list(loop_var_names | create_var_names)
|
||||
nonlocal_names.sort()
|
||||
# TODO(dev): Need a better way to deal this.
|
||||
if ARGS_NAME in nonlocal_names:
|
||||
nonlocal_names.remove(ARGS_NAME)
|
||||
|
||||
nonlocal_stmt_node = create_nonlocal_stmt_nodes(nonlocal_names)
|
||||
|
||||
# Python can create variable in loop and use it out of loop, E.g.
|
||||
#
|
||||
# while x < 10:
|
||||
# x += 1
|
||||
# y = x
|
||||
# z = y
|
||||
#
|
||||
# We don't need to create static variable for those variables, because
|
||||
# we do this in CreateUndefinedVarTransformer
|
||||
|
||||
condition_func_node = gast.FunctionDef(
|
||||
name=unique_name.generate(WHILE_CONDITION_PREFIX),
|
||||
args=gast.arguments(
|
||||
args=[],
|
||||
posonlyargs=[],
|
||||
vararg=None,
|
||||
kwonlyargs=[],
|
||||
kw_defaults=None,
|
||||
kwarg=None,
|
||||
defaults=[],
|
||||
),
|
||||
body=[*nonlocal_stmt_node, gast.Return(value=node.test)],
|
||||
decorator_list=[],
|
||||
returns=None,
|
||||
type_comment=None,
|
||||
type_params=[],
|
||||
)
|
||||
|
||||
new_stmts.append(condition_func_node)
|
||||
|
||||
new_body = node.body
|
||||
body_func_node = gast.FunctionDef(
|
||||
name=unique_name.generate(WHILE_BODY_PREFIX),
|
||||
args=gast.arguments(
|
||||
args=[],
|
||||
posonlyargs=[],
|
||||
vararg=None,
|
||||
kwonlyargs=[],
|
||||
kw_defaults=None,
|
||||
kwarg=None,
|
||||
defaults=[],
|
||||
),
|
||||
body=nonlocal_stmt_node + new_body,
|
||||
decorator_list=[],
|
||||
returns=None,
|
||||
type_comment=None,
|
||||
type_params=[],
|
||||
)
|
||||
new_stmts.append(body_func_node)
|
||||
|
||||
helper = GetterSetterHelper(None, None, nonlocal_names, push_pop_names)
|
||||
get_args_node = create_get_args_node(helper.union())
|
||||
set_args_node = create_set_args_node(helper.union())
|
||||
|
||||
while_loop_nodes = create_while_nodes(
|
||||
condition_func_node.name,
|
||||
body_func_node.name,
|
||||
nonlocal_names,
|
||||
push_pop_names,
|
||||
get_args_node.name,
|
||||
set_args_node.name,
|
||||
)
|
||||
new_stmts.extend([get_args_node, set_args_node])
|
||||
new_stmts.extend(while_loop_nodes)
|
||||
return new_stmts
|
||||
@@ -0,0 +1,130 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import ast_to_source_code
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class NameloadJstTransformer(BaseTransformer):
|
||||
"""
|
||||
change name and attribute load to __jst.Ld(name) pattern.
|
||||
for example:
|
||||
a.dtype --> __jst.Ld(__jst.Ld(a).dtype)
|
||||
|
||||
In paddle science and deepxde, we have to support changing tensor into variable
|
||||
in arbitrary occasion such as global tensor.
|
||||
|
||||
NOTE: we only deal with ctx=Load() case.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
return self.root
|
||||
|
||||
def _surround_with_ld(self, node):
|
||||
node = (
|
||||
gast.parse(f"_jst.Ld({ast_to_source_code(node).strip()})")
|
||||
.body[0]
|
||||
.value
|
||||
)
|
||||
return node
|
||||
|
||||
def visit_Call(self, node):
|
||||
"""
|
||||
Can't convert name of function call, because this will affect CallTransformer.
|
||||
"""
|
||||
node.args = [self.visit(arg) for arg in node.args]
|
||||
for keyword in node.keywords:
|
||||
keyword.value = self.visit(keyword.value)
|
||||
node.func = self.visit(node.func)
|
||||
return node
|
||||
|
||||
def create_visit_with_convert_load(self, node_type, skip_fn=None):
|
||||
def visit(node):
|
||||
assert isinstance(node, node_type)
|
||||
if skip_fn and skip_fn(node):
|
||||
return node
|
||||
self.generic_visit(node)
|
||||
if isinstance(node.ctx, gast.Load):
|
||||
node = self._surround_with_ld(node)
|
||||
return node
|
||||
|
||||
return visit
|
||||
|
||||
def visit_Attribute(self, node):
|
||||
def skip_fn(node):
|
||||
if isinstance(node.value, gast.Name) and node.value.id == "_jst":
|
||||
return True
|
||||
return False
|
||||
|
||||
return self.create_visit_with_convert_load(gast.Attribute, skip_fn)(
|
||||
node
|
||||
)
|
||||
|
||||
def visit_Subscript(self, node):
|
||||
return self.create_visit_with_convert_load(gast.Subscript)(node)
|
||||
|
||||
def visit_Name(self, node):
|
||||
return self.create_visit_with_convert_load(gast.Name)(node)
|
||||
|
||||
|
||||
class AttributeJstTransformer(BaseTransformer):
|
||||
"""
|
||||
change some special attribute into __jst.XXX(obj, "attr_name") format.
|
||||
for example:
|
||||
a.size --> __jst.attr(a, "size")
|
||||
|
||||
because `size` have different behavior when in dygraph / static graph mode
|
||||
NOTE: we only deal with ctx=Load() case.
|
||||
"""
|
||||
|
||||
def __init__(self, node):
|
||||
assert isinstance(node, gast.AST), (
|
||||
"Input non-gast.AST node for the initialization of ToTensorTransformer."
|
||||
)
|
||||
self.interested_name = {
|
||||
'size',
|
||||
}
|
||||
self.root = node
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
return self.root
|
||||
|
||||
def visit_Attribute(self, node):
|
||||
assert isinstance(node, gast.Attribute)
|
||||
assert isinstance(node.attr, str)
|
||||
if (
|
||||
isinstance(node.ctx, gast.Load)
|
||||
and node.attr in self.interested_name
|
||||
):
|
||||
attr = node.attr
|
||||
value = node.value
|
||||
node = (
|
||||
gast.parse(
|
||||
f'_jst.Attr({ast_to_source_code(value).strip()}, "{attr}")'
|
||||
)
|
||||
.body[0]
|
||||
.value
|
||||
)
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
@@ -0,0 +1,415 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.base import unique_name
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import (
|
||||
ORIGIN_INFO,
|
||||
Dygraph2StaticException,
|
||||
ast_to_source_code,
|
||||
)
|
||||
from .base import BaseTransformer
|
||||
from .break_continue_transformer import ForToWhileTransformer
|
||||
from .utils import create_bool_node, index_in_list
|
||||
|
||||
__all__ = []
|
||||
|
||||
# Constant for the name of the variable which stores the boolean state that we
|
||||
# should return
|
||||
RETURN_PREFIX = '__return'
|
||||
|
||||
# Constant for the name of the variable which stores the final return value
|
||||
RETURN_VALUE_PREFIX = '__return_value'
|
||||
|
||||
# Constant for the name of variables to initialize the __return_value
|
||||
RETURN_VALUE_INIT_NAME = '__return_value_init'
|
||||
|
||||
# Constant magic number representing returning no value. This constant amis to
|
||||
# support returning various lengths of variables. Static graph must have fixed
|
||||
# size of fetched output while dygraph can have flexible lengths of output, to
|
||||
# solve it in dy2stat, we put float64 value with this magic number at Static
|
||||
# graph as a place holder to indicate the returning placeholder means no value
|
||||
# should return.
|
||||
|
||||
|
||||
def get_return_size(return_node):
|
||||
assert isinstance(return_node, gast.Return), "Input is not gast.Return node"
|
||||
return_length = 0
|
||||
if return_node.value is not None:
|
||||
if isinstance(return_node.value, gast.Tuple):
|
||||
return_length = len(return_node.value.elts)
|
||||
else:
|
||||
return_length = 1
|
||||
return return_length
|
||||
|
||||
|
||||
class ReplaceReturnNoneTransformer(BaseTransformer):
|
||||
"""
|
||||
Replace 'return None' to 'return' because 'None' cannot be a valid input
|
||||
in control flow. In ReturnTransformer single 'Return' will be appended no
|
||||
value placeholder
|
||||
"""
|
||||
|
||||
def __init__(self, root_node):
|
||||
self.root = root_node
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_Return(self, node):
|
||||
if isinstance(node.value, gast.Name) and node.value.id == 'None':
|
||||
node.value = None
|
||||
return node
|
||||
if isinstance(node.value, gast.Constant) and node.value.value is None:
|
||||
node.value = None
|
||||
return node
|
||||
return node
|
||||
|
||||
|
||||
class ReturnAnalysisVisitor(gast.NodeVisitor):
|
||||
"""
|
||||
Visits gast Tree and analyze the information about 'return'.
|
||||
"""
|
||||
|
||||
def __init__(self, root_node):
|
||||
self.root = root_node
|
||||
assert isinstance(self.root, gast.FunctionDef), (
|
||||
"Input is not gast.FunctionDef node"
|
||||
)
|
||||
|
||||
# the number of return statements
|
||||
self.count_return = 0
|
||||
|
||||
# maximum number of variables
|
||||
self.max_return_length = 0
|
||||
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
"""
|
||||
don't analysis closure, just analyze current func def level.
|
||||
"""
|
||||
if node == self.root:
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Return(self, node):
|
||||
self.count_return += 1
|
||||
|
||||
return_length = get_return_size(node)
|
||||
self.max_return_length = max(self.max_return_length, return_length)
|
||||
|
||||
self.generic_visit(node)
|
||||
|
||||
def get_func_return_count(self):
|
||||
return self.count_return
|
||||
|
||||
def get_func_max_return_length(self):
|
||||
return self.max_return_length
|
||||
|
||||
|
||||
class ReturnTransformer(BaseTransformer):
|
||||
"""
|
||||
Transforms return statements into equivalent python statements containing
|
||||
only one return statement at last. The basics idea is using a return value
|
||||
variable to store the early return statements and boolean states with
|
||||
if-else to skip the statements after the return.
|
||||
|
||||
Go through all the function definition and call SingleReturnTransformer for each function.
|
||||
SingleReturnTransformer don't care the nested function def.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
pre_transformer = ReplaceReturnNoneTransformer(self.root)
|
||||
pre_transformer.transform()
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
node = self.generic_visit(node)
|
||||
node = SingleReturnTransformer(node).transform()
|
||||
return node
|
||||
|
||||
|
||||
class SingleReturnTransformer(BaseTransformer):
|
||||
"""
|
||||
This function only apply to single function. don't care the nested function_def
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
assert isinstance(self.root, gast.FunctionDef), (
|
||||
"Input is not gast.FunctionDef node"
|
||||
)
|
||||
|
||||
self.ancestor_nodes = []
|
||||
|
||||
# The name of return placeholder
|
||||
self.return_value_name = None
|
||||
|
||||
# Every return stmt corresponds to a bool value variable, and return name is the name of the boolean variable
|
||||
self.return_name = []
|
||||
|
||||
self.pre_analysis = None
|
||||
|
||||
def assert_parent_is_not_while(self, parent_node_of_return):
|
||||
if isinstance(parent_node_of_return, (gast.While, gast.For)):
|
||||
raise Dygraph2StaticException(
|
||||
"Found return statement in While or For body and loop "
|
||||
"is meaningless, please check you code and remove return in while/for."
|
||||
)
|
||||
|
||||
def generic_visit(self, node):
|
||||
# Because we change ancestor nodes during visit_Return, not current
|
||||
# node, original generic_visit of NodeTransformer will visit node
|
||||
# which may be deleted. To prevent that node being added into
|
||||
# transformed AST, We self-write a generic_visit and visit
|
||||
for field, value in gast.iter_fields(node):
|
||||
if isinstance(value, list):
|
||||
for item in value:
|
||||
if isinstance(item, gast.AST):
|
||||
self.visit(item)
|
||||
elif isinstance(value, gast.AST):
|
||||
self.visit(value)
|
||||
|
||||
def visit(self, node):
|
||||
"""
|
||||
Self-defined visit for appending ancestor
|
||||
"""
|
||||
self.ancestor_nodes.append(node)
|
||||
ret = super().visit(node)
|
||||
self.ancestor_nodes.pop()
|
||||
return ret
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
"""
|
||||
don't analysis closure, just analyze current func def level.
|
||||
"""
|
||||
if node == self.root:
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def append_assign_to_return_node(
|
||||
self, value, parent_node_of_return, return_name, assign_nodes
|
||||
):
|
||||
self.assert_parent_is_not_while(parent_node_of_return)
|
||||
assert value in [True, False], "value must be True or False."
|
||||
if isinstance(parent_node_of_return, gast.If):
|
||||
# Prepend control flow boolean nodes such as '__return@1 = True'
|
||||
node_str = f"{return_name} = _jst.create_bool_as_type({ast_to_source_code(parent_node_of_return.test).strip()}, {value})"
|
||||
|
||||
assign_node = gast.parse(node_str).body[0]
|
||||
assign_nodes.append(assign_node)
|
||||
|
||||
def transform(self):
|
||||
node = self.root
|
||||
self.pre_analysis = ReturnAnalysisVisitor(node)
|
||||
max_return_length = self.pre_analysis.get_func_max_return_length()
|
||||
while self.pre_analysis.get_func_return_count() > 0:
|
||||
# every visit will decrease the number of returns.
|
||||
# so we need a while.
|
||||
self.visit(node)
|
||||
self.pre_analysis = ReturnAnalysisVisitor(node)
|
||||
|
||||
if max_return_length == 0:
|
||||
return node
|
||||
|
||||
# Prepend initialization of final return and append final return statement
|
||||
return_flag_names = self.return_name
|
||||
value_name = self.return_value_name
|
||||
if value_name is not None:
|
||||
node.body.append(
|
||||
gast.Return(
|
||||
value=gast.Name(
|
||||
id=value_name,
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
)
|
||||
)
|
||||
)
|
||||
assign_return_value_node = gast.Assign(
|
||||
targets=[
|
||||
gast.Name(
|
||||
id=value_name,
|
||||
ctx=gast.Store(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
)
|
||||
],
|
||||
value=gast.Constant(kind=None, value=None),
|
||||
type_comment=None,
|
||||
)
|
||||
node.body.insert(0, assign_return_value_node)
|
||||
|
||||
for return_flag_name in return_flag_names:
|
||||
assign_return_flag_node = create_bool_node(return_flag_name, False)
|
||||
node.body.insert(0, assign_return_flag_node)
|
||||
|
||||
# Prepend no value placeholders
|
||||
return node
|
||||
|
||||
def visit_Return(self, node):
|
||||
return_name = unique_name.generate(RETURN_PREFIX)
|
||||
self.return_name.append(return_name)
|
||||
max_return_length = self.pre_analysis.get_func_max_return_length()
|
||||
parent_node_of_return = self.ancestor_nodes[-2]
|
||||
|
||||
for ancestor_index in reversed(range(len(self.ancestor_nodes) - 1)):
|
||||
ancestor = self.ancestor_nodes[ancestor_index]
|
||||
cur_node = self.ancestor_nodes[ancestor_index + 1]
|
||||
|
||||
def _deal_branches(branch_name):
|
||||
if hasattr(ancestor, branch_name):
|
||||
branch_node = getattr(ancestor, branch_name)
|
||||
if index_in_list(branch_node, cur_node) != -1:
|
||||
if cur_node == node:
|
||||
self._replace_return_in_stmt_list(
|
||||
branch_node,
|
||||
cur_node,
|
||||
return_name,
|
||||
max_return_length,
|
||||
parent_node_of_return,
|
||||
)
|
||||
self._replace_after_node_to_if_in_stmt_list(
|
||||
branch_node,
|
||||
cur_node,
|
||||
return_name,
|
||||
parent_node_of_return,
|
||||
)
|
||||
|
||||
_deal_branches("body")
|
||||
_deal_branches("orelse")
|
||||
# If return node in while loop, add `not return_name` in gast.While.test
|
||||
if isinstance(ancestor, gast.While):
|
||||
cond_var_node = gast.UnaryOp(
|
||||
op=gast.Not(),
|
||||
operand=gast.Name(
|
||||
id=return_name,
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
)
|
||||
ancestor.test = gast.BoolOp(
|
||||
op=gast.And(), values=[ancestor.test, cond_var_node]
|
||||
)
|
||||
continue
|
||||
|
||||
# If return node in for loop, add `not return_name` in gast.While.test
|
||||
if isinstance(ancestor, gast.For):
|
||||
cond_var_node = gast.UnaryOp(
|
||||
op=gast.Not(),
|
||||
operand=gast.Name(
|
||||
id=return_name,
|
||||
ctx=gast.Load(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
)
|
||||
parent_node = self.ancestor_nodes[ancestor_index - 1]
|
||||
for_to_while = ForToWhileTransformer(
|
||||
parent_node, ancestor, cond_var_node
|
||||
)
|
||||
new_stmts = for_to_while.transform()
|
||||
while_node = new_stmts[-1]
|
||||
self.ancestor_nodes[ancestor_index] = while_node
|
||||
|
||||
if ancestor == self.root:
|
||||
break
|
||||
# return_node is replaced so we shouldn't return here
|
||||
|
||||
def _replace_return_in_stmt_list(
|
||||
self,
|
||||
stmt_list,
|
||||
return_node,
|
||||
return_name,
|
||||
max_return_length,
|
||||
parent_node_of_return,
|
||||
):
|
||||
assert max_return_length >= 0, "Input illegal max_return_length"
|
||||
i = index_in_list(stmt_list, return_node)
|
||||
if i == -1:
|
||||
return False
|
||||
|
||||
assign_nodes = []
|
||||
self.append_assign_to_return_node(
|
||||
True, parent_node_of_return, return_name, assign_nodes
|
||||
)
|
||||
|
||||
return_length = get_return_size(return_node)
|
||||
# In this case we should NOT append RETURN_NO_VALUE placeholder
|
||||
if return_node.value is not None:
|
||||
if self.return_value_name is None:
|
||||
self.return_value_name = unique_name.generate(
|
||||
RETURN_VALUE_PREFIX
|
||||
)
|
||||
|
||||
assign_nodes.append(
|
||||
gast.Assign(
|
||||
targets=[
|
||||
gast.Name(
|
||||
id=self.return_value_name,
|
||||
ctx=gast.Store(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
)
|
||||
],
|
||||
value=return_node.value,
|
||||
type_comment=None,
|
||||
)
|
||||
)
|
||||
return_origin_info = getattr(return_node, ORIGIN_INFO, None)
|
||||
setattr(assign_nodes[-1], ORIGIN_INFO, return_origin_info)
|
||||
|
||||
# If there is a return in the body or else of if, the remaining statements
|
||||
# will not be executed, so they can be properly replaced.
|
||||
stmt_list[i:] = assign_nodes
|
||||
return True
|
||||
|
||||
def _replace_after_node_to_if_in_stmt_list(
|
||||
self, stmt_list, node, return_name, parent_node_of_return
|
||||
):
|
||||
i = index_in_list(stmt_list, node)
|
||||
if i < 0 or i >= len(stmt_list):
|
||||
return False
|
||||
if i == len(stmt_list) - 1:
|
||||
# No need to add, we consider this as added successfully
|
||||
return True
|
||||
|
||||
if_stmt = gast.If(
|
||||
test=gast.UnaryOp(
|
||||
op=gast.Not(),
|
||||
operand=gast.Name(
|
||||
id=return_name,
|
||||
ctx=gast.Store(),
|
||||
annotation=None,
|
||||
type_comment=None,
|
||||
),
|
||||
),
|
||||
body=stmt_list[i + 1 :],
|
||||
orelse=[],
|
||||
)
|
||||
|
||||
stmt_list[i + 1 :] = [if_stmt]
|
||||
|
||||
# Here assume that the parent node of return is gast.If
|
||||
assign_nodes = []
|
||||
self.append_assign_to_return_node(
|
||||
False, parent_node_of_return, return_name, assign_nodes
|
||||
)
|
||||
stmt_list[i:i] = assign_nodes
|
||||
return True
|
||||
@@ -0,0 +1,60 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class SuperTransformer(BaseTransformer):
|
||||
"""
|
||||
This class transforms super() into super(__class__, <first argument>).
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
self.first_arg = None
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
if self.first_arg is not None:
|
||||
return self.generic_visit(node)
|
||||
|
||||
positional_args = node.args.posonlyargs + node.args.args
|
||||
if not positional_args:
|
||||
return self.generic_visit(node)
|
||||
|
||||
self.first_arg = positional_args[0].id
|
||||
|
||||
return self.generic_visit(node)
|
||||
|
||||
def visit_Call(self, node):
|
||||
# super() -> _jst.WrapSuper(super)(x.__class__, x)
|
||||
self.generic_visit(node)
|
||||
if self.first_arg is None:
|
||||
return node
|
||||
if not isinstance(node.func, gast.Name):
|
||||
return node
|
||||
if node.func.id != "super":
|
||||
return node
|
||||
if node.args:
|
||||
return node
|
||||
|
||||
new_fn_call_str = f"_jst.WrapSuper(super)({self.first_arg}.__class__, {self.first_arg})"
|
||||
new_fn_call_ast = gast.parse(new_fn_call_str).body[0]
|
||||
return new_fn_call_ast
|
||||
@@ -0,0 +1,45 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import ast_to_source_code
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class TensorShapeTransformer(BaseTransformer):
|
||||
"""
|
||||
This class transforms variable.shape into Static Graph Ast.
|
||||
All 'xxx.shape' will be converted int '_jst.Shape(x)'.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_Attribute(self, node):
|
||||
self.generic_visit(node)
|
||||
if node.attr == 'shape':
|
||||
args = ast_to_source_code(node.value).strip()
|
||||
# NOTE(dev): we can deal with paddle.shape in this case, but it's
|
||||
# not pretty to modify into 'convert_shape(paddle)(x)[0]'.
|
||||
if args != 'paddle':
|
||||
convert_shape_func = f"_jst.Shape({args})"
|
||||
shape_node = gast.parse(convert_shape_func).body[0].value
|
||||
return shape_node
|
||||
return node
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import ast_to_source_code
|
||||
from .base import BaseTransformer
|
||||
|
||||
|
||||
def get_loads(node: gast.AST):
|
||||
for child in gast.walk(node):
|
||||
if isinstance(
|
||||
child, (gast.Name, gast.Attribute, gast.Subscript)
|
||||
) and isinstance(child.ctx, gast.Load):
|
||||
yield child
|
||||
|
||||
|
||||
class RegisterHookTransformer(BaseTransformer):
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
"""
|
||||
Main function to transform AST.
|
||||
"""
|
||||
self.visit(self.root)
|
||||
|
||||
def reorder_block_statements(self, stmts):
|
||||
register_hook_nodes = [
|
||||
n
|
||||
for n in stmts
|
||||
for stmt in gast.walk(n)
|
||||
if isinstance(stmt, gast.Attribute) and stmt.attr == "register_hook"
|
||||
]
|
||||
# Analyze the register_hook nodes name dependency
|
||||
dependents = {}
|
||||
for n in register_hook_nodes:
|
||||
if n not in stmts:
|
||||
continue
|
||||
for load_node in get_loads(n):
|
||||
load_name = ast_to_source_code(load_node)
|
||||
if load_name not in dependents:
|
||||
dependents[load_name] = []
|
||||
dependents[load_name].append(n)
|
||||
|
||||
# Reorder the register_hook nodes, insert it before the dependent nodes
|
||||
idx = 0
|
||||
reordered_stmts = list(stmts)
|
||||
while idx < len(reordered_stmts):
|
||||
stmt = reordered_stmts[idx]
|
||||
loads = get_loads(stmt)
|
||||
for load_node in loads:
|
||||
load_name = ast_to_source_code(load_node)
|
||||
if load_name in dependents:
|
||||
dep_nodes = dependents[load_name]
|
||||
for dep_node in dep_nodes:
|
||||
dep_idx = reordered_stmts.index(dep_node)
|
||||
if dep_idx <= idx:
|
||||
continue
|
||||
reordered_stmts.remove(dep_node)
|
||||
reordered_stmts.insert(idx, dep_node)
|
||||
idx += 1
|
||||
idx += 1
|
||||
return reordered_stmts
|
||||
|
||||
def visit_FunctionDef(self, node: gast.FunctionDef):
|
||||
node.body = self.reorder_block_statements(node.body)
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def visit_For(self, node: gast.For):
|
||||
node.body = self.reorder_block_statements(node.body)
|
||||
node.orelse = self.reorder_block_statements(node.orelse)
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def visit_While(self, node: gast.While):
|
||||
node.body = self.reorder_block_statements(node.body)
|
||||
node.orelse = self.reorder_block_statements(node.orelse)
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def visit_If(self, node: gast.If):
|
||||
node.body = self.reorder_block_statements(node.body)
|
||||
node.orelse = self.reorder_block_statements(node.orelse)
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def visit_With(self, node: gast.With):
|
||||
node.body = self.reorder_block_statements(node.body)
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def visit_Try(self, node: gast.Try):
|
||||
node.body = self.reorder_block_statements(node.body)
|
||||
node.orelse = self.reorder_block_statements(node.orelse)
|
||||
node.finalbody = self.reorder_block_statements(node.finalbody)
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
@@ -0,0 +1,144 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
# gast is a generic AST to represent Python2 and Python3's Abstract Syntax Tree(AST).
|
||||
# It provides a compatibility layer between the AST of various Python versions,
|
||||
# as produced by ast.parse from the standard ast module.
|
||||
# See details in https://github.com/serge-sans-paille/gast/
|
||||
|
||||
import os
|
||||
|
||||
from paddle.framework import use_pir_api
|
||||
|
||||
from .. import logging_utils
|
||||
from ..utils import ast_to_source_code
|
||||
from .assert_transformer import AssertTransformer
|
||||
from .base import BaseTransformer
|
||||
from .break_continue_transformer import (
|
||||
BreakContinueTransformer,
|
||||
BreakTransformOptimizer,
|
||||
)
|
||||
from .call_transformer import CallTransformer
|
||||
from .cast_transformer import CastTransformer
|
||||
from .create_variable_transformer import CreateVariableTransformer
|
||||
from .decorator_transformer import DecoratorTransformer
|
||||
from .early_return_transformer import EarlyReturnTransformer
|
||||
from .ifelse_transformer import IfElseTransformer
|
||||
from .logical_transformer import LogicalTransformer
|
||||
from .loop_transformer import LoopTransformer
|
||||
from .name_load_transformer import (
|
||||
AttributeJstTransformer,
|
||||
NameloadJstTransformer,
|
||||
)
|
||||
from .return_transformer import ReturnTransformer
|
||||
from .super_transformer import SuperTransformer
|
||||
from .tensor_shape_transformer import TensorShapeTransformer
|
||||
from .tensorhook_transformer import RegisterHookTransformer
|
||||
from .typehint_transformer import TypeHintTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def apply_optimization(transformers):
|
||||
"""
|
||||
Judge whether to apply optimized transformation, such as BreakTransformOptimizer.
|
||||
And not all optimized transformations are applied by default. It's controlled by
|
||||
'export FLAGS_optim_transformation=1'
|
||||
"""
|
||||
flag = str(os.environ.get('FLAGS_optim_transformation')) in [
|
||||
'1',
|
||||
'True',
|
||||
'true',
|
||||
]
|
||||
if flag:
|
||||
transformers.insert(3, BreakTransformOptimizer)
|
||||
|
||||
|
||||
class DygraphToStaticAst(BaseTransformer):
|
||||
"""
|
||||
Main class to transform Dygraph to Static Graph
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.translator_logger = logging_utils.TranslatorLogger()
|
||||
|
||||
def get_static_ast(self, root):
|
||||
self.root = root
|
||||
self.decorate_func_name = None
|
||||
|
||||
# inplace transfer
|
||||
self.transfer_from_node_type(self.root)
|
||||
return self.root
|
||||
|
||||
def _apply(self, transformer, node, log_level):
|
||||
transformer(node).transform()
|
||||
self.translator_logger.log_transformed_code(
|
||||
log_level, self.root, transformer.__name__
|
||||
)
|
||||
|
||||
def transfer_from_node_type(self, node):
|
||||
self.translator_logger.log(
|
||||
1, f"Source code: \n{ast_to_source_code(self.root)}"
|
||||
)
|
||||
# Generic transformation
|
||||
self.visit(node)
|
||||
|
||||
transformers = [
|
||||
TypeHintTransformer, # remove all typehint
|
||||
SuperTransformer, # super() -> super(__class__, <first argument>)
|
||||
RegisterHookTransformer,
|
||||
EarlyReturnTransformer,
|
||||
AttributeJstTransformer, # Tensor.size -> Tensor.size(), it's unnecessary in PIR mode
|
||||
TensorShapeTransformer, # Tensor.shape -> paddle.shape(Tensor)
|
||||
BreakContinueTransformer, # break/continue in loops
|
||||
ReturnTransformer, # return in functions
|
||||
LogicalTransformer, # logical and/or/not
|
||||
CreateVariableTransformer, # create undefined var for if / while / for
|
||||
LoopTransformer, # for/while -> while_op
|
||||
IfElseTransformer, # if/else -> if_op
|
||||
AssertTransformer, # assert statement
|
||||
CallTransformer, # transform call recursively
|
||||
CastTransformer, # type casting statement
|
||||
DecoratorTransformer, # transform decorators to function call
|
||||
NameloadJstTransformer,
|
||||
]
|
||||
|
||||
if use_pir_api():
|
||||
# It's unnecessary in PIR mode
|
||||
transformers.remove(AttributeJstTransformer)
|
||||
|
||||
apply_optimization(transformers)
|
||||
|
||||
for index, transformer in enumerate(transformers):
|
||||
self._apply(transformer, node, log_level=index + 1)
|
||||
|
||||
self.translator_logger.log_transformed_code(
|
||||
logging_utils.LOG_AllTransformer, self.root, "All Transformers"
|
||||
)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
if self.decorate_func_name is None:
|
||||
self.decorate_func_name = node.name
|
||||
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def get_module_name(self):
|
||||
"""
|
||||
Return the main function name which will be used as module name
|
||||
in ast_to_func.
|
||||
"""
|
||||
# Should consider BaseAPITransformer which add new module name in Yamei's PR.
|
||||
assert self.decorate_func_name, "decorate_func_name shall not be None."
|
||||
return self.decorate_func_name
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
|
||||
from paddle.utils import gast
|
||||
|
||||
from .base import BaseTransformer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class TypeHintTransformer(BaseTransformer):
|
||||
"""
|
||||
A class remove all the typehint in gast.Name(annotation).
|
||||
Please put it behind other transformers because other transformer may relay on typehints.
|
||||
"""
|
||||
|
||||
def __init__(self, root):
|
||||
self.root = root
|
||||
|
||||
def transform(self):
|
||||
self.visit(self.root)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
node.returns = None
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def visit_Name(self, node):
|
||||
node.annotation = None
|
||||
self.generic_visit(node)
|
||||
return node
|
||||
|
||||
def visit_AnnAssign(self, node):
|
||||
if node.value is None:
|
||||
return None
|
||||
assign_node = gast.Assign(
|
||||
targets=[node.target],
|
||||
value=node.value,
|
||||
type_comment=None,
|
||||
)
|
||||
return assign_node
|
||||
@@ -0,0 +1,622 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import textwrap
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
from paddle.base import unique_name
|
||||
from paddle.jit.dy2static.ast_utils import ast_to_source_code
|
||||
from paddle.utils import gast
|
||||
|
||||
from ..utils import PADDLE_MODULE_PREFIX, is_api_in_module_helper
|
||||
|
||||
GET_ARGS_FUNC_PREFIX = 'get_args'
|
||||
SET_ARGS_FUNC_PREFIX = 'set_args'
|
||||
ARGS_NAME = '__args'
|
||||
|
||||
TRUE_FUNC_PREFIX = 'true_fn'
|
||||
FALSE_FUNC_PREFIX = 'false_fn'
|
||||
|
||||
FOR_ITER_INDEX_PREFIX = '__for_loop_var_index'
|
||||
FOR_ITER_TUPLE_PREFIX = '__for_loop_iter_tuple'
|
||||
FOR_ITER_TARGET_PREFIX = '__for_loop_iter_target'
|
||||
FOR_ITER_ITERATOR_PREFIX = '__for_loop_iter_iterator'
|
||||
FOR_ITER_TUPLE_INDEX_PREFIX = '__for_loop_iter_tuple_index'
|
||||
FOR_ITER_VAR_LEN_PREFIX = '__for_loop_var_len'
|
||||
FOR_ITER_VAR_NAME_PREFIX = '__for_loop_iter_var'
|
||||
FOR_ITER_ZIP_TO_LIST_PREFIX = '__for_loop_iter_zip'
|
||||
|
||||
WHILE_CONDITION_PREFIX = 'while_condition'
|
||||
WHILE_BODY_PREFIX = 'while_body'
|
||||
FOR_CONDITION_PREFIX = 'for_loop_condition'
|
||||
FOR_BODY_PREFIX = 'for_loop_body'
|
||||
|
||||
|
||||
def index_in_list(array_list, item):
|
||||
try:
|
||||
return array_list.index(item)
|
||||
except ValueError:
|
||||
# Item not in array_list
|
||||
return -1
|
||||
|
||||
|
||||
class BaseNodeVisitor(gast.NodeVisitor):
|
||||
"""
|
||||
Implement customized NodeVisitor inherited from gast.NodeVisitor.
|
||||
Ancestor nodes are traced to easily support more operations of currently
|
||||
visited node.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.ancestor_nodes = []
|
||||
|
||||
def visit(self, node):
|
||||
"""Visit a node."""
|
||||
self.ancestor_nodes.append(node)
|
||||
|
||||
method = 'visit_' + node.__class__.__name__
|
||||
visitor = getattr(self, method, self.generic_visit)
|
||||
ret = visitor(node)
|
||||
self.ancestor_nodes.pop()
|
||||
return ret
|
||||
|
||||
|
||||
def create_undefined_var(name):
|
||||
func_code = f"{name} = _jst.UndefinedVar('{name}')"
|
||||
return gast.parse(func_code).body[0]
|
||||
|
||||
|
||||
def create_bool_node(name, value):
|
||||
'''
|
||||
Create a assign stmt for name = value .
|
||||
'''
|
||||
assert isinstance(value, bool)
|
||||
node = f"{name} = {value}"
|
||||
return gast.parse(node).body[0]
|
||||
|
||||
|
||||
def get_parent_mapping(root):
|
||||
to_parent: dict[gast.AST, gast.AST] = {}
|
||||
for node in gast.walk(root):
|
||||
for child in gast.iter_child_nodes(node):
|
||||
to_parent[child] = node
|
||||
return to_parent
|
||||
|
||||
|
||||
def create_name_str(name_ids):
|
||||
"""
|
||||
Return "('x', 'y')" for [x, y]
|
||||
"""
|
||||
if not name_ids:
|
||||
return 'None'
|
||||
|
||||
names_str = ["'{}'".format(name.replace("'", "\\'")) for name in name_ids]
|
||||
return "({}, )".format(','.join(names_str))
|
||||
|
||||
|
||||
def create_function_def_node(nodes, name, input_args, return_name_ids):
|
||||
"""
|
||||
Wrapper all statements of nodes into one ast.FunctionDef, which can be
|
||||
called by ast.Call.
|
||||
"""
|
||||
nodes = copy.copy(nodes)
|
||||
# add return statement
|
||||
if return_name_ids:
|
||||
nodes.append(gast.Return(value=generate_name_node(return_name_ids)))
|
||||
else:
|
||||
nodes.append(gast.Return(value=None))
|
||||
func_def_node = gast.FunctionDef(
|
||||
name=name,
|
||||
args=input_args,
|
||||
body=nodes,
|
||||
decorator_list=[],
|
||||
returns=None,
|
||||
type_comment=None,
|
||||
type_params=[],
|
||||
)
|
||||
return func_def_node
|
||||
|
||||
|
||||
def create_assign_node(name, node):
|
||||
"""
|
||||
Creates a `gast.Assign` node by given name_id as target and node as value.
|
||||
"""
|
||||
targets = generate_name_node(name, ctx=gast.Store())
|
||||
assign_node = gast.Assign(
|
||||
targets=[targets],
|
||||
value=node,
|
||||
type_comment=None,
|
||||
)
|
||||
return targets, assign_node
|
||||
|
||||
|
||||
def create_get_args_node(names):
|
||||
"""
|
||||
Create get_args function as follows:
|
||||
|
||||
def get_args_0():
|
||||
nonlocal x, y
|
||||
return x, y
|
||||
"""
|
||||
|
||||
def empty_node():
|
||||
func_def = f"""
|
||||
def {unique_name.generate(GET_ARGS_FUNC_PREFIX)}():
|
||||
return
|
||||
"""
|
||||
return gast.parse(textwrap.dedent(func_def)).body[0]
|
||||
|
||||
assert isinstance(names, (list, tuple))
|
||||
node = create_nonlocal_stmt_nodes(names)
|
||||
if not names:
|
||||
return empty_node()
|
||||
if node == []:
|
||||
nonlocal_vars = "\n"
|
||||
else:
|
||||
nonlocal_vars = ast_to_source_code(node[0])
|
||||
template = """
|
||||
def {func_name}():
|
||||
{nonlocal_vars}
|
||||
return {vars},
|
||||
"""
|
||||
func_def = template.format(
|
||||
func_name=unique_name.generate(GET_ARGS_FUNC_PREFIX),
|
||||
nonlocal_vars=nonlocal_vars,
|
||||
vars=",".join(names),
|
||||
)
|
||||
return gast.parse(textwrap.dedent(func_def)).body[0]
|
||||
|
||||
|
||||
def create_set_args_node(names):
|
||||
"""
|
||||
Create set_args function as follows:
|
||||
|
||||
def set_args_0(__args):
|
||||
nonlocal x, y
|
||||
x, y = __args
|
||||
"""
|
||||
|
||||
def empty_node():
|
||||
func_def = f"""
|
||||
def {unique_name.generate(SET_ARGS_FUNC_PREFIX)}({ARGS_NAME}):
|
||||
pass
|
||||
"""
|
||||
return gast.parse(textwrap.dedent(func_def)).body[0]
|
||||
|
||||
assert isinstance(names, (list, tuple))
|
||||
node = create_nonlocal_stmt_nodes(names)
|
||||
if not names:
|
||||
return empty_node()
|
||||
if node == []:
|
||||
nonlocal_vars = "\n"
|
||||
else:
|
||||
nonlocal_vars = ast_to_source_code(node[0])
|
||||
template = """
|
||||
def {func_name}({args}):
|
||||
{nonlocal_vars}
|
||||
{vars}, = {args}
|
||||
"""
|
||||
func_def = template.format(
|
||||
func_name=unique_name.generate(SET_ARGS_FUNC_PREFIX),
|
||||
args=ARGS_NAME,
|
||||
nonlocal_vars=nonlocal_vars,
|
||||
vars=",".join(names),
|
||||
)
|
||||
return gast.parse(textwrap.dedent(func_def)).body[0]
|
||||
|
||||
|
||||
def create_nonlocal_stmt_nodes(names):
|
||||
assert isinstance(names, (list, tuple))
|
||||
|
||||
mapped = list(filter(lambda n: '.' not in n, names))
|
||||
mapped = list(filter(lambda n: '[' not in n, mapped))
|
||||
names = sorted(
|
||||
mapped, key=mapped.index
|
||||
) # to keep the order, we can't use set() to unique
|
||||
if not names:
|
||||
return []
|
||||
func_code = "nonlocal {}".format(','.join(names))
|
||||
return [gast.parse(func_code).body[0]]
|
||||
|
||||
|
||||
def generate_name_node(name_ids, ctx=gast.Load(), gen_tuple_if_single=False):
|
||||
"""
|
||||
If name_ids is list or tuple or set with multiple strings, this function
|
||||
generates gast.Tuple of gast.Name.
|
||||
If the name_ids is single string or contains only 1 string, this function
|
||||
returns gast.Name if gen_tuple_if_single==False else returns gast.Tuple
|
||||
with only one gast.Name
|
||||
|
||||
This function is used at several gast.Return statements.
|
||||
"""
|
||||
if isinstance(name_ids, str):
|
||||
name_ids = [name_ids]
|
||||
if not isinstance(name_ids, (list, tuple, set)):
|
||||
raise TypeError(
|
||||
f'name_ids must be list or tuple or set, but received {type(name_ids)}'
|
||||
)
|
||||
|
||||
def create_node_for_name(name):
|
||||
if '.' not in name:
|
||||
return gast.Name(
|
||||
id=name, ctx=ctx, annotation=None, type_comment=None
|
||||
)
|
||||
return gast.parse(name).body[0].value
|
||||
|
||||
gast_names = [create_node_for_name(name_id) for name_id in name_ids]
|
||||
if len(gast_names) == 1 and not gen_tuple_if_single:
|
||||
name_node = gast_names[0]
|
||||
else:
|
||||
name_node = gast.Tuple(elts=gast_names, ctx=ctx)
|
||||
return name_node
|
||||
|
||||
|
||||
def get_attribute_full_name(node):
|
||||
assert isinstance(node, gast.Attribute), (
|
||||
"Input non-Attribute node to get attribute full name"
|
||||
)
|
||||
return ast_to_source_code(node).strip()
|
||||
|
||||
|
||||
def is_api_in_module(node, module_prefix):
|
||||
assert isinstance(node, gast.Call), (
|
||||
"Input non-Call node for is_api_in_module"
|
||||
)
|
||||
|
||||
# Python can have gast.Call as function, for example: convert_call(func)(x)
|
||||
# We only check the most outside function
|
||||
func_node = node.func
|
||||
while isinstance(func_node, gast.Call):
|
||||
func_node = func_node.func
|
||||
|
||||
func_str = ast_to_source_code(func_node).strip()
|
||||
try:
|
||||
import paddle
|
||||
import paddle.jit.dy2static as _jst
|
||||
from paddle import to_tensor
|
||||
|
||||
globals = {
|
||||
'np': np,
|
||||
'paddle': paddle,
|
||||
'_jst': _jst,
|
||||
'to_tensor': to_tensor,
|
||||
}
|
||||
|
||||
fn = eval(func_str, globals)
|
||||
return is_api_in_module_helper(fn, module_prefix)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def is_paddle_api(node):
|
||||
return is_api_in_module(node, PADDLE_MODULE_PREFIX)
|
||||
|
||||
|
||||
class NameScope:
|
||||
def __init__(self):
|
||||
"""
|
||||
A NameScope is a object which manager all the variable names.
|
||||
only FunctionDef and Controlflow node will have a namescope property.
|
||||
|
||||
type can be "function" and "controlflow"
|
||||
|
||||
we don't analyze the read only variable because they don't affect the analysis.
|
||||
"""
|
||||
self.globals = set()
|
||||
self.nonlocals = set()
|
||||
self.args = set()
|
||||
self.father = None # point to the nearest function name scope.
|
||||
self.w_vars = set() # all qualified + normal names been stored
|
||||
self.created = set() # useful for control flow compatibility
|
||||
# only valid in control_flow nodes
|
||||
# may be remove later.
|
||||
self.push_pop_vars = set() # we call push and pop in the vars
|
||||
|
||||
def set_father(self, father):
|
||||
self.father = father
|
||||
|
||||
def existed_vars(self):
|
||||
"""vars existing in current scope.
|
||||
they must not contain qualified names.
|
||||
"""
|
||||
local_vars = self.w_vars - self.globals - self.nonlocals - self.args
|
||||
return set(filter(lambda x: '.' not in x, local_vars))
|
||||
|
||||
def created_vars(self):
|
||||
return self.created
|
||||
|
||||
def modified_vars(self):
|
||||
# may be globals / non-locals / args / qualified names and created_vars
|
||||
return self.w_vars
|
||||
|
||||
def variadic_length_vars(self):
|
||||
"""
|
||||
At present, we do not support global append, such as
|
||||
|
||||
import numpy as np
|
||||
a = []
|
||||
def func():
|
||||
a.append() # global names `a`, we will raise a warning.
|
||||
p.append(a, 1) # global names `np`, we will raise a warning.
|
||||
"""
|
||||
non_global_push_pop_names = []
|
||||
for var in self.push_pop_vars:
|
||||
if self._is_simple_name(var) and self.is_global_var(var):
|
||||
warnings.warn(
|
||||
f"Find variable `{var}` defined in global scope"
|
||||
f" and call `{var}.append() or {var}.pop()`"
|
||||
f", which will be ignored and never be transferred into"
|
||||
f" tensor array."
|
||||
)
|
||||
else:
|
||||
non_global_push_pop_names.append(var)
|
||||
return set(non_global_push_pop_names)
|
||||
|
||||
def control_flow_vars(self):
|
||||
valid_names = self.w_vars
|
||||
tmp = (self.father.global_vars & valid_names,)
|
||||
return {"global": tmp, "nonlocal": self.w_vars - tmp}
|
||||
|
||||
def _is_simple_name(self, name):
|
||||
if '.' in name or '[' in name:
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_global_var(self, name):
|
||||
"""
|
||||
Return whether the name is a var created in global scope.
|
||||
Search from bottom to top. If it is not created or modified,
|
||||
it means global vars; otherwise, it means local vars.
|
||||
Only valid after FunctionNameLivenessAnalysis visitor.
|
||||
"""
|
||||
assert self._is_simple_name(name), (
|
||||
"is_global_var accept a simple name, but get `{name}`."
|
||||
)
|
||||
ancestor = self
|
||||
while ancestor is not None:
|
||||
if name in ancestor.globals:
|
||||
return True
|
||||
if name in (ancestor.nonlocals | ancestor.w_vars):
|
||||
return False
|
||||
ancestor = ancestor.father
|
||||
return True
|
||||
|
||||
def is_local_var(self, name):
|
||||
return not self.is_global_var(name)
|
||||
|
||||
def merge_from(self, name_scope):
|
||||
self.globals |= name_scope.globals
|
||||
self.nonlocals |= name_scope.nonlocals
|
||||
self.args |= name_scope.args
|
||||
self.w_vars |= name_scope.w_vars
|
||||
self.push_pop_vars |= name_scope.push_pop_vars
|
||||
|
||||
|
||||
class FunctionNameLivenessAnalysis(gast.NodeVisitor):
|
||||
"""analyze the liveness of a function.
|
||||
|
||||
every variables stored in this scope will be collected,
|
||||
in addition with global/nonlocal information and
|
||||
push_pop information.
|
||||
|
||||
1. global variable is stored in node.var_globals.
|
||||
2. nonlocal variable is stored in node.var_nonlocals.
|
||||
3. arguments is stored in node.var_args.
|
||||
4. if a variable's push and pop attribute is called,
|
||||
it will be collected in push_pop_vars. They are
|
||||
used for transformation to tensor_array.
|
||||
NOTE: push_pop_vars **may not** in w_vars.
|
||||
a.push(0) don't modify the variable a, but the content
|
||||
of a.
|
||||
|
||||
For example:
|
||||
|
||||
def func(*args, **kargs):
|
||||
a = 12
|
||||
global i,j
|
||||
nonlocal x,y
|
||||
print(a)
|
||||
i = k
|
||||
b = []
|
||||
c = [1,2,3]
|
||||
for m in range(10):
|
||||
q = 12
|
||||
b.push(1)
|
||||
c.pop()
|
||||
|
||||
After this visitor we have:
|
||||
# node is the FunctionDef node with name: "func"
|
||||
node.pd_scope = NameScope(
|
||||
globals = ['i', 'j'],
|
||||
nonlocals = ['x', 'y'],
|
||||
args = ['args', 'kargs'],
|
||||
wr_vars = ['a', 'i', 'q', 'm', 'c', 'b']
|
||||
push_pop_vars = ['b', 'c']
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(self, root_node):
|
||||
self.scope_node_stack = [] # controlflow, functiondef node
|
||||
self.visit(root_node)
|
||||
|
||||
def _reset_name_scope(self, node):
|
||||
# always reset the node as empty namescope.
|
||||
node.pd_scope = NameScope()
|
||||
|
||||
def _get_name_scope(self, node):
|
||||
if not hasattr(node, "pd_scope"):
|
||||
node.pd_scope = NameScope()
|
||||
return node.pd_scope
|
||||
|
||||
def _current_name_scope(self):
|
||||
return self._get_name_scope(self.scope_node_stack[-1])
|
||||
|
||||
def _father_name_scope(self):
|
||||
if len(self.scope_node_stack) == 1:
|
||||
return None
|
||||
return self._get_name_scope(self.scope_node_stack[-2])
|
||||
|
||||
def _nearest_function_scope(self):
|
||||
if len(self.scope_node_stack) == 1:
|
||||
return None
|
||||
for node in self.scope_node_stack[-2::-1]:
|
||||
if isinstance(node, gast.FunctionDef):
|
||||
return self._get_name_scope(node)
|
||||
|
||||
def visit_ListComp(self, node):
|
||||
"""[ i for i in range(10) ]
|
||||
In this case, `i` will not created in FunctionScope.
|
||||
We don't collect `i` by not calling generic_visit.
|
||||
"""
|
||||
pass
|
||||
|
||||
def visit_DictComp(self, node):
|
||||
"""the same as ListComp."""
|
||||
pass
|
||||
|
||||
def visit_Name(self, node):
|
||||
self.generic_visit(node)
|
||||
write_context = (gast.Store, gast.AugStore, gast.Del)
|
||||
if isinstance(node.ctx, write_context):
|
||||
self._current_name_scope().w_vars.add(node.id)
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
def pre_func():
|
||||
self._current_name_scope().args |= set(
|
||||
self._get_argument_names(node)
|
||||
)
|
||||
|
||||
def post_func():
|
||||
"""NOTE: why we need merge w_vars and push_pop_vars here ?
|
||||
because we do ifelse_transformer after loop_transformer. Loops will changed into functions. but we know this function will be called in if. so we add w_vars to father function scope.
|
||||
"""
|
||||
control_flow_function_def = [
|
||||
WHILE_BODY_PREFIX,
|
||||
WHILE_BODY_PREFIX,
|
||||
FOR_CONDITION_PREFIX,
|
||||
FOR_BODY_PREFIX,
|
||||
TRUE_FUNC_PREFIX,
|
||||
FALSE_FUNC_PREFIX,
|
||||
]
|
||||
|
||||
def is_control_flow_def_node():
|
||||
for prefix in control_flow_function_def:
|
||||
if node.name.startswith(prefix):
|
||||
return True
|
||||
return False
|
||||
|
||||
if self._father_name_scope() and is_control_flow_def_node():
|
||||
self._father_name_scope().w_vars |= (
|
||||
self._current_name_scope().w_vars
|
||||
)
|
||||
self._father_name_scope().push_pop_vars |= (
|
||||
self._current_name_scope().push_pop_vars
|
||||
)
|
||||
|
||||
self._visit_scope_node(node, pre_func, post_func)
|
||||
|
||||
def _visit_scope_node(self, node, pre_func, post_func):
|
||||
"""scope node main visit logic.
|
||||
pre_func and post_func is callbacks
|
||||
"""
|
||||
self._reset_name_scope(node)
|
||||
self.scope_node_stack.append(node)
|
||||
self._current_name_scope().set_father(self._nearest_function_scope())
|
||||
if pre_func:
|
||||
pre_func()
|
||||
self.generic_visit(node)
|
||||
if post_func:
|
||||
post_func()
|
||||
self.scope_node_stack.pop()
|
||||
|
||||
def _visit_controlflow_node(self, node):
|
||||
def post_func():
|
||||
self._father_name_scope().merge_from(self._current_name_scope())
|
||||
self._nearest_function_scope().merge_from(
|
||||
self._current_name_scope()
|
||||
)
|
||||
self._current_name_scope().created = (
|
||||
self._nearest_function_scope().existed_vars()
|
||||
- node.before_created
|
||||
)
|
||||
# gather created vars into father and used in CreateUndefinedVarTransform
|
||||
self._nearest_function_scope().created |= (
|
||||
self._current_name_scope().created
|
||||
)
|
||||
|
||||
def pre_func():
|
||||
node.before_created = self._nearest_function_scope().existed_vars()
|
||||
|
||||
self._visit_scope_node(node, pre_func, post_func)
|
||||
|
||||
def visit_For(self, node):
|
||||
self._visit_controlflow_node(node)
|
||||
|
||||
def visit_While(self, node):
|
||||
self._visit_controlflow_node(node)
|
||||
|
||||
def visit_If(self, node):
|
||||
self._visit_controlflow_node(node)
|
||||
|
||||
def visit_Global(self, node):
|
||||
self._current_name_scope().globals |= set(node.names)
|
||||
|
||||
def visit_Nonlocal(self, node):
|
||||
self._current_name_scope().nonlocals |= set(node.names)
|
||||
|
||||
def visit_Attribute(self, node):
|
||||
self.generic_visit(node)
|
||||
write_context = (gast.Store, gast.AugStore, gast.Del)
|
||||
if isinstance(node.ctx, write_context):
|
||||
name = ast_to_source_code(node).strip()
|
||||
self._current_name_scope().w_vars.add(name)
|
||||
|
||||
def visit_Subscript(self, node):
|
||||
self.generic_visit(node)
|
||||
write_context = (gast.Store, gast.AugStore, gast.Del)
|
||||
if isinstance(node.ctx, write_context):
|
||||
while isinstance(node.value, gast.Subscript):
|
||||
node = node.value
|
||||
if isinstance(node.value, gast.Name):
|
||||
self._current_name_scope().w_vars.add(node.value.id)
|
||||
|
||||
def visit_Call(self, node):
|
||||
self.generic_visit(node)
|
||||
if not isinstance(node.func, gast.Attribute):
|
||||
return
|
||||
variadic_length_method = ['append', 'pop']
|
||||
if node.func.attr not in variadic_length_method:
|
||||
return
|
||||
# we don't treat push and pop as a write operator. such as a[i]=10 is not modify a.
|
||||
name = ast_to_source_code(node.func.value).strip()
|
||||
self._current_name_scope().push_pop_vars.add(name)
|
||||
|
||||
def _get_argument_names(self, node):
|
||||
"""get all arguments name in the functiondef node.
|
||||
this node is local to the function and shouldn't
|
||||
be created.
|
||||
"""
|
||||
assert isinstance(node, gast.FunctionDef), (
|
||||
"Input node is not function define node"
|
||||
)
|
||||
names = list(node.args.args)
|
||||
names.append(node.args.vararg)
|
||||
names.append(node.args.kwarg)
|
||||
names = [i.id for i in names if i is not None]
|
||||
return names
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,50 @@
|
||||
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
||||
# Copyright (c) 2021 NVIDIA Corporation. 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.
|
||||
|
||||
from paddle.base import core
|
||||
from paddle.base.core import Load
|
||||
|
||||
|
||||
class Layer:
|
||||
def __init__(self):
|
||||
self.cpp_layer = None
|
||||
# {name: Function}
|
||||
self.functions = {}
|
||||
|
||||
def load(self, load_path, place):
|
||||
self.cpp_layer = Load(load_path, place)
|
||||
|
||||
for name in self.cpp_layer.function_names():
|
||||
function = self.cpp_layer.function(name)
|
||||
info = self.cpp_layer.function_info(name)
|
||||
self.functions[name] = Function(function, info)
|
||||
setattr(self, name, self.functions[name])
|
||||
|
||||
|
||||
class Function:
|
||||
def __init__(self, function, info):
|
||||
self.function = function
|
||||
self.info = FunctionInfo(info)
|
||||
|
||||
def __call__(self, *args):
|
||||
return core.eager.jit_function_call(self.function, args)
|
||||
|
||||
|
||||
class FunctionInfo:
|
||||
def __init__(self, info):
|
||||
self.info = info
|
||||
|
||||
def name(self):
|
||||
return self.info.name()
|
||||
@@ -0,0 +1,188 @@
|
||||
# Copyright (c) 2025 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from typing import TYPE_CHECKING, Protocol, TypeVar, overload
|
||||
|
||||
from typing_extensions import (
|
||||
ParamSpec,
|
||||
)
|
||||
|
||||
import paddle
|
||||
from paddle.jit.dy2static.utils import DYNAMIC_DIMS_ATTR_NAME
|
||||
|
||||
from .dy2static.utils import (
|
||||
TransformOptions,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
|
||||
_RetT = TypeVar("_RetT")
|
||||
_InputT = ParamSpec("_InputT")
|
||||
|
||||
|
||||
class _NotToStaticDecorator(Protocol):
|
||||
@overload
|
||||
def __call__(
|
||||
self, func: Callable[_InputT, _RetT]
|
||||
) -> Callable[_InputT, _RetT]: ...
|
||||
|
||||
@overload
|
||||
def __call__(self, func: None = ...) -> _NotToStaticDecorator: ...
|
||||
|
||||
|
||||
@overload
|
||||
def not_to_static(
|
||||
func: Callable[_InputT, _RetT],
|
||||
) -> Callable[_InputT, _RetT]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def not_to_static(func: None = ...) -> _NotToStaticDecorator: ...
|
||||
|
||||
|
||||
# Legacy decorator only for AST
|
||||
def not_to_static(func=None):
|
||||
"""
|
||||
A Decorator to suppresses the convention of a function.
|
||||
|
||||
Args:
|
||||
func(callable): The function to decorate.
|
||||
|
||||
Returns:
|
||||
callable: A function which won't be converted in Dynamic-to-Static.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
|
||||
>>> import paddle
|
||||
|
||||
>>> @paddle.jit.not_to_static
|
||||
... def func_not_to_static(x):
|
||||
... res = x - 1
|
||||
... return res
|
||||
|
||||
>>> @paddle.jit.to_static
|
||||
... def func(x):
|
||||
... if paddle.mean(x) < 0:
|
||||
... out = func_not_to_static(x)
|
||||
... else:
|
||||
... out = x + 1
|
||||
... return out
|
||||
>>> x = paddle.ones([1, 2], dtype='float32')
|
||||
>>> out = func(x)
|
||||
>>> print(out)
|
||||
Tensor(shape=[1, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
|
||||
[[2., 2.]])
|
||||
"""
|
||||
return unified(func, for_sot=False, for_ast=True)
|
||||
|
||||
|
||||
def unified(
|
||||
fn: Callable[_InputT, _RetT] | type[paddle.nn.Layer] | None = None,
|
||||
*,
|
||||
for_sot: bool = True,
|
||||
for_ast: bool = True,
|
||||
) -> Callable[_InputT, _RetT]:
|
||||
"""
|
||||
Mark a function already unified in dygraph and static mode. So
|
||||
that it won't be transformed again in SOT or AST mode.
|
||||
|
||||
Args:
|
||||
fn(callable): The function to decorate.
|
||||
for_sot(bool): Whether to mark the function as unified in SOT mode.
|
||||
for_ast(bool): Whether to mark the function as unified in AST mode.
|
||||
"""
|
||||
|
||||
def _mark_as_unified(fn, *, for_sot: bool, for_ast: bool):
|
||||
mode = TransformOptions.ToStaticMode.Nil()
|
||||
if for_sot:
|
||||
mode |= TransformOptions.ToStaticMode.SOT
|
||||
if for_ast:
|
||||
mode |= TransformOptions.ToStaticMode.AST
|
||||
options = TransformOptions().with_skip_transform_mode(mode)
|
||||
options.attach(fn)
|
||||
return fn
|
||||
|
||||
if fn is None:
|
||||
return lambda fn: _mark_as_unified(fn, for_sot=for_sot, for_ast=for_ast)
|
||||
return _mark_as_unified(fn, for_sot=for_sot, for_ast=for_ast)
|
||||
|
||||
|
||||
def capture_control_flow(
|
||||
fn: Callable[_InputT, _RetT] | None = None,
|
||||
) -> Callable[_InputT, _RetT]:
|
||||
def _mark_as_need_capture_control_flow(fn):
|
||||
options = TransformOptions().with_need_capture_control_flow(True)
|
||||
options.attach(fn)
|
||||
return fn
|
||||
|
||||
if fn is None:
|
||||
return _mark_as_need_capture_control_flow
|
||||
return _mark_as_need_capture_control_flow(fn)
|
||||
|
||||
|
||||
def force_dynamic(
|
||||
fn: Callable[_InputT, _RetT] | type[paddle.nn.Layer] | None = None,
|
||||
) -> Callable[_InputT, _RetT]:
|
||||
"""
|
||||
Mark a function or paddle.nn.Layer to be executed in dynamic mode, it will
|
||||
break the graph and prevent it from being converted to static mode.
|
||||
"""
|
||||
from paddle.jit import sot
|
||||
|
||||
if inspect.isclass(fn) and issubclass(fn, paddle.nn.Layer):
|
||||
sot.utils.paddle_api_config.add_break_graph_layer_class(fn)
|
||||
return fn
|
||||
if inspect.isfunction(fn):
|
||||
sot.utils.paddle_api_config.add_break_graph_function(fn)
|
||||
return fn
|
||||
|
||||
raise TypeError(
|
||||
f"Expected a callable or paddle.nn.Layer, but got {type(fn).__name__}."
|
||||
)
|
||||
|
||||
|
||||
def dynamic_dims(
|
||||
tensor: paddle.Tensor,
|
||||
dims: int | Sequence[int],
|
||||
) -> None:
|
||||
"""
|
||||
Mark a tensor as having dynamic dimensions.
|
||||
This is used to indicate that the tensor's shape may change dynamically
|
||||
during execution, which is particularly useful in dynamic-to-static
|
||||
conversion scenarios.
|
||||
Args:
|
||||
tensor (paddle.Tensor): The tensor to mark as dynamic.
|
||||
dims (int | Sequence[int]): The dimensions to mark as dynamic.
|
||||
"""
|
||||
if not isinstance(tensor, paddle.Tensor):
|
||||
raise TypeError(
|
||||
f"Expected a paddle.Tensor, but got {type(tensor).__name__}."
|
||||
)
|
||||
|
||||
if isinstance(dims, int):
|
||||
dims = (dims,)
|
||||
|
||||
if not isinstance(dims, (list, tuple)):
|
||||
raise TypeError("Dimensions must be a list or tuple.")
|
||||
if not all(isinstance(dim, int) for dim in dims):
|
||||
raise TypeError("All dimensions must be integers.")
|
||||
|
||||
setattr(tensor, DYNAMIC_DIMS_ATTR_NAME, tuple(dims))
|
||||
@@ -0,0 +1,797 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.base import core, framework, unique_name
|
||||
from paddle.base.dygraph.base import switch_to_static_graph
|
||||
from paddle.framework import in_dynamic_mode
|
||||
from paddle.nn.layer import layers
|
||||
from paddle.pir.core import datatype_to_vartype
|
||||
|
||||
__all__ = []
|
||||
|
||||
PIR_INFER_MODEL_SUFFIX = ".json"
|
||||
|
||||
from .translated_layer import (
|
||||
BUFFER_NAME_PREFIX,
|
||||
INFER_PARAMS_SUFFIX,
|
||||
PARAMETER_NAME_PREFIX,
|
||||
)
|
||||
|
||||
|
||||
def _load_pir_program(model_file_path):
|
||||
program = paddle.static.Program()
|
||||
trainable = paddle.base.core.deserialize_pir_program(
|
||||
model_file_path, program
|
||||
)
|
||||
|
||||
return program, trainable
|
||||
|
||||
|
||||
@switch_to_static_graph
|
||||
def _generate_unique_var_name(prefix):
|
||||
return unique_name.generate_with_ignorable_key(prefix)
|
||||
|
||||
|
||||
@switch_to_static_graph
|
||||
def _generate_unique_var_name(prefix):
|
||||
return unique_name.generate(prefix)
|
||||
|
||||
|
||||
from paddle.static.pir_io import get_pir_parameters
|
||||
|
||||
|
||||
def _get_pir_parameters_var_names(program):
|
||||
persistable_vars = []
|
||||
persistable_names = []
|
||||
rename_new_old_dict = {}
|
||||
param, opt = get_pir_parameters(program)
|
||||
vars = param + opt
|
||||
for var in vars:
|
||||
persistable_vars.append(var)
|
||||
origin_name = var.name
|
||||
var.name = _generate_unique_var_name(var.name)
|
||||
rename_new_old_dict[var.name] = origin_name
|
||||
persistable_names.append(var.name)
|
||||
|
||||
return (
|
||||
persistable_vars,
|
||||
persistable_names,
|
||||
rename_new_old_dict,
|
||||
)
|
||||
|
||||
|
||||
class _PirProgramHolder:
|
||||
def __init__(self, program, trainable):
|
||||
super().__init__()
|
||||
|
||||
# input, output, persistable,
|
||||
self._input_vars = []
|
||||
self._output_vars = []
|
||||
self._parameter_vars = []
|
||||
self._parameter_names = []
|
||||
|
||||
self.support_train = trainable
|
||||
# append suffix var name dict
|
||||
self._suffix_varname_dict = None
|
||||
self._infer_program = program
|
||||
|
||||
self._preprocess()
|
||||
|
||||
def _preprocess(self):
|
||||
(
|
||||
self._parameter_vars,
|
||||
self._parameter_names,
|
||||
self._suffix_varname_dict,
|
||||
) = _get_pir_parameters_var_names(self._infer_program)
|
||||
block = self._infer_program.global_block()
|
||||
for op in block.ops:
|
||||
if op.name() == 'pd_op.data':
|
||||
self._input_vars.append(op.result(0))
|
||||
elif op.name() == 'pd_op.feed':
|
||||
var_name = op.attr()["name"]
|
||||
org_value = op.result(0)
|
||||
with block:
|
||||
value = paddle._pir_ops.data(
|
||||
name=var_name,
|
||||
shape=org_value.shape,
|
||||
dtype=org_value.dtype,
|
||||
place=paddle.base.core.Place(),
|
||||
)
|
||||
org_value.replace_all_uses_with(value)
|
||||
value.get_defining_op().move_before(op)
|
||||
block.remove_op(op)
|
||||
if op.name() == 'pd_op.fetch':
|
||||
self._output_vars.append(op.operand_source(0))
|
||||
with block:
|
||||
paddle._pir_ops.set_persistable_value(
|
||||
op.operand_source(0),
|
||||
"output_" + str(len(self._output_vars) - 1),
|
||||
)
|
||||
block.remove_op(op)
|
||||
|
||||
@property
|
||||
def infer_program(self):
|
||||
return self._infer_program
|
||||
|
||||
@property
|
||||
def input_vars(self):
|
||||
return self._input_vars
|
||||
|
||||
@property
|
||||
def output_vars(self):
|
||||
return self._output_vars
|
||||
|
||||
@property
|
||||
def persistable_names(self):
|
||||
return self._parameter_names
|
||||
|
||||
@property
|
||||
def persistable_vars(self):
|
||||
return self._parameter_vars
|
||||
|
||||
|
||||
# [ PirTranslatedLayer : Run program in dygraph mode ]
|
||||
#
|
||||
# DESIGN IDEA: using an special operator `PirRunProgram`, execute program inside operator.
|
||||
#
|
||||
# Op's Inputs:
|
||||
# - the input variable of the user feed
|
||||
# - the necessary parameters of the network
|
||||
# Op's Outputs:
|
||||
# - the output variable of fetch
|
||||
#
|
||||
# This op receives a complete program desc, internally creates scope
|
||||
# and executor, executes this program. Key points:
|
||||
#
|
||||
# 1. Data Sharing:
|
||||
# The variable/parameter of the dynamic graph is not in the scope, so before the op
|
||||
# executes the program internally, create persistent variables with the
|
||||
# same name as feed, parameters, and fetch in the scope, and share the
|
||||
# DenseTensor of the op input.
|
||||
#
|
||||
# 2. Forward and Backward Separation:
|
||||
# Because the dynamic graph op performs the forward and backward separately,
|
||||
# in the forward op RunProgram, we only execute the forward part of whole program,
|
||||
# and in the backward op RunProgramGrad, we execute the backward part of program.
|
||||
# We can not separate the program into forward and backward part, which will
|
||||
# make some control flow execution logic wrong.
|
||||
|
||||
|
||||
# NOTE: [compatible] deal with model saved by save_inference_model,
|
||||
# which need get var info from program desc
|
||||
|
||||
|
||||
def _load_pir_parameter_vars(model_path, program_holder, params_filename):
|
||||
# construct var dict
|
||||
load_var_dict = {}
|
||||
load_var_list = []
|
||||
other_var_dict = {}
|
||||
load_densetensor_list = []
|
||||
persistable_var = program_holder.persistable_vars
|
||||
persistable_var_name = program_holder.persistable_names
|
||||
origin_persistable_var_name = [
|
||||
program_holder._suffix_varname_dict[var_name]
|
||||
for var_name in persistable_var_name
|
||||
]
|
||||
for name, var in sorted(zip(origin_persistable_var_name, persistable_var)):
|
||||
if var.persistable:
|
||||
# use default shape and dtype
|
||||
new_var = framework.EagerParamBase(
|
||||
shape=var.shape, # only to pass check, this shape is not meaningful
|
||||
dtype=core.VarDesc.VarType.FP32,
|
||||
name=var.name,
|
||||
persistable=True,
|
||||
)
|
||||
|
||||
new_var.stop_gradient = var.stop_gradient
|
||||
load_var_dict[name] = new_var
|
||||
load_var_list.append(new_var)
|
||||
load_densetensor_list.append(new_var.get_tensor())
|
||||
|
||||
else:
|
||||
new_var = core.eager.Tensor(
|
||||
dtype=datatype_to_vartype[var.dtype],
|
||||
dims=var.shape,
|
||||
name=var.name,
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
place=framework._current_expected_place(),
|
||||
persistable=False,
|
||||
)
|
||||
other_var_dict[name] = new_var
|
||||
|
||||
# load all vars
|
||||
assert params_filename is not None, "params_filename should not be None."
|
||||
var_file_path = os.path.join(model_path, params_filename)
|
||||
if os.path.exists(var_file_path):
|
||||
core.load_combine_func(
|
||||
var_file_path,
|
||||
list(load_var_dict.keys()),
|
||||
load_densetensor_list,
|
||||
False,
|
||||
framework._current_expected_place(),
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The file {var_file_path} does not exist. Please check the model path."
|
||||
)
|
||||
|
||||
load_var_dict.update(other_var_dict)
|
||||
return load_var_dict
|
||||
|
||||
|
||||
def _construct_program_holders(model_path, model_filename=None):
|
||||
# make sure the path has been checked
|
||||
program_holder_dict = {}
|
||||
|
||||
if model_filename is not None:
|
||||
# [compatible] if assign model_filename, only can load one program as Layer.forward
|
||||
model_filename = os.path.basename(model_filename)
|
||||
model_file_path = os.path.join(model_path, model_filename)
|
||||
model_name = model_filename[: -len(PIR_INFER_MODEL_SUFFIX)]
|
||||
# Load every file that meets the requirements in the directory model_path.
|
||||
for filename in os.listdir(model_path):
|
||||
if model_filename == filename:
|
||||
func_name = 'forward'
|
||||
model_file_path = os.path.join(model_path, model_filename)
|
||||
elif filename.endswith(
|
||||
PIR_INFER_MODEL_SUFFIX
|
||||
) and filename.startswith(model_name):
|
||||
parsing_names = filename[
|
||||
len(model_name) : -len(PIR_INFER_MODEL_SUFFIX) + 1
|
||||
].split('.')
|
||||
if len(parsing_names) == 3 and len(parsing_names[1]) > 0:
|
||||
func_name = parsing_names[1]
|
||||
model_file_path = os.path.join(model_path, filename)
|
||||
else:
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
program, trainable = _load_pir_program(model_file_path)
|
||||
program_holder_dict[func_name] = _PirProgramHolder(
|
||||
program, trainable
|
||||
)
|
||||
else:
|
||||
for _, _, file_names in os.walk(model_path):
|
||||
for name in file_names:
|
||||
if 'model' in name:
|
||||
model_file_path = os.path.join(model_path, name)
|
||||
method_name = name.strip('_')
|
||||
if method_name == 'model':
|
||||
method_name = 'forward'
|
||||
else:
|
||||
method_name.replace('model', '')
|
||||
program, trainable = _load_pir_program(model_file_path)
|
||||
program_holder_dict[method_name] = _PirProgramHolder(
|
||||
program, trainable
|
||||
)
|
||||
|
||||
return program_holder_dict
|
||||
|
||||
|
||||
def _construct_params_and_buffers(model_path, programs, params_filename=None):
|
||||
params_path = os.path.join(model_path, str(params_filename))
|
||||
|
||||
if params_filename is not None and not os.path.exists(params_path):
|
||||
# When saving XX, there is only '*.pdmodel'
|
||||
return {}
|
||||
else:
|
||||
var_dict = _load_pir_parameter_vars(
|
||||
model_path, programs['forward'], params_filename
|
||||
)
|
||||
model_name = params_filename[: -len(INFER_PARAMS_SUFFIX)]
|
||||
# Load every file that meets the requirements in the directory model_path.
|
||||
for file_name in os.listdir(model_path):
|
||||
if file_name.startswith(model_name) and file_name.endswith(
|
||||
INFER_PARAMS_SUFFIX
|
||||
):
|
||||
parsing_names = file_name[
|
||||
len(model_name) : -len(INFER_PARAMS_SUFFIX) + 1
|
||||
].split('.')
|
||||
if len(parsing_names) == 3 and len(parsing_names[1]) > 0:
|
||||
func_name = parsing_names[1]
|
||||
else:
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
|
||||
var_dict.update(
|
||||
_load_pir_parameter_vars(
|
||||
model_path, programs[func_name], file_name
|
||||
)
|
||||
)
|
||||
|
||||
return var_dict
|
||||
|
||||
|
||||
def _run_dygraph(instance, input, program_holder, method_name):
|
||||
# 1. prepare inputs, outputs, attrs
|
||||
input_tensors = []
|
||||
input_tensor_names = []
|
||||
|
||||
for i, value in enumerate(input):
|
||||
if not isinstance(value, (np.ndarray, core.eager.Tensor)):
|
||||
raise TypeError(
|
||||
f"The type of input in PirTranslatedLayer must be numpy array or Variable(Tensor), but received {type(value)}."
|
||||
)
|
||||
# NOTE: In order to unify the API, firstly convert the input to Tensor
|
||||
if isinstance(value, np.ndarray):
|
||||
tensor = core.eager.Tensor(
|
||||
value=value,
|
||||
name=program_holder.input_vars[i].name,
|
||||
persistable=False,
|
||||
place=framework._current_expected_place(),
|
||||
zero_copy=True,
|
||||
)
|
||||
else:
|
||||
tensor = value
|
||||
# NOTE: we changed var name here,
|
||||
# but it may be an important name set by user
|
||||
tensor.name = program_holder.input_vars[i].name
|
||||
input_tensor_names.append(tensor.name)
|
||||
input_tensors.append(tensor)
|
||||
|
||||
if instance._get_partial_program_layer(method_name) is None:
|
||||
persistable_tensors = []
|
||||
origin_persistable_var_name = [
|
||||
program_holder._suffix_varname_dict[var_name]
|
||||
for var_name in program_holder.persistable_names
|
||||
]
|
||||
for var_name in origin_persistable_var_name:
|
||||
dy_var_name = instance._persistable_var_name_dict[var_name]
|
||||
if dy_var_name in instance._parameters:
|
||||
persistable_tensors.append(instance._parameters[dy_var_name])
|
||||
elif dy_var_name in instance._buffers:
|
||||
persistable_tensors.append(instance._buffers[dy_var_name])
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The persistable variable {var_name} does not exist in current PirTranslatedLayer."
|
||||
)
|
||||
|
||||
from paddle.jit.dy2static.pir_partial_program import PartialProgramLayer
|
||||
|
||||
inputs = program_holder.input_vars
|
||||
outputs = program_holder.output_vars
|
||||
parameters = (persistable_tensors, program_holder.persistable_vars)
|
||||
|
||||
layer = PartialProgramLayer(
|
||||
program_holder.infer_program,
|
||||
inputs,
|
||||
outputs,
|
||||
parameters,
|
||||
)
|
||||
instance._set_partial_program_layer(method_name, layer)
|
||||
layer = instance._get_partial_program_layer(method_name)
|
||||
if instance._is_test:
|
||||
layer.training = False
|
||||
else:
|
||||
if not program_holder.support_train:
|
||||
raise ValueError(
|
||||
"The model is not trainable, please check model_file of jit.save."
|
||||
)
|
||||
else:
|
||||
layer.training = True
|
||||
|
||||
return layer(input_tensors)
|
||||
|
||||
|
||||
def _run_static_graph(inputs, program_holder, src_program):
|
||||
'''
|
||||
This function is used when the pirTranslatedLayer is
|
||||
applied for dy_to_static conversion.
|
||||
'''
|
||||
dst_program = paddle.static.default_main_program()
|
||||
value_map = paddle.pir.IrMapping()
|
||||
# Establish a mapping relationship between existing parameters
|
||||
# and corresponding parameters in the program to be copied
|
||||
len_dst_op = len(dst_program.global_block().ops)
|
||||
for dst_op in dst_program.global_block().ops:
|
||||
if dst_op.name() == "builtin.parameter":
|
||||
for src_op in src_program.global_block().ops[:len_dst_op]:
|
||||
if (
|
||||
src_op.name() == dst_op.name()
|
||||
and src_op.result(0).name == dst_op.result(0).name
|
||||
):
|
||||
for i in range(src_op.num_results()):
|
||||
value_map.add(src_op.result(i), dst_op.result(i))
|
||||
# Establish a mapping relationship between truly inputs
|
||||
# and corresponding inputs in the program to be copied
|
||||
src_inputs = program_holder.input_vars
|
||||
if len(src_inputs) != len(inputs):
|
||||
raise ValueError(
|
||||
f"The number of input is invalid, expected {len(src_inputs)}, but received {len(inputs)}."
|
||||
)
|
||||
for src_input, input_ in zip(src_inputs, inputs):
|
||||
value_map.add(src_input, input_)
|
||||
|
||||
# find the insert point for copy
|
||||
current_insert_point = paddle.pir.get_current_insertion_point()
|
||||
current_block = current_insert_point.block()
|
||||
src_program.copy_to_block(value_map, current_block)
|
||||
|
||||
output = [value_map.look_up(v) for v in program_holder.output_vars]
|
||||
return output[0] if len(output) == 1 else output
|
||||
|
||||
|
||||
def _collect_current_and_parent_var(program, block_idx):
|
||||
'''
|
||||
Get variables in current block and its parent block.
|
||||
|
||||
Args:
|
||||
program(Program): The program containing the current block.
|
||||
block_idx(int): index of current block.
|
||||
|
||||
Returns:
|
||||
List: list of variables.
|
||||
'''
|
||||
vars = []
|
||||
if block_idx < 0:
|
||||
return vars
|
||||
for var in program.block(block_idx).vars:
|
||||
vars.append(var)
|
||||
parent_idx = program.block(block_idx).parent_idx
|
||||
if parent_idx > -1:
|
||||
vars += _collect_current_and_parent_var(program, parent_idx)
|
||||
return vars
|
||||
|
||||
|
||||
class PirTranslatedLayer(layers.Layer):
|
||||
"""
|
||||
PirTranslatedLayer is a ``paddle.nn.Layer`` for holding the model
|
||||
loaded by :ref:`api_paddle_jit_load` . It can be used like a
|
||||
general Layer object in eval or train mode.
|
||||
|
||||
.. note:
|
||||
The PirTranslatedLayer objects should not be created by constructor, it only can be loaded and constructed by :ref:`api_paddle_jit_load` .
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> import paddle.nn as nn
|
||||
>>> import paddle.optimizer as opt
|
||||
|
||||
>>> BATCH_SIZE = 16
|
||||
>>> BATCH_NUM = 4
|
||||
>>> EPOCH_NUM = 4
|
||||
|
||||
>>> IMAGE_SIZE = 784
|
||||
>>> CLASS_NUM = 10
|
||||
|
||||
>>> # define a random dataset
|
||||
>>> class RandomDataset(paddle.io.Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([IMAGE_SIZE]).astype('float32')
|
||||
... label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> class LinearNet(nn.Layer):
|
||||
... def __init__(self):
|
||||
... super().__init__()
|
||||
... self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
|
||||
...
|
||||
... @paddle.jit.to_static
|
||||
... def forward(self, x):
|
||||
... return self._linear(x)
|
||||
>>> def train(layer, loader, loss_fn, opt):
|
||||
... for epoch_id in range(EPOCH_NUM):
|
||||
... for batch_id, (image, label) in enumerate(loader()):
|
||||
... out = layer(image)
|
||||
... loss = loss_fn(out, label)
|
||||
... loss.backward()
|
||||
... opt.step()
|
||||
... opt.clear_grad()
|
||||
... print("Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))
|
||||
>>> # 1. train & save model.
|
||||
>>> # create network
|
||||
>>> layer = LinearNet()
|
||||
>>> loss_fn = nn.CrossEntropyLoss()
|
||||
>>> adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())
|
||||
|
||||
>>> # create data loader
|
||||
>>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
|
||||
>>> loader = paddle.io.DataLoader(
|
||||
... dataset,
|
||||
... batch_size=BATCH_SIZE,
|
||||
... shuffle=True,
|
||||
... drop_last=True,
|
||||
... num_workers=2,
|
||||
... )
|
||||
>>> # train
|
||||
>>> train(layer, loader, loss_fn, adam)
|
||||
|
||||
>>> # save
|
||||
>>> model_path = "linear.example.model"
|
||||
>>> paddle.jit.save(layer, model_path)
|
||||
|
||||
>>> # 2. load model as PirTranslatedLayer
|
||||
>>> # load
|
||||
>>> translated_layer = paddle.jit.load(model_path)
|
||||
|
||||
>>> # inference
|
||||
>>> translated_layer.eval()
|
||||
>>> x = paddle.randn([1, IMAGE_SIZE], 'float32')
|
||||
>>> pred = translated_layer(x)
|
||||
|
||||
>>> # fine-tune
|
||||
>>> translated_layer.train()
|
||||
>>> adam = opt.Adam(learning_rate=0.001, parameters=translated_layer.parameters())
|
||||
>>> train(translated_layer, loader, loss_fn, adam)
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
programs: dict[str, paddle.static.Program],
|
||||
persistable_vars: dict[str, paddle.Tensor],
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if not isinstance(programs, dict):
|
||||
raise TypeError(
|
||||
"PirTranslatedLayer need to use _ProgramHolder's dict for initialization."
|
||||
)
|
||||
if not isinstance(persistable_vars, dict):
|
||||
raise TypeError(
|
||||
"PirTranslatedLayer need to use persistable variable dict for initialization."
|
||||
)
|
||||
|
||||
self._program_holder_dict = programs
|
||||
|
||||
# NOTE(chenweihang): [ why not use var name directly? ]
|
||||
# When add parameter or buffer to Layer by follow apis,
|
||||
# the variable name can't contain `.`, because which may cause
|
||||
# AttributeError when access the newly added parameter or buffer
|
||||
# in the form of `self.**.**``, but the EagerParamBase or BarBase
|
||||
# name contains `.` originally, such as `linear_0.w_0`, so here
|
||||
# need to generate new var name for each var
|
||||
self._persistable_var_name_dict = {}
|
||||
# the PirTranslatedLayer object held var names count started from 0
|
||||
with unique_name.guard():
|
||||
for name, var in persistable_vars.items():
|
||||
if isinstance(var, framework.EagerParamBase):
|
||||
dy_name = _generate_unique_var_name(PARAMETER_NAME_PREFIX)
|
||||
self._persistable_var_name_dict[name] = dy_name
|
||||
self.add_parameter(dy_name, var)
|
||||
elif isinstance(var, core.eager.Tensor):
|
||||
dy_name = _generate_unique_var_name(BUFFER_NAME_PREFIX)
|
||||
self._persistable_var_name_dict[name] = dy_name
|
||||
self.register_buffer(dy_name, var)
|
||||
else:
|
||||
raise TypeError(
|
||||
"Adding persistent variable which to layer is not supported now"
|
||||
)
|
||||
|
||||
self._is_test = True
|
||||
self._input_args_names = None
|
||||
self._partial_program_layers = {}
|
||||
|
||||
@staticmethod
|
||||
@framework.dygraph_only
|
||||
def _construct(model_path, configs=None):
|
||||
# 0. dir and filename check
|
||||
model_path = os.path.normpath(model_path)
|
||||
if not os.path.isdir(model_path):
|
||||
raise ValueError(f"There is no directory named '{model_path}'")
|
||||
model_filename = None
|
||||
params_filename = None
|
||||
if configs is not None:
|
||||
model_filename = configs.model_filename
|
||||
params_filename = configs.params_filename
|
||||
|
||||
# 1. load program desc & construct _ProgramHolder
|
||||
programs = _construct_program_holders(model_path, model_filename)
|
||||
|
||||
# 2. load layer parameters & buffers
|
||||
persistable_vars = _construct_params_and_buffers(
|
||||
model_path, programs, params_filename
|
||||
)
|
||||
|
||||
# 3. construct PirTranslatedLayer object
|
||||
translated_layer = PirTranslatedLayer(programs, persistable_vars)
|
||||
|
||||
# 4. create PirTranslatedLayer's execution method
|
||||
for method_name, program_holder in programs.items():
|
||||
if translated_layer._input_args_names is None:
|
||||
translated_layer._input_args_names = [
|
||||
ins.name for ins in program_holder.input_vars
|
||||
]
|
||||
setattr(
|
||||
PirTranslatedLayer,
|
||||
method_name,
|
||||
PirTranslatedLayer._execution_method_creator(
|
||||
method_name, program_holder
|
||||
),
|
||||
)
|
||||
|
||||
# 5. set PirTranslatedLayer's default mode to eval
|
||||
translated_layer.eval()
|
||||
|
||||
return translated_layer
|
||||
|
||||
@staticmethod
|
||||
def _execution_method_creator(method_name, program_holder):
|
||||
def __i_m_p_l__(self, *input):
|
||||
program_holder = self._program_holder_dict[__i_m_p_l__.__name__]
|
||||
# When using jit.save, it runs in static graph mode.
|
||||
# Run in dynamic graph mode when the model is inferring.
|
||||
if in_dynamic_mode():
|
||||
return _run_dygraph(self, input, program_holder, method_name)
|
||||
else:
|
||||
return _run_static_graph(
|
||||
input, program_holder, program_holder.infer_program
|
||||
)
|
||||
|
||||
__i_m_p_l__.__name__ = method_name
|
||||
return __i_m_p_l__
|
||||
|
||||
def train(self):
|
||||
self._is_test = False
|
||||
self.training = True
|
||||
|
||||
def eval(self):
|
||||
self._is_test = True
|
||||
self.training = False
|
||||
|
||||
def program(self, method_name='forward'):
|
||||
"""
|
||||
Gets translated program of specified method.
|
||||
|
||||
Args:
|
||||
- method_name (string): method name corresponding to the program
|
||||
to be obtained. Default: 'forward'.
|
||||
|
||||
Returns:
|
||||
Program
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> from paddle import nn
|
||||
>>> import paddle.optimizer as opt
|
||||
|
||||
>>> BATCH_SIZE = 16
|
||||
>>> BATCH_NUM = 4
|
||||
>>> EPOCH_NUM = 4
|
||||
|
||||
>>> IMAGE_SIZE = 784
|
||||
>>> CLASS_NUM = 10
|
||||
|
||||
>>> # define a random dataset
|
||||
>>> class RandomDataset(paddle.io.Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([IMAGE_SIZE]).astype('float32')
|
||||
... label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> class LinearNet(nn.Layer):
|
||||
... def __init__(self):
|
||||
... super().__init__()
|
||||
... self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
|
||||
...
|
||||
... @paddle.jit.to_static
|
||||
... def forward(self, x):
|
||||
... return self._linear(x)
|
||||
>>> def train(layer, loader, loss_fn, opt):
|
||||
... for epoch_id in range(EPOCH_NUM):
|
||||
... for batch_id, (image, label) in enumerate(loader()):
|
||||
... out = layer(image)
|
||||
... loss = loss_fn(out, label)
|
||||
... loss.backward()
|
||||
... opt.step()
|
||||
... opt.clear_grad()
|
||||
... print("Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))
|
||||
>>> # create network
|
||||
>>> layer = LinearNet()
|
||||
>>> loss_fn = nn.CrossEntropyLoss()
|
||||
>>> adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())
|
||||
>>> # create data loader
|
||||
>>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
|
||||
>>> loader = paddle.io.DataLoader(
|
||||
... dataset,
|
||||
... batch_size=BATCH_SIZE,
|
||||
... shuffle=True,
|
||||
... drop_last=True,
|
||||
... num_workers=2,
|
||||
... )
|
||||
>>> # train
|
||||
>>> train(layer, loader, loss_fn, adam)
|
||||
|
||||
>>> # save
|
||||
>>> model_path = "linear.example.model"
|
||||
>>> paddle.jit.save(layer, model_path)
|
||||
|
||||
>>> # load
|
||||
>>> translated_layer = paddle.jit.load(model_path)
|
||||
|
||||
>>> # get program
|
||||
>>> program = translated_layer.program()
|
||||
"""
|
||||
# 1. get program holder
|
||||
program_holder = self._get_program_holder(method_name)
|
||||
|
||||
# 2. get inference program desc
|
||||
program = program_holder.infer_program
|
||||
|
||||
return program
|
||||
|
||||
def _get_program_holder(self, method_name='forward'):
|
||||
program_holder = self._program_holder_dict.get(method_name, None)
|
||||
if program_holder is None:
|
||||
raise ValueError(
|
||||
f"The method `{method_name}` does not exist in loaded PirTranslatedLayer."
|
||||
)
|
||||
return program_holder
|
||||
|
||||
def _input_spec(self, method_name='forward'):
|
||||
# 1. get program holder
|
||||
program_holder = self._get_program_holder(method_name)
|
||||
|
||||
# 2. build input spec by input desc
|
||||
input_spec = []
|
||||
for var in program_holder.input_vars:
|
||||
spec = paddle.static.InputSpec(
|
||||
shape=var.shape,
|
||||
dtype=var.dtype,
|
||||
name=var.name,
|
||||
)
|
||||
input_spec.append(spec)
|
||||
|
||||
return input_spec
|
||||
|
||||
def _output_spec(self, method_name='forward'):
|
||||
# 1. get program holder
|
||||
program_holder = self._get_program_holder(method_name)
|
||||
|
||||
# 2. build output spec by output desc
|
||||
output_spec = []
|
||||
for var in program_holder.output_vars:
|
||||
# NOTE(chenweihang): InputSpec describes a tensor, not just input.
|
||||
# Maybe the name is not good enough. Here we use InputSpec to
|
||||
# construct the description of Output tensor
|
||||
spec = paddle.static.InputSpec(
|
||||
shape=var.shape,
|
||||
dtype=var.dtype,
|
||||
name=var.name,
|
||||
)
|
||||
output_spec.append(spec)
|
||||
|
||||
return output_spec
|
||||
|
||||
def _get_partial_program_layer(self, method_name):
|
||||
return self._partial_program_layers.get(method_name, None)
|
||||
|
||||
def _set_partial_program_layer(self, method_name, layer):
|
||||
self._partial_program_layers[method_name] = layer
|
||||
@@ -0,0 +1,97 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import wraps
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
from paddle.framework import core
|
||||
|
||||
from .dy2static.utils import ENV_SOT_EVENT_LEVEL
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
P = ParamSpec("P")
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class SotProfiler:
|
||||
def __enter__(self):
|
||||
self.enable()
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.disable()
|
||||
|
||||
def enable(self, tag=None):
|
||||
core.nvprof_start()
|
||||
core.nvprof_enable_record_event()
|
||||
|
||||
def disable(self):
|
||||
core.nvprof_stop()
|
||||
|
||||
|
||||
class EventGuard:
|
||||
def __init__(self, event_name, event_level=1):
|
||||
self.event_name = event_name
|
||||
self.event_level = event_level
|
||||
self.need_pop = False
|
||||
|
||||
def __call__(self, fn):
|
||||
if ENV_SOT_EVENT_LEVEL.get() < self.event_level:
|
||||
return fn
|
||||
|
||||
@wraps(fn)
|
||||
def wrapped(*args, **kwargs):
|
||||
with self:
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return wrapped
|
||||
|
||||
def __enter__(self):
|
||||
if ENV_SOT_EVENT_LEVEL.get() >= self.event_level:
|
||||
core.nvprof_nvtx_push(self.event_name)
|
||||
self.need_pop = True
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
if self.need_pop:
|
||||
core.nvprof_nvtx_pop()
|
||||
|
||||
|
||||
def event_register(
|
||||
event_name_formatter: Callable[P, str] | str, event_level=1
|
||||
) -> Callable[[Callable[P, T]], Callable[P, T]]:
|
||||
def event_wrapper(func: Callable[P, T]) -> Callable[P, T]:
|
||||
@wraps(func)
|
||||
def call_with_event(*args: P.args, **kwargs: P.kwargs):
|
||||
event_name = (
|
||||
event_name_formatter(*args, **kwargs)
|
||||
if callable(event_name_formatter)
|
||||
else event_name_formatter
|
||||
)
|
||||
with EventGuard(event_name, event_level=event_level):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return call_with_event
|
||||
|
||||
def do_nothing(func: Callable[P, T]) -> Callable[P, T]:
|
||||
return func
|
||||
|
||||
if ENV_SOT_EVENT_LEVEL.get() >= event_level:
|
||||
return event_wrapper
|
||||
else:
|
||||
return do_nothing
|
||||
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from . import (
|
||||
profiler as profiler,
|
||||
psdb, # noqa: F401
|
||||
)
|
||||
from .opcode_translator.breakpoint import ( # noqa: F401
|
||||
BM,
|
||||
add_breakpoint,
|
||||
add_event,
|
||||
)
|
||||
from .translate import symbolic_translate # noqa: F401
|
||||
@@ -0,0 +1,656 @@
|
||||
# Copyright (c) 2023 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.
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from contextlib import nullcontext
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
import paddle
|
||||
from paddle.base.data_feeder import convert_dtype
|
||||
from paddle.base.framework import convert_nptype_to_datatype_or_vartype
|
||||
from paddle.base.unique_name import (
|
||||
UniqueNameGenerator,
|
||||
guard as UniqueNameGuard,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.placement_type import (
|
||||
get_shard_spec,
|
||||
to_placements,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.dist_input_spec import (
|
||||
DistributedInputSpec,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.utils import (
|
||||
convert_to_dims_mapping,
|
||||
)
|
||||
from paddle.jit.dy2static.utils import (
|
||||
ALREADY_D2S,
|
||||
extract_tensor_dynamic_dims,
|
||||
graph_tracing_guard,
|
||||
)
|
||||
from paddle.pir import is_fake_value
|
||||
from paddle.static import InputSpec
|
||||
from paddle.utils import flatten, is_sequence
|
||||
|
||||
from .symbolic_shape.symbolic_value import SymbolicInt
|
||||
from .utils import (
|
||||
Cache,
|
||||
Singleton,
|
||||
get_min_non_specialized_number,
|
||||
map_if_extend,
|
||||
meta_str,
|
||||
)
|
||||
from .utils.exceptions import BreakGraphError, NullMetaBreak
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import numpy.typing as npt
|
||||
|
||||
DynamicSymbolT = TypeVar("DynamicSymbolT")
|
||||
SOT_INFER_META_INNER_VAR = "___SOT_INFER_META_INNER_VAR"
|
||||
|
||||
|
||||
class DistInfo:
|
||||
def __init__(self, mesh=None, dims_mapping=None, local_shape=None):
|
||||
self.mesh = mesh
|
||||
self.dims_mapping = dims_mapping
|
||||
self.local_shape = local_shape
|
||||
|
||||
@staticmethod
|
||||
def from_tensor(tensor: paddle.Tensor) -> DistInfo:
|
||||
assert isinstance(tensor, paddle.Tensor) and tensor.is_dist(), (
|
||||
f"Expect a Tensor, but got a {type(tensor)}."
|
||||
)
|
||||
|
||||
mesh = tensor.process_mesh
|
||||
sharding_specs = get_shard_spec(
|
||||
mesh, tensor.placements, len(tensor.shape)
|
||||
)
|
||||
dims_mapping = convert_to_dims_mapping(sharding_specs, mesh)
|
||||
local_shape = tensor._local_value().shape
|
||||
return DistInfo(mesh, dims_mapping, local_shape)
|
||||
|
||||
@staticmethod
|
||||
def from_value(value: paddle.pir.Value) -> DistInfo:
|
||||
assert isinstance(value, paddle.pir.Value) and value.is_dist(), (
|
||||
f"Expect a Value, but got a {type(value)}."
|
||||
)
|
||||
return DistInfo(
|
||||
value.dist_attr().process_mesh,
|
||||
value.dist_attr().dims_mapping,
|
||||
value._local_shape,
|
||||
)
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
return DistInfo(
|
||||
mesh=copy.deepcopy(self.mesh),
|
||||
dims_mapping=copy.deepcopy(self.dims_mapping),
|
||||
local_shape=copy.deepcopy(self.local_shape),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"DistInfo(mesh={self.mesh}, dims_mapping={self.dims_mapping}, local_shape={self.local_shape})"
|
||||
|
||||
|
||||
class MetaInfoOrNull:
|
||||
def __init__(self, meta: MetaInfo | None):
|
||||
self.meta = meta
|
||||
|
||||
@staticmethod
|
||||
def null():
|
||||
return MetaInfoOrNull(None)
|
||||
|
||||
def is_null(self):
|
||||
return self.meta is None
|
||||
|
||||
def unwrap_or_breakgraph(self):
|
||||
if self.meta is None:
|
||||
raise BreakGraphError(NullMetaBreak())
|
||||
return self.meta
|
||||
|
||||
def unwrap_unsafe(self):
|
||||
assert self.meta is not None, "MetaInfo is None"
|
||||
return self.meta
|
||||
|
||||
def with_dynamic_axes(
|
||||
self, name: str, dynamic_axes: list[int]
|
||||
) -> MetaInfoOrNull:
|
||||
if self.meta is None:
|
||||
return MetaInfoOrNull.null()
|
||||
return self.meta.with_dynamic_axes(name, dynamic_axes).wrap()
|
||||
|
||||
def to_input_spec(self):
|
||||
if self.meta is None:
|
||||
return None
|
||||
return self.meta.to_input_spec()
|
||||
|
||||
def guard_str(self):
|
||||
if self.meta is None:
|
||||
return "(Null)"
|
||||
return self.meta.guard_str()
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
if self.meta is None:
|
||||
return MetaInfoOrNull(None)
|
||||
return MetaInfoOrNull(copy.deepcopy(self.meta))
|
||||
|
||||
@staticmethod
|
||||
def mix_axes(axes1: list[int], axes2: list[int]) -> list[int]:
|
||||
return sorted(set(axes1 + axes2))
|
||||
|
||||
@staticmethod
|
||||
def from_tensor(
|
||||
tensor: paddle.Tensor, *, dynamic_axes: list[int] | None = None
|
||||
) -> MetaInfoOrNull:
|
||||
if not tensor._is_dense_tensor_hold_allocation():
|
||||
return MetaInfoOrNull.null()
|
||||
assert isinstance(tensor, paddle.Tensor), (
|
||||
"Expect a Tensor, but got a Value."
|
||||
)
|
||||
|
||||
assert -1 not in tensor.shape, (
|
||||
"Tensor shape should not contain -1, maybe you pass a Value to from_tensor"
|
||||
)
|
||||
user_specified_dynamic_axes = extract_tensor_dynamic_dims(tensor)
|
||||
dynamic_axes = dynamic_axes or []
|
||||
dynamic_axes = MetaInfoOrNull.mix_axes(
|
||||
dynamic_axes, list(user_specified_dynamic_axes)
|
||||
)
|
||||
shape = [
|
||||
SymbolicInt(dim) if i in dynamic_axes else dim
|
||||
for i, dim in enumerate(tensor.shape)
|
||||
]
|
||||
if tensor.is_dist():
|
||||
dist_info = DistInfo.from_tensor(tensor)
|
||||
else:
|
||||
dist_info = None
|
||||
return MetaInfo(
|
||||
shape,
|
||||
tensor.dtype,
|
||||
tensor.stop_gradient,
|
||||
tensor.name,
|
||||
tensor.persistable,
|
||||
tensor.type,
|
||||
tensor.place,
|
||||
None,
|
||||
dist_info=dist_info,
|
||||
).wrap()
|
||||
|
||||
@staticmethod
|
||||
def from_numpy(
|
||||
nparray: npt.NDArray[Any], *, dynamic_axes: list[int] | None = None
|
||||
) -> MetaInfoOrNull:
|
||||
dtype = convert_nptype_to_datatype_or_vartype(nparray.dtype)
|
||||
dynamic_axes = dynamic_axes or []
|
||||
shape = [
|
||||
SymbolicInt() if i in dynamic_axes else dim
|
||||
for i, dim in enumerate(nparray.shape)
|
||||
]
|
||||
return MetaInfo(
|
||||
shape,
|
||||
dtype,
|
||||
True, # stop_gradient
|
||||
None,
|
||||
None, # persistable
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
dist_info=None,
|
||||
).wrap()
|
||||
|
||||
@staticmethod
|
||||
def from_value(value) -> MetaInfoOrNull:
|
||||
if is_fake_value(value):
|
||||
return MetaInfoOrNull.null()
|
||||
name = SOT_INFER_META_INNER_VAR
|
||||
shape = [SymbolicInt() if dim == -1 else dim for dim in value.shape]
|
||||
for dim in shape:
|
||||
if isinstance(dim, int):
|
||||
assert dim >= 0, (
|
||||
"Dimensions must be non-negative integers or SymbolicInt. "
|
||||
f"Encountered value {dim} in shape {shape}."
|
||||
)
|
||||
|
||||
if isinstance(value, paddle.pir.Value) and value.is_dist():
|
||||
dist_info = DistInfo.from_value(value)
|
||||
else:
|
||||
dist_info = None
|
||||
return MetaInfo(
|
||||
shape,
|
||||
value.dtype,
|
||||
value.stop_gradient,
|
||||
name,
|
||||
value.persistable,
|
||||
None, # type is not a unified attribute in dygraph and static mode.
|
||||
None, # We can't infer the right place in compile time.
|
||||
None, # there's no spec_name specified when from_value.
|
||||
dist_info=dist_info,
|
||||
).wrap()
|
||||
|
||||
def __repr__(self):
|
||||
if self.meta is None:
|
||||
return "MetaInfoOrNull(None)"
|
||||
return f"MetaInfoOrNull({self.meta})"
|
||||
|
||||
def __eq__(self, other):
|
||||
if self.meta is None:
|
||||
return other.meta is None
|
||||
if other.meta is None:
|
||||
return False
|
||||
return self.meta == other.meta
|
||||
|
||||
def __hash__(self):
|
||||
if self.meta is None:
|
||||
return hash(None)
|
||||
return hash(self.meta)
|
||||
|
||||
|
||||
class MetaInfo:
|
||||
shape: list[int | SymbolicInt]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
shape,
|
||||
dtype,
|
||||
stop_gradient,
|
||||
name,
|
||||
persistable,
|
||||
type,
|
||||
place,
|
||||
spec_name=None,
|
||||
dist_info=None,
|
||||
):
|
||||
assert -1 not in shape, (
|
||||
"NOTE: Shape should not contain -1, consider convert it to SymbolicInt."
|
||||
)
|
||||
self.name = name
|
||||
self.persistable = persistable
|
||||
self.type = type
|
||||
self.place = place
|
||||
self.shape = shape
|
||||
self.dtype = dtype
|
||||
self.stop_gradient = stop_gradient
|
||||
self.dist_info = dist_info
|
||||
self.spec_name = spec_name
|
||||
|
||||
def wrap(self):
|
||||
return MetaInfoOrNull(self)
|
||||
|
||||
def shape_with_special_symbol(
|
||||
self, dynamic_symbol: DynamicSymbolT = -1
|
||||
) -> list[int | DynamicSymbolT]:
|
||||
return [
|
||||
dynamic_symbol if isinstance(dim, SymbolicInt) else dim
|
||||
for dim in self.shape
|
||||
]
|
||||
|
||||
def with_dynamic_axes(self, name: str, dynamic_axes: list[int]) -> MetaInfo:
|
||||
mixed_dynamic_axes = MetaInfoOrNull.mix_axes(
|
||||
self.dynamic_axes, dynamic_axes
|
||||
)
|
||||
# NOTE(SigureMo): Make sure create a new shape list with dynamic axes.
|
||||
# We will create a new shape list variable lazily in the future.
|
||||
shape = [
|
||||
(
|
||||
SymbolicInt(dim)
|
||||
if (
|
||||
i in mixed_dynamic_axes and not isinstance(dim, SymbolicInt)
|
||||
)
|
||||
else dim
|
||||
)
|
||||
for i, dim in enumerate(self.shape)
|
||||
]
|
||||
return MetaInfo(
|
||||
shape,
|
||||
self.dtype,
|
||||
self.stop_gradient,
|
||||
self.name,
|
||||
self.persistable,
|
||||
self.type,
|
||||
self.place,
|
||||
spec_name=name,
|
||||
dist_info=self.dist_info,
|
||||
)
|
||||
|
||||
@property
|
||||
def dynamic_axes(self):
|
||||
return [
|
||||
i
|
||||
for i, dim in enumerate(self.shape)
|
||||
if isinstance(dim, SymbolicInt)
|
||||
]
|
||||
|
||||
def is_inner_var(self):
|
||||
return self.name == SOT_INFER_META_INNER_VAR
|
||||
|
||||
def is_dynamic_shape(self):
|
||||
"""
|
||||
if SymbolicInt in shape, return True
|
||||
else: return False
|
||||
"""
|
||||
return len(self.dynamic_axes) > 0
|
||||
|
||||
def to_input_spec(self) -> DistributedInputSpec | ConstrainedInputSpec:
|
||||
shape = self.shape_with_special_symbol(None)
|
||||
if self.dist_info is not None:
|
||||
placements = to_placements(
|
||||
self.dist_info.dims_mapping, self.dist_info.mesh
|
||||
)
|
||||
return DistributedInputSpec(
|
||||
shape,
|
||||
dtype=self.dtype,
|
||||
stop_gradient=self.stop_gradient,
|
||||
mesh=self.dist_info.mesh,
|
||||
placements=placements,
|
||||
local_shape=self.dist_info.local_shape,
|
||||
)
|
||||
else:
|
||||
return ConstrainedInputSpec(
|
||||
self.dynamic_axes,
|
||||
shape,
|
||||
dtype=self.dtype,
|
||||
name=self.spec_name,
|
||||
stop_gradient=self.stop_gradient,
|
||||
)
|
||||
|
||||
def guard_str(self):
|
||||
shape = self.shape_with_special_symbol(SymbolicInt())
|
||||
return f"({shape}, {self.dtype}, {self.stop_gradient})"
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
return MetaInfo(
|
||||
list(self.shape),
|
||||
self.dtype,
|
||||
self.stop_gradient,
|
||||
self.name,
|
||||
self.persistable,
|
||||
self.type,
|
||||
self.place,
|
||||
self.spec_name,
|
||||
dist_info=copy.deepcopy(self.dist_info),
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return meta_str(self.shape, self.dtype, self.stop_gradient)
|
||||
|
||||
def __eq__(self, meta):
|
||||
return (
|
||||
self.shape == meta.shape
|
||||
and self.dtype == meta.dtype
|
||||
and self.stop_gradient == meta.stop_gradient
|
||||
)
|
||||
|
||||
def __hash__(self):
|
||||
return hash((tuple(self.shape), self.dtype, self.stop_gradient))
|
||||
|
||||
|
||||
class VariableCreator(metaclass=Singleton):
|
||||
"""
|
||||
We use the static graph Variable to infer the meta information of Tensor.
|
||||
This singleton class is used to create Variable for infer meta.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.var_name_generator = UniqueNameGenerator(SOT_INFER_META_INNER_VAR)
|
||||
self.var_cache = {}
|
||||
self.main_program = paddle.static.Program()
|
||||
self.startup_program = paddle.static.Program()
|
||||
|
||||
def gen_name(self, meta_or_null: MetaInfoOrNull):
|
||||
if meta_or_null.is_null():
|
||||
return "null"
|
||||
meta = meta_or_null.unwrap_unsafe()
|
||||
name = f"{meta.dtype}_{meta.stop_gradient}_"
|
||||
name += "_".join(map(str, meta.shape))
|
||||
return name
|
||||
|
||||
def create_var(self, meta_or_null: MetaInfoOrNull):
|
||||
if meta_or_null.is_null():
|
||||
return None
|
||||
meta = meta_or_null.unwrap_unsafe()
|
||||
shape = meta.shape_with_special_symbol(-1)
|
||||
|
||||
with paddle.static.program_guard(
|
||||
self.main_program, self.startup_program
|
||||
):
|
||||
var = paddle.static.input.data(
|
||||
name=self.gen_name(meta.wrap()),
|
||||
shape=shape,
|
||||
dtype=convert_dtype(meta.dtype),
|
||||
)
|
||||
var.stop_gradient = meta.stop_gradient
|
||||
|
||||
if meta.dist_info is not None:
|
||||
mesh = meta.dist_info.mesh
|
||||
placements = to_placements(meta.dist_info.dims_mapping, mesh)
|
||||
var = paddle._pir_ops.shard_tensor(var, mesh, placements)
|
||||
var.stop_gradient = meta.stop_gradient
|
||||
assert not isinstance(var, paddle.Tensor), (
|
||||
"Expect a Variable, but got a Tensor."
|
||||
)
|
||||
return var
|
||||
|
||||
def get_variable(self, meta: MetaInfoOrNull, without_cache=False):
|
||||
var_feature_name = self.gen_name(meta)
|
||||
if without_cache:
|
||||
return self.create_var(meta)
|
||||
if var_feature_name not in self.var_cache:
|
||||
self.var_cache[var_feature_name] = self.create_var(meta)
|
||||
return self.var_cache[var_feature_name]
|
||||
|
||||
def infer_meta(self, func, *args, **kwargs):
|
||||
with (
|
||||
paddle.base.framework._dygraph_guard(None),
|
||||
UniqueNameGuard(self.var_name_generator),
|
||||
):
|
||||
if func is paddle.distributed.shard_tensor:
|
||||
args, kwargs = (
|
||||
convert_meta_to_variable(args, without_cache=True),
|
||||
convert_meta_to_variable(kwargs, without_cache=True),
|
||||
)
|
||||
else:
|
||||
args, kwargs = (
|
||||
convert_meta_to_variable(args),
|
||||
convert_meta_to_variable(kwargs),
|
||||
)
|
||||
|
||||
graph_tracing_context_manager = nullcontext()
|
||||
with paddle.static.program_guard(
|
||||
self.main_program, self.startup_program
|
||||
):
|
||||
if isinstance(func, str):
|
||||
# TODO(Aurelius84): Is length of args always greater than 0?
|
||||
# Do we need add condition check here?
|
||||
func = getattr(args[0], func)
|
||||
args = args[1:]
|
||||
if hasattr(func, ALREADY_D2S):
|
||||
graph_tracing_context_manager = graph_tracing_guard(
|
||||
self.main_program
|
||||
)
|
||||
with graph_tracing_context_manager:
|
||||
out = func(*args, **kwargs)
|
||||
return convert_variable_to_meta_info(out)
|
||||
|
||||
|
||||
def convert_meta_to_variable(args, without_cache=False):
|
||||
return map_if_extend(
|
||||
args,
|
||||
pred=lambda x: isinstance(x, MetaInfoOrNull),
|
||||
true_fn=lambda x: VariableCreator().get_variable(
|
||||
x, without_cache=without_cache
|
||||
),
|
||||
false_fn=lambda x: x,
|
||||
)
|
||||
|
||||
|
||||
def convert_meta_to_input_spec(args):
|
||||
return map_if_extend(
|
||||
args,
|
||||
pred=lambda x: isinstance(x, MetaInfoOrNull),
|
||||
true_fn=lambda x: x.to_input_spec(),
|
||||
# TODO(xiongkun): can x be tensor ?
|
||||
false_fn=lambda x: (
|
||||
paddle.static.InputSpec.from_tensor(x)
|
||||
if isinstance(x, paddle.Tensor)
|
||||
else x
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def convert_variable_to_meta_info(args):
|
||||
return map_if_extend(
|
||||
args,
|
||||
pred=lambda x: isinstance(x, paddle.pir.Value),
|
||||
true_fn=lambda x: MetaInfoOrNull.from_value(x),
|
||||
false_fn=lambda x: x,
|
||||
)
|
||||
|
||||
|
||||
def infer_meta(func, *args, **kwargs):
|
||||
fn = SpecialInferMeta().get_infermeta_fn(func)
|
||||
if fn:
|
||||
return fn(*args, **kwargs)
|
||||
return VariableCreator().infer_meta(func, *args, **kwargs)
|
||||
|
||||
|
||||
def infer_meta_for_layer(layer, *args, **kwargs):
|
||||
assert isinstance(layer, paddle.nn.Layer), (
|
||||
f"Expect a Layer, but got {layer}."
|
||||
)
|
||||
layer = paddle.jit.to_static(layer, full_graph=True)
|
||||
|
||||
args_, kwargs_ = convert_meta_to_input_spec((args, kwargs))
|
||||
|
||||
(
|
||||
concrete_program,
|
||||
partial_program_layer,
|
||||
) = layer.forward.get_concrete_program(*args_, **kwargs_)
|
||||
|
||||
output_values = partial_program_layer._outputs.var_list
|
||||
|
||||
out = partial_program_layer._restore_out(
|
||||
[
|
||||
x
|
||||
for x in paddle.utils.flatten(
|
||||
convert_variable_to_meta_info(output_values)
|
||||
)
|
||||
if isinstance(x, MetaInfoOrNull)
|
||||
]
|
||||
)
|
||||
layer.forward.rollback()
|
||||
return out
|
||||
|
||||
|
||||
def ast_infer_meta(static_function, *args, **kwargs):
|
||||
args_, kwargs_ = convert_meta_to_input_spec((args, kwargs))
|
||||
|
||||
(
|
||||
concrete_program,
|
||||
partial_program_layer,
|
||||
) = static_function.get_concrete_program(*args_, **kwargs_)
|
||||
|
||||
out = partial_program_layer._restore_out(
|
||||
[
|
||||
x
|
||||
for x in paddle.utils.flatten(
|
||||
convert_variable_to_meta_info(concrete_program.outputs)
|
||||
)
|
||||
if isinstance(x, MetaInfoOrNull)
|
||||
]
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class SpecialInferMeta(metaclass=Singleton):
|
||||
"""
|
||||
There are some functions that cannot be inferred directly through static graph,
|
||||
and need to be implemented manually. This class is used to implement infer meta
|
||||
for these functions.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def get_infermeta_fn(self, fn):
|
||||
try:
|
||||
funcname = fn.__name__
|
||||
return getattr(self, f"infermeta_{funcname}")
|
||||
except:
|
||||
pass
|
||||
return None
|
||||
|
||||
def infermeta_grad(
|
||||
self,
|
||||
outputs,
|
||||
inputs,
|
||||
grad_outputs=None,
|
||||
retain_graph=None,
|
||||
create_graph=False,
|
||||
only_inputs=True,
|
||||
allow_unused=False,
|
||||
no_grad_vars=None,
|
||||
):
|
||||
if not is_sequence(inputs):
|
||||
inputs = [inputs]
|
||||
return inputs
|
||||
|
||||
|
||||
class InferMetaCache(Cache, metaclass=Singleton):
|
||||
def __init__(self):
|
||||
super().__init__(copy=True)
|
||||
|
||||
def key_fn(
|
||||
self, func, *args, **kwargs
|
||||
): # args & kwargs have transformed to MetaInfo
|
||||
return (
|
||||
func,
|
||||
tuple(flatten(args)),
|
||||
tuple(kwargs.keys()),
|
||||
tuple(flatten(kwargs)),
|
||||
)
|
||||
|
||||
def value_fn(self, func, *args, **kwargs):
|
||||
return infer_meta(func, *args, **kwargs)
|
||||
|
||||
|
||||
class LayerInferMetaCache(Cache, metaclass=Singleton):
|
||||
def __init__(self):
|
||||
super().__init__(copy=True)
|
||||
|
||||
def key_fn(self, layer, *args, **kwargs):
|
||||
params = [
|
||||
MetaInfoOrNull.from_tensor(x)
|
||||
for x in layer.parameters(include_sublayers=True)
|
||||
]
|
||||
return (
|
||||
layer,
|
||||
tuple(params),
|
||||
tuple(flatten(args)),
|
||||
tuple(kwargs.keys()),
|
||||
tuple(flatten(kwargs)),
|
||||
)
|
||||
|
||||
def value_fn(self, layer, *args, **kwargs):
|
||||
return infer_meta_for_layer(layer, *args, **kwargs)
|
||||
|
||||
|
||||
class ConstrainedInputSpec(InputSpec):
|
||||
def __init__(self, dynamic_axes: list[int], *args, **kwargs):
|
||||
self.ranges: list[
|
||||
tuple[int, int | None, int | None]
|
||||
] = [] # (idx of dim, min, max)
|
||||
super().__init__(*args, **kwargs)
|
||||
min_non_specialized_number = get_min_non_specialized_number()
|
||||
for i in dynamic_axes:
|
||||
self.ranges.append((i, min_non_specialized_number, None))
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from .eval_frame_callback import eval_frame_callback # noqa: F401
|
||||
from .skip_files import setup_skip_files
|
||||
|
||||
setup_skip_files()
|
||||
@@ -0,0 +1,180 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
import inspect
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ..opcode_translator.instruction_utils import instrs_info
|
||||
from ..utils import Singleton, log
|
||||
from .executor.opcode_executor import OpcodeExecutorBase
|
||||
|
||||
# this file is a debug utils files for quick debug
|
||||
# >>> sot.add_breakpoint(file, line)
|
||||
# >>> sot.remove_breakpoint(file, line)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Breakpoint:
|
||||
file: str
|
||||
line: int
|
||||
co_name: str
|
||||
offset: int
|
||||
|
||||
def __hash__(self):
|
||||
return hash((self.file, self.line, self.co_name, self.offset))
|
||||
|
||||
|
||||
class BreakpointManager(metaclass=Singleton):
|
||||
def __init__(self):
|
||||
self.breakpoints = set()
|
||||
self.executors = OpcodeExecutorBase.call_stack
|
||||
self.activate = 0
|
||||
self.record_event = []
|
||||
|
||||
def clear_event(self, event):
|
||||
self.record_event.clear()
|
||||
|
||||
def add_event(self, event):
|
||||
"""
|
||||
event in ['All' ,'FallbackError', 'BreakGraphError', 'InnerError']
|
||||
"""
|
||||
self.record_event.append(event)
|
||||
|
||||
def add(self, file, line, co_name=None, offset=None):
|
||||
log(1, f"add breakpoint at {file}:{line}\n")
|
||||
self.breakpoints.add(Breakpoint(file, line, co_name, offset))
|
||||
|
||||
def addn(self, *lines):
|
||||
"""
|
||||
called inside a executor. add a list of line number in current file.
|
||||
"""
|
||||
if not isinstance(lines, (list, tuple)):
|
||||
lines = [lines]
|
||||
for line in lines:
|
||||
file = self.cur_exe.vframe.code.co_filename
|
||||
self.add(file, line)
|
||||
|
||||
def clear(self):
|
||||
self.breakpoints.clear()
|
||||
|
||||
def hit(self, file, line, co_name, offset):
|
||||
if Breakpoint(file, line, None, None) in self.breakpoints:
|
||||
return True
|
||||
if Breakpoint(file, line, co_name, offset) in self.breakpoints:
|
||||
return True
|
||||
return False
|
||||
|
||||
def locate(self, exe):
|
||||
for i, _e in enumerate(self.executors):
|
||||
if _e is exe:
|
||||
self.activate = i
|
||||
return
|
||||
raise RuntimeError("Not found executor.")
|
||||
|
||||
def up(self):
|
||||
if self.activate == 0:
|
||||
return
|
||||
self.activate -= 1
|
||||
print("current function is: ", self.cur_exe.vframe.code.co_name)
|
||||
|
||||
def down(self):
|
||||
if self.activate >= len(self.executors) - 1:
|
||||
return
|
||||
self.activate += 1
|
||||
print("current function is: ", self.cur_exe.vframe.code.co_name)
|
||||
|
||||
def opcode(self, cur_exe=None):
|
||||
if cur_exe is None:
|
||||
cur_exe = self.cur_exe
|
||||
instr = cur_exe._instructions[cur_exe.vframe.lasti - 1]
|
||||
message = f"[Translate {cur_exe}] (line {cur_exe._current_line:>3}) {instr.opname:<12} {instr.argval}, stack is {cur_exe._stack}\n"
|
||||
return message
|
||||
|
||||
def bt(self):
|
||||
"""
|
||||
display all inline calls: backtrace.
|
||||
"""
|
||||
for exe in self.executors:
|
||||
lines, _ = inspect.getsourcelines(exe.vframe.code)
|
||||
print(
|
||||
" "
|
||||
+ exe.vframe.code.co_filename
|
||||
+ f"({exe._current_line})"
|
||||
+ f"{exe.vframe.code.co_name}()"
|
||||
)
|
||||
print(f"-> {lines[0].strip()}")
|
||||
print(f"-> {self.opcode(exe)}")
|
||||
pass
|
||||
|
||||
def on_event(self, event):
|
||||
if "All" in self.record_event or event in self.record_event:
|
||||
print("event captured.")
|
||||
self.activate = len(self.executors) - 1
|
||||
breakpoint() # noqa: T100
|
||||
|
||||
def _dis_source_code(self):
|
||||
cur_exe = self.executors[self.activate]
|
||||
lines, start_line = inspect.getsourcelines(cur_exe.vframe.code)
|
||||
cur_line = cur_exe._current_line
|
||||
lines[cur_line - start_line + 1 : cur_line - start_line + 1] = (
|
||||
" ^^^^^ HERE \n"
|
||||
)
|
||||
print("\033[31mSource Code is: \033[0m")
|
||||
print("".join(lines))
|
||||
|
||||
def dis(self, range=5):
|
||||
"""
|
||||
display all instruction code and source code.
|
||||
"""
|
||||
print("displaying debug info...")
|
||||
cur_exe = self.cur_exe
|
||||
print(self._dis_source_code())
|
||||
|
||||
print(f"\n{cur_exe.vframe.code}")
|
||||
lasti = cur_exe.vframe.lasti
|
||||
instr_str = instrs_info(
|
||||
cur_exe._instructions, lasti - 1, range, want_str=True
|
||||
)
|
||||
print(instr_str)
|
||||
|
||||
@property
|
||||
def cur_exe(self):
|
||||
exe = self.executors[self.activate]
|
||||
return exe
|
||||
|
||||
def sir(self):
|
||||
"""
|
||||
display sir in a page.
|
||||
"""
|
||||
print("displaying sir...")
|
||||
self.cur_exe.print_sir()
|
||||
|
||||
def pe(self, e):
|
||||
"""
|
||||
print exception.
|
||||
"""
|
||||
lines = traceback.format_tb(e.__traceback__)
|
||||
print("".join(lines))
|
||||
|
||||
|
||||
def add_breakpoint(file, line, co_name=None, offset=None):
|
||||
BM.add(file, line, co_name, offset)
|
||||
|
||||
|
||||
def add_event(event):
|
||||
BM.add_event(event)
|
||||
|
||||
|
||||
BM = BreakpointManager()
|
||||
@@ -0,0 +1,25 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, NamedTuple
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from types import CodeType
|
||||
|
||||
|
||||
class CustomCode(NamedTuple):
|
||||
code: CodeType | None
|
||||
disable_eval_frame: bool
|
||||
@@ -0,0 +1,92 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dis
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ..profiler import EventGuard
|
||||
from ..utils import log_do, log_enabled, log_format
|
||||
from .executor.executor_cache import OpcodeExecutorCache
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .custom_code import CustomCode
|
||||
|
||||
|
||||
def print_locals(frame):
|
||||
local_key = [
|
||||
key for key in frame.f_locals.keys() if not key.startswith("__")
|
||||
]
|
||||
print(
|
||||
f"[eval_frame_callback] {frame.f_code.co_name} with locals {local_key}"
|
||||
)
|
||||
print(
|
||||
f"[eval_frame_callback] {' ' * len(frame.f_code.co_name)} with cellvars + freevars: {frame.f_code.co_cellvars + frame.f_code.co_freevars}"
|
||||
)
|
||||
|
||||
def convert_obj(obj):
|
||||
import paddle
|
||||
|
||||
if isinstance(obj, paddle.Tensor):
|
||||
return "Tensor(" + str(obj.shape) + ")"
|
||||
if isinstance(obj, list):
|
||||
return [convert_obj(i) for i in obj]
|
||||
return obj
|
||||
|
||||
for key in local_key:
|
||||
print(
|
||||
f"[eval_frame_callback] {' ' * len(frame.f_code.co_name)} {key} = {convert_obj(frame.f_locals[key])}"
|
||||
)
|
||||
|
||||
|
||||
def eval_frame_callback(frame, **kwargs) -> CustomCode:
|
||||
with EventGuard(
|
||||
f"eval_frame_callback: {frame.f_code.co_name}", event_level=2
|
||||
):
|
||||
# A quick way to check if the log is enabled
|
||||
need_log = log_enabled(1)
|
||||
if need_log:
|
||||
log_format(
|
||||
2,
|
||||
"[eval_frame_callback] start to translate: {}\n",
|
||||
frame.f_code,
|
||||
)
|
||||
log_do(4, lambda: print_locals(frame))
|
||||
|
||||
log_format(
|
||||
3,
|
||||
"[eval_frame_callback] OriginCode: {}\n",
|
||||
frame.f_code.co_name,
|
||||
)
|
||||
log_do(3, lambda: dis.dis(frame.f_code))
|
||||
|
||||
custom_code = OpcodeExecutorCache()(frame, **kwargs)
|
||||
|
||||
if need_log:
|
||||
if custom_code.code is None:
|
||||
log_format(
|
||||
3,
|
||||
"[eval_frame_callback] NewCode (same as origin code): {}\n",
|
||||
frame.f_code.co_name,
|
||||
)
|
||||
else:
|
||||
log_format(
|
||||
3,
|
||||
"[eval_frame_callback] NewCode: {}\n",
|
||||
custom_code.code.co_name,
|
||||
)
|
||||
log_do(3, lambda: dis.dis(custom_code.code))
|
||||
|
||||
return custom_code
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from . import variable_dispatch # noqa: F401
|
||||
@@ -0,0 +1,71 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
# This file stores the customized function that will be called by the dispatch mechanism.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ...utils import BreakGraphError, BreakGraphReasonBase, FallbackError
|
||||
|
||||
|
||||
def create_raise_break_graph_handler(reason: BreakGraphReasonBase):
|
||||
def raise_break_graph_fn(*args, **kwarg):
|
||||
raise BreakGraphError(reason)
|
||||
|
||||
return raise_break_graph_fn
|
||||
|
||||
|
||||
def raise_not_implement_fn(*args, **kwarg):
|
||||
raise FallbackError("raise by raise_break_graph_fn.")
|
||||
|
||||
|
||||
# just a function for operator.in
|
||||
def operator_in(left, right):
|
||||
return left in right
|
||||
|
||||
|
||||
def operator_not_in(left, right):
|
||||
return left not in right
|
||||
|
||||
|
||||
def operator_exception_match(left, right):
|
||||
pass
|
||||
|
||||
|
||||
def operator_BAD(left, right):
|
||||
pass
|
||||
|
||||
|
||||
def operator_is_none(val):
|
||||
pass
|
||||
|
||||
|
||||
def operator_is_not_none(val):
|
||||
pass
|
||||
|
||||
|
||||
def tensor_dim(x):
|
||||
pass
|
||||
|
||||
|
||||
def generator_send(x):
|
||||
pass
|
||||
|
||||
|
||||
def place_get_device_id():
|
||||
pass
|
||||
|
||||
|
||||
def place_get_device_type():
|
||||
pass
|
||||
@@ -0,0 +1,298 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import inspect
|
||||
import operator
|
||||
from functools import cached_property, reduce
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
from ...utils import InnerError, NameGenerator, hashable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
T = TypeVar("T")
|
||||
Args = tuple[T, ...]
|
||||
Kwargs = dict[str, T]
|
||||
|
||||
|
||||
def format_type(type_: type[Any] | tuple[type[Any], ...]) -> str:
|
||||
if not isinstance(type_, tuple):
|
||||
type_ = (type_,)
|
||||
return " | ".join([t.__name__ for t in type_])
|
||||
|
||||
|
||||
def format_param(param: Parameter) -> str:
|
||||
kind = param.kind
|
||||
if kind == inspect.Parameter.VAR_POSITIONAL:
|
||||
return f"*{format_type(param.type)}"
|
||||
elif kind == inspect.Parameter.VAR_KEYWORD:
|
||||
return f"**{format_type(param.type)}"
|
||||
else:
|
||||
return format_type(param.type)
|
||||
|
||||
|
||||
def convert_annotation_to_type(type_str: str) -> tuple[type[Any], ...]:
|
||||
"""
|
||||
Convert type annotation to runtime value. Because we are using :pep:`563`
|
||||
to use the future annotation syntax, we cannot use `get_type_hints <https://docs.python.org/3.8/library/typing.html#typing.get_type_hints>`_
|
||||
directly. Currently, only the builtins and variables namespaces are supported.
|
||||
|
||||
Returns:
|
||||
tuple: The converted type.
|
||||
"""
|
||||
|
||||
import builtins
|
||||
|
||||
from . import variables
|
||||
|
||||
type_str = type_str.strip()
|
||||
if type_str == "Any":
|
||||
type_str = "object"
|
||||
|
||||
if "|" in type_str:
|
||||
return reduce(
|
||||
operator.add, map(convert_annotation_to_type, type_str.split("|"))
|
||||
)
|
||||
|
||||
search_namespaces = [variables, builtins]
|
||||
for namespace in search_namespaces:
|
||||
if hasattr(namespace, type_str):
|
||||
return (getattr(namespace, type_str),)
|
||||
raise InnerError(f"Cannot find type {type_str} in {search_namespaces}")
|
||||
|
||||
|
||||
class Parameter:
|
||||
name_gen = NameGenerator("param_")
|
||||
annotation: str
|
||||
name: str
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
annotation: str,
|
||||
*,
|
||||
kind: inspect._ParameterKind = inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
name: str | None = None,
|
||||
default: Any = inspect._empty,
|
||||
):
|
||||
self.name = name if name is not None else Parameter.name_gen.next()
|
||||
self.annotation = annotation
|
||||
self.kind = kind
|
||||
self.default = default
|
||||
|
||||
def to_parameter(self) -> inspect.Parameter:
|
||||
return inspect.Parameter(
|
||||
self.name,
|
||||
kind=self.kind,
|
||||
annotation=self.annotation,
|
||||
default=copy.copy(self.default),
|
||||
)
|
||||
|
||||
@cached_property
|
||||
def type(self) -> tuple[type[Any], ...]:
|
||||
return convert_annotation_to_type(self.annotation)
|
||||
|
||||
def match_arg(self, arg: Any) -> bool:
|
||||
if self.kind == inspect.Parameter.VAR_POSITIONAL:
|
||||
is_tuple = isinstance(arg, tuple)
|
||||
return is_tuple and all(isinstance(a, self.type) for a in arg)
|
||||
elif self.kind == inspect.Parameter.VAR_KEYWORD:
|
||||
is_dict = isinstance(arg, dict)
|
||||
return is_dict and all(
|
||||
isinstance(a, self.type) for a in arg.values()
|
||||
)
|
||||
else:
|
||||
return isinstance(arg, self.type)
|
||||
|
||||
@staticmethod
|
||||
def from_str(annotation: str) -> Parameter:
|
||||
return Parameter(annotation)
|
||||
|
||||
@staticmethod
|
||||
def from_parameter(parameter: inspect.Parameter) -> Parameter:
|
||||
if parameter.annotation != parameter.empty and not isinstance(
|
||||
parameter.annotation, str
|
||||
):
|
||||
raise InnerError(
|
||||
f"Parameter {parameter} has annotation {parameter.annotation} "
|
||||
"which is not a string. Please add `from __future__ import annotations` "
|
||||
"to the top of your file."
|
||||
)
|
||||
annotation = (
|
||||
parameter.annotation
|
||||
if parameter.annotation != parameter.empty
|
||||
else "Any"
|
||||
)
|
||||
|
||||
return Parameter(
|
||||
annotation,
|
||||
kind=parameter.kind,
|
||||
name=parameter.name,
|
||||
default=parameter.default,
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
default_repr = f"= {self.default!r}"
|
||||
return f"Parameter({', '.join([self.annotation, default_repr])})"
|
||||
|
||||
|
||||
def optional(annotation: str, default: Any = None) -> Parameter:
|
||||
return Parameter(annotation, default=default)
|
||||
|
||||
|
||||
class Pattern:
|
||||
parameters: dict[str, Parameter]
|
||||
signature: inspect.Signature
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*parameters: Parameter,
|
||||
):
|
||||
self.parameters = {
|
||||
parameter.name: parameter for parameter in parameters
|
||||
}
|
||||
self.signature = inspect.Signature(
|
||||
[parameter.to_parameter() for parameter in self.parameters.values()]
|
||||
)
|
||||
|
||||
def match_inputs(self, /, *args: Any, **kwargs: Any) -> bool:
|
||||
"""
|
||||
Match the input parameters of the function.
|
||||
|
||||
Returns:
|
||||
bool: Whether the input parameters match the pattern.
|
||||
"""
|
||||
try:
|
||||
bound_args = self.signature.bind(*args, **kwargs)
|
||||
except TypeError:
|
||||
return False
|
||||
for arg_name, arg_value in bound_args.arguments.items():
|
||||
if arg_name not in self.parameters:
|
||||
continue
|
||||
if not self.parameters[arg_name].match_arg(arg_value):
|
||||
return False
|
||||
return True
|
||||
|
||||
def __repr__(self) -> str:
|
||||
types_repr = ", ".join(
|
||||
[format_param(param) for param in self.parameters.values()]
|
||||
)
|
||||
return f"Pattern({types_repr})"
|
||||
|
||||
|
||||
class Dispatcher:
|
||||
"""
|
||||
Used for pattern registration and distribution.
|
||||
|
||||
For more design ideas, refer to the `Builtin dispatcher <https://github.com/PaddlePaddle/PaddleSOT/blob/develop/docs/design/builtin-dispatcher.md>`_ for details.
|
||||
|
||||
Examples:
|
||||
|
||||
>>> def builtin_add(a: int, b: int) -> int: ...
|
||||
>>> Dispatcher.register(builtin_add, ("int", "int"), lambda a, b: a + b)
|
||||
>>> handler = Dispatcher.dispatch(builtin_add, 1, 2)
|
||||
>>> handler(1, 2)
|
||||
3
|
||||
"""
|
||||
|
||||
handlers: dict[
|
||||
Callable[..., Any], list[tuple[Pattern, Callable[..., Any]]]
|
||||
] = {}
|
||||
graph: Any = None
|
||||
|
||||
@classmethod
|
||||
def register(
|
||||
cls,
|
||||
fn: Callable[..., Any],
|
||||
parameters: tuple[str | Parameter, ...],
|
||||
handler: Callable[..., Any],
|
||||
):
|
||||
"""
|
||||
Registering function signature.
|
||||
|
||||
Args:
|
||||
fn: The function to be registered.
|
||||
parameters: The parameters of the function to be registered.
|
||||
handler: The handler function.
|
||||
"""
|
||||
_parameters = tuple(
|
||||
(
|
||||
Parameter.from_str(parameter)
|
||||
if isinstance(parameter, str)
|
||||
else parameter
|
||||
)
|
||||
for parameter in parameters
|
||||
)
|
||||
if fn not in cls.handlers:
|
||||
cls.handlers[fn] = []
|
||||
cls.handlers[fn].append((Pattern(*_parameters), handler))
|
||||
|
||||
@classmethod
|
||||
def register_decorator(cls, fn: Callable[..., Any]):
|
||||
"""
|
||||
Decorator mode of register, Used to register some complex functions.
|
||||
|
||||
Args:
|
||||
fn: The function to be registered.
|
||||
|
||||
Examples:
|
||||
>>> def builtin_add(a: int, b: int) -> int: ...
|
||||
>>> @Dispatcher.register_decorator(builtin_add)
|
||||
... def builtin_add_dispatcher(a: int, b: int) -> int:
|
||||
... return a + b
|
||||
>>> handler = Dispatcher.dispatch(builtin_add, 1, 2)
|
||||
>>> handler(1, 2)
|
||||
3
|
||||
"""
|
||||
|
||||
def decorator(handler: Callable[..., Any]):
|
||||
signature = inspect.signature(handler)
|
||||
parameters = tuple(
|
||||
Parameter.from_parameter(parameter)
|
||||
for parameter in signature.parameters.values()
|
||||
)
|
||||
cls.register(fn, parameters, handler)
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def call(cls, fn, *args, **kwargs):
|
||||
func = cls.dispatch(fn, *args, **kwargs)
|
||||
if func is None:
|
||||
raise InnerError(
|
||||
f"Cannot find handler for {fn} with args {args} and kwargs {kwargs}"
|
||||
)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def dispatch(
|
||||
cls, fn: Callable[..., Any], *args: Any, **kwargs: Any
|
||||
) -> Callable[..., Any] | None:
|
||||
"""
|
||||
Find the matching handler from the registered functions.
|
||||
|
||||
Args:
|
||||
fn: The function to be dispatched.
|
||||
args: The args of the function.
|
||||
kwargs: The kwargs of the function.
|
||||
"""
|
||||
if not hashable(fn) or fn not in cls.handlers:
|
||||
return None
|
||||
for pattern, handler in cls.handlers[fn]:
|
||||
if pattern.match_inputs(*args, **kwargs):
|
||||
return handler
|
||||
return None
|
||||
@@ -0,0 +1,129 @@
|
||||
# Copyright (c) 2025 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ...utils import InnerError
|
||||
from .variables import ConstantVariable, ExceptionVariable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .function_graph import FunctionGraph
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ExceptionStack:
|
||||
# This data structure manages exceptions as in CPython, primarily handling
|
||||
# the __context__ attribute of SotCapturedException.
|
||||
|
||||
_exception_stack: list[ExceptionVariable] = dataclasses.field(
|
||||
default_factory=list
|
||||
)
|
||||
_current_exception: ExceptionVariable | None = dataclasses.field(
|
||||
default=None
|
||||
)
|
||||
|
||||
def clear_current_exception(self):
|
||||
self._current_exception = None
|
||||
|
||||
def set_current_exception(
|
||||
self, val: ExceptionVariable, graph: FunctionGraph
|
||||
):
|
||||
self._set_context_and_break_context_reference_cycle(val, graph)
|
||||
self._current_exception = val
|
||||
|
||||
def move_current_exception_to_stack(self):
|
||||
self.push(self.get_current_exception())
|
||||
self.clear_current_exception()
|
||||
|
||||
def get_current_exception(self):
|
||||
if self._current_exception is None:
|
||||
raise InnerError("Current exception should not be None")
|
||||
return self._current_exception
|
||||
|
||||
def _set_context_and_break_context_reference_cycle(
|
||||
self, val: ExceptionVariable, graph: FunctionGraph
|
||||
):
|
||||
# set Exception.__context__
|
||||
self._set_context_recursive(val, len(self._exception_stack) - 1)
|
||||
self._break_context_reference_cycle(val, graph)
|
||||
|
||||
def _set_context_recursive(self, val: ExceptionVariable, prev_idx: int):
|
||||
# Recursively sets the __context__ attribute for ExceptionVariable objects
|
||||
# in self._exception_stack. Ensures that __context__ is properly linked
|
||||
# to the previous exception in the stack.
|
||||
if (ctx := val.__context__) and not isinstance(ctx, ConstantVariable):
|
||||
return val
|
||||
if (
|
||||
len(self._exception_stack) + prev_idx > 0
|
||||
): # Prevent invalid negative indexing
|
||||
prev = self._exception_stack[prev_idx]
|
||||
self._set_context_recursive(prev, prev_idx - 1)
|
||||
val.setattr("__context__", prev)
|
||||
return val
|
||||
|
||||
def _break_context_reference_cycle(
|
||||
self, val: ExceptionVariable, graph: FunctionGraph
|
||||
):
|
||||
# Detects and breaks cycles in exception __context__ chains using Floyd's algorithm,
|
||||
# following CPython's implementation.
|
||||
|
||||
fast = slow = val
|
||||
slow_update_toggle = False
|
||||
while True:
|
||||
context = fast.__context__
|
||||
if isinstance(
|
||||
context, ConstantVariable
|
||||
): # End of the chain; no context set
|
||||
break
|
||||
|
||||
if context is val:
|
||||
# The chain loops back to the original exception; break the cycle.
|
||||
fast.setattr(
|
||||
"__context__", ConstantVariable.wrap_literal(None, graph)
|
||||
)
|
||||
break
|
||||
|
||||
fast = context
|
||||
if fast is slow:
|
||||
# Cycle detected; all exceptions on the path have been visited and checked.
|
||||
break
|
||||
|
||||
if slow_update_toggle:
|
||||
slow = slow.__context__
|
||||
slow_update_toggle = not slow_update_toggle
|
||||
|
||||
def pop(self) -> ExceptionVariable:
|
||||
return self._exception_stack.pop()
|
||||
|
||||
def push(self, val: ExceptionVariable) -> None:
|
||||
self._exception_stack.append(val)
|
||||
|
||||
def empty(self) -> bool:
|
||||
return len(self._exception_stack) == 0
|
||||
|
||||
def __len__(self):
|
||||
return len(self._exception_stack)
|
||||
|
||||
def __repr__(self):
|
||||
return f"ExceptionStack({self._exception_stack})"
|
||||
|
||||
def __getitem__(self, idx: int) -> ExceptionVariable:
|
||||
return self._exception_stack[idx]
|
||||
|
||||
def cleanup(self) -> None:
|
||||
self._exception_stack[:] = []
|
||||
self._current_exception = None
|
||||
@@ -0,0 +1,483 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import traceback
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
|
||||
from ...profiler import EventGuard, event_register
|
||||
from ...psdb import NO_FALLBACK_CODES
|
||||
from ...utils import (
|
||||
ENV_SOT_ALLOW_DYNAMIC_SHAPE,
|
||||
ENV_SOT_ENABLE_COMPILE_TIME_LIMIT,
|
||||
ENV_SOT_ENABLE_GUARD_TREE,
|
||||
ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
|
||||
ENV_SOT_UNSAFE_CACHE_FASTPATH,
|
||||
BreakGraphError,
|
||||
CompileCountInfo,
|
||||
ConditionalFallbackError,
|
||||
FallbackError,
|
||||
InfoCollector,
|
||||
InnerError,
|
||||
Singleton,
|
||||
SotCapturedException,
|
||||
is_strict_mode,
|
||||
log,
|
||||
log_do,
|
||||
log_once,
|
||||
)
|
||||
from ..custom_code import CustomCode
|
||||
from .function_graph import FunctionGraph
|
||||
from .guard import Guard
|
||||
from .opcode_executor import OpcodeExecutor, OpcodeExecutorBase
|
||||
from .virtual_frame import VirtualFrame
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import types
|
||||
|
||||
GuardedFunction = tuple[CustomCode, Guard]
|
||||
GuardedFunctions = list[GuardedFunction]
|
||||
GuardChain = list[paddle.framework.core.GuardNodeBase]
|
||||
GuardChainList = list[GuardChain]
|
||||
|
||||
dummy_guard: Guard = lambda frame: True
|
||||
dummy_guard.expr = "lambda frame: True"
|
||||
dummy_guard.inlined_expr = "lambda frame: True"
|
||||
if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
|
||||
dummy_guard.mirror_guard = lambda frame: True
|
||||
|
||||
|
||||
class OpcodeExecutorCache(metaclass=Singleton):
|
||||
"""
|
||||
A singleton class that implements a cache for translated instructions.
|
||||
This cache is used to store previously translated instructions along with their corresponding guard functions.
|
||||
|
||||
Attributes:
|
||||
cache (dict): A dictionary that maps code objects to tuples of a cache getter function and a list of guarded functions.
|
||||
translate_count (int): The count of how many instructions have been translated. It is used to test whether the cache hits.
|
||||
"""
|
||||
|
||||
MAX_CACHE_SIZE = 20
|
||||
MAX_COMPILE_TIME_PER_CODE = 40
|
||||
MAX_COMPILE_TIME_TOTAL = 15 * 60
|
||||
CACHE_HIT_FASTPATH_THRESHOLD = 32
|
||||
cache: dict[
|
||||
types.CodeType, tuple[GuardedFunctions, paddle.framework.core.GuardTree]
|
||||
]
|
||||
translate_count: int
|
||||
code_symbolic_inputs: dict[types.CodeType, dict[str, None | dict[int, int]]]
|
||||
compile_time_stats: dict[types.CodeType, float]
|
||||
consecutive_cache_hit_count: defaultdict[types.CodeType, int]
|
||||
|
||||
def __init__(self):
|
||||
self.cache = {}
|
||||
self.translate_count = 0
|
||||
self.code_symbolic_inputs = {}
|
||||
self.compile_time_stats = {}
|
||||
self.consecutive_cache_hit_count = defaultdict(int)
|
||||
|
||||
def get_symbolic_inputs(
|
||||
self, code: types.CodeType
|
||||
) -> dict[str, dict[int, int] | None]:
|
||||
self.code_symbolic_inputs.setdefault(code, {})
|
||||
return self.code_symbolic_inputs[code]
|
||||
|
||||
def clear(self):
|
||||
"""
|
||||
Clears the cache and resets the translate count.
|
||||
"""
|
||||
self.cache.clear()
|
||||
self.translate_count = 0
|
||||
self.code_symbolic_inputs.clear()
|
||||
self.compile_time_stats.clear()
|
||||
|
||||
def dump_state(self):
|
||||
return {
|
||||
"cache": self.cache,
|
||||
"translate_count": self.translate_count,
|
||||
"code_symbolic_inputs": self.code_symbolic_inputs,
|
||||
"compile_time_stats": self.compile_time_stats,
|
||||
}
|
||||
|
||||
def load_state(self, state):
|
||||
self.cache = state["cache"]
|
||||
self.translate_count = state["translate_count"]
|
||||
self.code_symbolic_inputs = state["code_symbolic_inputs"]
|
||||
self.compile_time_stats = state["compile_time_stats"]
|
||||
|
||||
def __call__(self, frame: types.FrameType, **kwargs) -> CustomCode:
|
||||
code: types.CodeType = frame.f_code
|
||||
if code not in self.cache:
|
||||
log(2, f"[Cache] Firstly call {code}\n")
|
||||
new_custom_code, guard_fn, guard_chain = self.translate(
|
||||
frame, **kwargs
|
||||
)
|
||||
assert guard_fn is not None
|
||||
assert guard_chain is not None
|
||||
self.cache[code] = (
|
||||
[(new_custom_code, guard_fn)],
|
||||
paddle.framework.core.GuardTree([guard_chain]),
|
||||
)
|
||||
return new_custom_code
|
||||
guarded_fns, guard_tree = self.cache[code]
|
||||
compile_time_for_code = self.compile_time_stats.get(code, 0)
|
||||
compile_time_total = sum(self.compile_time_stats.values())
|
||||
return self.lookup(
|
||||
frame,
|
||||
guarded_fns,
|
||||
guard_tree,
|
||||
compile_time_for_code,
|
||||
compile_time_total,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def is_fastpath_threshold_reached(self, code):
|
||||
# Returns True if the number of consecutive cache hits for the given code
|
||||
# exceeds the UNSAFE_CACHE_FASTPATH threshold.
|
||||
return (
|
||||
self.consecutive_cache_hit_count.get(code, 0)
|
||||
>= self.CACHE_HIT_FASTPATH_THRESHOLD
|
||||
)
|
||||
|
||||
@event_register("lookup")
|
||||
def lookup(
|
||||
self,
|
||||
frame: types.FrameType,
|
||||
guarded_fns: GuardedFunctions,
|
||||
guard_tree: paddle.framework.core.GuardTree,
|
||||
compile_time_for_code: float,
|
||||
compile_time_total: float,
|
||||
**kwargs,
|
||||
) -> CustomCode:
|
||||
"""
|
||||
Looks up the cache for a matching code object and returns a custom code object if a matching guard function is found, otherwise None.
|
||||
|
||||
Args:
|
||||
frame (types.FrameType): The frame whose code object needs to be looked up in the cache.
|
||||
guarded_fns (GuardedFunctions): The list of guarded functions associated with the code object.
|
||||
|
||||
Returns:
|
||||
CustomCode: The custom code object if a matching guard function is found, otherwise None.
|
||||
"""
|
||||
code: types.CodeType = frame.f_code
|
||||
|
||||
if len(guarded_fns) >= self.MAX_CACHE_SIZE:
|
||||
log(2, "[Cache] Exceed max cache size, skip it\n")
|
||||
return CustomCode(None, False)
|
||||
|
||||
enable_strict_guard = ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get()
|
||||
enable_guard_tree = ENV_SOT_ENABLE_GUARD_TREE.get()
|
||||
enable_unsafe_cache_fastpath = ENV_SOT_UNSAFE_CACHE_FASTPATH.get()
|
||||
enable_compile_time_limit = ENV_SOT_ENABLE_COMPILE_TIME_LIMIT.get()
|
||||
|
||||
if enable_unsafe_cache_fastpath and (
|
||||
self.is_fastpath_threshold_reached(code)
|
||||
):
|
||||
# NOTE: In inference scenarios, cache misses are generally rare, so we can enable this unsafe short path.
|
||||
log(
|
||||
2,
|
||||
"[Cache] The CACHE_HIT_FASTPATH_THRESHOLD has been reached, so fast path is now enabled\n",
|
||||
)
|
||||
return guarded_fns[0][0]
|
||||
|
||||
cache_index = None
|
||||
if enable_strict_guard or enable_guard_tree:
|
||||
log(4, f"[Cache] Guard tree: \n{guard_tree.stringify()}")
|
||||
cache_index = guard_tree.lookup(frame)
|
||||
|
||||
if not enable_strict_guard and enable_guard_tree:
|
||||
if cache_index is not None:
|
||||
# TODO(zrr1999): add a mapping between custom_code and cache_index
|
||||
return guarded_fns[cache_index][0]
|
||||
else:
|
||||
log(2, "[Cache] all guards missed (guard tree mode)\n")
|
||||
if (
|
||||
enable_compile_time_limit
|
||||
and compile_time_for_code >= self.MAX_COMPILE_TIME_PER_CODE
|
||||
):
|
||||
log(
|
||||
2,
|
||||
"[Cache] Exceed max compile time per code, skip it\n",
|
||||
)
|
||||
return CustomCode(None, False)
|
||||
if (
|
||||
enable_compile_time_limit
|
||||
and compile_time_total >= self.MAX_COMPILE_TIME_TOTAL
|
||||
):
|
||||
log_once(
|
||||
f"[SOT] Current total compile time is {compile_time_total}, exceed max compile time total {self.MAX_COMPILE_TIME_TOTAL}, fallback new function to dygraph"
|
||||
)
|
||||
log(
|
||||
2,
|
||||
"[Cache] Exceed max compile time total, skip it\n",
|
||||
)
|
||||
return CustomCode(None, False)
|
||||
new_custom_code, guard_fn, guard_chain = self.translate(
|
||||
frame, **kwargs
|
||||
)
|
||||
if guard_fn is not None:
|
||||
assert guard_chain is not None
|
||||
guarded_fns.append((new_custom_code, guard_fn))
|
||||
guard_tree.add_guard_chain(guard_chain)
|
||||
return new_custom_code
|
||||
|
||||
for index, (custom_code, guard_fn) in enumerate(guarded_fns):
|
||||
if enable_strict_guard:
|
||||
mirror_guard_error = None
|
||||
try:
|
||||
with EventGuard("try mirror guard"):
|
||||
mirror_guard_result = guard_fn.mirror_guard(frame)
|
||||
except Exception as e:
|
||||
log(2, f"[Cache] Mirror guard error: {e}\n")
|
||||
mirror_guard_error = e
|
||||
|
||||
try:
|
||||
with EventGuard("try guard"):
|
||||
guard_result = guard_fn(frame)
|
||||
if enable_strict_guard and (not enable_unsafe_cache_fastpath):
|
||||
assert mirror_guard_result == guard_result, (
|
||||
"faster guard result is not equal to guard result, "
|
||||
f"guard_expr: {getattr(guard_fn, 'expr', 'None')} \n"
|
||||
f"faster_guard_expr: {getattr(guard_fn.mirror_guard, 'expr', 'None')},"
|
||||
)
|
||||
if guard_result:
|
||||
log(
|
||||
2,
|
||||
f"[Cache] Cache hit, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
|
||||
)
|
||||
if not enable_unsafe_cache_fastpath:
|
||||
# TODO(zrr1999): cache_index should be equal to index when enable_strict_guard.
|
||||
assert cache_index is None or index == cache_index, (
|
||||
f"cache_index({cache_index}) is not equal to index({index})"
|
||||
)
|
||||
|
||||
if enable_unsafe_cache_fastpath:
|
||||
if index == 0:
|
||||
self.consecutive_cache_hit_count[code] += 1
|
||||
else:
|
||||
# Move the current hit to the front
|
||||
# Note: Be cautious when modifying the order of elements in a list during iteration,
|
||||
# as it can lead to unexpected behavior.
|
||||
guarded_fns[:] = [
|
||||
guarded_fns[index],
|
||||
*guarded_fns[:index],
|
||||
*guarded_fns[index + 1 :],
|
||||
]
|
||||
self.consecutive_cache_hit_count[code] = 0
|
||||
|
||||
return custom_code
|
||||
else:
|
||||
log_do(
|
||||
4,
|
||||
self.analyse_guard_global_object(guard_fn),
|
||||
)
|
||||
log(
|
||||
2,
|
||||
f"[Cache] Cache miss, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
|
||||
)
|
||||
log_do(
|
||||
2,
|
||||
self.analyse_guard_error(guard_fn, frame),
|
||||
)
|
||||
except Exception as e:
|
||||
log(2, f"[Cache] Guard function error: {e}\n")
|
||||
log(
|
||||
2,
|
||||
f"[Cache] Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
|
||||
)
|
||||
log_do(
|
||||
2,
|
||||
self.analyse_guard_error(guard_fn, frame),
|
||||
)
|
||||
if enable_strict_guard and (not enable_unsafe_cache_fastpath):
|
||||
assert type(e) == type(mirror_guard_error) and str(
|
||||
e
|
||||
) == str(mirror_guard_error), (
|
||||
"mirror guard error is not equal to guard error, "
|
||||
f"guard_error: {e} \n"
|
||||
f"mirror_guard_error: {mirror_guard_error},"
|
||||
)
|
||||
|
||||
log(2, "[Cache] all guards missed\n")
|
||||
if (
|
||||
enable_compile_time_limit
|
||||
and compile_time_for_code >= self.MAX_COMPILE_TIME_PER_CODE
|
||||
):
|
||||
log(2, "[Cache] Exceed max compile time per code, skip it\n")
|
||||
return CustomCode(None, False)
|
||||
if (
|
||||
enable_compile_time_limit
|
||||
and compile_time_total >= self.MAX_COMPILE_TIME_TOTAL
|
||||
):
|
||||
log_once(
|
||||
f"[SOT] Current compile time total is {compile_time_total}, exceed max compile time total {self.MAX_COMPILE_TIME_TOTAL}, fallback new function to dygraph"
|
||||
)
|
||||
log(
|
||||
2,
|
||||
"[Cache] Exceed max compile time total, skip it\n",
|
||||
)
|
||||
return CustomCode(None, False)
|
||||
new_custom_code, guard_fn, guard_chain = self.translate(frame, **kwargs)
|
||||
if guard_fn is not None:
|
||||
assert guard_chain is not None
|
||||
guarded_fns.append((new_custom_code, guard_fn))
|
||||
guard_tree.add_guard_chain(guard_chain)
|
||||
return new_custom_code
|
||||
|
||||
def before_translate_hook(self, frame: types.FrameType):
|
||||
if not ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():
|
||||
return
|
||||
|
||||
def translate(
|
||||
self, frame: types.FrameType, **kwargs
|
||||
) -> tuple[CustomCode, Guard | None, GuardChain | None]:
|
||||
"""
|
||||
Translates the given frame's code object and returns the cache getter function and a guarded function for the translated code object.
|
||||
|
||||
Args:
|
||||
frame (types.FrameType): The frame whose code object needs to be translated.
|
||||
|
||||
Returns:
|
||||
tuple[CustomCode, Guard]: The cache getter function and a guarded function for the translated code object.
|
||||
"""
|
||||
self.before_translate_hook(frame)
|
||||
self.translate_count += 1
|
||||
custom_new_code, guard_fn, guard_chain = start_translate(
|
||||
frame, **kwargs
|
||||
)
|
||||
return custom_new_code, guard_fn, guard_chain
|
||||
|
||||
def analyse_guard_global_object(self, guard_fn):
|
||||
def inner():
|
||||
for key in guard_fn.__globals__.keys():
|
||||
if key.startswith("__object"):
|
||||
print(
|
||||
f"[Cache] meet global object: {key} : {guard_fn.__globals__[key]}",
|
||||
)
|
||||
|
||||
return inner
|
||||
|
||||
def analyse_guard_error(self, guard_fn, frame):
|
||||
def inner():
|
||||
guard_expr = guard_fn.inlined_expr
|
||||
lambda_head = "lambda frame: "
|
||||
guard_expr = guard_expr.replace(lambda_head, "")
|
||||
guards = guard_expr.split(" and ")
|
||||
for guard_str in guards:
|
||||
guard = eval(lambda_head + guard_str, guard_fn.__globals__)
|
||||
result = False
|
||||
try:
|
||||
result = guard(frame)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"[Cache] Error occurred when checking guard {guard_str}: {e}"
|
||||
)
|
||||
return
|
||||
if result is False:
|
||||
print(f"[Cache] missed at {guard_str}")
|
||||
return
|
||||
print("[Cache] missed guard not found.")
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
def start_translate(
|
||||
frame: types.FrameType,
|
||||
**kwargs,
|
||||
) -> tuple[CustomCode, Guard | None, GuardChain | None]:
|
||||
"""
|
||||
Starts the translation process for the given frame and returns the translated code object, its guard function and its guard tree node, or None if translation fails.
|
||||
|
||||
Args:
|
||||
frame: The frame to be translated.
|
||||
|
||||
Returns:
|
||||
tuple[CustomCode, Guard | None, GuardChain | None]: The translated code object, its guard function and its guard tree node, or None if translation fails.
|
||||
"""
|
||||
simulator = None
|
||||
graph = FunctionGraph(frame.f_code, frame.f_globals, **kwargs)
|
||||
try:
|
||||
vframe = VirtualFrame.from_real_frame(frame, graph)
|
||||
simulator = OpcodeExecutor(vframe, graph)
|
||||
simulator.check_code_simulatable()
|
||||
InfoCollector().attach(CompileCountInfo, frame.f_code)
|
||||
|
||||
new_custom_code, guard_fn = simulator.transform(frame)
|
||||
if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
|
||||
assert guard_fn(frame)
|
||||
assert guard_fn.mirror_guard(frame)
|
||||
|
||||
if not simulator._graph.need_cache:
|
||||
return (
|
||||
CustomCode(None, True),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
guard_chain = simulator.guard_chain
|
||||
if len(guard_chain) == 0:
|
||||
guard_chain: GuardChain = [paddle.framework.core.DummyGuardNode()]
|
||||
return new_custom_code, guard_fn, guard_chain
|
||||
# TODO(0x45f): handle BreakGraphError to trigger fallback
|
||||
except BreakGraphError as e:
|
||||
raise RuntimeError(
|
||||
f"Found BreakGraphError raised, it should not be catch at start_translate!\n{e}"
|
||||
)
|
||||
except FallbackError as e:
|
||||
if frame.f_code in NO_FALLBACK_CODES:
|
||||
raise InnerError(
|
||||
f"{frame.f_code.co_name} should not fallback, but got '{e}'"
|
||||
)
|
||||
if is_strict_mode():
|
||||
raise
|
||||
log(
|
||||
2,
|
||||
f"Unsupported Frame is {frame.f_code}, error message is: \n"
|
||||
+ "".join(traceback.format_exception(type(e), e, e.__traceback__)),
|
||||
)
|
||||
dummy_guard_chain: GuardChain = [paddle.framework.core.DummyGuardNode()]
|
||||
guard, guard_chain = dummy_guard, dummy_guard_chain
|
||||
|
||||
if isinstance(e, ConditionalFallbackError):
|
||||
# Guard global variables only
|
||||
graph.input_variables.clear()
|
||||
guard = graph.guard_fn
|
||||
guard_chain = graph.guard_chain
|
||||
|
||||
return (
|
||||
CustomCode(None, e.disable_eval_frame),
|
||||
guard,
|
||||
guard_chain,
|
||||
)
|
||||
except SotCapturedException as e:
|
||||
log(
|
||||
1,
|
||||
"Note: This fallback may be triggered by user code, or it could result from an internal "
|
||||
"SOT exception being incorrectly captured. Please investigate carefully.\n",
|
||||
)
|
||||
if is_strict_mode():
|
||||
raise
|
||||
dummy_guard_chain: GuardChain = [paddle.framework.core.DummyGuardNode()]
|
||||
return (CustomCode(None, True), dummy_guard, dummy_guard_chain)
|
||||
except Exception as e:
|
||||
raise InnerError(OpcodeExecutorBase.error_message_summary(e)) from e
|
||||
finally:
|
||||
if simulator is not None:
|
||||
simulator.cleanup()
|
||||
del simulator
|
||||
gc.collect()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,327 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import types
|
||||
import weakref
|
||||
from collections.abc import Callable
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
import paddle
|
||||
|
||||
from ...profiler import EventGuard
|
||||
from ...utils import (
|
||||
ENV_SOT_ENABLE_FASTER_GUARD,
|
||||
ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
|
||||
current_symbol_registry,
|
||||
log,
|
||||
log_do,
|
||||
)
|
||||
|
||||
Guard = Callable[[types.FrameType], bool]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .variables import VariableBase
|
||||
|
||||
GuardBase = paddle.framework.core.GuardBase
|
||||
CheckGuardInputT = TypeVar("CheckGuardInputT", bound=VariableBase)
|
||||
|
||||
# NOTE(SigureMo): [How to write Stringified Guard?]
|
||||
# 1. we should capture free variables manually, the string cannot capture free
|
||||
# variables automatically.
|
||||
# 2. Be aware that the comparison logic before and after stringify may be different.
|
||||
# 3. we should compute as much as possible at "compile time" and encode the
|
||||
# computation in the Guard string, rather than passing it to runtime to minimize
|
||||
# runtime overhead.
|
||||
|
||||
|
||||
class StringifiedExpression:
|
||||
"""
|
||||
Used to store string based expressions for generating Guard.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
expr_template: str,
|
||||
sub_exprs: list[StringifiedExpression],
|
||||
free_vars: dict[str, Any],
|
||||
):
|
||||
self.expr_template = expr_template
|
||||
expr = self.expr_template.format(
|
||||
*[sub_expr.symbol for sub_expr in sub_exprs]
|
||||
)
|
||||
self.registered_expr = expr
|
||||
self.symbol = current_symbol_registry().request_symbol(expr)
|
||||
self.sub_exprs = sub_exprs
|
||||
self.free_vars = free_vars
|
||||
|
||||
@cached_property
|
||||
def inlined_expr(self):
|
||||
return self.expr_template.format(
|
||||
*[sub_expr.inlined_expr for sub_expr in self.sub_exprs]
|
||||
)
|
||||
|
||||
def gen_expr(self):
|
||||
def gen_expr_fn():
|
||||
return self.expr_template.format(
|
||||
*[sub_expr.gen_expr() for sub_expr in self.sub_exprs]
|
||||
)
|
||||
|
||||
return current_symbol_registry().gen_expr(
|
||||
self.registered_expr, gen_expr_fn
|
||||
)
|
||||
|
||||
def __hash__(self):
|
||||
if self.free_vars:
|
||||
return hash((self.inlined_expr, id(self)))
|
||||
else:
|
||||
return hash(self.inlined_expr)
|
||||
|
||||
|
||||
class FasterStringifiedExpression(StringifiedExpression):
|
||||
def __init__(
|
||||
self,
|
||||
expr_template: str,
|
||||
faster_guard: GuardBase,
|
||||
sub_exprs: list[StringifiedExpression],
|
||||
free_vars: dict[str, Any],
|
||||
):
|
||||
self.faster_guard = faster_guard
|
||||
if ENV_SOT_ENABLE_FASTER_GUARD.get():
|
||||
if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
|
||||
self.py_guard_expr_template = original_expr_template = (
|
||||
expr_template
|
||||
)
|
||||
else:
|
||||
original_expr_template = expr_template
|
||||
expr_template, free_vars = gen_faster_guard_expr_template(
|
||||
faster_guard, sub_exprs, free_vars
|
||||
)
|
||||
log(
|
||||
3,
|
||||
f"[FasterGuard] transform {original_expr_template} to {expr_template}\n",
|
||||
)
|
||||
|
||||
super().__init__(expr_template, sub_exprs, free_vars)
|
||||
|
||||
def gen_mirror_guard(
|
||||
self, enable_faster_gurad: bool
|
||||
) -> StringifiedExpression:
|
||||
if not enable_faster_gurad:
|
||||
# gen faster_guard_expr
|
||||
expr_template, expr_free_vars = gen_faster_guard_expr_template(
|
||||
self.faster_guard,
|
||||
self.sub_exprs,
|
||||
self.free_vars,
|
||||
)
|
||||
return StringifiedExpression(
|
||||
expr_template, self.sub_exprs, expr_free_vars
|
||||
)
|
||||
# gen pyGuard_expr
|
||||
return StringifiedExpression(
|
||||
self.py_guard_expr_template, self.sub_exprs, self.free_vars
|
||||
)
|
||||
|
||||
|
||||
def gen_faster_guard_expr_template(
|
||||
faster_guard: GuardBase,
|
||||
sub_exprs: list[StringifiedExpression],
|
||||
free_vars: dict[str, Any],
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
guard_cls_name = faster_guard.__class__.__name__
|
||||
guard_name = f"{guard_cls_name}_{id(faster_guard)}"
|
||||
expr_template = guard_name + "(" + ", ".join(["{}"] * len(sub_exprs)) + ")"
|
||||
free_vars = union_free_vars(free_vars, {guard_name: faster_guard.check})
|
||||
return expr_template, free_vars
|
||||
|
||||
|
||||
def union_free_vars(*free_vars: dict[str, Any]):
|
||||
return {k: v for d in free_vars for k, v in d.items()}
|
||||
|
||||
|
||||
def make_guard(stringified_guards: list[StringifiedExpression]) -> Guard:
|
||||
"""
|
||||
Make a guard from a list of StringifiedExpression.
|
||||
|
||||
For more design ideas, refer to the `Stringified guard <https://github.com/PaddlePaddle/PaddleSOT/blob/develop/docs/design/stringify-guard.md>`_ for details.
|
||||
|
||||
Args:
|
||||
stringified_guards: a list of StringifiedExpression.
|
||||
"""
|
||||
with EventGuard("make_guard"):
|
||||
num_guards = len(stringified_guards)
|
||||
if not num_guards:
|
||||
guard = lambda frame: True
|
||||
guard.expr = "lambda frame: True"
|
||||
guard.original_guard = guard
|
||||
if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
|
||||
guard.mirror_guard = lambda frame: True
|
||||
return guard
|
||||
|
||||
free_vars = union_free_vars(
|
||||
*(expr.free_vars for expr in stringified_guards)
|
||||
)
|
||||
inlined_guard_expr = "lambda frame: " + " and ".join(
|
||||
[expr.inlined_expr for expr in stringified_guards]
|
||||
)
|
||||
guard_expr: str = "lambda frame: " + " and ".join(
|
||||
[expr.gen_expr() for expr in stringified_guards]
|
||||
)
|
||||
|
||||
guard = eval(guard_expr, free_vars)
|
||||
|
||||
log(3, f"[Guard] {inlined_guard_expr}\n")
|
||||
guard.inlined_expr = inlined_guard_expr
|
||||
guard.expr = guard_expr
|
||||
|
||||
def check_guard_callable(guard: GuardBase):
|
||||
assert callable(guard), "guard must be callable."
|
||||
|
||||
if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
|
||||
mirror_guard_expr_list: list[str] = []
|
||||
mirror_guard_temp_free_vars: dict[str, Any] = {}
|
||||
enable_faster_gurad = ENV_SOT_ENABLE_FASTER_GUARD.get()
|
||||
for expr in stringified_guards:
|
||||
if isinstance(expr, FasterStringifiedExpression):
|
||||
expr = expr.gen_mirror_guard(enable_faster_gurad)
|
||||
mirror_guard_expr_list.append(expr.inlined_expr)
|
||||
mirror_guard_temp_free_vars.update(expr.free_vars)
|
||||
mirror_guard_expr = "lambda frame: " + " and ".join(
|
||||
mirror_guard_expr_list
|
||||
)
|
||||
mirror_guard_free_vars = union_free_vars(
|
||||
mirror_guard_temp_free_vars
|
||||
)
|
||||
guard.mirror_guard = eval(mirror_guard_expr, mirror_guard_free_vars)
|
||||
guard.mirror_guard.expr = mirror_guard_expr
|
||||
check_guard_callable(guard.mirror_guard)
|
||||
|
||||
check_guard_callable(guard)
|
||||
|
||||
return guard
|
||||
|
||||
|
||||
def support_weak_ref(obj):
|
||||
if isinstance(obj, types.FunctionType):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# TODO(zrr1999): unify check_guard and check_faster_guard
|
||||
def check_guard(
|
||||
fn: Callable[[CheckGuardInputT], list[StringifiedExpression]],
|
||||
) -> Callable[[CheckGuardInputT], list[StringifiedExpression]]:
|
||||
def wrapper(self: CheckGuardInputT) -> list[StringifiedExpression]:
|
||||
assert self.tracker.is_traceable(), (
|
||||
"Cannot make guard from a non-tracable guard variable."
|
||||
)
|
||||
|
||||
def guard_log():
|
||||
frame_value_tracer = self.tracker.trace_value_from_frame()
|
||||
print(
|
||||
f"[Guard] guard_fn for {self}, tracker={self.tracker.__class__.__name__}, value={frame_value_tracer.registered_expr}"
|
||||
)
|
||||
|
||||
log_do(4, guard_log)
|
||||
return fn(self)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def check_faster_guard(
|
||||
fn: Callable[[CheckGuardInputT], list[paddle.framework.core.GuardNodeBase]],
|
||||
) -> Callable[[CheckGuardInputT], list[paddle.framework.core.GuardNodeBase]]:
|
||||
def wrapper(
|
||||
self: CheckGuardInputT,
|
||||
) -> list[paddle.framework.core.GuardNodeBase]:
|
||||
assert self.tracker.is_traceable(), (
|
||||
"Cannot make guard from a non-tracable guard variable."
|
||||
)
|
||||
|
||||
def guard_log():
|
||||
frame_value_tracer = self.tracker.trace_value_from_frame()
|
||||
print(
|
||||
f"[Guard Tree] guard_fn for {self}, tracker={self.tracker.__class__.__name__}, value={frame_value_tracer.registered_expr}"
|
||||
)
|
||||
|
||||
log_do(4, guard_log)
|
||||
return fn(self)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@check_guard
|
||||
def object_equal_stringified_guard(self) -> list[StringifiedExpression]:
|
||||
frame_value_tracer = self.tracker.trace_value_from_frame()
|
||||
|
||||
obj_free_var_name = f"__{self.id}"
|
||||
weak_ref_obj = self.get_py_value()
|
||||
if support_weak_ref(weak_ref_obj):
|
||||
weak_ref_obj = weakref.ref(self.get_py_value())
|
||||
return [
|
||||
FasterStringifiedExpression(
|
||||
f"{obj_free_var_name}() is not None and {{}} == {obj_free_var_name}()",
|
||||
paddle.framework.core.WeakRefMatchGuard(self.get_py_value()),
|
||||
[frame_value_tracer],
|
||||
union_free_vars(
|
||||
frame_value_tracer.free_vars,
|
||||
{obj_free_var_name: weak_ref_obj},
|
||||
),
|
||||
)
|
||||
]
|
||||
return [
|
||||
FasterStringifiedExpression(
|
||||
f"{{}} == {obj_free_var_name}",
|
||||
paddle.framework.core.ValueMatchGuard(weak_ref_obj),
|
||||
[frame_value_tracer],
|
||||
union_free_vars(
|
||||
frame_value_tracer.free_vars,
|
||||
{obj_free_var_name: self.get_py_value()},
|
||||
),
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@check_faster_guard
|
||||
def object_equal_faster_guard(
|
||||
self,
|
||||
) -> list[paddle.framework.core.GuardNodeBase]:
|
||||
expr_node = self.tracker.guard_tree_expr_node()
|
||||
|
||||
weak_ref_obj = self.get_py_value()
|
||||
if support_weak_ref(weak_ref_obj):
|
||||
weak_ref_obj = weakref.ref(self.get_py_value())
|
||||
return [
|
||||
paddle.framework.core.GuardNode(
|
||||
paddle.framework.core.WeakRefMatchGuard(self.get_py_value()),
|
||||
[expr_node],
|
||||
)
|
||||
]
|
||||
return [
|
||||
paddle.framework.core.GuardNode(
|
||||
paddle.framework.core.ValueMatchGuard(weak_ref_obj),
|
||||
[expr_node],
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def stringify_pyobject(obj: object) -> tuple[str, dict[str, Any]]:
|
||||
if isinstance(obj, paddle.core.VarDesc.VarType):
|
||||
return f"paddle.core.VarDesc.VarType({obj.value})", {"paddle": paddle}
|
||||
elif isinstance(obj, paddle.core.DataType):
|
||||
return f"paddle.core.DataType({obj.value})", {"paddle": paddle}
|
||||
# For builtin values
|
||||
return f"{obj!r}", {}
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
# flags for instructions
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class FORMAT_VALUE_FLAG:
|
||||
FVC_MASK = 0x3
|
||||
FVC_NONE = 0x0
|
||||
FVC_STR = 0x1
|
||||
FVC_REPR = 0x2
|
||||
FVC_ASCII = 0x3
|
||||
FVS_MASK = 0x4
|
||||
FVS_HAVE_SPEC = 0x4
|
||||
|
||||
|
||||
class CONVERT_VALUE_FLAG:
|
||||
CV_STR = 1
|
||||
CV_REPR = 2
|
||||
CV_ASCII = 3
|
||||
|
||||
|
||||
# https://github.com/python/cpython/blob/3.14/Include/internal/pycore_opcode_utils.h#L63-L68
|
||||
class MAKE_FUNCTION_FLAG:
|
||||
MF_HAS_ANNOTATE = 0x10
|
||||
MF_HAS_CLOSURE = 0x08
|
||||
MF_HAS_ANNOTATION = 0x04
|
||||
MF_HAS_KWDEFAULTS = 0x02
|
||||
MF_HAS_DEFAULTS = 0x01
|
||||
|
||||
|
||||
class CALL_FUNCTION_EX_FLAG:
|
||||
CFE_HAS_KWARGS = 0x01
|
||||
|
||||
|
||||
# see https://github.com/python/cpython/blob/3.12/Python/intrinsics.c#L211-L225
|
||||
class IntrinsicsUnaryFunctions(Enum):
|
||||
INTRINSIC_1_INVALID = 0
|
||||
INTRINSIC_PRINT = 1 # no support, only non-interactive mode
|
||||
INTRINSIC_IMPORT_STAR = 2 # no support, `from module import *`
|
||||
INTRINSIC_STOPITERATION_ERROR = 3 # no support, generator or coroutine
|
||||
INTRINSIC_ASYNC_GEN_WRAP = 4 # no support, async
|
||||
INTRINSIC_UNARY_POSITIVE = 5
|
||||
INTRINSIC_LIST_TO_TUPLE = 6
|
||||
INTRINSIC_TYPEVAR = 7 # no support, PEP 695
|
||||
INTRINSIC_PARAMSPEC = 8 # no support, PEP 695
|
||||
INTRINSIC_TYPEVARTUPLE = 9 # no support, PEP 695
|
||||
INTRINSIC_SUBSCRIPT_GENERIC = 10 # no support, PEP 695
|
||||
INTRINSIC_TYPEALIAS = 11 # no support, PEP 695
|
||||
|
||||
|
||||
# https://github.com/python/cpython/blob/3.14/Include/internal/pycore_opcode_utils.h#L70-L76
|
||||
# All are attributes of 'builtins'
|
||||
LOAD_COMMON_CONSTANT_FLAG = (
|
||||
"AssertionError",
|
||||
"NotImplementedError",
|
||||
"tuple",
|
||||
"all",
|
||||
"any",
|
||||
)
|
||||
@@ -0,0 +1,306 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Concatenate, Generic, TypeAlias, TypeVar
|
||||
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
|
||||
MutableDataT = TypeVar("MutableDataT", bound="MutableData")
|
||||
DataGetter: TypeAlias = Callable[[MutableDataT, Any], Any]
|
||||
|
||||
InnerMutableDataT = TypeVar(
|
||||
"InnerMutableDataT", bound="dict[str, Any] | list[Any]"
|
||||
)
|
||||
|
||||
|
||||
class Mutation:
|
||||
ABBR: str
|
||||
|
||||
|
||||
class MutationSet(Mutation):
|
||||
"""
|
||||
Setting a value.
|
||||
This mutation is used for MutableDictLikeData and MutableListLikeData.
|
||||
"""
|
||||
|
||||
ABBR = "S"
|
||||
|
||||
def __init__(self, key, value):
|
||||
self.key = key
|
||||
self.value = value
|
||||
|
||||
def __repr__(self):
|
||||
return f"MutationSet({self.key}, {self.value})"
|
||||
|
||||
|
||||
class MutationDel(Mutation):
|
||||
"""
|
||||
Deleting a value.
|
||||
This mutation is used for MutableDictLikeData and MutableListLikeData.
|
||||
"""
|
||||
|
||||
ABBR = "D"
|
||||
|
||||
def __init__(self, key):
|
||||
self.key = key
|
||||
|
||||
def __repr__(self):
|
||||
return f"MutationDel({self.key})"
|
||||
|
||||
|
||||
class MutationNew(Mutation):
|
||||
"""
|
||||
Adding a new value.
|
||||
This mutation is only used for MutableDictLikeData.
|
||||
"""
|
||||
|
||||
ABBR = "N"
|
||||
|
||||
def __init__(self, key, value):
|
||||
self.key = key
|
||||
self.value = value
|
||||
|
||||
def __repr__(self):
|
||||
return f"MutationNew({self.key}, {self.value})"
|
||||
|
||||
|
||||
class MutationInsert(Mutation):
|
||||
"""
|
||||
Inserting a value.
|
||||
This mutation is only used for MutableListLikeData.
|
||||
"""
|
||||
|
||||
ABBR = "I"
|
||||
|
||||
def __init__(self, index, value):
|
||||
self.index = index
|
||||
self.value = value
|
||||
|
||||
def __repr__(self):
|
||||
return f"MutationInsert({self.index}, {self.value})"
|
||||
|
||||
|
||||
class MutationPermutate(Mutation):
|
||||
"""
|
||||
Permutating all the values.
|
||||
This mutation is only used for MutableListLikeData.
|
||||
"""
|
||||
|
||||
ABBR = "P"
|
||||
|
||||
def __init__(self, permutation):
|
||||
self.permutation = permutation
|
||||
|
||||
def __repr__(self):
|
||||
return f"MutationPermutate({self.permutation})"
|
||||
|
||||
|
||||
def record_mutation(
|
||||
mutation_fn: Callable[Concatenate[MutableDataT, P], Mutation],
|
||||
) -> Callable[Concatenate[MutableDataT, P], None]:
|
||||
def wrapper(self, *args: P.args, **kwargs: P.kwargs):
|
||||
mutation = mutation_fn(self, *args, **kwargs)
|
||||
self.records.append(mutation)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class MutableData(Generic[InnerMutableDataT]):
|
||||
"""
|
||||
An intermediate data structure between data and variable, it records all the mutations.
|
||||
"""
|
||||
|
||||
read_cache: InnerMutableDataT
|
||||
|
||||
class Empty:
|
||||
def __repr__(self):
|
||||
return "Empty()"
|
||||
|
||||
def __init__(self, data: Any, getter: DataGetter):
|
||||
self.original_data = data
|
||||
self.getter = getter
|
||||
self.records: list[Mutation] = []
|
||||
|
||||
def is_empty(self, value):
|
||||
return isinstance(value, MutableData.Empty)
|
||||
|
||||
@property
|
||||
def version(self):
|
||||
return len(self.records)
|
||||
|
||||
@property
|
||||
def has_changed(self):
|
||||
return self.version != 0
|
||||
|
||||
def check_changed(self, key: Any) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def rollback(self, version: int):
|
||||
assert version <= self.version
|
||||
self.records[:] = self.records[:version]
|
||||
|
||||
def get(self, key):
|
||||
raise NotImplementedError
|
||||
|
||||
def set(self, key, value):
|
||||
raise NotImplementedError
|
||||
|
||||
def apply(self, mutation: Mutation, write_cache: InnerMutableDataT):
|
||||
raise NotImplementedError
|
||||
|
||||
def reproduce(self, version: int | None = None) -> InnerMutableDataT:
|
||||
if version is None:
|
||||
version = self.version
|
||||
write_cache = self.read_cache.copy()
|
||||
for mutation in self.records[:version]:
|
||||
self.apply(mutation, write_cache)
|
||||
return write_cache
|
||||
|
||||
def __repr__(self) -> str:
|
||||
records_abbrs = "".join([mutation.ABBR for mutation in self.records])
|
||||
return f"{self.__class__.__name__}({records_abbrs})"
|
||||
|
||||
|
||||
class MutableDictLikeData(MutableData["dict[str, Any]"]):
|
||||
def __init__(self, data: Any, getter: DataGetter):
|
||||
super().__init__(data, getter)
|
||||
self.read_cache = {}
|
||||
|
||||
def clear_read_cache(self):
|
||||
self.read_cache.clear()
|
||||
|
||||
def check_changed(self, key: Any) -> bool:
|
||||
if not self.has_changed:
|
||||
return False
|
||||
for mutation in self.records:
|
||||
if (
|
||||
isinstance(mutation, (MutationNew, MutationDel, MutationSet))
|
||||
and mutation.key == key
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
def get(self, key: Any):
|
||||
# TODO(SigureMo): Optimize performance of this.
|
||||
write_cache = self.reproduce(self.version)
|
||||
if key not in write_cache:
|
||||
self.read_cache[key] = self.getter(self, key)
|
||||
return self.reproduce(self.version)[key]
|
||||
|
||||
def get_all(self):
|
||||
original_keys = list(self.original_data.keys())
|
||||
for mutation in self.records:
|
||||
if isinstance(mutation, MutationNew):
|
||||
original_keys.append(mutation.key)
|
||||
elif isinstance(mutation, MutationDel):
|
||||
original_keys.remove(mutation.key)
|
||||
return {key: self.get(key) for key in original_keys}
|
||||
|
||||
@record_mutation
|
||||
def set(self, key: Any, value: Any) -> Mutation:
|
||||
is_new = False
|
||||
if self.is_empty(self.get(key)):
|
||||
is_new = True
|
||||
return (
|
||||
MutationSet(key, value) if not is_new else MutationNew(key, value)
|
||||
)
|
||||
|
||||
@record_mutation
|
||||
def delete(self, key):
|
||||
return MutationDel(key)
|
||||
|
||||
def apply(self, mutation: Mutation, write_cache: dict[str, Any]):
|
||||
if isinstance(mutation, MutationNew):
|
||||
write_cache[mutation.key] = mutation.value
|
||||
elif isinstance(mutation, MutationSet):
|
||||
write_cache[mutation.key] = mutation.value
|
||||
elif isinstance(mutation, MutationDel):
|
||||
write_cache[mutation.key] = MutableData.Empty()
|
||||
else:
|
||||
raise ValueError(f"Unknown mutation type {mutation}")
|
||||
|
||||
def reproduce(self, version: int | None = None):
|
||||
if version is None:
|
||||
version = self.version
|
||||
write_cache = self.read_cache.copy()
|
||||
for mutation in self.records[:version]:
|
||||
self.apply(mutation, write_cache)
|
||||
return write_cache
|
||||
|
||||
|
||||
class MutableListLikeData(MutableData["list[Any]"]):
|
||||
def __init__(self, data: Any, getter: DataGetter):
|
||||
super().__init__(data, getter)
|
||||
self.read_cache = [
|
||||
self.getter(self, idx) for idx in range(len(self.original_data))
|
||||
]
|
||||
|
||||
def clear_read_cache(self):
|
||||
self.read_cache[:] = []
|
||||
|
||||
def check_changed(self, key: Any) -> bool:
|
||||
return self.has_changed
|
||||
|
||||
@property
|
||||
def length(self):
|
||||
return len(self.reproduce())
|
||||
|
||||
def get(self, key):
|
||||
write_cache = self.reproduce(self.version)
|
||||
return write_cache[key]
|
||||
|
||||
def get_all(self) -> list[Any]:
|
||||
items = self.reproduce(self.version)
|
||||
return items
|
||||
|
||||
@record_mutation
|
||||
def set(self, key: int, value: Any):
|
||||
return MutationSet(self._regularize_index(key), value)
|
||||
|
||||
@record_mutation
|
||||
def delete(self, key: int):
|
||||
return MutationDel(self._regularize_index(key))
|
||||
|
||||
@record_mutation
|
||||
def insert(self, index: int, value: Any):
|
||||
return MutationInsert(self._regularize_index(index), value)
|
||||
|
||||
@record_mutation
|
||||
def permutate(self, permutation: list[int]):
|
||||
return MutationPermutate(permutation)
|
||||
|
||||
def _regularize_index(self, index: int):
|
||||
if index < 0:
|
||||
index += self.length
|
||||
return index
|
||||
|
||||
def apply(self, mutation: Mutation, write_cache: list[Any]):
|
||||
if isinstance(mutation, MutationSet):
|
||||
write_cache[mutation.key] = mutation.value
|
||||
elif isinstance(mutation, MutationDel):
|
||||
write_cache[:] = (
|
||||
write_cache[: mutation.key] + write_cache[mutation.key + 1 :]
|
||||
)
|
||||
elif isinstance(mutation, MutationInsert):
|
||||
write_cache.insert(mutation.index, mutation.value)
|
||||
elif isinstance(mutation, MutationPermutate):
|
||||
write_cache[:] = [write_cache[i] for i in mutation.permutation]
|
||||
else:
|
||||
raise ValueError(f"Unknown mutation type {mutation}")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,228 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import sys
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ...utils import (
|
||||
BreakGraphError,
|
||||
DataDependencyControlFlowBreak,
|
||||
FallbackError,
|
||||
UnsupportedIteratorBreak,
|
||||
)
|
||||
from ...utils.exceptions import SotCapturedStopIteration
|
||||
from ..instruction_utils import Instruction
|
||||
from .dispatch_functions import generator_send
|
||||
from .opcode_executor import OpcodeExecutorBase, Stop
|
||||
from .tracker import DanglingTracker
|
||||
from .variables import (
|
||||
BuiltinVariable,
|
||||
ConstantVariable,
|
||||
GeneratorVariable,
|
||||
IterVariable,
|
||||
ObjectVariable,
|
||||
UserDefinedIterVariable,
|
||||
VariableBase,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .function_graph import FunctionGraph
|
||||
from .virtual_frame import VirtualFrame
|
||||
|
||||
|
||||
def inline_for_iter_impl(exe: OpcodeExecutorBase, instr: Instruction):
|
||||
iterator = exe.stack.top
|
||||
assert isinstance(iterator, IterVariable)
|
||||
|
||||
exe._graph.add_global_guarded_variable(iterator)
|
||||
|
||||
# simply get next
|
||||
if not isinstance(iterator, UserDefinedIterVariable):
|
||||
try:
|
||||
exe.stack.push(iterator.next())
|
||||
except SotCapturedStopIteration:
|
||||
exe.stack.pop()
|
||||
assert isinstance(instr.jump_to, Instruction)
|
||||
exe.vframe.lasti = exe.indexof(instr.jump_to)
|
||||
if sys.version_info >= (3, 12):
|
||||
assert exe._instructions[exe.vframe.lasti].opname == "END_FOR"
|
||||
skip_n_instrs = 2 if sys.version_info >= (3, 13) else 1
|
||||
exe.vframe.lasti += skip_n_instrs
|
||||
|
||||
else:
|
||||
exe._graph.remove_global_guarded_variable(iterator)
|
||||
raise BreakGraphError(
|
||||
UnsupportedIteratorBreak(
|
||||
reason_str=f"Found {iterator.__class__.__name__} as iterator."
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class OpcodeInlineExecutor(OpcodeExecutorBase):
|
||||
"""
|
||||
A class that represents an executor for inlined opcode operations.
|
||||
|
||||
Args:
|
||||
fn_variable: The function variable.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vframe: VirtualFrame,
|
||||
code_var: VariableBase,
|
||||
graph: FunctionGraph,
|
||||
):
|
||||
super().__init__(vframe, graph)
|
||||
self.return_value: VariableBase | None = None
|
||||
self._code_var = code_var
|
||||
self._name = "InlineFn"
|
||||
|
||||
def inline_call(self) -> VariableBase:
|
||||
"""
|
||||
Execute the inline call of the function.
|
||||
"""
|
||||
self._graph.add_global_guarded_variable(self._code_var)
|
||||
self.run()
|
||||
assert self.return_value is not None
|
||||
return self.return_value
|
||||
|
||||
def RETURN_VALUE(self, instr: Instruction):
|
||||
assert len(self.stack) == 1, (
|
||||
f"Stack must have one element, but get {len(self.stack)} elements."
|
||||
)
|
||||
self.return_value = self.stack.pop()
|
||||
return Stop(state="Return")
|
||||
|
||||
def RETURN_CONST(self, instr: Instruction):
|
||||
self.return_value = self.vframe.consts[instr.arg]
|
||||
return Stop(state="Return")
|
||||
|
||||
def _break_graph_when_if(self, result, instr: Instruction):
|
||||
"""
|
||||
Helper method to raise a BreakGraphError when breaking the graph in a jump operation.
|
||||
|
||||
Args:
|
||||
result: The result of the operation.
|
||||
instr (Instruction): The jump instruction.
|
||||
"""
|
||||
|
||||
raise BreakGraphError(DataDependencyControlFlowBreak())
|
||||
|
||||
def FOR_ITER(self, instr: Instruction):
|
||||
return inline_for_iter_impl(self, instr)
|
||||
|
||||
|
||||
class OpcodeInlineGeneratorExecutor(OpcodeExecutorBase):
|
||||
def __init__(
|
||||
self,
|
||||
vframe: VirtualFrame,
|
||||
code_var: VariableBase,
|
||||
graph: FunctionGraph,
|
||||
):
|
||||
super().__init__(vframe, graph)
|
||||
self.return_value: VariableBase | None = None
|
||||
self._code_var = code_var
|
||||
self._name = "InlineGen"
|
||||
|
||||
def inline_call(self) -> VariableBase:
|
||||
self._graph.add_global_guarded_variable(self._code_var)
|
||||
self.run()
|
||||
assert self.return_value is not None
|
||||
return self.return_value
|
||||
|
||||
def RETURN_GENERATOR(self, instr: Instruction):
|
||||
vframe = self.vframe
|
||||
code_var = self._code_var
|
||||
# NOTE: we set the real tracker in calling function
|
||||
self.return_value = GeneratorVariable(
|
||||
code_var, vframe, self._graph, DanglingTracker()
|
||||
)
|
||||
return Stop(state="Return")
|
||||
|
||||
def SEND(self, instr: Instruction):
|
||||
assert len(self.stack) >= 2
|
||||
recv = self.stack.pop()
|
||||
source_obj = self.stack.top
|
||||
if not isinstance(source_obj, IterVariable):
|
||||
raise FallbackError(
|
||||
"Yield from for non-generator object is not supported."
|
||||
)
|
||||
self.stack.push(
|
||||
BuiltinVariable(generator_send, self._graph, DanglingTracker())(
|
||||
source_obj, recv
|
||||
)
|
||||
)
|
||||
|
||||
def END_SEND(self, instr: Instruction):
|
||||
value = self.stack.pop()
|
||||
receiver = self.stack.pop() # pop the receiver
|
||||
self.stack.push(value)
|
||||
|
||||
def GEN_START(self, instr: Instruction):
|
||||
tos = self.stack.pop()
|
||||
assert isinstance(tos, ConstantVariable)
|
||||
assert tos.value is None
|
||||
|
||||
def YIELD_VALUE(self, instr: Instruction):
|
||||
assert len(self.stack) >= 1
|
||||
self.return_value = self.stack.pop()
|
||||
return Stop(state="Yield")
|
||||
|
||||
def GET_YIELD_FROM_ITER(self, instr: Instruction):
|
||||
source_obj = self.stack.top
|
||||
if isinstance(source_obj, ObjectVariable) and inspect.iscoroutine(
|
||||
source_obj.value
|
||||
):
|
||||
raise FallbackError(
|
||||
"Get yield from iter for coroutine object is not supported."
|
||||
)
|
||||
if isinstance(source_obj, GeneratorVariable):
|
||||
return
|
||||
source_obj = self.stack.pop()
|
||||
iter_variable = BuiltinVariable(iter, self._graph, DanglingTracker())(
|
||||
source_obj
|
||||
)
|
||||
self.stack.push(iter_variable)
|
||||
|
||||
def YIELD_FROM(self, instr: Instruction):
|
||||
recv = self.stack.pop()
|
||||
source_obj = self.stack.top
|
||||
if not isinstance(source_obj, IterVariable):
|
||||
raise FallbackError(
|
||||
"Yield from for non-generator object is not supported."
|
||||
)
|
||||
self.return_value = BuiltinVariable(
|
||||
generator_send, self._graph, DanglingTracker()
|
||||
)(source_obj, recv)
|
||||
assert self.vframe.lasti > 0
|
||||
self.vframe.lasti -= 1
|
||||
return Stop(state="Yield")
|
||||
|
||||
def FOR_ITER(self, instr: Instruction):
|
||||
return inline_for_iter_impl(self, instr)
|
||||
|
||||
def RETURN_VALUE(self, instr: Instruction):
|
||||
assert len(self.stack) == 1, (
|
||||
f"Stack must have one element, but get {len(self.stack)} elements."
|
||||
)
|
||||
self.return_value = self.stack.pop()
|
||||
return Stop(state="Return")
|
||||
|
||||
def RETURN_CONST(self, instr: Instruction):
|
||||
self.return_value = self.vframe.consts[instr.arg]
|
||||
return Stop(state="Return")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,237 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
from typing import TypeAlias
|
||||
|
||||
from .mutable_data import DataGetter, MutableData
|
||||
from .pycode_generator import PyCodeGen
|
||||
from .variables import VariableBase
|
||||
|
||||
IdGetter: TypeAlias = Callable[[Any], int]
|
||||
MutableDataT = TypeVar("MutableDataT", bound=MutableData)
|
||||
|
||||
|
||||
class SideEffectsState(NamedTuple):
|
||||
data_id_to_proxy: dict[int, MutableData]
|
||||
proxy_variables: list[VariableBase]
|
||||
mutable_variables: list[VariableBase]
|
||||
proxy_versions: list[int]
|
||||
mutable_attrs: list[dict[str, Any]]
|
||||
|
||||
|
||||
class SideEffects:
|
||||
def __init__(self):
|
||||
self.data_id_to_proxy: dict[int, MutableData] = {}
|
||||
self.proxy_variables: list[VariableBase] = []
|
||||
self.mutable_variables: list[VariableBase] = []
|
||||
|
||||
def record_proxy_variable(self, variable: VariableBase):
|
||||
if variable not in self.proxy_variables:
|
||||
self.proxy_variables.append(variable)
|
||||
|
||||
def record_mutable_variable(self, variable: VariableBase):
|
||||
if variable not in self.mutable_variables:
|
||||
self.mutable_variables.append(variable)
|
||||
|
||||
def get_proxy(
|
||||
self,
|
||||
proxy_type: type[MutableDataT],
|
||||
data: Any,
|
||||
getter: DataGetter,
|
||||
id_getter: IdGetter = id,
|
||||
) -> MutableDataT:
|
||||
data_id = id_getter(data)
|
||||
if data_id not in self.data_id_to_proxy:
|
||||
self.data_id_to_proxy[data_id] = proxy_type(data, getter)
|
||||
return self.data_id_to_proxy[data_id] # type: ignore
|
||||
|
||||
def get_state(self):
|
||||
return SideEffectsState(
|
||||
self.data_id_to_proxy.copy(),
|
||||
self.proxy_variables.copy(),
|
||||
self.mutable_variables.copy(),
|
||||
[proxy.version for proxy in self.data_id_to_proxy.values()],
|
||||
[
|
||||
{attr: getattr(var, attr)}
|
||||
for var in self.mutable_variables
|
||||
for attr in var.mutable_attrs
|
||||
],
|
||||
)
|
||||
|
||||
def restore_state(self, state: SideEffectsState):
|
||||
self.data_id_to_proxy = state.data_id_to_proxy
|
||||
self.proxy_variables = state.proxy_variables
|
||||
self.mutable_variables = state.mutable_variables
|
||||
# NOTE(SigureMo): We can use the `strict=True` option in zip after
|
||||
# Python 3.10.
|
||||
assert len(self.data_id_to_proxy.values()) == len(
|
||||
state.proxy_versions
|
||||
), "proxy_versions length not match"
|
||||
assert sum(
|
||||
len(var.mutable_attrs) for var in self.mutable_variables
|
||||
) == len(state.mutable_attrs), "mutable_attrs length not match"
|
||||
|
||||
for proxy, version in zip(
|
||||
self.data_id_to_proxy.values(), state.proxy_versions
|
||||
):
|
||||
proxy.rollback(version)
|
||||
|
||||
for (variable, attr), attr_dict in zip(
|
||||
(
|
||||
(var, attr)
|
||||
for var in self.mutable_variables
|
||||
for attr in var.mutable_attrs
|
||||
),
|
||||
(attr_dict for attr_dict in state.mutable_attrs),
|
||||
):
|
||||
setattr(variable, attr, attr_dict[attr])
|
||||
|
||||
|
||||
class SideEffectRestorer:
|
||||
def pre_gen(self, codegen: PyCodeGen):
|
||||
raise NotImplementedError
|
||||
|
||||
def post_gen(self, codegen: PyCodeGen):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class DictSideEffectRestorer(SideEffectRestorer):
|
||||
"""
|
||||
old_dict.clear()
|
||||
old_dict.update(new_dict)
|
||||
"""
|
||||
|
||||
def __init__(self, var: VariableBase):
|
||||
super().__init__()
|
||||
self.var = var
|
||||
|
||||
def pre_gen(self, codegen: PyCodeGen):
|
||||
# Reference to the original dict.
|
||||
# load old_dict.update and new_dict to stack.
|
||||
self.var.reconstruct(codegen)
|
||||
codegen.gen_load_method("update")
|
||||
# Generate dict by each key-value pair.
|
||||
self.var.reconstruct(codegen, use_tracker=False)
|
||||
# load old_dict.clear to stack.
|
||||
self.var.reconstruct(codegen)
|
||||
codegen.gen_load_method("clear")
|
||||
|
||||
def post_gen(self, codegen: PyCodeGen):
|
||||
# Call methods to apply side effects.
|
||||
codegen.gen_call_method(0) # call clear
|
||||
codegen.gen_pop_top()
|
||||
codegen.gen_call_method(1) # call update
|
||||
codegen.gen_pop_top()
|
||||
|
||||
|
||||
class ListSideEffectRestorer(SideEffectRestorer):
|
||||
"""
|
||||
old_list[:] = new_list
|
||||
"""
|
||||
|
||||
def __init__(self, var: VariableBase):
|
||||
super().__init__()
|
||||
self.var = var
|
||||
|
||||
def pre_gen(self, codegen: PyCodeGen):
|
||||
# Reference to the original list.
|
||||
# load new_list to stack.
|
||||
self.var.reconstruct(codegen, use_tracker=False)
|
||||
# load old_list[:] to stack.
|
||||
self.var.reconstruct(codegen)
|
||||
codegen.gen_load_const(None)
|
||||
codegen.gen_load_const(None)
|
||||
codegen.gen_build_slice(2)
|
||||
|
||||
def post_gen(self, codegen: PyCodeGen):
|
||||
# Call STORE_SUBSCR to apply side effects.
|
||||
codegen.gen_store_subscr()
|
||||
|
||||
|
||||
class GlobalSetSideEffectRestorer(SideEffectRestorer):
|
||||
"""
|
||||
global_var = new_value
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, var: VariableBase):
|
||||
super().__init__()
|
||||
self.name = name
|
||||
self.var = var
|
||||
|
||||
def pre_gen(self, codegen: PyCodeGen):
|
||||
self.var.reconstruct(codegen)
|
||||
|
||||
def post_gen(self, codegen: PyCodeGen):
|
||||
codegen.gen_store_global(self.name)
|
||||
|
||||
|
||||
class GlobalDelSideEffectRestorer(SideEffectRestorer):
|
||||
"""
|
||||
del global_var
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__()
|
||||
self.name = name
|
||||
|
||||
def pre_gen(self, codegen: PyCodeGen):
|
||||
# do nothing
|
||||
...
|
||||
|
||||
def post_gen(self, codegen: PyCodeGen):
|
||||
codegen.gen_delete_global(self.name)
|
||||
|
||||
|
||||
class ObjSetSideEffectRestorer(SideEffectRestorer):
|
||||
"""
|
||||
obj.attr = new_value
|
||||
"""
|
||||
|
||||
def __init__(self, obj: VariableBase, name: str, var: VariableBase):
|
||||
super().__init__()
|
||||
self.obj = obj
|
||||
self.name = name
|
||||
self.var = var
|
||||
|
||||
def pre_gen(self, codegen: PyCodeGen):
|
||||
# value
|
||||
self.var.reconstruct(codegen)
|
||||
# obj
|
||||
self.obj.reconstruct(codegen)
|
||||
|
||||
def post_gen(self, codegen: PyCodeGen):
|
||||
codegen.gen_store_attr(self.name)
|
||||
|
||||
|
||||
class ObjDelSideEffectRestorer(SideEffectRestorer):
|
||||
"""
|
||||
del obj.attr
|
||||
"""
|
||||
|
||||
def __init__(self, obj: VariableBase, name: str):
|
||||
super().__init__()
|
||||
self.obj = obj
|
||||
self.name = name
|
||||
|
||||
def pre_gen(self, codegen: PyCodeGen):
|
||||
self.obj.reconstruct(codegen)
|
||||
|
||||
def post_gen(self, codegen: PyCodeGen):
|
||||
codegen.gen_delete_attr(self.name)
|
||||
@@ -0,0 +1,619 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import sys
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
|
||||
from ...utils import InnerError, NameGenerator
|
||||
from .guard import StringifiedExpression, stringify_pyobject, union_free_vars
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from ...utils.magic_methods import BinaryOp, UnaryOp
|
||||
from .pycode_generator import PyCodeGen
|
||||
from .variables import FunctionVariable, VariableBase
|
||||
|
||||
|
||||
class Tracker:
|
||||
"""
|
||||
Tracker is a base class responsible for tracking variables or objects in Python code.
|
||||
It is used to identify how a variable is derived from the initial state of the frame.
|
||||
|
||||
Args:
|
||||
inputs: The list of variables to be tracked.
|
||||
|
||||
Note:
|
||||
It serves as an abstract class and should not be instantiated directly.
|
||||
"""
|
||||
|
||||
inputs: Sequence[VariableBase]
|
||||
name_generator = NameGenerator("tracker_")
|
||||
|
||||
def __init__(self, inputs: Sequence[VariableBase], changed: bool = False):
|
||||
self.inputs = inputs
|
||||
self.changed = changed
|
||||
self.id = Tracker.name_generator.next()
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen) -> None:
|
||||
"""
|
||||
Generate instructions based on the tracked variables.
|
||||
|
||||
Args:
|
||||
codegen (PyCodeGen): An instance of PyCodeGen to generate instructions.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
raise NotImplementedError(
|
||||
f"{self.__class__.__name__} has no guard_tree_expr_node"
|
||||
)
|
||||
|
||||
# TODO(xiongkun): trace_value_from_frame is not a good name, it should be more related to guard but not traceable.
|
||||
def trace_value_from_frame(self) -> StringifiedExpression:
|
||||
"""
|
||||
Trace the value of the tracked variables from the frame. It used for generating the guard.
|
||||
|
||||
Returns:
|
||||
The value of the tracked variables.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def is_traceable(self) -> bool:
|
||||
"""
|
||||
Determine if all the tracked variables can be traced from the frame.
|
||||
|
||||
Returns:
|
||||
bool: True if all tracked variables are traceable, False otherwise.
|
||||
"""
|
||||
if self.changed:
|
||||
return False
|
||||
for input in self.inputs:
|
||||
if not input.tracker.is_traceable():
|
||||
return False
|
||||
return True
|
||||
|
||||
def need_guard(self) -> bool:
|
||||
return self.is_traceable()
|
||||
|
||||
|
||||
class DummyTracker(Tracker):
|
||||
"""
|
||||
DummyTracker is a subclass of Tracker that specifically tracks variables cannot be reproduced from the frame.
|
||||
It is mostly generated by complex operations (instructions).
|
||||
|
||||
Args:
|
||||
inputs (list[VariableBase]): The input variables associated with the generated variables.
|
||||
"""
|
||||
|
||||
def __init__(self, inputs: Sequence[VariableBase]):
|
||||
super().__init__(inputs)
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
raise InnerError("DummyTracker has no instructions")
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
raise InnerError("DummyTracker can't trace value from frame")
|
||||
|
||||
def is_traceable(self):
|
||||
return False
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"DummyTracker(num_inputs={len(self.inputs)})"
|
||||
|
||||
def need_guard(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class SymbolicOperationTracker(Tracker):
|
||||
"""
|
||||
SymbolicOperationTracker is a subclass of Tracker that specifically tracks variables cannot be reproduced from the frame.
|
||||
It is mostly generated by complex operations of symbolic variables.
|
||||
|
||||
Args:
|
||||
inputs (list[VariableBase]): The input variables associated with the generated variables.
|
||||
"""
|
||||
|
||||
def __init__(self, inputs: Sequence[VariableBase], op: UnaryOp | BinaryOp):
|
||||
super().__init__(inputs)
|
||||
self.op = op
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
raise InnerError("SymbolicOperationTracker has no instructions")
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
raise InnerError(
|
||||
"SymbolicOperationTracker can't trace value from frame"
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"SymbolicOperationTracker(num_inputs={len(self.inputs)})"
|
||||
|
||||
def is_traceable(self):
|
||||
return False
|
||||
|
||||
|
||||
class DanglingTracker(Tracker):
|
||||
"""
|
||||
DanglingTracker is a subclass of Tracker that specifically tracks variables that are not in the frame.
|
||||
Variables whose tracker is DanglingTracker should not be placed on the stack, except for NullVariable.
|
||||
DanglingTracker is often used in conjunction with BuiltinVariable to reuse the dispatch mechanism.
|
||||
|
||||
Examples:
|
||||
>>> import operator
|
||||
>>> from sot.opcode_translator.executor.variables import (
|
||||
... BuiltinVariable,
|
||||
... ConstantVariable,
|
||||
... )
|
||||
>>> a = ConstantVariable.wrap_literal(1, None)
|
||||
>>> b = ConstantVariable.wrap_literal(2, None)
|
||||
>>> c = BuiltinVariable(operator.add, None, DanglingTracker())(a, b)
|
||||
>>> c.value
|
||||
3
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__([])
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
raise InnerError("DanglingTracker has no instructions")
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
raise InnerError("DanglingTracker can't trace value from frame")
|
||||
|
||||
def is_traceable(self):
|
||||
return False
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "DanglingTracker()"
|
||||
|
||||
|
||||
class LocalTracker(Tracker):
|
||||
"""
|
||||
LocalTracker is a subclass of Tracker that specifically tracks variables from f_locals of frame.
|
||||
|
||||
Args:
|
||||
name (str): The name of the variable in f_locals to be tracked.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__([])
|
||||
self.name = name
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen) -> None:
|
||||
codegen.gen_load_fast(self.name)
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
return paddle.framework.core.LocalVarExprNode(self.name)
|
||||
|
||||
def trace_value_from_frame(self) -> StringifiedExpression:
|
||||
return StringifiedExpression(f"frame.f_locals['{self.name}']", [], {})
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"LocalTracker(name={self.name})"
|
||||
|
||||
|
||||
class CellTracker(LocalTracker):
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
codegen.gen_load_deref(self.name)
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
return paddle.framework.core.LocalVarExprNode(self.name)
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
return StringifiedExpression(f"frame.f_locals['{self.name}']", [], {})
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"CellTracker(name={self.name})"
|
||||
|
||||
|
||||
class GlobalTracker(Tracker):
|
||||
"""
|
||||
GlobalTracker is a subclass of Tracker that specifically tracks variables from f_globals of frame.
|
||||
|
||||
Args:
|
||||
name (str): The name of the variable in f_globals to be tracked.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__([])
|
||||
self.name = name
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen) -> None:
|
||||
codegen.gen_load_global(self.name, push_null=False)
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
return paddle.framework.core.GlobalVarExprNode(self.name)
|
||||
|
||||
def trace_value_from_frame(self) -> StringifiedExpression:
|
||||
return StringifiedExpression(f"frame.f_globals['{self.name}']", [], {})
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"GlobalTracker(name={self.name})"
|
||||
|
||||
|
||||
class BuiltinTracker(Tracker):
|
||||
"""
|
||||
BuiltinTracker is a subclass of Tracker that specifically tracks variables from f_builtins of frame.
|
||||
|
||||
Args:
|
||||
name (str): The name of the variable in f_builtins to be tracked.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__([])
|
||||
self.name = name
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen) -> None:
|
||||
codegen.gen_load_global(self.name, push_null=False)
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
return paddle.framework.core.ConstantExprNode(
|
||||
getattr(builtins, self.name)
|
||||
)
|
||||
|
||||
def trace_value_from_frame(self) -> StringifiedExpression:
|
||||
return StringifiedExpression(
|
||||
f"builtins.__dict__['{self.name}']", [], {"builtins": builtins}
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"BuiltinTracker(name={self.name})"
|
||||
|
||||
|
||||
class ConstTracker(Tracker):
|
||||
"""
|
||||
ConstTracker is a subclass of Tracker that specifically tracks a constant value.
|
||||
|
||||
Args:
|
||||
value (Any): The value of the constant.
|
||||
"""
|
||||
|
||||
def __init__(self, value):
|
||||
super().__init__([])
|
||||
self.value = value
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
codegen.gen_load_const(self.value)
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
return paddle.framework.core.ConstantExprNode(self.value)
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
value_str, value_free_vars = stringify_pyobject(self.value)
|
||||
return StringifiedExpression(
|
||||
value_str, [], union_free_vars(value_free_vars)
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"ConstTracker(value={self.value})"
|
||||
|
||||
def need_guard(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class GetAttrTracker(Tracker):
|
||||
"""
|
||||
GetAttrTracker is a subclass of Tracker that specifically tracks the attribute access of an variable.
|
||||
|
||||
Args:
|
||||
obj (VariableBase): The object whose attribute is to be tracked.
|
||||
attr (str): The attribute to be tracked.
|
||||
"""
|
||||
|
||||
def __init__(self, obj: VariableBase, attr: str, changed: bool = False):
|
||||
super().__init__([obj], changed)
|
||||
self.obj = obj
|
||||
self.attr = attr
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
self.obj.tracker.gen_instructions(codegen)
|
||||
codegen.gen_load_attr(self.attr)
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
obj_tracer = self.obj.tracker.guard_tree_expr_node()
|
||||
return paddle.framework.core.AttributeExprNode(
|
||||
obj_tracer,
|
||||
self.attr,
|
||||
)
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
obj_tracer = self.obj.tracker.trace_value_from_frame()
|
||||
if self.attr.isidentifier():
|
||||
expr = f"{{}}.{self.attr}"
|
||||
else:
|
||||
expr = f"getattr({{}}, '{self.attr}')"
|
||||
return StringifiedExpression(
|
||||
expr,
|
||||
[obj_tracer],
|
||||
union_free_vars(obj_tracer.free_vars),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"GetAttrTracker(attr={self.attr})"
|
||||
|
||||
def need_guard(self) -> bool:
|
||||
return self.is_traceable() and self.obj.tracker.need_guard()
|
||||
|
||||
|
||||
class GetItemTracker(Tracker):
|
||||
"""
|
||||
GetItemTracker is a subclass of Tracker that specifically tracks item access of a container variable.
|
||||
|
||||
It generates instructions and traces the item value from the frame.
|
||||
|
||||
Args:
|
||||
container_var (VariableBase): The container object whose item is to be tracked.
|
||||
key: The key/index of the item to be tracked.
|
||||
"""
|
||||
|
||||
def __init__(self, container_var: VariableBase, key: object, changed=False):
|
||||
super().__init__([container_var], changed)
|
||||
self.container = container_var
|
||||
self.key = key
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
self.container.tracker.gen_instructions(codegen)
|
||||
if isinstance(self.key, slice):
|
||||
codegen.gen_load_const(self.key.start)
|
||||
codegen.gen_load_const(self.key.stop)
|
||||
codegen.gen_load_const(self.key.step)
|
||||
codegen.gen_build_slice(3)
|
||||
else:
|
||||
codegen.gen_load_const(self.key)
|
||||
codegen.gen_subscribe()
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
container_tracer = self.container.tracker.guard_tree_expr_node()
|
||||
return paddle.framework.core.ItemExprNode(
|
||||
container_tracer,
|
||||
paddle.framework.core.ConstantExprNode(self.key),
|
||||
)
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
container_tracer = self.container.tracker.trace_value_from_frame()
|
||||
key_string, key_free_vars = stringify_pyobject(self.key)
|
||||
return StringifiedExpression(
|
||||
f"{{}}[{key_string}]",
|
||||
[container_tracer],
|
||||
union_free_vars(container_tracer.free_vars, key_free_vars),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"GetItemTracker(key={self.key!r})"
|
||||
|
||||
def need_guard(self) -> bool:
|
||||
return self.is_traceable() and self.container.tracker.need_guard()
|
||||
|
||||
|
||||
class GetIterTracker(Tracker):
|
||||
"""
|
||||
GetIterTracker is a subclass of Tracker that specifically tracks iteration of a variable.
|
||||
|
||||
It generates instructions and traces the iterator from the frame.
|
||||
|
||||
Args:
|
||||
iter_source (VariableBase): The source variable to be iterated.
|
||||
"""
|
||||
|
||||
def __init__(self, iter_source: VariableBase):
|
||||
super().__init__([iter_source])
|
||||
self.iter_source = iter_source
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
self.iter_source.tracker.gen_instructions(codegen)
|
||||
codegen.add_instr("GET_ITER")
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
# TODO(zrr1999): implement IterExprNode
|
||||
raise NotImplementedError("IterExprNode is not implemented")
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
iter_source_tracer = self.iter_source.tracker.trace_value_from_frame()
|
||||
return StringifiedExpression(
|
||||
"iter({})",
|
||||
[iter_source_tracer],
|
||||
union_free_vars(iter_source_tracer.free_vars),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "GetIterTracker()"
|
||||
|
||||
|
||||
class CreateLayerTracker(Tracker):
|
||||
def __init__(self, layer_class, args, kwargs):
|
||||
super().__init__([layer_class, *list(args), *list(kwargs.values())])
|
||||
self.layer_class = layer_class
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
if sys.version_info >= (3, 11):
|
||||
codegen.gen_push_null()
|
||||
|
||||
self.layer_class.reconstruct(codegen)
|
||||
for variable in self.args:
|
||||
variable.reconstruct(codegen)
|
||||
|
||||
if len(self.kwargs) == 0:
|
||||
codegen.gen_call_function(argc=len(self.args))
|
||||
else:
|
||||
codegen.gen_build_tuple(len(self.args))
|
||||
for k, v in self.kwargs.items():
|
||||
codegen.gen_load_const(k)
|
||||
v.reconstruct(codegen)
|
||||
codegen.gen_build_map(len(self.kwargs))
|
||||
codegen.gen_call_function_ex(has_kwargs=True)
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
# TODO(zrr1999): implement LayerExprNode.guard_tree_expr_node
|
||||
raise NotImplementedError("LayerExprNode is not implemented")
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
class_tracer = self.layer_class.tracker.trace_value_from_frame()
|
||||
arg_tracers = [
|
||||
arg.tracker.trace_value_from_frame() for arg in self.args
|
||||
]
|
||||
kwarg_tracers_dict = {
|
||||
k: v.tracker.trace_value_from_frame()
|
||||
for k, v in self.kwargs.items()
|
||||
}
|
||||
kwarg_tracers = list(kwarg_tracers_dict.values())
|
||||
|
||||
expr = "{}("
|
||||
expr += ", ".join(["{}"] * len(arg_tracers))
|
||||
if len(arg_tracers) and len(kwarg_tracers) > 0:
|
||||
expr += ", "
|
||||
expr += ", ".join(f"{k}={{}}" for k in kwarg_tracers_dict.keys())
|
||||
expr += ")"
|
||||
|
||||
return StringifiedExpression(
|
||||
expr,
|
||||
[class_tracer, *arg_tracers, *kwarg_tracers],
|
||||
union_free_vars(
|
||||
*(
|
||||
tracer.free_vars
|
||||
for tracer in chain(
|
||||
[class_tracer], arg_tracers, kwarg_tracers
|
||||
)
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"CreateLayerTracker(Layer={self.layer_class}, args={self.args}, kwargs={self.kwargs})"
|
||||
|
||||
|
||||
class FunctionClosureTracker(Tracker):
|
||||
"""
|
||||
A tracker class that represents a function closure variable.
|
||||
|
||||
Args:
|
||||
fn: The FunctionVariable object.
|
||||
idx: The index of the closure variable.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, fn: FunctionVariable, idx: int):
|
||||
super().__init__([fn])
|
||||
self.fn = fn
|
||||
self.idx = idx
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
"""
|
||||
Generate bytecode instructions to trace the value of the function closure variable.
|
||||
|
||||
Args:
|
||||
codegen: The PyCodeGen object used to generate bytecode.
|
||||
|
||||
"""
|
||||
self.fn.tracker.gen_instructions(codegen)
|
||||
codegen.gen_load_attr("__closure__")
|
||||
codegen.gen_load_const(self.idx)
|
||||
codegen.gen_subscribe()
|
||||
codegen.gen_load_attr("cell_contents")
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
fn_tracer = self.fn.tracker.guard_tree_expr_node()
|
||||
return paddle.framework.core.AttributeExprNode(
|
||||
paddle.framework.core.ItemExprNode(
|
||||
paddle.framework.core.AttributeExprNode(
|
||||
fn_tracer,
|
||||
"__closure__",
|
||||
),
|
||||
paddle.framework.core.ConstantExprNode(self.idx),
|
||||
),
|
||||
"cell_contents",
|
||||
)
|
||||
|
||||
def trace_value_from_frame(self):
|
||||
"""
|
||||
Trace the value of the function closure variable from the frame.
|
||||
|
||||
Returns:
|
||||
The traced value of the function closure variable.
|
||||
|
||||
"""
|
||||
fn_tracer = self.fn.tracker.trace_value_from_frame()
|
||||
return StringifiedExpression(
|
||||
f"{{}}.__closure__[{self.idx}].cell_contents",
|
||||
[fn_tracer],
|
||||
union_free_vars(fn_tracer.free_vars),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"FunctionClosureTracker(fn={self.fn}, idx={self.idx})"
|
||||
|
||||
|
||||
class FunctionGlobalTracker(Tracker):
|
||||
"""
|
||||
A tracker class that represents a function global variable.
|
||||
|
||||
Args:
|
||||
fn: FunctionVariable object.
|
||||
name: The name of the global variable.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, fn: FunctionVariable, name: str):
|
||||
super().__init__([fn])
|
||||
self.fn = fn
|
||||
self.name = name
|
||||
|
||||
def gen_instructions(self, codegen: PyCodeGen):
|
||||
"""
|
||||
Generate bytecode instructions in order to put the variables at the top of the stack.
|
||||
|
||||
Args:
|
||||
codegen: The PyCodeGen object used to generate bytecode.
|
||||
|
||||
"""
|
||||
self.fn.tracker.gen_instructions(codegen)
|
||||
codegen.gen_load_attr("__globals__")
|
||||
codegen.gen_load_const(self.name)
|
||||
codegen.gen_subscribe()
|
||||
|
||||
def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
|
||||
fn_tracer = self.fn.tracker.guard_tree_expr_node()
|
||||
return paddle.framework.core.ItemExprNode(
|
||||
paddle.framework.core.AttributeExprNode(
|
||||
fn_tracer,
|
||||
"__globals__",
|
||||
),
|
||||
paddle.framework.core.ConstantExprNode(self.name),
|
||||
)
|
||||
|
||||
def trace_value_from_frame(self) -> StringifiedExpression:
|
||||
"""
|
||||
Trace the value of the function global variable from the frame.
|
||||
|
||||
Returns:
|
||||
StringifiedExpression: The traced value of the function global variable.
|
||||
|
||||
"""
|
||||
fn_tracer = self.fn.tracker.trace_value_from_frame()
|
||||
return StringifiedExpression(
|
||||
f"{{}}.__globals__['{self.name}']",
|
||||
[fn_tracer],
|
||||
union_free_vars(fn_tracer.free_vars),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"FunctionGlobalTracker(fn={self.fn}, name={self.name})"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,217 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
ValidateValueFunc = Callable[[Any], None]
|
||||
|
||||
|
||||
StackDataT = TypeVar("StackDataT")
|
||||
|
||||
|
||||
class VariableStack(Generic[StackDataT]):
|
||||
"""
|
||||
A stack class for storing variables.
|
||||
|
||||
Examples:
|
||||
>>> var1, var2, var3, var4 = range(1, 5)
|
||||
>>> stack = VariableStack()
|
||||
>>> stack.push(var1)
|
||||
>>> stack.push(var3)
|
||||
>>> stack.insert(1, var2)
|
||||
>>> stack
|
||||
[1, 2, 3]
|
||||
>>> stack.pop()
|
||||
3
|
||||
>>> stack.pop_n(2)
|
||||
[1, 2]
|
||||
>>> stack.push(var1)
|
||||
>>> stack.push(var2)
|
||||
>>> stack.push(var3)
|
||||
>>> stack
|
||||
[1, 2, 3]
|
||||
>>> stack.top
|
||||
3
|
||||
>>> stack.peek[1]
|
||||
3
|
||||
>>> stack.peek[:1]
|
||||
[3]
|
||||
>>> stack.peek[:2]
|
||||
[2, 3]
|
||||
>>> stack.peek[1] = var4
|
||||
>>> stack
|
||||
[1, 2, 4]
|
||||
|
||||
"""
|
||||
|
||||
class VariablePeeker:
|
||||
@overload
|
||||
def __getitem__(self, index: int) -> StackDataT: ...
|
||||
|
||||
@overload
|
||||
def __getitem__(self, index: slice) -> list[StackDataT]: ...
|
||||
|
||||
@overload
|
||||
def __call__(self, index: int = 1) -> StackDataT: ...
|
||||
|
||||
@overload
|
||||
def __call__(self, index: slice) -> list[StackDataT]: ...
|
||||
|
||||
def __init__(
|
||||
self, data: list[StackDataT], validate_value_func: ValidateValueFunc
|
||||
):
|
||||
self._data = data
|
||||
self.validate_value_func = validate_value_func
|
||||
|
||||
def __getitem__(
|
||||
self, index: int | slice
|
||||
) -> StackDataT | list[StackDataT]:
|
||||
if isinstance(index, int):
|
||||
assert 0 < index <= len(self._data)
|
||||
return self._data[-index]
|
||||
if isinstance(index, slice):
|
||||
assert index.start is None and index.step is None, (
|
||||
"slice which has start or step not supported"
|
||||
)
|
||||
assert 0 < index.stop <= len(self._data)
|
||||
return self._data[-index.stop :]
|
||||
raise NotImplementedError(f"index type {type(index)} not supported")
|
||||
|
||||
def __setitem__(self, index: int, value: Any):
|
||||
assert isinstance(index, int), (
|
||||
f"index type {type(index)} not supported"
|
||||
)
|
||||
assert 0 < index <= len(self._data), (
|
||||
f"index should be in [1, {len(self._data)}], but get {index}"
|
||||
)
|
||||
self.validate_value_func(value)
|
||||
self._data[-index] = value
|
||||
|
||||
def __call__(
|
||||
self, index: int | slice = 1
|
||||
) -> StackDataT | list[StackDataT]:
|
||||
return self[index]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data: list[StackDataT] | None = None,
|
||||
*,
|
||||
validate_value_func: ValidateValueFunc | None = None,
|
||||
):
|
||||
if data is None:
|
||||
data = []
|
||||
else:
|
||||
data = data.copy()
|
||||
self.validate_value_func = (
|
||||
(lambda _: None)
|
||||
if validate_value_func is None
|
||||
else validate_value_func
|
||||
)
|
||||
self._data = data
|
||||
self._peeker = VariableStack.VariablePeeker(
|
||||
self._data, self.validate_value_func
|
||||
)
|
||||
|
||||
def copy(self):
|
||||
return VariableStack(
|
||||
self._data, validate_value_func=self.validate_value_func
|
||||
)
|
||||
|
||||
def push(self, val: StackDataT):
|
||||
"""
|
||||
Pushes a variable onto the stack.
|
||||
|
||||
Args:
|
||||
val: The variable to be pushed.
|
||||
|
||||
"""
|
||||
self.validate_value_func(val)
|
||||
self._data.append(val)
|
||||
|
||||
def insert(self, index: int, val: StackDataT):
|
||||
"""
|
||||
Inserts a variable onto the stack.
|
||||
|
||||
Args:
|
||||
index: The index at which the variable is to be inserted, the top of the stack is at index 0.
|
||||
val: The variable to be inserted.
|
||||
|
||||
"""
|
||||
assert 0 <= index <= len(self), (
|
||||
f"index should be in [0, {len(self)}], but get {index}"
|
||||
)
|
||||
self.validate_value_func(val)
|
||||
self._data.insert(len(self) - index, val)
|
||||
|
||||
def pop(self) -> StackDataT:
|
||||
"""
|
||||
Pops the top value from the stack.
|
||||
|
||||
Returns:
|
||||
The popped value.
|
||||
|
||||
"""
|
||||
assert len(self) > 0, "stack is empty"
|
||||
return self._data.pop()
|
||||
|
||||
def pop_n(self, n: int) -> list[StackDataT]:
|
||||
"""
|
||||
Pops the top n values from the stack.
|
||||
|
||||
Args:
|
||||
n: The number of values to pop.
|
||||
|
||||
Returns:
|
||||
A list of the popped values.
|
||||
|
||||
"""
|
||||
assert len(self) >= n >= 0, (
|
||||
f"n should be in [0, {len(self)}], but get {n}"
|
||||
)
|
||||
if n == 0:
|
||||
return []
|
||||
retval = self._data[-n:]
|
||||
self._data[-n:] = []
|
||||
return retval
|
||||
|
||||
@property
|
||||
def peek(self) -> VariablePeeker:
|
||||
return self._peeker
|
||||
|
||||
@property
|
||||
def top(self) -> StackDataT:
|
||||
assert len(self) > 0, "stack is empty"
|
||||
return self.peek[1]
|
||||
|
||||
@top.setter
|
||||
def top(self, value):
|
||||
assert len(self) > 0, "stack is empty"
|
||||
self.peek[1] = value
|
||||
|
||||
def __contains__(self, value):
|
||||
return value in self._data
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._data)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._data)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return str(self._data)
|
||||
@@ -0,0 +1,80 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from .base import ( # noqa: F401
|
||||
VariableBase,
|
||||
VariableFactory,
|
||||
find_traceable_vars,
|
||||
map_variables,
|
||||
)
|
||||
from .basic import ( # noqa: F401
|
||||
CellVariable,
|
||||
ConstantVariable,
|
||||
DataClassInstanceVariable,
|
||||
DataVariable,
|
||||
DygraphTracerVariable,
|
||||
EnumVariable,
|
||||
ExceptionVariable,
|
||||
FunctionGlobalVariable,
|
||||
GlobalVariable,
|
||||
InterpolationVariable,
|
||||
ModuleVariable,
|
||||
NullVariable,
|
||||
NumPyArrayVariable,
|
||||
NumPyNumberVariable,
|
||||
NumPyVariable,
|
||||
ObjectVariable,
|
||||
ParameterVariable,
|
||||
PlaceVariable,
|
||||
SliceVariable,
|
||||
SuperVariable,
|
||||
SymbolicVariable,
|
||||
TemplateVariable,
|
||||
TensorVariable,
|
||||
)
|
||||
from .callable import ( # noqa: F401
|
||||
BuiltinVariable,
|
||||
CallableVariable,
|
||||
ClassVariable,
|
||||
ContainerLayerVariable,
|
||||
DataClassVariable,
|
||||
FunctionVariable,
|
||||
LayerVariable,
|
||||
MethodVariable,
|
||||
NumPyApiVariable,
|
||||
PaddleApiVariable,
|
||||
PaddleLayerVariable,
|
||||
PartialVariable,
|
||||
UserCodeVariable,
|
||||
UserDefinedFunctionVariable,
|
||||
UserDefinedGeneratorFunctionVariable,
|
||||
UserDefinedLayerVariable,
|
||||
)
|
||||
from .container import ( # noqa: F401
|
||||
ContainerVariable,
|
||||
DictVariable,
|
||||
ListVariable,
|
||||
RangeVariable,
|
||||
SizeVariable,
|
||||
TupleVariable,
|
||||
)
|
||||
from .iter import ( # noqa: F401
|
||||
EnumerateVariable,
|
||||
GeneratorVariable,
|
||||
IterVariable,
|
||||
MapVariable,
|
||||
SequenceIterVariable,
|
||||
UserDefinedIterVariable,
|
||||
ZipVariable,
|
||||
)
|
||||
@@ -0,0 +1,731 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import operator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import fields
|
||||
from functools import cached_property
|
||||
from queue import Queue
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import paddle
|
||||
from paddle.jit.dy2static.utils import (
|
||||
dataclass_from_dict,
|
||||
)
|
||||
|
||||
from ....profiler import event_register
|
||||
from ....utils import (
|
||||
NameGenerator,
|
||||
get_unbound_method,
|
||||
log,
|
||||
)
|
||||
from ....utils.exceptions import FallbackError, HasNoAttributeError
|
||||
from ..dispatcher import Dispatcher
|
||||
from ..guard import (
|
||||
FasterStringifiedExpression,
|
||||
StringifiedExpression,
|
||||
check_faster_guard,
|
||||
check_guard,
|
||||
union_free_vars,
|
||||
)
|
||||
from ..mutable_data import MutableDictLikeData
|
||||
from ..tracker import (
|
||||
BuiltinTracker,
|
||||
ConstTracker,
|
||||
DummyTracker,
|
||||
GetAttrTracker,
|
||||
GetItemTracker,
|
||||
GetIterTracker,
|
||||
GlobalTracker,
|
||||
LocalTracker,
|
||||
Tracker,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
from typing import TypeAlias
|
||||
|
||||
from ..function_graph import FunctionGraph
|
||||
from ..pycode_generator import PyCodeGen
|
||||
|
||||
# Each variable object should implement a method called `from_value`,
|
||||
# which should adhere to the FromValueFunc signature.
|
||||
FromValueFunc: TypeAlias = Callable[
|
||||
[Any, FunctionGraph, Tracker], "VariableBase | None"
|
||||
]
|
||||
|
||||
|
||||
@event_register("find_traceable_vars")
|
||||
def find_traceable_vars(
|
||||
root_vars: list[VariableBase],
|
||||
) -> list[VariableBase]:
|
||||
"""
|
||||
This function is used to find all traceable variables in the given list of variables.
|
||||
|
||||
Args:
|
||||
root_vars (list[VariableBase]): A list of root variables from which the ordering starts.
|
||||
|
||||
Returns:
|
||||
list[VariableBase]: A list of variables that are traceable.
|
||||
"""
|
||||
results: list[VariableBase] = []
|
||||
visited: set[VariableBase] = set()
|
||||
queue: Queue[VariableBase] = Queue()
|
||||
|
||||
for root in root_vars:
|
||||
queue.put(root)
|
||||
|
||||
while not queue.empty():
|
||||
var = queue.get()
|
||||
if var in visited:
|
||||
continue
|
||||
|
||||
visited.add(var)
|
||||
if var.tracker.need_guard():
|
||||
results.append(var)
|
||||
continue
|
||||
|
||||
# Pruning traceable variable, if the variable is traceable, we don't need to
|
||||
# trace its inputs.
|
||||
inputs = var.get_inputs()
|
||||
|
||||
for var in inputs:
|
||||
if var not in visited and var not in queue.queue:
|
||||
queue.put(var)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def map_variables(
|
||||
map_func,
|
||||
variables: list[VariableBase],
|
||||
*,
|
||||
restore_variable=False,
|
||||
) -> list[VariableBase]:
|
||||
"""
|
||||
This function maps the given map_func to the given list of variables in a recursive manner.
|
||||
Args:
|
||||
map_func (Callable[[VariableBase], Any]): The function to be mapped to each variable.
|
||||
variables (list[VariableBase]): A list of variables to which the map_func is to be applied.
|
||||
|
||||
Returns:
|
||||
tuple: The result of applying the map_func to the variables.
|
||||
"""
|
||||
from .basic import DataClassInstanceVariable, SliceVariable
|
||||
from .container import ContainerVariable
|
||||
|
||||
def _map_container_variable(variable: VariableBase | object):
|
||||
if not isinstance(variable, ContainerVariable):
|
||||
return variable
|
||||
new_container = paddle.utils.map_structure(
|
||||
_map_variable, variable.get_wrapped_items()
|
||||
)
|
||||
if not restore_variable:
|
||||
return new_container
|
||||
return VariableFactory.from_value(
|
||||
new_container,
|
||||
variable.graph,
|
||||
DummyTracker(paddle.utils.flatten(new_container)),
|
||||
)
|
||||
|
||||
def _map_slice_variable(variable: VariableBase | object):
|
||||
if not isinstance(variable, SliceVariable):
|
||||
return variable
|
||||
new_slice = slice(
|
||||
map_func(variable.getattr("start")),
|
||||
map_func(variable.getattr("stop")),
|
||||
map_func(variable.getattr("step")),
|
||||
)
|
||||
if not restore_variable:
|
||||
return new_slice
|
||||
return VariableFactory.from_value(
|
||||
new_slice,
|
||||
variable.graph,
|
||||
DummyTracker([new_slice.start, new_slice.stop, new_slice.step]),
|
||||
)
|
||||
|
||||
def _map_dataclass_variable(variable: VariableBase | object):
|
||||
if not isinstance(variable, DataClassInstanceVariable):
|
||||
return variable
|
||||
new_dataclass = dataclass_from_dict(
|
||||
variable.get_py_type(),
|
||||
{
|
||||
fd.name: _map_variable(variable.getattr(fd.name))
|
||||
for fd in fields(variable.get_py_type())
|
||||
},
|
||||
)
|
||||
if not restore_variable:
|
||||
return new_dataclass
|
||||
return VariableFactory.from_value(
|
||||
new_dataclass,
|
||||
variable.graph,
|
||||
DummyTracker(
|
||||
[
|
||||
variable.getattr(fd.name)
|
||||
for fd in fields(variable.get_py_type())
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
def _map_variable(variable: VariableBase | object):
|
||||
variable = _map_container_variable(variable)
|
||||
variable = _map_slice_variable(variable)
|
||||
variable = _map_dataclass_variable(variable)
|
||||
return map_func(variable)
|
||||
|
||||
return paddle.utils.map_structure(_map_variable, variables)
|
||||
|
||||
|
||||
class VariableFactory:
|
||||
"""
|
||||
A factory class for creating variables from arbitrary values.
|
||||
|
||||
This class provides a set of registration and factory methods for creating variables
|
||||
of different types based on the type of the input value.
|
||||
|
||||
"""
|
||||
|
||||
registered_funcs: dict[str, list[str]] = {"default": []}
|
||||
mapping_str_func: dict[str, FromValueFunc] = {}
|
||||
|
||||
@staticmethod
|
||||
def default_from_value(value, graph, tracker):
|
||||
"""
|
||||
A default factory function that creates an ObjectVariable from the given value.
|
||||
|
||||
Args:
|
||||
value: The input value.
|
||||
graph: The FunctionGraph object that this variable is associated with.
|
||||
tracker: The Tracker object that tracks the information of this variable.
|
||||
|
||||
Returns:
|
||||
ObjectVariable: A new ObjectVariable representing the input value.
|
||||
"""
|
||||
from .basic import ObjectVariable
|
||||
|
||||
return ObjectVariable(value, graph, tracker)
|
||||
|
||||
@staticmethod
|
||||
def register_from_value(*, successor: str | None = None):
|
||||
"""
|
||||
A decorator function that registers a function for creating a Variable from a value.
|
||||
|
||||
Args:
|
||||
successor (str | None, optional): The name of the successor function that will be called after this function when creating a Variable. If None, the function is added to a default list of functions.
|
||||
|
||||
Returns:
|
||||
The _register_from_value decorator function, which takes the function to be registered as an argument.
|
||||
"""
|
||||
registered_funcs = VariableFactory.registered_funcs
|
||||
mapping_str_func = VariableFactory.mapping_str_func
|
||||
|
||||
def _register_from_value(func: FromValueFunc):
|
||||
"""
|
||||
Function to register a function for creating a Variable from a value
|
||||
"""
|
||||
# Get the name of the function
|
||||
name = func.__qualname__.split(".")[0]
|
||||
# Map the name of the function to the function
|
||||
mapping_str_func[name] = func
|
||||
if successor is None:
|
||||
registered_funcs["default"].append(
|
||||
name
|
||||
) # If successor is None, add the function to the "default" list
|
||||
elif successor not in registered_funcs.keys():
|
||||
registered_funcs[successor] = [
|
||||
name
|
||||
] # If the successor is not in the registered_funcs dictionary, set the value to a list containing only name
|
||||
else:
|
||||
registered_funcs[successor].append(
|
||||
name
|
||||
) # If the successor is in the registered_funcs dictionary, append name to the existing list of functions for that successor
|
||||
|
||||
log(
|
||||
4, VariableFactory.registered_funcs
|
||||
) # Print the registered_funcs dictionary if the logging level is at least 4
|
||||
return _register_from_value
|
||||
|
||||
@staticmethod
|
||||
def from_value(
|
||||
value: Any,
|
||||
graph: FunctionGraph,
|
||||
tracker: Tracker,
|
||||
) -> VariableBase:
|
||||
"""
|
||||
Create a new variable object from the given value.
|
||||
|
||||
This method searches through the registered from_value functions to find one
|
||||
that can create a variable object from the given value. If no matching function
|
||||
is found, the default_from_value function is used.
|
||||
|
||||
Args:
|
||||
value (Any): The input value.
|
||||
graph (FunctionGraph): The FunctionGraph object that this variable is associated with.
|
||||
tracker (Tracker): The Tracker object that tracks the information of this variable.
|
||||
|
||||
Returns:
|
||||
VariableBase: A new variable object representing the input value.
|
||||
"""
|
||||
registered_funcs = VariableFactory.registered_funcs
|
||||
|
||||
def _find_var(key: str = "default") -> VariableBase | None:
|
||||
for name in registered_funcs[key]:
|
||||
if name in registered_funcs.keys():
|
||||
# If the function name is a key in the registered_funcs dictionary, recursively find a Variable using that function
|
||||
var = _find_var(name)
|
||||
if var is not None:
|
||||
return var
|
||||
# Get the function corresponding to the name from the mapping_str_func dictionary
|
||||
func = VariableFactory.mapping_str_func[name]
|
||||
var = func(
|
||||
value, graph, tracker
|
||||
) # Call the function to create a Variable from the value
|
||||
if var is not None:
|
||||
return var
|
||||
|
||||
var = _find_var()
|
||||
if var is None:
|
||||
var = VariableFactory.default_from_value(
|
||||
value, graph, tracker
|
||||
) # If a Variable could not be found using the registered functions, use the default function to create a new Variable
|
||||
return var
|
||||
|
||||
|
||||
def infer_debug_name_from_tracker(tracker: Tracker) -> str | None:
|
||||
res = None
|
||||
if isinstance(tracker, (LocalTracker, GlobalTracker, BuiltinTracker)):
|
||||
res = f"{tracker.name}"
|
||||
elif isinstance(tracker, ConstTracker):
|
||||
res = f"{tracker.value}"
|
||||
elif isinstance(tracker, GetItemTracker) and tracker.container.debug_name:
|
||||
res = f"{tracker.container.debug_name}[{tracker.key}]"
|
||||
elif isinstance(tracker, GetAttrTracker) and tracker.obj.debug_name:
|
||||
res = f"{tracker.obj.debug_name}.{tracker.attr}"
|
||||
return res
|
||||
|
||||
|
||||
class VariableBase:
|
||||
"""
|
||||
VariableBase is a basic concept and each symbols in VM stack is regarded as
|
||||
an Variable Object in symbolic tracing process.
|
||||
|
||||
There are two key data structures during Python runtime:
|
||||
PyFrameObject, which provides the instance for function logical lock usage,
|
||||
and PyCodeObject, which provides the bytecode for the corresponding function.
|
||||
With these data, the Python virtual machine executes the bytecode sequentially on a stack to complete function logic.
|
||||
|
||||
Args:
|
||||
tracker(Tracker): The Tracker object that tracks the information of this variable.
|
||||
|
||||
Note:
|
||||
We should push an object of a subclass of VariableBase instead of an object of VariableBase onto the VM stack.
|
||||
It serves as an abstract class and should not be instantiated directly.
|
||||
"""
|
||||
|
||||
tracker: Tracker # An attribute to store the Tracker object associated with the variable
|
||||
value: Any
|
||||
name_generator = NameGenerator(
|
||||
"object_"
|
||||
) # A class-level attribute to generate names for new variables
|
||||
mutable_attrs = []
|
||||
|
||||
def __init__(self, graph: FunctionGraph, tracker: Tracker):
|
||||
self.graph = graph
|
||||
self.tracker = tracker
|
||||
self.id = VariableBase.name_generator.next()
|
||||
self.debug_name = infer_debug_name_from_tracker(tracker)
|
||||
|
||||
@property
|
||||
def main_info(self) -> dict[str, Any]:
|
||||
"""
|
||||
Property method to return a dictionary of main information about the variable
|
||||
|
||||
Returns:
|
||||
main_info: Main information of the variable.
|
||||
"""
|
||||
return {}
|
||||
|
||||
@property
|
||||
def debug_info(self) -> dict[str, Any]:
|
||||
"""
|
||||
Property method to return a dictionary of debug information about the variable
|
||||
"""
|
||||
info = {
|
||||
"id": self.id,
|
||||
}
|
||||
if self.debug_name:
|
||||
info["debug_name"] = self.debug_name
|
||||
return info
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self.id)
|
||||
|
||||
@check_faster_guard
|
||||
def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
|
||||
expr_node = self.tracker.guard_tree_expr_node()
|
||||
return [
|
||||
paddle.framework.core.GuardNode(
|
||||
paddle.framework.core.ValueMatchGuard(self.get_py_value()),
|
||||
[expr_node],
|
||||
)
|
||||
]
|
||||
|
||||
@check_guard
|
||||
def make_stringified_guard(self) -> list[StringifiedExpression]:
|
||||
"""
|
||||
Create a StringifiedExpression object that represents a guard expression for this variable.
|
||||
|
||||
Returns:
|
||||
StringifiedExpression: An object that contains the guard expression and the free variables used in the expression.
|
||||
"""
|
||||
|
||||
# Get a ValueTracer object from the Tracker object associated with the variable
|
||||
frame_value_tracer = self.tracker.trace_value_from_frame()
|
||||
return [
|
||||
FasterStringifiedExpression(
|
||||
f"id(type({{0}})) == {id(self.get_py_type())} and {{0}} == {self.get_py_value()!r}",
|
||||
paddle.framework.core.ValueMatchGuard(self.get_py_value()),
|
||||
[frame_value_tracer],
|
||||
union_free_vars(frame_value_tracer.free_vars),
|
||||
)
|
||||
]
|
||||
|
||||
def get_py_value(self, allow_tensor=False) -> Any:
|
||||
"""
|
||||
Abstract method to get the value of the variable
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_py_type(self):
|
||||
"""
|
||||
Method to get the type of the variable's value
|
||||
"""
|
||||
return type(self.get_py_value())
|
||||
|
||||
def is_none(self) -> bool:
|
||||
"""
|
||||
Method to check if the variable's value is None
|
||||
"""
|
||||
return self.get_py_value() is None
|
||||
|
||||
def reconstruct(
|
||||
self,
|
||||
codegen: PyCodeGen,
|
||||
*,
|
||||
use_tracker: bool = True,
|
||||
add_to_global_guarded_vars: bool = True,
|
||||
):
|
||||
if self.tracker.is_traceable() and use_tracker:
|
||||
self.tracker.gen_instructions(codegen)
|
||||
else:
|
||||
if add_to_global_guarded_vars:
|
||||
self.graph.add_global_guarded_variable(self)
|
||||
self._reconstruct(codegen)
|
||||
|
||||
def _reconstruct(self, codegen: PyCodeGen) -> None:
|
||||
"""
|
||||
Abstract method to construct an opcode and append it into codegen.instructions
|
||||
"""
|
||||
raise FallbackError(
|
||||
f'{self.__class__.__name__} does not implement "_reconstruct" method'
|
||||
)
|
||||
|
||||
def flatten_inner_vars(self) -> list[VariableBase]:
|
||||
"""
|
||||
Recursively flatten the items in this container variable to a list of Variable objects.
|
||||
|
||||
Returns:
|
||||
list[VariableBase]: Flattened items of a container variable.
|
||||
"""
|
||||
return [self]
|
||||
|
||||
def get_inputs(self) -> list[VariableBase]:
|
||||
"""
|
||||
This method is used to get the inputs for the current variable.
|
||||
|
||||
Returns:
|
||||
list[VariableBase]: Inputs for the current variable.
|
||||
"""
|
||||
return self.tracker.inputs
|
||||
|
||||
def get_traceable_inputs(self) -> list[VariableBase]:
|
||||
"""
|
||||
This method is used to get the traceable inputs for the current variable.
|
||||
|
||||
Returns:
|
||||
list[VariableBase]: Traceable inputs for the current variable.
|
||||
"""
|
||||
return list(
|
||||
filter(lambda x: x.tracker.is_traceable(), self.tracker.inputs)
|
||||
)
|
||||
|
||||
def call_function(self, /, *args, **kwargs):
|
||||
pass
|
||||
|
||||
@cached_property
|
||||
def attr_proxy(self):
|
||||
return self.graph.side_effects.get_proxy(
|
||||
MutableDictLikeData, self.get_py_value(), self.attr_proxy_getter
|
||||
)
|
||||
|
||||
def attr_proxy_getter(self, proxy: MutableDictLikeData, name: str):
|
||||
if not hasattr(proxy.original_data, name): # can't true.
|
||||
return MutableDictLikeData.Empty()
|
||||
|
||||
attr = getattr(proxy.original_data, name)
|
||||
if inspect.ismethod(attr) or (
|
||||
hasattr(attr, "__self__")
|
||||
and inspect.ismethoddescriptor(
|
||||
getattr(attr.__self__.__class__, name, None)
|
||||
)
|
||||
):
|
||||
from .callable import MethodVariable
|
||||
|
||||
fn = None
|
||||
instance = self
|
||||
if inspect.ismethoddescriptor(
|
||||
getattr(attr.__self__.__class__, name, None)
|
||||
):
|
||||
class_var = VariableFactory.from_value(
|
||||
self.get_py_type(),
|
||||
self.graph,
|
||||
GetAttrTracker(self, "__class__"),
|
||||
)
|
||||
fn = VariableFactory.from_value(
|
||||
getattr(attr.__self__.__class__, name),
|
||||
self.graph,
|
||||
GetAttrTracker(class_var, name),
|
||||
)
|
||||
if not hasattr(self.get_py_type(), name):
|
||||
instance = None
|
||||
return MethodVariable.wrap_method(
|
||||
value=attr,
|
||||
instance=instance,
|
||||
fn=fn,
|
||||
graph=self.graph,
|
||||
tracker=GetAttrTracker(self, name),
|
||||
)
|
||||
|
||||
return VariableFactory.from_value(
|
||||
attr, self.graph, tracker=GetAttrTracker(self, name)
|
||||
)
|
||||
|
||||
def hasattr(self, name: str):
|
||||
from .basic import ConstantVariable
|
||||
|
||||
try:
|
||||
self.getattr(name)
|
||||
return ConstantVariable(
|
||||
True, graph=self.graph, tracker=DummyTracker([self])
|
||||
)
|
||||
except HasNoAttributeError:
|
||||
# NOTE(SigureMo): Only the HasNoAttributeError is raised, we can
|
||||
# ensure that the attribute does not exist. Otherwise, we should
|
||||
# raise the error.
|
||||
return ConstantVariable(
|
||||
False, graph=self.graph, tracker=DummyTracker([self])
|
||||
)
|
||||
|
||||
def getattr(self, name: str, default=None):
|
||||
result = self.attr_proxy.get(name)
|
||||
if isinstance(result, MutableDictLikeData.Empty):
|
||||
if default is not None:
|
||||
assert isinstance(default, VariableBase)
|
||||
return default
|
||||
raise HasNoAttributeError(
|
||||
f"{self.__class__.__name__} {self} has no attribute {name}"
|
||||
)
|
||||
return result
|
||||
|
||||
def setattr(self, key: str, value):
|
||||
from .basic import ConstantVariable
|
||||
|
||||
self.attr_proxy.set(key, value)
|
||||
self.graph.side_effects.record_proxy_variable(self)
|
||||
return ConstantVariable.wrap_literal(None, self.graph)
|
||||
|
||||
def delattr(self, key: str):
|
||||
from .basic import ConstantVariable
|
||||
|
||||
self.attr_proxy.delete(key)
|
||||
self.graph.side_effects.record_proxy_variable(self)
|
||||
return ConstantVariable.wrap_literal(None, self.graph)
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
return self.setitem(key, value)
|
||||
|
||||
def setitem(self, key, value):
|
||||
raise FallbackError(f"{self} is not support setitem.")
|
||||
|
||||
def __repr__(self):
|
||||
info = self.main_info | self.debug_info
|
||||
info_str = ", ".join([f"{value}" for value in info.values()])
|
||||
return f"{self.__class__.__name__}({info_str})"
|
||||
|
||||
def __str__(self):
|
||||
return self.__repr__()
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return Dispatcher.call(operator.getitem, self, idx)
|
||||
|
||||
def getitem(self, item):
|
||||
class_var = VariableFactory.from_value(
|
||||
self.get_py_value().__class__,
|
||||
self.graph,
|
||||
GetAttrTracker(self, '__class__'),
|
||||
)
|
||||
fn_var = VariableFactory.from_value(
|
||||
get_unbound_method(self.get_py_value(), '__getitem__'),
|
||||
self.graph,
|
||||
GetAttrTracker(class_var, '__getitem__'),
|
||||
)
|
||||
self.graph.add_global_guarded_variable(item)
|
||||
item = item.get_py_value()
|
||||
output = fn_var(self, item)
|
||||
return output
|
||||
|
||||
def __call__(self, /, *args, **kwargs):
|
||||
"""
|
||||
Call the object represented by this variable with the given arguments.
|
||||
|
||||
Args:
|
||||
*args: Positional arguments to pass to the object's __call__ method.
|
||||
**kwargs: Keyword arguments to pass to the object's __call__ method.
|
||||
|
||||
Returns:
|
||||
VariableBase: A new variable representing the result of calling the object's __call__ method.
|
||||
"""
|
||||
from .callable import BuiltinVariable, UserDefinedFunctionVariable
|
||||
|
||||
class_var = VariableFactory.from_value(
|
||||
self.get_py_value().__class__,
|
||||
self.graph,
|
||||
GetAttrTracker(self, '__class__'),
|
||||
)
|
||||
assert class_var is not None
|
||||
# if __call__ is a method, we should add self to arguments.
|
||||
if inspect.ismethod(self.get_py_value().__call__):
|
||||
args = (self, *args)
|
||||
unbound_method = get_unbound_method(self.get_py_value(), '__call__')
|
||||
if hasattr(unbound_method, "__code__"):
|
||||
fn_var = UserDefinedFunctionVariable(
|
||||
unbound_method,
|
||||
self.graph,
|
||||
GetAttrTracker(class_var, '__call__'),
|
||||
)
|
||||
else:
|
||||
fn_var = BuiltinVariable(
|
||||
self.value,
|
||||
self.graph,
|
||||
GetAttrTracker(class_var, '__call__'),
|
||||
)
|
||||
output = fn_var(*args, **kwargs)
|
||||
return output
|
||||
|
||||
def get_iter(self):
|
||||
from . import (
|
||||
BuiltinVariable,
|
||||
ConstantVariable,
|
||||
SequenceIterVariable,
|
||||
UserDefinedFunctionVariable,
|
||||
UserDefinedIterVariable,
|
||||
)
|
||||
|
||||
if not hasattr(self.value, "__iter__"):
|
||||
return UserDefinedIterVariable(
|
||||
self, self.graph, GetIterTracker(self)
|
||||
)
|
||||
iter_name_var = ConstantVariable.wrap_literal("__iter__", self.graph)
|
||||
iter_method = BuiltinVariable(
|
||||
getattr, graph=self.graph, tracker=DummyTracker([self])
|
||||
)(self, iter_name_var)
|
||||
# If the target object is a builtin object like list_iterator, the iter_method's fn will be a ObjectVariable instead of UserDefinedFunctionVariable.
|
||||
if not isinstance(iter_method.fn, UserDefinedFunctionVariable):
|
||||
return UserDefinedIterVariable(
|
||||
self, self.graph, GetIterTracker(self)
|
||||
)
|
||||
iter_result = iter_method()
|
||||
|
||||
if not isinstance(iter_result, SequenceIterVariable):
|
||||
return UserDefinedIterVariable(
|
||||
self, self.graph, GetIterTracker(self)
|
||||
)
|
||||
|
||||
return iter_result
|
||||
|
||||
@VariableFactory.register_from_value()
|
||||
def from_value(
|
||||
value: Any,
|
||||
graph: FunctionGraph | None,
|
||||
tracker: Tracker,
|
||||
) -> VariableBase | None:
|
||||
"""
|
||||
Create a new variable from a given value, or return None if the value cannot be converted to a variable.
|
||||
Args:
|
||||
value (Any): The value to create a variable from.
|
||||
graph (FunctionGraph | None): The graph in which the variable will be used.
|
||||
tracker (Tracker): The variable tracker to put the new variable in if created.
|
||||
|
||||
Returns:
|
||||
VariableBase | None: A new variable if one can be created from the given value, or None if the value cannot be converted to a variable.
|
||||
"""
|
||||
if isinstance(value, VariableBase):
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def signature_clear_guard(fn, name):
|
||||
if not hasattr(fn, name):
|
||||
yield
|
||||
else:
|
||||
saved_attr = getattr(fn, name)
|
||||
delattr(fn, name)
|
||||
yield
|
||||
setattr(fn, name, saved_attr)
|
||||
|
||||
|
||||
def fn_bind_inputs(
|
||||
fn: Callable[..., Any],
|
||||
graph: FunctionGraph,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
):
|
||||
# temparay clear the fn.__signature__ to avoid signature check error
|
||||
with (
|
||||
signature_clear_guard(fn, "__signature__"),
|
||||
signature_clear_guard(fn, "__wrapped__"),
|
||||
):
|
||||
sig = inspect.signature(fn)
|
||||
bound_args = sig.bind(*args, **kwargs)
|
||||
bound_args.apply_defaults()
|
||||
parameters = {}
|
||||
for name, value in bound_args.arguments.items():
|
||||
assert name in sig.parameters
|
||||
# Convert varargs and kwargs to Variable
|
||||
if sig.parameters[name].kind == inspect.Parameter.VAR_POSITIONAL:
|
||||
tracker = DummyTracker(value)
|
||||
elif sig.parameters[name].kind == inspect.Parameter.VAR_KEYWORD:
|
||||
tracker = DummyTracker(list(value.values()))
|
||||
# Convert default args to Variable
|
||||
elif not isinstance(value, VariableBase):
|
||||
tracker = ConstTracker(value)
|
||||
else:
|
||||
tracker = value.tracker
|
||||
value = VariableFactory.from_value(value, graph, tracker)
|
||||
parameters[name] = value
|
||||
return parameters
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,460 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import types
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from paddle._typing import unreached
|
||||
|
||||
from ....profiler import EventGuard
|
||||
from ....utils import do_until_stop_iteration
|
||||
from ....utils.exceptions import (
|
||||
BreakGraphError,
|
||||
BreakGraphInlineCallBreak,
|
||||
FallbackError,
|
||||
FallbackInlineCallBreak,
|
||||
OtherInlineCallBreak,
|
||||
SotCapturedExceptionFactory,
|
||||
SotCapturedStopIteration,
|
||||
SotErrorBase,
|
||||
UnsupportedOperationBreak,
|
||||
)
|
||||
from ..guard import check_faster_guard
|
||||
from ..tracker import ConstTracker, DanglingTracker, DummyTracker
|
||||
from .base import (
|
||||
VariableBase,
|
||||
VariableFactory,
|
||||
)
|
||||
from .basic import ConstantVariable
|
||||
from .callable import BuiltinVariable
|
||||
from .container import TupleVariable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
import paddle
|
||||
|
||||
from ..function_graph import FunctionGraph
|
||||
from ..pycode_generator import PyCodeGen
|
||||
from ..tracker import Tracker
|
||||
from ..virtual_frame import VirtualFrame
|
||||
|
||||
|
||||
class IterVariable(VariableBase):
|
||||
"""
|
||||
This Variable (include subclasses) should be generated only when simulate GET_ITER opcode
|
||||
"""
|
||||
|
||||
def __init__(self, graph: FunctionGraph, tracker: Tracker):
|
||||
super().__init__(graph, tracker)
|
||||
|
||||
def next(self):
|
||||
raise NotImplementedError(f"Can not simulate `next` for {type(self)}")
|
||||
|
||||
def to_list(self):
|
||||
raise NotImplementedError(
|
||||
f"Can not simulate `to_list` for {type(self)}"
|
||||
)
|
||||
|
||||
def send(self, value: VariableBase):
|
||||
return self.next()
|
||||
|
||||
def get_iter(self):
|
||||
return self
|
||||
|
||||
|
||||
class SequenceIterVariable(IterVariable):
|
||||
"""
|
||||
The basic SequenceIterVariable wraps iterators which can be simulated by call getitem
|
||||
Currently includes: List | Tuple | Dict (keys) | Range | Tensor | nn.LayerList
|
||||
|
||||
these interfaces is needed:
|
||||
- next
|
||||
- to_list
|
||||
- has_side_effect
|
||||
- _reconstruct
|
||||
"""
|
||||
|
||||
mutable_attrs = ["idx"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
held: VariableBase | list[VariableBase],
|
||||
graph: FunctionGraph,
|
||||
tracker: Tracker,
|
||||
):
|
||||
if not isinstance(held, list):
|
||||
held = [held]
|
||||
super().__init__(graph, tracker)
|
||||
self.holds = held
|
||||
self.idx = 0
|
||||
self.graph.side_effects.record_mutable_variable(self)
|
||||
|
||||
@check_faster_guard
|
||||
def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
|
||||
return [
|
||||
guard for held in self.holds for guard in held.make_faster_guard()
|
||||
]
|
||||
|
||||
def make_stringified_guard(self):
|
||||
return [
|
||||
guard
|
||||
for held in self.holds
|
||||
for guard in held.make_stringified_guard()
|
||||
]
|
||||
|
||||
def next(self):
|
||||
held = self.holds[0]
|
||||
if self.idx < len(held):
|
||||
val = held[self.idx]
|
||||
self.idx += 1
|
||||
return val
|
||||
else:
|
||||
raise SotCapturedExceptionFactory.create(StopIteration())
|
||||
|
||||
def to_list(self) -> list:
|
||||
if self.has_side_effect():
|
||||
raise FallbackError("Can not convert an used iterator into list")
|
||||
held = self.holds[0]
|
||||
self.idx = len(held)
|
||||
retval = []
|
||||
for i in range(len(held)):
|
||||
retval.append(held[i])
|
||||
return retval
|
||||
|
||||
def has_side_effect(self) -> bool:
|
||||
return self.idx != 0
|
||||
|
||||
def _reconstruct(self, codegen: PyCodeGen):
|
||||
if self.has_side_effect():
|
||||
super()._reconstruct(codegen)
|
||||
else:
|
||||
self.holds[0].reconstruct(codegen)
|
||||
codegen.gen_get_iter()
|
||||
|
||||
@property
|
||||
def main_info(self) -> dict[str, Any]:
|
||||
return {
|
||||
"idx": self.idx,
|
||||
}
|
||||
|
||||
def flatten_inner_vars(self) -> list[VariableBase]:
|
||||
held = self.holds
|
||||
return [
|
||||
inner_var for obj in held for inner_var in obj.flatten_inner_vars()
|
||||
]
|
||||
|
||||
|
||||
class EnumerateVariable(SequenceIterVariable):
|
||||
"""
|
||||
EnumerateVariable holds a SequenceIterVariable and return additional index
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, val_iterator: IterVariable, graph: FunctionGraph, tracker: Tracker
|
||||
):
|
||||
super().__init__(val_iterator, graph, tracker)
|
||||
|
||||
def next(self):
|
||||
val = self.holds[0].next()
|
||||
idx_var = ConstantVariable(self.idx, self.graph, ConstTracker(self.idx))
|
||||
self.idx += 1
|
||||
return TupleVariable(
|
||||
(idx_var, val), self.graph, DummyTracker([idx_var, val])
|
||||
)
|
||||
|
||||
def to_list(self):
|
||||
values = self.holds[0].to_list()
|
||||
idx = [
|
||||
ConstantVariable(i, self.graph, ConstTracker(i))
|
||||
for i in range(len(values))
|
||||
]
|
||||
return list(zip(idx, values))
|
||||
|
||||
def has_side_effect(self) -> bool:
|
||||
return self.holds[0].has_side_effect()
|
||||
|
||||
def _reconstruct(self, codegen: PyCodeGen):
|
||||
if self.has_side_effect():
|
||||
super()._reconstruct(codegen)
|
||||
else:
|
||||
codegen.gen_load_global("enumerate", push_null=True)
|
||||
self.holds[0].reconstruct(codegen)
|
||||
codegen.gen_call_function(1)
|
||||
|
||||
@staticmethod
|
||||
def from_iterator(value, graph: FunctionGraph | None, tracker: Tracker):
|
||||
iter_variable = value.get_iter()
|
||||
if isinstance(iter_variable, UserDefinedIterVariable):
|
||||
return UserDefinedIterVariable(value, graph, tracker)
|
||||
else:
|
||||
return EnumerateVariable(iter_variable, graph, tracker)
|
||||
|
||||
|
||||
class ZipVariable(SequenceIterVariable):
|
||||
"""
|
||||
ZipVariable holds a list of SequenceIterVariable
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, iters: list[IterVariable], graph: FunctionGraph, tracker: Tracker
|
||||
):
|
||||
super().__init__(iters, graph, tracker)
|
||||
|
||||
def next(self):
|
||||
# can not use <listcomp> here, because it will raise a RuntimeError("StopIteration")
|
||||
# but we want a StopIteration Exception
|
||||
values = []
|
||||
for iter_var in self.holds:
|
||||
next_var = iter_var.next()
|
||||
values.append(next_var)
|
||||
|
||||
return VariableFactory.from_value(
|
||||
tuple(values), self.graph, DummyTracker(values)
|
||||
)
|
||||
|
||||
def to_list(self):
|
||||
lists = [iter_vars.to_list() for iter_vars in self.holds]
|
||||
min_len = min(len(l) for l in lists)
|
||||
result = []
|
||||
for i in range(min_len):
|
||||
result.append(
|
||||
VariableFactory.from_value(
|
||||
tuple(l[i] for l in lists),
|
||||
self.graph,
|
||||
DummyTracker(list(self.holds)),
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
def has_side_effect(self) -> bool:
|
||||
return any(iter_var.has_side_effect() for iter_var in self.holds)
|
||||
|
||||
def _reconstruct(self, codegen: PyCodeGen):
|
||||
if self.has_side_effect():
|
||||
super()._reconstruct(codegen)
|
||||
else:
|
||||
codegen.gen_load_global("zip", push_null=True)
|
||||
for iter_var in self.holds:
|
||||
iter_var.reconstruct(codegen)
|
||||
codegen.gen_call_function(len(self.holds))
|
||||
|
||||
@staticmethod
|
||||
def from_iterator(
|
||||
value: Sequence[VariableBase],
|
||||
graph: FunctionGraph | None,
|
||||
tracker: Tracker,
|
||||
):
|
||||
assert isinstance(value, (list, tuple))
|
||||
zip_targets = []
|
||||
|
||||
for variable in value:
|
||||
iter_variable = variable.get_iter()
|
||||
if isinstance(iter_variable, UserDefinedIterVariable):
|
||||
return UserDefinedIterVariable(value, graph, tracker)
|
||||
zip_targets.append(iter_variable)
|
||||
|
||||
return ZipVariable(zip_targets, graph, tracker)
|
||||
|
||||
|
||||
class MapVariable(SequenceIterVariable):
|
||||
"""
|
||||
MapVariable holds a SequenceIterVariable and return a Iterable Variable after map function
|
||||
"""
|
||||
|
||||
def __init__(self, fn, iters: list[IterVariable], graph, tracker):
|
||||
super().__init__(iters, graph, tracker)
|
||||
self.fn = fn
|
||||
|
||||
def next(self):
|
||||
return self.fn(*[iter_var.next() for iter_var in self.holds])
|
||||
|
||||
def to_list(self) -> list:
|
||||
lists = [iter_var.to_list() for iter_var in self.holds]
|
||||
min_len = min(len(l) for l in lists)
|
||||
result = []
|
||||
for i in range(min_len):
|
||||
result.append(self.fn(*(l[i] for l in lists)))
|
||||
return result
|
||||
|
||||
def has_side_effect(self) -> bool:
|
||||
return any(iter_var.has_side_effect() for iter_var in self.holds)
|
||||
|
||||
def _reconstruct(self, codegen: PyCodeGen):
|
||||
if self.has_side_effect():
|
||||
super()._reconstruct(codegen)
|
||||
else:
|
||||
codegen.gen_load_global("map", push_null=True)
|
||||
self.fn.reconstruct(codegen)
|
||||
for iter_var in self.holds:
|
||||
iter_var.reconstruct(codegen)
|
||||
codegen.gen_call_function(len(self.holds) + 1)
|
||||
|
||||
@staticmethod
|
||||
def from_iterator(
|
||||
fn,
|
||||
value: Sequence[VariableBase],
|
||||
graph: FunctionGraph | None,
|
||||
tracker: Tracker,
|
||||
):
|
||||
map_targets = []
|
||||
|
||||
for variable in value:
|
||||
iter_variable = variable.get_iter()
|
||||
if isinstance(iter_variable, UserDefinedIterVariable):
|
||||
return UserDefinedIterVariable(value, graph, tracker)
|
||||
map_targets.append(iter_variable)
|
||||
|
||||
return MapVariable(fn, map_targets, graph, tracker)
|
||||
|
||||
|
||||
class GeneratorVariable(IterVariable):
|
||||
def __init__(
|
||||
self,
|
||||
code_var: VariableBase,
|
||||
vframe: VirtualFrame,
|
||||
graph: FunctionGraph,
|
||||
tracker: Tracker,
|
||||
):
|
||||
self.code_var = code_var
|
||||
self.vframe = vframe
|
||||
self.shared_stack = []
|
||||
super().__init__(graph, tracker)
|
||||
|
||||
def send(self, /, value: VariableBase):
|
||||
from ..opcode_inline_executor import OpcodeInlineGeneratorExecutor
|
||||
|
||||
checkpoint = self.graph.save_memo()
|
||||
frame_state = self.vframe.get_state()
|
||||
try:
|
||||
inline_gen_executor = OpcodeInlineGeneratorExecutor(
|
||||
self.vframe, self.code_var, self.graph
|
||||
)
|
||||
self.vframe.stack.push(value)
|
||||
with EventGuard(
|
||||
f"Inline Gen Call: {inline_gen_executor.vframe.code.co_name}, file {inline_gen_executor.vframe.code.co_filename}, line {int(inline_gen_executor.vframe.code.co_firstlineno)}"
|
||||
):
|
||||
output: VariableBase = inline_gen_executor.inline_call()
|
||||
if inline_gen_executor.stop_state == "Return":
|
||||
raise SotCapturedExceptionFactory.create(StopIteration())
|
||||
except SotCapturedStopIteration:
|
||||
raise
|
||||
except SotErrorBase as error:
|
||||
self.graph.restore_memo(checkpoint)
|
||||
self.vframe.restore_state(frame_state)
|
||||
filename = self.code_var.value.co_filename
|
||||
lineno = self.code_var.value.co_firstlineno
|
||||
code_name = self.code_var.value.co_name
|
||||
location_info = f'File "{filename}", line {lineno}, in {code_name}'
|
||||
|
||||
exception_class = OtherInlineCallBreak
|
||||
if isinstance(error, BreakGraphError):
|
||||
exception_class = BreakGraphInlineCallBreak
|
||||
elif isinstance(error, FallbackError):
|
||||
exception_class = FallbackInlineCallBreak
|
||||
|
||||
raise BreakGraphError(
|
||||
exception_class(
|
||||
f"{location_info} encountered breakgraph error caused by\n {error}"
|
||||
)
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
def getattr(self, name: str, default=None):
|
||||
from ..dispatch_functions import generator_send
|
||||
|
||||
known_generator_attrs = {"send"}
|
||||
if name not in known_generator_attrs:
|
||||
raise BreakGraphError(
|
||||
UnsupportedOperationBreak(
|
||||
reason_str=f"Get attribute {name} from generator is not supported."
|
||||
)
|
||||
)
|
||||
if name == "send":
|
||||
return BuiltinVariable(
|
||||
generator_send, self.graph, DanglingTracker()
|
||||
).bind_dangling_fn(self, "send")
|
||||
unreached()
|
||||
|
||||
def get_py_value(self, allow_tensor=False):
|
||||
raise BreakGraphError(
|
||||
UnsupportedOperationBreak(
|
||||
reason_str="Get real value from generator is not supported."
|
||||
)
|
||||
)
|
||||
|
||||
def get_py_type(self):
|
||||
return types.GeneratorType
|
||||
|
||||
def next(self):
|
||||
return self.send(ConstantVariable.wrap_literal(None, self.graph))
|
||||
|
||||
def to_list(self):
|
||||
return do_until_stop_iteration(lambda: self.next())
|
||||
|
||||
@property
|
||||
def main_info(self) -> dict[str, Any]:
|
||||
return {
|
||||
"co_name": self.code_var.value.co_name,
|
||||
}
|
||||
|
||||
# @VariableFactory.register_from_value()
|
||||
# def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
|
||||
# if inspect.isgenerator(value):
|
||||
# return GeneratorVariable()
|
||||
# return None
|
||||
|
||||
|
||||
# what UserDefinedIterVariable holds doesn't matter, because use user defined iterator will trigger break graph
|
||||
class UserDefinedIterVariable(IterVariable):
|
||||
def __init__(
|
||||
self,
|
||||
held: VariableBase | list[VariableBase],
|
||||
graph: FunctionGraph,
|
||||
tracker: Tracker,
|
||||
):
|
||||
if not isinstance(held, list):
|
||||
held = [held]
|
||||
self.holds = held
|
||||
super().__init__(graph, tracker)
|
||||
|
||||
def to_list(self):
|
||||
raise BreakGraphError(
|
||||
UnsupportedOperationBreak(
|
||||
reason_str="Break graph when iterating user defined iterator"
|
||||
)
|
||||
)
|
||||
|
||||
def next(self):
|
||||
raise BreakGraphError(
|
||||
UnsupportedOperationBreak(
|
||||
reason_str="Break graph when iterating user defined iterator"
|
||||
)
|
||||
)
|
||||
|
||||
@check_faster_guard
|
||||
def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
|
||||
return [
|
||||
guard for held in self.holds for guard in held.make_faster_guard()
|
||||
]
|
||||
|
||||
def make_stringified_guard(self):
|
||||
return [
|
||||
guard
|
||||
for held in self.holds
|
||||
for guard in held.make_stringified_guard()
|
||||
]
|
||||
@@ -0,0 +1,238 @@
|
||||
# Copyright (c) 2025 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, NamedTuple
|
||||
|
||||
from ...utils import log
|
||||
from .tracker import (
|
||||
BuiltinTracker,
|
||||
CellTracker,
|
||||
ConstTracker,
|
||||
DanglingTracker,
|
||||
FunctionClosureTracker,
|
||||
LocalTracker,
|
||||
)
|
||||
from .variable_stack import VariableStack
|
||||
from .variables.base import VariableBase, VariableFactory, fn_bind_inputs
|
||||
from .variables.basic import (
|
||||
CellVariable,
|
||||
FunctionGlobalVariable,
|
||||
GlobalVariable,
|
||||
NullVariable,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import types
|
||||
from typing import TypeAlias
|
||||
|
||||
from ..instruction_utils import Instruction
|
||||
from .function_graph import FunctionGraph
|
||||
from .variables.callable import FunctionVariable
|
||||
|
||||
# The type to represent the (*args, **kwargs) pack in the call.
|
||||
CallArgsPack: TypeAlias = tuple[tuple[Any, ...], dict[str, Any]]
|
||||
|
||||
|
||||
def validate_value(value):
|
||||
assert isinstance(value, VariableBase), (
|
||||
f"value: {value}, type should be VariableBase(or derived), but get {type(value)}"
|
||||
)
|
||||
assert not isinstance(value.tracker, DanglingTracker) or isinstance(
|
||||
value, (NullVariable, CellVariable)
|
||||
), f"dangling variable {value} should not be pushed into stack."
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlockStackItem:
|
||||
# `PyTryBlock` in CPython source code
|
||||
type: str
|
||||
inst: Instruction
|
||||
handler: Instruction
|
||||
level: int
|
||||
|
||||
|
||||
class VirtualFrameState(NamedTuple):
|
||||
locals: dict[str, VariableBase]
|
||||
builtins: dict[str, VariableBase]
|
||||
cells: dict[str, VariableBase]
|
||||
lasti: int
|
||||
stack_data: list[VariableBase]
|
||||
block_stack: list[BlockStackItem]
|
||||
|
||||
|
||||
class VirtualFrame:
|
||||
code: types.CodeType
|
||||
locals: dict[str, Any] # TODO: should we use DictVariable instead of dict?
|
||||
globals: GlobalVariable
|
||||
builtins: dict[str, Any]
|
||||
consts: list[Any]
|
||||
cells: dict[str, Any]
|
||||
lasti: int
|
||||
stack: VariableStack
|
||||
block_stack: list[BlockStackItem]
|
||||
|
||||
def __init__(self, code: types.CodeType):
|
||||
self.code = code
|
||||
self.locals = {}
|
||||
self.globals = None # type: ignore
|
||||
self.builtins = {}
|
||||
self.cells = {}
|
||||
self.lasti = 0
|
||||
self.consts = []
|
||||
self.stack = VariableStack(validate_value_func=validate_value)
|
||||
self.block_stack: list[BlockStackItem] = []
|
||||
|
||||
@staticmethod
|
||||
def from_real_frame(frame: types.FrameType, graph: FunctionGraph):
|
||||
code = frame.f_code
|
||||
locals = frame.f_locals
|
||||
vframe = VirtualFrame(code)
|
||||
|
||||
# convert locals
|
||||
free_or_cell_vars = code.co_cellvars + code.co_freevars
|
||||
for name, value in locals.items():
|
||||
tracker = (
|
||||
CellTracker(name)
|
||||
if name in free_or_cell_vars
|
||||
else LocalTracker(name)
|
||||
)
|
||||
vframe.locals[name] = VariableFactory.from_value(
|
||||
value, graph, tracker
|
||||
)
|
||||
|
||||
for name in free_or_cell_vars:
|
||||
# create a cell for each variable.
|
||||
vframe.cells[name] = CellVariable() # put in cells.
|
||||
if name in vframe.locals:
|
||||
vframe.cells[name].set_value(vframe.locals[name])
|
||||
|
||||
# convert globals
|
||||
vframe.globals = GlobalVariable(
|
||||
frame.f_globals,
|
||||
graph,
|
||||
DanglingTracker(),
|
||||
)
|
||||
|
||||
# convert builtins
|
||||
for name, value in builtins.__dict__.items():
|
||||
vframe.builtins[name] = VariableFactory.from_value(
|
||||
value, graph, BuiltinTracker(name)
|
||||
)
|
||||
# Temporarily use the builtins from the graph to avoid the conversion overhead.
|
||||
graph.builtins = vframe.builtins
|
||||
|
||||
# prepare consts
|
||||
for value in code.co_consts:
|
||||
vframe.consts.append(
|
||||
VariableFactory.from_value(value, graph, ConstTracker(value))
|
||||
)
|
||||
return vframe
|
||||
|
||||
@staticmethod
|
||||
def from_inline_call(
|
||||
code: types.CodeType,
|
||||
fn_var: FunctionVariable,
|
||||
fn_value: types.FunctionType,
|
||||
graph: FunctionGraph,
|
||||
call_args_pack: CallArgsPack,
|
||||
):
|
||||
call_args, call_kwargs = call_args_pack
|
||||
vframe = VirtualFrame(code)
|
||||
vframe.globals = FunctionGlobalVariable(
|
||||
fn_var,
|
||||
fn_value.__globals__,
|
||||
graph,
|
||||
DanglingTracker(),
|
||||
)
|
||||
|
||||
# convert builtins
|
||||
# NOTE(SigureMo): inline call should inherit the builtins from the caller to reduce the conversion overhead.
|
||||
vframe.builtins = graph.builtins
|
||||
|
||||
# prepare consts
|
||||
for value in code.co_consts:
|
||||
vframe.consts.append(
|
||||
VariableFactory.from_value(value, graph, ConstTracker(value))
|
||||
)
|
||||
|
||||
# convert locals
|
||||
vframe.locals.update(
|
||||
fn_bind_inputs(fn_value, graph, *call_args, **call_kwargs)
|
||||
)
|
||||
|
||||
log(
|
||||
5,
|
||||
f"[INLINE CALL] {code.co_name} with locals: ",
|
||||
vframe.locals,
|
||||
)
|
||||
|
||||
# handle implicit variables in comprehensions
|
||||
vframe.handle_comps(fn_value)
|
||||
|
||||
# convert closures
|
||||
closure = fn_var.get_py_value().__closure__
|
||||
for name in code.co_cellvars + code.co_freevars:
|
||||
# create a cell for each variable.
|
||||
vframe.cells[name] = CellVariable() # put in cells.
|
||||
if name in vframe.locals:
|
||||
vframe.cells[name].set_value(vframe.locals[name])
|
||||
|
||||
if closure is None:
|
||||
return vframe
|
||||
assert len(closure) == len(code.co_freevars)
|
||||
for idx, (name, cell) in enumerate(zip(code.co_freevars, closure)):
|
||||
value = cell.cell_contents
|
||||
value = VariableFactory.from_value(
|
||||
value, graph, FunctionClosureTracker(fn_var, idx)
|
||||
)
|
||||
# wrapped by a CellVariable
|
||||
if not isinstance(value, CellVariable):
|
||||
value = CellVariable(value)
|
||||
vframe.cells[name] = value
|
||||
return vframe
|
||||
|
||||
def handle_comps(self, fn_value):
|
||||
is_comp = any(
|
||||
x in fn_value.__name__
|
||||
for x in ['<listcomp>', '<dictcomp>', '<setcomp>', '<genexpr>']
|
||||
)
|
||||
if not is_comp:
|
||||
return
|
||||
pattern = r'implicit\d+'
|
||||
for name in list(self.locals.keys()):
|
||||
if re.match(pattern, name):
|
||||
self.locals[name.replace('implicit', '.')] = self.locals[name]
|
||||
|
||||
def get_state(self):
|
||||
return VirtualFrameState(
|
||||
locals=self.locals.copy(),
|
||||
builtins=self.builtins.copy(),
|
||||
cells=self.cells.copy(),
|
||||
lasti=self.lasti,
|
||||
stack_data=list(self.stack._data),
|
||||
block_stack=self.block_stack.copy(),
|
||||
)
|
||||
|
||||
def restore_state(self, state: VirtualFrameState):
|
||||
self.locals = state.locals
|
||||
self.builtins = state.builtins
|
||||
self.cells = state.cells
|
||||
self.lasti = state.lasti
|
||||
self.stack._data = state.stack_data
|
||||
self.block_stack = state.block_stack
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from .instruction_pass import apply_instr_pass # noqa: F401
|
||||
from .instruction_utils import ( # noqa: F401
|
||||
Instruction,
|
||||
Space,
|
||||
calc_offset_from_bytecode_offset,
|
||||
calc_stack_effect,
|
||||
convert_instruction,
|
||||
gen_instr,
|
||||
get_instruction_size,
|
||||
get_instructions,
|
||||
instrs_info,
|
||||
modify_extended_args,
|
||||
modify_instrs,
|
||||
modify_vars,
|
||||
relocate_jump_target,
|
||||
replace_instr,
|
||||
reset_offset,
|
||||
)
|
||||
from .opcode_analysis import ( # noqa: F401
|
||||
analysis_used_names,
|
||||
)
|
||||
@@ -0,0 +1,334 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dis
|
||||
import sys
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from paddle.jit.sot.utils import log, log_do
|
||||
|
||||
from ...utils import InnerError
|
||||
from .instruction_utils import instrs_info
|
||||
from .opcode_info import TO_FUSED_INSTS
|
||||
from .stack_analyse import StackAnalyser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .instruction_utils import Instruction
|
||||
|
||||
|
||||
def apply_instr_pass(instrs: list[Instruction], code_options):
|
||||
log(4, f"[Opcode Pass]: Original New Code {code_options['co_name']}:\n")
|
||||
log_do(4, lambda: print(instrs_info(instrs)))
|
||||
supported_passes = [
|
||||
remove_load_store_pass,
|
||||
remove_duplicate_resume,
|
||||
check_precall_followed_by_call,
|
||||
]
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
supported_passes.append(check_for_iter_jump_to)
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
supported_passes.append(fuse_double_super_instrs)
|
||||
|
||||
for instr_pass in supported_passes:
|
||||
instr_pass(instrs, code_options)
|
||||
|
||||
log(
|
||||
4,
|
||||
f"[Opcode Pass]: New Code After Opcode Pass {code_options['co_name']}:\n",
|
||||
)
|
||||
log_do(4, lambda: print(instrs_info(instrs)))
|
||||
|
||||
|
||||
def find_stored_once_local_vars(instrs: list[Instruction], code_options):
|
||||
"""
|
||||
find out the local var names which is only stored once
|
||||
"""
|
||||
stored_vars = {}
|
||||
|
||||
# The input vars are considered as stored at the beginning
|
||||
input_names = code_options['co_varnames'][: code_options['co_argcount']]
|
||||
|
||||
for name in input_names:
|
||||
stored_vars[name] = 1
|
||||
|
||||
for instr in instrs:
|
||||
if instr.opname == "STORE_FAST":
|
||||
if instr.argval in stored_vars:
|
||||
stored_vars[instr.argval] += 1
|
||||
else:
|
||||
stored_vars[instr.argval] = 1
|
||||
|
||||
stored_once = {name for name, count in stored_vars.items() if count == 1}
|
||||
return stored_once
|
||||
|
||||
|
||||
def find_loaded_once_local_vars(instrs: list[Instruction], code_options):
|
||||
"""
|
||||
find out the local var names which is only stored once
|
||||
"""
|
||||
loaded_vars = {}
|
||||
for instr in instrs:
|
||||
if instr.opname in ["LOAD_FAST", "LOAD_FAST_BORROW", "LOAD_FAST_CHECK"]:
|
||||
if instr.argval in loaded_vars:
|
||||
loaded_vars[instr.argval] += 1
|
||||
else:
|
||||
loaded_vars[instr.argval] = 1
|
||||
|
||||
loaded_once = {name for name, count in loaded_vars.items() if count == 1}
|
||||
return loaded_once
|
||||
|
||||
|
||||
def find_related_local_opcodes(instrs: list[Instruction], code_options):
|
||||
"""
|
||||
Find opcode pairs consisting of LOAD_FAST, LOAD_FAST_BORROW, STORE_FAST, and LOAD_FAST_CHECK.
|
||||
"""
|
||||
stack = []
|
||||
opcode_pairs = []
|
||||
for instr in instrs:
|
||||
if instr.opname in ["LOAD_FAST", "LOAD_FAST_BORROW", "LOAD_FAST_CHECK"]:
|
||||
stack.append(instr)
|
||||
elif instr.opname == "STORE_FAST":
|
||||
if len(stack) > 0 and stack[-1] is not None:
|
||||
opcode_pairs.append((stack[-1], instr))
|
||||
stack.pop()
|
||||
elif "ROT" in instr.opname or "DUP" in instr.opname:
|
||||
return []
|
||||
else:
|
||||
try:
|
||||
pop_n, push_n = StackAnalyser().stack_effect(instr)
|
||||
if pop_n == 0:
|
||||
stack.extend([None] * push_n)
|
||||
else:
|
||||
stack = stack[:-pop_n] + [None] * push_n
|
||||
except AttributeError:
|
||||
break
|
||||
|
||||
return opcode_pairs
|
||||
|
||||
|
||||
def remove_load_store_pass(instrs: list[Instruction], code_options):
|
||||
"""
|
||||
This question is extremely complex, so we just simplify it as
|
||||
'remove renames which is between var names who only stored once'
|
||||
and we only consider the local vars.
|
||||
"""
|
||||
|
||||
def stored_from(load_instr, instrs):
|
||||
idx = instrs.index(load_instr) - 1
|
||||
while idx >= 0:
|
||||
instr = instrs[idx]
|
||||
if (
|
||||
instr.opname == "STORE_FAST"
|
||||
and instr.argval == load_instr.argval
|
||||
):
|
||||
return instr
|
||||
idx -= 1
|
||||
return None
|
||||
|
||||
def code_exist(opname, argval, instrs):
|
||||
for instr in instrs:
|
||||
if instr.opname == opname and instr.argval == argval:
|
||||
return True
|
||||
return False
|
||||
|
||||
# remove rename and load store
|
||||
jump_target = {
|
||||
instr.jump_to for instr in instrs if instr.jump_to is not None
|
||||
}
|
||||
|
||||
modified = True
|
||||
while modified:
|
||||
modified = False
|
||||
stored_once = find_stored_once_local_vars(instrs, code_options)
|
||||
|
||||
# find out all LOAD_FAST -> STORE_FAST pair
|
||||
opcode_pairs = find_related_local_opcodes(instrs, code_options)
|
||||
|
||||
for load_a, store_b in opcode_pairs:
|
||||
if load_a in jump_target or store_b in jump_target:
|
||||
continue
|
||||
a_name = load_a.argval
|
||||
b_name = store_b.argval
|
||||
|
||||
# if these two names are only stored once
|
||||
# it means these two name only have one value all the time
|
||||
# so we can just rename them, to delete some codes
|
||||
if a_name in stored_once and b_name in stored_once:
|
||||
instrs.remove(load_a)
|
||||
instrs.remove(store_b)
|
||||
if a_name != b_name:
|
||||
for instr in instrs:
|
||||
if (
|
||||
instr.opname
|
||||
in (
|
||||
"LOAD_FAST_CHECK",
|
||||
"LOAD_FAST",
|
||||
"LOAD_FAST_BORROW",
|
||||
"STORE_FAST",
|
||||
)
|
||||
and instr.argval == b_name
|
||||
):
|
||||
instr.argval = a_name
|
||||
instr.arg = load_a.arg
|
||||
modified = True
|
||||
|
||||
# if
|
||||
# LOAD A
|
||||
# STORE B
|
||||
# A or B is not stored only once (maybe it is input)
|
||||
# we give a more general way to simplify the codes
|
||||
#
|
||||
# if A will not be loaded again after (6)STORE B, it means we can move (6)STORE B ahead to (1)STORE A
|
||||
# TIP: there is no more STORE A between (1) and (5)
|
||||
# (1) STORE A -> STORE B
|
||||
# ... ...
|
||||
# (2) LOAD A -> LOAD B
|
||||
# ...
|
||||
# (3) LOAD B -> not support
|
||||
# ...
|
||||
# (4) STORE B -> not support
|
||||
# ... ...
|
||||
# (5) LOAD A -> ---- (rm)
|
||||
# (6) STORE B ---- (rm)
|
||||
# ...
|
||||
# (7) STORE B
|
||||
# (8) LOAD A
|
||||
# so we can rename the rest LOAD A below as LOAD B
|
||||
#
|
||||
# What changed:
|
||||
# 1. if (4) exist, B changed:
|
||||
# (1) ~ (4), (6) ~
|
||||
# 2. if (4) not exist, B changed:
|
||||
# (1), (6)
|
||||
# 3. A changed:
|
||||
# (1) ~
|
||||
#
|
||||
# To do this transform, we should make sure
|
||||
# 1. (4) is not exist in (1) ~ (5): it is too complex
|
||||
# 2. (3) is not exist in (1) ~ (5): load B in the range that B value is changed
|
||||
# 3. (7) (8) is not exist in (6)~: load A in range that A value is changed, if we load B instead, but B also changed
|
||||
# we can simplify this as "no more LOAD A after (6)"
|
||||
else:
|
||||
last_store_a = stored_from(load_a, instrs)
|
||||
if last_store_a is None:
|
||||
# if last store a just not exist, we can not do this transform
|
||||
continue
|
||||
|
||||
last_store_idx = instrs.index(last_store_a)
|
||||
code_range = instrs[last_store_idx : instrs.index(store_b)]
|
||||
if (
|
||||
not code_exist("STORE_FAST", b_name, code_range)
|
||||
and not code_exist("LOAD_FAST_CHECK", b_name, code_range)
|
||||
and not code_exist("LOAD_FAST", b_name, code_range)
|
||||
and not code_exist("LOAD_FAST_BORROW", b_name, code_range)
|
||||
and not code_exist(
|
||||
"LOAD_FAST_CHECK",
|
||||
a_name,
|
||||
instrs[instrs.index(store_b) :],
|
||||
)
|
||||
and not code_exist(
|
||||
"LOAD_FAST", a_name, instrs[instrs.index(store_b) :]
|
||||
)
|
||||
and not code_exist(
|
||||
"LOAD_FAST_BORROW",
|
||||
a_name,
|
||||
instrs[instrs.index(store_b) :],
|
||||
)
|
||||
):
|
||||
last_store_a.argval = b_name
|
||||
last_store_a.arg = store_b.arg
|
||||
instrs.remove(load_a)
|
||||
instrs.remove(store_b)
|
||||
for instr in instrs[last_store_idx:]:
|
||||
if (
|
||||
instr.opname
|
||||
in (
|
||||
"LOAD_FAST_CHECK",
|
||||
"LOAD_FAST",
|
||||
"LOAD_FAST_BORROW",
|
||||
"STORE_FAST",
|
||||
)
|
||||
and instr.argval == a_name
|
||||
):
|
||||
instr.argval = b_name
|
||||
instr.arg = store_b.arg
|
||||
|
||||
|
||||
def remove_duplicate_resume(instrs: list[Instruction], code_options):
|
||||
resumes = list(filter(lambda instr: instr.opname == "RESUME", instrs))
|
||||
if not resumes:
|
||||
return
|
||||
for resume in resumes[1:]:
|
||||
instrs.remove(resume)
|
||||
|
||||
|
||||
def check_precall_followed_by_call(instrs: list[Instruction], code_options):
|
||||
"""
|
||||
PRECALL should be followed by CALL, otherwise it will cause a segmentation fault
|
||||
"""
|
||||
for instr, next_instr in zip(instrs[:-1], instrs[1:]):
|
||||
if instr.opname == "PRECALL" and next_instr.opname != "CALL":
|
||||
raise InnerError(
|
||||
f"PRECALL is not followed by CALL in {code_options['co_name']}"
|
||||
)
|
||||
|
||||
|
||||
def check_for_iter_jump_to(instrs: list[Instruction], code_options):
|
||||
"""
|
||||
Check if the `jump_to` of FOR_ITER is END_FOR, in Python3.12+
|
||||
"""
|
||||
for instr in instrs:
|
||||
if instr.opname == "FOR_ITER":
|
||||
assert instr.jump_to is not None
|
||||
if instr.jump_to.opname != "END_FOR":
|
||||
raise InnerError("FOR_ITER jump_to is not END_FOR")
|
||||
|
||||
|
||||
def fuse_double_super_instrs(instrs: list[Instruction], code_options):
|
||||
"""
|
||||
Fuse two consecutive LOAD_FAST or STORE_FAST instructions into one.
|
||||
"""
|
||||
co_varnames = code_options['co_varnames']
|
||||
|
||||
def able_to_merge(idx: int):
|
||||
return (
|
||||
idx > 0
|
||||
and (instrs[idx - 1].opname, instrs[idx].opname)
|
||||
in TO_FUSED_INSTS.keys()
|
||||
and not instrs[idx].is_jump_target
|
||||
and not instrs[idx - 1].is_jump_target
|
||||
and co_varnames.index(instrs[idx - 1].argval) < 16
|
||||
and co_varnames.index(instrs[idx].argval) < 16
|
||||
)
|
||||
|
||||
def merge_two_op(prev_instr: Instruction, instr: Instruction):
|
||||
merge_key = (instrs[idx - 1].opname, instrs[idx].opname)
|
||||
prev_instr.opname = TO_FUSED_INSTS[merge_key]
|
||||
prev_instr.opcode = dis.opmap[prev_instr.opname]
|
||||
prev_instr.is_generated = True
|
||||
prev_instr.argval = (prev_instr.argval, instr.argval)
|
||||
instrs.remove(instr)
|
||||
|
||||
idx = 0
|
||||
# We must manually control the indices, so we cannot use a for loop.
|
||||
while idx < len(instrs):
|
||||
if able_to_merge(idx):
|
||||
merge_two_op(instrs[idx - 1], instrs[idx])
|
||||
continue
|
||||
|
||||
idx += 1
|
||||
@@ -0,0 +1,586 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import dis
|
||||
import sys
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ...utils import InnerError
|
||||
from .opcode_info import (
|
||||
ABS_JUMP,
|
||||
ALL_JUMP,
|
||||
FUSED_INSTS,
|
||||
PYOPCODE_CACHE_SIZE,
|
||||
REL_BWD_JUMP,
|
||||
REL_JUMP,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import types
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Instruction:
|
||||
opcode: int
|
||||
opname: str
|
||||
arg: int | None
|
||||
argval: Any
|
||||
offset: int | None = None
|
||||
starts_line: int | None = None
|
||||
is_jump_target: bool = False
|
||||
jump_to: Instruction | None = None
|
||||
is_generated: bool = True
|
||||
|
||||
# for analysis EXTENDED_ARG
|
||||
first_ex_arg: Instruction | None = None
|
||||
ex_arg_for: Instruction | None = None
|
||||
|
||||
# used in modify_extended_args
|
||||
def __hash__(self):
|
||||
return id(self)
|
||||
|
||||
def __eq__(self, instr):
|
||||
return id(self) == id(instr)
|
||||
|
||||
|
||||
def get_instruction_size(instr: Instruction) -> int:
|
||||
cache_size = 0
|
||||
if sys.version_info >= (3, 11):
|
||||
cache_size = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
|
||||
return 2 * (cache_size + 1)
|
||||
|
||||
|
||||
def gen_instr(name, arg=None, argval=None, gened=True, jump_to=None):
|
||||
return Instruction(
|
||||
opcode=dis.opmap[name],
|
||||
opname=name,
|
||||
arg=arg,
|
||||
argval=argval,
|
||||
is_generated=gened,
|
||||
jump_to=jump_to,
|
||||
)
|
||||
|
||||
|
||||
def convert_instruction(instr: dis.Instruction) -> Instruction:
|
||||
"""
|
||||
Converts a disassembled instruction to a customized Instruction object.
|
||||
|
||||
Args:
|
||||
instr (dis.Instruction): The disassembled instruction.
|
||||
|
||||
Returns:
|
||||
Instruction: A customized Instruction object.
|
||||
"""
|
||||
return Instruction(
|
||||
instr.opcode,
|
||||
instr.opname,
|
||||
instr.arg,
|
||||
instr.argval,
|
||||
instr.offset,
|
||||
instr.line_number if sys.version_info >= (3, 13) else instr.starts_line,
|
||||
instr.is_jump_target,
|
||||
jump_to=None,
|
||||
is_generated=False,
|
||||
)
|
||||
|
||||
|
||||
def replace_jump_target(
|
||||
instrs: list[Instruction],
|
||||
replacements: dict[Instruction, Instruction],
|
||||
) -> None:
|
||||
"""Replace jump targets based on the replacements dictionary.
|
||||
|
||||
Args:
|
||||
instrs (list[Instruction]): The list of instructions to modify.
|
||||
replacements (dict[Instruction, Instruction]): Mapping from old jump targets to new ones.
|
||||
"""
|
||||
for instr in instrs:
|
||||
if instr.jump_to in replacements:
|
||||
instr.jump_to = replacements[instr.jump_to]
|
||||
|
||||
|
||||
def expand_super_instrs(instructions: list[Instruction]) -> list[Instruction]:
|
||||
expanded_instrs = []
|
||||
replacements = {}
|
||||
|
||||
def copy_instruction(
|
||||
instr, opname, argval, arg, is_jump_target, is_generated
|
||||
):
|
||||
return Instruction(
|
||||
opcode=dis.opmap[opname],
|
||||
opname=opname,
|
||||
arg=arg,
|
||||
argval=argval,
|
||||
is_jump_target=is_jump_target,
|
||||
is_generated=is_generated,
|
||||
jump_to=instr.jump_to,
|
||||
)
|
||||
|
||||
for instr in instructions:
|
||||
if instr.opname in FUSED_INSTS:
|
||||
instr1 = copy_instruction(
|
||||
instr,
|
||||
FUSED_INSTS[instr.opname][0],
|
||||
instr.argval[0],
|
||||
instr.arg >> 4,
|
||||
instr.is_jump_target,
|
||||
True,
|
||||
)
|
||||
instr2 = copy_instruction(
|
||||
instr,
|
||||
FUSED_INSTS[instr.opname][1],
|
||||
instr.argval[1],
|
||||
instr.arg & 15,
|
||||
False,
|
||||
False,
|
||||
)
|
||||
replacements[instr] = instr1
|
||||
expanded_instrs.append(instr1)
|
||||
expanded_instrs.append(instr2)
|
||||
# If the LOAD_ATTR opcode will lead to load_method in 3.13+, we manually split it into two instructions,
|
||||
# to avoid Uncontrollable specialization that changes the behavior of LOAD_ATTR,
|
||||
# which can lead to incorrect results when the current graph is smaller than the MIN_GRAPH_SIZE
|
||||
elif (
|
||||
sys.version_info >= (3, 13)
|
||||
and instr.opname == "LOAD_ATTR"
|
||||
and instr.arg & 1
|
||||
):
|
||||
instr1 = copy_instruction(
|
||||
instr,
|
||||
"LOAD_ATTR",
|
||||
instr.argval,
|
||||
instr.arg & ~1,
|
||||
instr.is_jump_target,
|
||||
True,
|
||||
)
|
||||
instr2 = Instruction(
|
||||
dis.opmap["PUSH_NULL"],
|
||||
"PUSH_NULL",
|
||||
None,
|
||||
None,
|
||||
is_generated=True,
|
||||
)
|
||||
replacements[instr] = instr1
|
||||
expanded_instrs.append(instr1)
|
||||
expanded_instrs.append(instr2)
|
||||
else:
|
||||
expanded_instrs.append(instr)
|
||||
|
||||
replace_jump_target(expanded_instrs, replacements)
|
||||
return expanded_instrs
|
||||
|
||||
|
||||
def replace_load_fast_borrow_with_strong_ref(
|
||||
instructions: list[Instruction],
|
||||
) -> list[Instruction]:
|
||||
"""
|
||||
Patch LOAD_FAST_BORROW to LOAD_FAST for Python 3.14+.
|
||||
|
||||
LOAD_FAST_BORROW loads a value using a borrowing reference and does not
|
||||
increment the reference count. In some cases this can cause subsequent
|
||||
STORE_FAST or other operations to retain a reference that becomes invalid
|
||||
once the borrowed value is released, leading to incorrect behavior or
|
||||
crashes when the variable is accessed later.
|
||||
To avoid these issues, we replace LOAD_FAST_BORROW with LOAD_FAST here.
|
||||
"""
|
||||
replacements = {}
|
||||
expanded_instrs = []
|
||||
for instr in instructions:
|
||||
if instr.opname == "LOAD_FAST_BORROW":
|
||||
instr1 = Instruction(
|
||||
dis.opmap["LOAD_FAST"],
|
||||
"LOAD_FAST",
|
||||
instr.arg,
|
||||
instr.argval,
|
||||
is_generated=instr.is_generated,
|
||||
is_jump_target=instr.is_jump_target,
|
||||
jump_to=instr.jump_to,
|
||||
)
|
||||
replacements[instr] = instr1
|
||||
expanded_instrs.append(instr1)
|
||||
else:
|
||||
expanded_instrs.append(instr)
|
||||
|
||||
replace_jump_target(expanded_instrs, replacements)
|
||||
return expanded_instrs
|
||||
|
||||
|
||||
def get_instructions(code: types.CodeType) -> list[Instruction]:
|
||||
"""
|
||||
Returns parsed instructions from the given code object and exclude
|
||||
any opcodes that contain `EXTENDED_ARG`.
|
||||
|
||||
Args:
|
||||
code (types.CodeType): The code object to extract instructions from.
|
||||
|
||||
Returns:
|
||||
list[Instruction]: A list of Instruction objects representing the
|
||||
bytecode instructions in the code object.
|
||||
"""
|
||||
# instrs do not contain EXTENDED_ARG
|
||||
instrs = list(map(convert_instruction, dis.get_instructions(code)))
|
||||
for instr in instrs:
|
||||
if instr.opname in ALL_JUMP:
|
||||
origin_jump_target = calc_offset_from_bytecode_offset(
|
||||
instr.argval, instrs
|
||||
)
|
||||
jump_offset = origin_jump_target
|
||||
|
||||
while instrs[jump_offset].opname == "EXTENDED_ARG":
|
||||
jump_offset += 1
|
||||
|
||||
if origin_jump_target != jump_offset:
|
||||
# copy infos from EXTENDED_ARG to other opcode
|
||||
|
||||
if instrs[origin_jump_target].is_jump_target:
|
||||
instrs[jump_offset].is_jump_target = instrs[
|
||||
origin_jump_target
|
||||
].is_jump_target
|
||||
if instrs[origin_jump_target].starts_line:
|
||||
instrs[jump_offset].starts_line = instrs[
|
||||
origin_jump_target
|
||||
].starts_line
|
||||
|
||||
instr.jump_to = instrs[jump_offset]
|
||||
|
||||
# if the origin opcode contains EXTENDED_ARG, it should be like:
|
||||
# >> EXTENDED_ARG 1
|
||||
# XX 388 <- 256 + 132
|
||||
# filter all EXTENDED_ARG here
|
||||
instrs = [x for x in instrs if x.opname != "EXTENDED_ARG"]
|
||||
prepare_passes = [expand_super_instrs]
|
||||
if sys.version_info >= (3, 14):
|
||||
prepare_passes.append(replace_load_fast_borrow_with_strong_ref)
|
||||
|
||||
for pass_fn in prepare_passes:
|
||||
instrs = pass_fn(instrs)
|
||||
return instrs
|
||||
|
||||
|
||||
def modify_instrs(instructions: list[Instruction]) -> None:
|
||||
"""
|
||||
Modifies the given list of instructions. It contains three steps:
|
||||
|
||||
1. reset offset
|
||||
2. relocate jump target
|
||||
3. add EXTENDED_ARG instruction if needed
|
||||
|
||||
Args:
|
||||
instructions (list): The list of Instruction objects representing bytecode instructions.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
modify_completed = False
|
||||
while not modify_completed:
|
||||
reset_offset(instructions)
|
||||
has_inverted_jump = relocate_jump_target(instructions)
|
||||
modify_completed = (
|
||||
modify_extended_args(instructions) and not has_inverted_jump
|
||||
)
|
||||
|
||||
|
||||
def reset_offset(instructions: list[Instruction]) -> None:
|
||||
"""
|
||||
Resets the offset for each instruction in the list.
|
||||
|
||||
Args:
|
||||
instructions (list): The list of Instruction objects representing bytecode instructions.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
from ..executor.pycode_generator import get_instruction_size
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
current_offset = 0
|
||||
for instr in instructions:
|
||||
instr.offset = current_offset
|
||||
current_offset += get_instruction_size(instr)
|
||||
return
|
||||
for idx, instr in enumerate(instructions):
|
||||
instr.offset = idx * 2
|
||||
|
||||
|
||||
def correct_jump_direction(
|
||||
instr: Instruction, arg: int
|
||||
) -> tuple[Instruction, bool]:
|
||||
"""
|
||||
Corrects the jump direction of the given instruction.
|
||||
NOTE(zrr1999): In Python 3.11, JUMP_ABSOLUTE is removed, so python generates JUMP_FORWARD or JUMP_BACKWARD instead,
|
||||
but in for loop breakgraph, we reuse JUMP_BACKWARD to jump forward, so we need to change it to JUMP_FORWARD.
|
||||
|
||||
Args:
|
||||
instr (Instruction): The instruction to be corrected.
|
||||
invert_jump (bool): Whether to invert the jump direction.
|
||||
"""
|
||||
if instr.opname in ABS_JUMP:
|
||||
instr.arg = arg
|
||||
return instr, False
|
||||
elif instr.opname in REL_JUMP:
|
||||
if arg < 0:
|
||||
if instr.opname in REL_BWD_JUMP:
|
||||
forward_op_name = instr.opname.replace("BACKWARD", "FORWARD")
|
||||
if forward_op_name not in dis.opmap:
|
||||
raise InnerError(f"Unknown jump type {instr.opname}")
|
||||
instr.opname = forward_op_name
|
||||
instr.opcode = dis.opmap[forward_op_name]
|
||||
else: # instr.opname in REL_FWD_JUMP
|
||||
backward_op_name = instr.opname.replace("FORWARD", "BACKWARD")
|
||||
if backward_op_name not in dis.opmap:
|
||||
raise InnerError(f"Unknown jump type {instr.opname}")
|
||||
instr.opname = backward_op_name
|
||||
instr.opcode = dis.opmap[backward_op_name]
|
||||
instr.arg = -arg
|
||||
invert_jump = True
|
||||
else:
|
||||
instr.arg = arg
|
||||
invert_jump = False
|
||||
return instr, invert_jump
|
||||
else:
|
||||
raise ValueError(f"unknown jump type: {instr.opname}")
|
||||
|
||||
|
||||
def relocate_jump_target(instructions: list[Instruction]) -> bool:
|
||||
"""
|
||||
If a jump instruction is found, this function will adjust the jump targets based on the presence of EXTENDED_ARG instructions.
|
||||
If an EXTENDED_ARG instruction exists for the jump target, use its offset as the new target.
|
||||
|
||||
Args:
|
||||
instructions (list): The list of Instruction objects representing bytecode instructions.
|
||||
|
||||
Returns:
|
||||
bool: True if the jump direction is inverted, False otherwise.
|
||||
"""
|
||||
has_inverted_jump = False
|
||||
extended_arg = []
|
||||
for instr in instructions:
|
||||
if instr.opname == "EXTENDED_ARG":
|
||||
extended_arg.append(instr)
|
||||
continue
|
||||
|
||||
if instr.opname in ALL_JUMP:
|
||||
assert instr.jump_to is not None
|
||||
assert instr.offset is not None
|
||||
# if jump target has extended_arg, should jump to the first extended_arg opcode
|
||||
jump_target = (
|
||||
instr.jump_to.offset
|
||||
if instr.jump_to.first_ex_arg is None
|
||||
else instr.jump_to.first_ex_arg.offset
|
||||
)
|
||||
assert jump_target is not None
|
||||
|
||||
if instr.opname in ABS_JUMP:
|
||||
new_arg = jump_target
|
||||
else: # instr.opname in REL_JUMP
|
||||
cache_size = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
|
||||
new_arg = jump_target - (2 * cache_size) - instr.offset - 2
|
||||
if instr.opname in REL_BWD_JUMP:
|
||||
new_arg = -new_arg
|
||||
|
||||
new_arg //= 2
|
||||
_, invert_jump = correct_jump_direction(instr, new_arg)
|
||||
has_inverted_jump = has_inverted_jump or invert_jump
|
||||
assert instr.arg is not None
|
||||
if extended_arg:
|
||||
instr.arg &= 0xFF
|
||||
new_arg = new_arg >> 8
|
||||
for ex in reversed(extended_arg):
|
||||
ex.arg = new_arg & 0xFF
|
||||
new_arg = new_arg >> 8
|
||||
|
||||
# need more extended_args instr
|
||||
# set arg in the first extended_arg
|
||||
if new_arg > 0:
|
||||
extended_arg[0].arg += new_arg << 8
|
||||
extended_arg.clear()
|
||||
return has_inverted_jump
|
||||
|
||||
|
||||
def modify_extended_args(instructions: list[Instruction]) -> bool:
|
||||
"""
|
||||
This function replaces any instruction with an argument greater than or equal to 256 with one or more EXTENDED_ARG instructions.
|
||||
|
||||
Args:
|
||||
instructions (list): The list of Instruction objects representing bytecode instructions.
|
||||
|
||||
Returns:
|
||||
bool: True if the modification is completed, False otherwise.
|
||||
"""
|
||||
|
||||
modify_completed = True
|
||||
extend_args_record = {}
|
||||
for instr in instructions:
|
||||
if instr.arg and instr.arg >= 256: # more than one byte
|
||||
_instrs = [
|
||||
instr
|
||||
] # replace instr with _instrs later (it is a set of instrs), all operations will be recorded in extend_args_record
|
||||
val = instr.arg
|
||||
instr.arg = val & 0xFF
|
||||
val = val >> 8
|
||||
while val > 0:
|
||||
_instrs.append(gen_instr("EXTENDED_ARG", arg=val & 0xFF))
|
||||
val = val >> 8
|
||||
|
||||
extend_args_record.update({instr: list(reversed(_instrs))})
|
||||
|
||||
if extend_args_record:
|
||||
# if new EXTENDED_ARG inserted, we need update offset and jump target
|
||||
modify_completed = False
|
||||
|
||||
def bind_ex_arg_with_instr(ex_arg, instr):
|
||||
# move opcode info to EXTENDED_ARG
|
||||
ex_arg.starts_line = instr.starts_line
|
||||
instr.starts_line = None
|
||||
ex_arg.is_jump_target = instr.is_jump_target
|
||||
instr.is_jump_target = False
|
||||
|
||||
if instr.ex_arg_for is not None:
|
||||
# instr is also an ex_arg for another instr
|
||||
instr.ex_arg_for.first_ex_arg = ex_arg
|
||||
ex_arg.ex_arg_for = instr.ex_arg_for
|
||||
instr.ex_arg_for = None
|
||||
else:
|
||||
instr.first_ex_arg = ex_arg
|
||||
ex_arg.ex_arg_for = instr
|
||||
|
||||
for key, val in extend_args_record.items():
|
||||
bind_ex_arg_with_instr(val[0], key)
|
||||
replace_instr(instructions, instr=key, new_instr=val)
|
||||
|
||||
return modify_completed
|
||||
|
||||
|
||||
def modify_vars(instructions: list[Instruction], code_options):
|
||||
co_varnames = code_options['co_varnames']
|
||||
co_freevars = code_options['co_freevars']
|
||||
for instrs in instructions:
|
||||
if instrs.opname in [
|
||||
'LOAD_FAST',
|
||||
'LOAD_FAST_BORROW',
|
||||
'LOAD_FAST_CHECK',
|
||||
'STORE_FAST',
|
||||
'DELETE_FAST',
|
||||
]:
|
||||
assert instrs.argval in co_varnames, (
|
||||
f"`{instrs.argval}` not in {co_varnames}"
|
||||
)
|
||||
instrs.arg = co_varnames.index(instrs.argval)
|
||||
elif instrs.opname == "LOAD_DEREF" or instrs.opname == "STORE_DEREF":
|
||||
if sys.version_info >= (3, 11):
|
||||
namemap = co_varnames + co_freevars
|
||||
assert instrs.argval in namemap, (
|
||||
f"`{instrs.argval}` not in {namemap}"
|
||||
)
|
||||
instrs.arg = namemap.index(instrs.argval)
|
||||
elif instrs.opname in FUSED_INSTS.keys():
|
||||
assert instrs.argval[0] in co_varnames, (
|
||||
f"`{instrs.argval[0]}` not in {co_varnames}"
|
||||
)
|
||||
assert instrs.argval[1] in co_varnames, (
|
||||
f"`{instrs.argval[1]}` not in {co_varnames}"
|
||||
)
|
||||
instrs.arg = (
|
||||
co_varnames.index(instrs.argval[0]) << 4
|
||||
) + co_varnames.index(instrs.argval[1])
|
||||
|
||||
|
||||
def calc_offset_from_bytecode_offset(
|
||||
bytecode_offset: int,
|
||||
instructions: list[dis.Instruction] | list[Instruction],
|
||||
) -> int:
|
||||
"""
|
||||
Calculate the index from bytecode offset, because it have 2 bytes per instruction (for Python <= 3.10).
|
||||
|
||||
Args:
|
||||
bytecode_offset (int): The bytecode offset of the instruction.
|
||||
|
||||
Returns:
|
||||
int: The index of the instruction in the instruction list.
|
||||
"""
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
instruction_offsets = [x.offset for x in instructions]
|
||||
return instruction_offsets.index(bytecode_offset)
|
||||
return bytecode_offset // 2
|
||||
|
||||
|
||||
def replace_instr(instructions, instr, new_instr):
|
||||
idx = instructions.index(instr)
|
||||
instructions[idx : idx + 1] = new_instr
|
||||
|
||||
|
||||
def instrs_info(instrs, mark=None, range=None, want_str=True):
|
||||
ret = []
|
||||
start = -1
|
||||
end = 1000000
|
||||
if mark is not None and range is not None:
|
||||
start = mark - range
|
||||
end = mark + range + 1
|
||||
for idx, instr in enumerate(instrs):
|
||||
if idx < start or idx >= end:
|
||||
continue
|
||||
if instr.starts_line is not None:
|
||||
ret.append("")
|
||||
ret.append(
|
||||
"{line:<8s}{is_jump_target:>2s}{offset:>4d} {opname:<30s}{arg:<4s}{argval:<40s}{mark}".format(
|
||||
line=str(instr.starts_line) if instr.starts_line else "",
|
||||
is_jump_target=">>" if instr.is_jump_target else " ",
|
||||
offset=(
|
||||
instr.offset if instr.offset or instr.offset == 0 else -1
|
||||
),
|
||||
opname=instr.opname,
|
||||
arg=str(instr.arg) if instr.arg is not None else "",
|
||||
argval=f"({instr.argval})" if instr.argval else "",
|
||||
mark="",
|
||||
)
|
||||
)
|
||||
if idx == mark:
|
||||
ret[-1] = "\033[31m" + ret[-1] + "\033[0m"
|
||||
if want_str:
|
||||
return "\n".join(ret)
|
||||
return ret
|
||||
|
||||
|
||||
def calc_stack_effect(instr: Instruction, *, jump: bool | None = None) -> int:
|
||||
"""
|
||||
Gets the stack effect of the given instruction. In Python 3.11, the stack effect of `CALL` is -1,
|
||||
refer to https://github.com/python/cpython/blob/3.11/Python/compile.c#L1123-L1124.
|
||||
|
||||
Args:
|
||||
instr: The instruction.
|
||||
|
||||
Returns:
|
||||
The stack effect of the instruction.
|
||||
|
||||
"""
|
||||
if sys.version_info[:2] == (3, 11):
|
||||
if instr.opname == "PRECALL":
|
||||
return 0
|
||||
elif instr.opname == "CALL":
|
||||
# NOTE(zrr1999): push_n = 1, pop_n = oparg + 2, stack_effect = push_n - pop_n = -oparg-1
|
||||
assert instr.arg is not None
|
||||
return -instr.arg - 1
|
||||
return dis.stack_effect(instr.opcode, instr.arg, jump=jump)
|
||||
|
||||
|
||||
class Space(Enum):
|
||||
locals = 1
|
||||
globals = 2
|
||||
cells = 3
|
||||
not_found = 4
|
||||
@@ -0,0 +1,161 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from paddle.jit.utils import OrderedSet
|
||||
|
||||
from .opcode_info import (
|
||||
ALL_JUMP,
|
||||
HAS_FREE,
|
||||
HAS_LOCAL,
|
||||
UNCONDITIONAL_JUMP,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .instruction_utils import Instruction
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class NameRecorder:
|
||||
reads: OrderedSet[str]
|
||||
writes: OrderedSet[str]
|
||||
|
||||
def __or__(self, other):
|
||||
reads = self.reads | other.reads
|
||||
writes = self.writes | other.writes
|
||||
return NameRecorder(reads, writes)
|
||||
|
||||
|
||||
def is_read_opcode(opname):
|
||||
if opname in [
|
||||
"LOAD_FAST",
|
||||
"LOAD_FAST_CHECK",
|
||||
"LOAD_DEREF",
|
||||
"LOAD_NAME",
|
||||
"LOAD_GLOBAL",
|
||||
"LOAD_CLOSURE",
|
||||
"LOAD_FAST_BORROW",
|
||||
]:
|
||||
return True
|
||||
if opname in (
|
||||
"DELETE_FAST",
|
||||
"DELETE_DEREF",
|
||||
"DELETE_NAME",
|
||||
"DELETE_GLOBAL",
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_write_opcode(opname):
|
||||
if opname in ["STORE_FAST", "STORE_NAME", "STORE_DEREF", "STORE_GLOBAL"]:
|
||||
return True
|
||||
if opname in (
|
||||
"DELETE_FAST",
|
||||
"DELETE_DEREF",
|
||||
"DELETE_NAME",
|
||||
"DELETE_GLOBAL",
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def analysis_used_names(
|
||||
instructions: list[Instruction],
|
||||
current_instr_idx: int,
|
||||
stop_instr_idx: int | None = None,
|
||||
) -> tuple[OrderedSet[str], OrderedSet[str]]:
|
||||
"""
|
||||
Analyze the inputs of the instructions from current_instr_idx to stop_instr_idx.
|
||||
|
||||
Args:
|
||||
instructions (list[Instruction]): The instructions to analyze.
|
||||
current_instr_idx (int): The index of the current instruction.
|
||||
stop_instr_idx (int | None, optional): The index of the instruction to stop. Defaults to None.
|
||||
If None, the analysis will stop at the end of the instructions.
|
||||
|
||||
Returns:
|
||||
State: The analysis result.
|
||||
"""
|
||||
name_recorder = NameRecorder(OrderedSet(), OrderedSet())
|
||||
|
||||
# start idx and writes names can decide the analysis result below
|
||||
# so, just check the pair of (idx, writes), to skip repeat simulation
|
||||
# (writes can decide if a name should be add to reads)
|
||||
# one idx can has multi writes for whom is not subset with each other
|
||||
# if A is subset of B, we just record A, simulate A might add more reads
|
||||
visited_states = {}
|
||||
|
||||
def check_and_update_visited_states(idx, writes):
|
||||
writes = set(writes)
|
||||
|
||||
if idx in visited_states:
|
||||
history = visited_states[idx]
|
||||
for record in history:
|
||||
if record.issubset(writes):
|
||||
return True
|
||||
elif writes.issubset(record):
|
||||
history.remove(record)
|
||||
history.append(writes)
|
||||
return False
|
||||
else:
|
||||
visited_states[idx] = [writes]
|
||||
|
||||
return False
|
||||
|
||||
def fork(
|
||||
name_recorder: NameRecorder, start: int, jump: bool, jump_target: int
|
||||
) -> NameRecorder:
|
||||
new_start = start + 1 if not jump else jump_target
|
||||
new_state = NameRecorder(
|
||||
OrderedSet(name_recorder.reads),
|
||||
OrderedSet(name_recorder.writes),
|
||||
)
|
||||
return walk(new_state, new_start)
|
||||
|
||||
def walk(name_recorder: NameRecorder, start: int) -> NameRecorder:
|
||||
end = len(instructions) if stop_instr_idx is None else stop_instr_idx
|
||||
for i in range(start, end):
|
||||
if check_and_update_visited_states(i, name_recorder.writes):
|
||||
return name_recorder
|
||||
|
||||
instr = instructions[i]
|
||||
if instr.opname in HAS_LOCAL | HAS_FREE:
|
||||
if is_read_opcode(instr.opname) and instr.argval not in (
|
||||
name_recorder.writes
|
||||
):
|
||||
name_recorder.reads.add(instr.argval)
|
||||
elif is_write_opcode(instr.opname):
|
||||
name_recorder.writes.add(instr.argval)
|
||||
elif instr.opname in ALL_JUMP:
|
||||
assert instr.jump_to is not None
|
||||
target_idx = instructions.index(instr.jump_to)
|
||||
# Fork to two branches, jump or not
|
||||
jump_branch = fork(name_recorder, i, True, target_idx)
|
||||
not_jump_branch = (
|
||||
fork(name_recorder, i, False, target_idx)
|
||||
if instr.opname not in UNCONDITIONAL_JUMP
|
||||
else NameRecorder(OrderedSet(), OrderedSet())
|
||||
)
|
||||
return jump_branch | not_jump_branch
|
||||
elif instr.opname == "RETURN_VALUE":
|
||||
return name_recorder
|
||||
return name_recorder
|
||||
|
||||
name_recorder = walk(name_recorder, current_instr_idx)
|
||||
return name_recorder.reads, name_recorder.writes
|
||||
@@ -0,0 +1,167 @@
|
||||
# Copyright (c) 2023 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.
|
||||
from __future__ import annotations
|
||||
|
||||
import opcode
|
||||
import sys
|
||||
from enum import Enum
|
||||
|
||||
REL_JUMP = {opcode.opname[x] for x in opcode.hasjrel}
|
||||
REL_BWD_JUMP = {opname for opname in REL_JUMP if "BACKWARD" in opname}
|
||||
REL_FWD_JUMP = REL_JUMP - REL_BWD_JUMP
|
||||
ABS_JUMP = {opcode.opname[x] for x in opcode.hasjabs}
|
||||
HAS_LOCAL = {opcode.opname[x] for x in opcode.haslocal}
|
||||
HAS_FREE = {opcode.opname[x] for x in opcode.hasfree}
|
||||
NEED_TO_BOOL = {"UNARY_NOT", "POP_JUMP_IF_FALSE", "POP_JUMP_IF_TRUE"}
|
||||
ALL_JUMP = REL_JUMP | ABS_JUMP
|
||||
UNCONDITIONAL_JUMP = {"JUMP_ABSOLUTE", "JUMP_FORWARD"}
|
||||
if sys.version_info >= (3, 11):
|
||||
UNCONDITIONAL_JUMP.add("JUMP_BACKWARD")
|
||||
RETURN = {"RETURN_VALUE"}
|
||||
if (3, 12) <= sys.version_info < (3, 14):
|
||||
RETURN.add("RETURN_CONST")
|
||||
|
||||
|
||||
class JumpDirection(Enum):
|
||||
FORWARD = "FORWARD"
|
||||
BACKWARD = "BACKWARD"
|
||||
|
||||
|
||||
class PopJumpCond(Enum):
|
||||
FALSE = "FALSE"
|
||||
TRUE = "TRUE"
|
||||
NONE = "NONE"
|
||||
NOT_NONE = "NOT_NONE"
|
||||
|
||||
|
||||
def _get_pyopcode_cache_size() -> dict[str, int]:
|
||||
if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
|
||||
# Cache for some opcodes, it's for Python 3.11
|
||||
# https://github.com/python/cpython/blob/3.11/Include/internal/pycore_opcode.h#L41-L53
|
||||
return {
|
||||
"BINARY_SUBSCR": 4,
|
||||
"STORE_SUBSCR": 1,
|
||||
"UNPACK_SEQUENCE": 1,
|
||||
"STORE_ATTR": 4,
|
||||
"LOAD_ATTR": 4,
|
||||
"COMPARE_OP": 2,
|
||||
"LOAD_GLOBAL": 5,
|
||||
"BINARY_OP": 1,
|
||||
"LOAD_METHOD": 10,
|
||||
"PRECALL": 1,
|
||||
"CALL": 4,
|
||||
}
|
||||
elif sys.version_info >= (3, 12) and sys.version_info < (3, 13):
|
||||
# Cache for some opcodes, it's for Python 3.12
|
||||
# https://github.com/python/cpython/blob/3.12/Include/internal/pycore_opcode.h#L34-L47
|
||||
return {
|
||||
"BINARY_SUBSCR": 1,
|
||||
"STORE_SUBSCR": 1,
|
||||
"UNPACK_SEQUENCE": 1,
|
||||
"FOR_ITER": 1,
|
||||
"STORE_ATTR": 4,
|
||||
"LOAD_ATTR": 9,
|
||||
"COMPARE_OP": 1,
|
||||
"LOAD_GLOBAL": 4,
|
||||
"BINARY_OP": 1,
|
||||
"SEND": 1,
|
||||
"LOAD_SUPER_ATTR": 1,
|
||||
"CALL": 3,
|
||||
}
|
||||
elif sys.version_info >= (3, 13) and sys.version_info < (3, 14):
|
||||
# Cache for some opcodes, it's for Python 3.13
|
||||
# https://github.com/python/cpython/blob/3.13/Include/internal/pycore_opcode_metadata.h#L1598-L1618
|
||||
return {
|
||||
"JUMP_BACKWARD": 1,
|
||||
"TO_BOOL": 3,
|
||||
"BINARY_SUBSCR": 1,
|
||||
"STORE_SUBSCR": 1,
|
||||
"SEND": 1,
|
||||
"UNPACK_SEQUENCE": 1,
|
||||
"STORE_ATTR": 4,
|
||||
"LOAD_GLOBAL": 4,
|
||||
"LOAD_SUPER_ATTR": 1,
|
||||
"LOAD_ATTR": 9,
|
||||
"COMPARE_OP": 1,
|
||||
"CONTAINS_OP": 1,
|
||||
"POP_JUMP_IF_TRUE": 1,
|
||||
"POP_JUMP_IF_FALSE": 1,
|
||||
"POP_JUMP_IF_NONE": 1,
|
||||
"POP_JUMP_IF_NOT_NONE": 1,
|
||||
"FOR_ITER": 1,
|
||||
"CALL": 3,
|
||||
"BINARY_OP": 1,
|
||||
}
|
||||
elif sys.version_info >= (3, 14) and sys.version_info < (3, 15):
|
||||
# Cache for some opcodes, it's for Python 3.14
|
||||
# https://github.com/python/cpython/blob/3.14/Include/internal/pycore_opcode_metadata.h#L1764-L1784
|
||||
return {
|
||||
"TO_BOOL": 3,
|
||||
"STORE_SUBSCR": 1,
|
||||
"SEND": 1,
|
||||
"UNPACK_SEQUENCE": 1,
|
||||
"STORE_ATTR": 4,
|
||||
"LOAD_GLOBAL": 4,
|
||||
"LOAD_SUPER_ATTR": 1,
|
||||
"LOAD_ATTR": 9,
|
||||
"COMPARE_OP": 1,
|
||||
"CONTAINS_OP": 1,
|
||||
"JUMP_BACKWARD": 1,
|
||||
"POP_JUMP_IF_TRUE": 1,
|
||||
"POP_JUMP_IF_FALSE": 1,
|
||||
"POP_JUMP_IF_NONE": 1,
|
||||
"POP_JUMP_IF_NOT_NONE": 1,
|
||||
"FOR_ITER": 1,
|
||||
"CALL": 3,
|
||||
"CALL_KW": 3,
|
||||
"BINARY_OP": 5,
|
||||
}
|
||||
elif sys.version_info >= (3, 15):
|
||||
raise NotImplementedError(
|
||||
f"Need to supplement cache operation code, for Python {sys.version_info}"
|
||||
)
|
||||
else:
|
||||
return {}
|
||||
|
||||
|
||||
PYOPCODE_CACHE_SIZE = _get_pyopcode_cache_size()
|
||||
|
||||
|
||||
class ExceptionHandler:
|
||||
opcode = 257
|
||||
opname = "EXCEPT_HANDLER"
|
||||
|
||||
|
||||
def _get_binary_op_arg_map() -> dict[str, int]:
|
||||
if sys.version_info < (3, 11):
|
||||
return {}
|
||||
res = {}
|
||||
for i, op in enumerate(opcode._nb_ops):
|
||||
res[op[0]] = i
|
||||
return res
|
||||
|
||||
|
||||
BINARY_OP_ARG_MAP: dict[str, int] = _get_binary_op_arg_map()
|
||||
|
||||
FUSED_INSTS: dict[str, tuple[str, str]] = {
|
||||
"LOAD_FAST_LOAD_FAST": ("LOAD_FAST", "LOAD_FAST"),
|
||||
"LOAD_FAST_BORROW_LOAD_FAST_BORROW": (
|
||||
"LOAD_FAST_BORROW",
|
||||
"LOAD_FAST_BORROW",
|
||||
),
|
||||
"STORE_FAST_STORE_FAST": ("STORE_FAST", "STORE_FAST"),
|
||||
"STORE_FAST_LOAD_FAST": ("STORE_FAST", "LOAD_FAST"),
|
||||
}
|
||||
|
||||
TO_FUSED_INSTS = {v: k for k, v in FUSED_INSTS.items()}
|
||||
@@ -0,0 +1,93 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from ...utils import Singleton
|
||||
|
||||
|
||||
class StackAnalyser(metaclass=Singleton):
|
||||
def stack_effect(self, instr):
|
||||
if "BINARY" in instr.opname or "INPLACE" in instr.opname:
|
||||
return 2, 1
|
||||
elif "UNARY" in instr.opname:
|
||||
return 1, 1
|
||||
return getattr(self, instr.opname)(instr)
|
||||
|
||||
def LOAD_GLOBAL(self, instr):
|
||||
return 0, 1
|
||||
|
||||
def LOAD_CONST(self, instr):
|
||||
return 0, 1
|
||||
|
||||
def LOAD_FAST(self, instr):
|
||||
return 0, 1
|
||||
|
||||
def LOAD_ATTR(self, instr):
|
||||
return 1, 1
|
||||
|
||||
def LOAD_METHOD(self, instr):
|
||||
return 1, 2
|
||||
|
||||
def STORE_FAST(self, instr):
|
||||
return 1, 0
|
||||
|
||||
def BUILD_TUPLE(self, instr):
|
||||
return instr.arg, 1
|
||||
|
||||
def BUILD_LIST(self, instr):
|
||||
return instr.arg, 1
|
||||
|
||||
def BUILD_SLICE(self, instr):
|
||||
if instr.arg == 3:
|
||||
return 3, 1
|
||||
else:
|
||||
return 2, 1
|
||||
|
||||
def UNPACK_SEQUENCE(self, instr):
|
||||
return 1, instr.arg
|
||||
|
||||
def CALL_FUNCTION(self, instr):
|
||||
return instr.arg + 1, 1
|
||||
|
||||
def DUP_TOP(self, instr):
|
||||
return 0, 1
|
||||
|
||||
def DUP_TOP_TWO(self, instr):
|
||||
return 0, 2
|
||||
|
||||
def ROT_N(self, instr):
|
||||
return 0, 0
|
||||
|
||||
def ROT_TWO(self, instr):
|
||||
return 0, 0
|
||||
|
||||
def ROT_THREE(self, instr):
|
||||
return 0, 0
|
||||
|
||||
def ROT_FOUR(self, instr):
|
||||
return 0, 0
|
||||
|
||||
def GET_ITER(self, instr):
|
||||
return 1, 1
|
||||
|
||||
def POP_TOP(self, instr):
|
||||
return 1, 0
|
||||
|
||||
def PUSH_NULL(self, instr):
|
||||
return 0, 1
|
||||
|
||||
def NOP(self, instr):
|
||||
return 0, 0
|
||||
|
||||
def EXTENDED_ARG(self, instr):
|
||||
return 0, 0
|
||||
@@ -0,0 +1,167 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
import _collections_abc
|
||||
import _weakrefset
|
||||
import abc
|
||||
import codecs
|
||||
import collections
|
||||
import contextlib
|
||||
import copy
|
||||
import copyreg
|
||||
import dataclasses
|
||||
import enum
|
||||
import functools
|
||||
import genericpath
|
||||
import importlib
|
||||
import inspect
|
||||
import linecache
|
||||
import logging
|
||||
import multiprocessing
|
||||
import operator
|
||||
import os
|
||||
import posixpath
|
||||
import random
|
||||
import re
|
||||
import selectors
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import tokenize
|
||||
import traceback
|
||||
import types
|
||||
import typing
|
||||
import unittest
|
||||
import uuid
|
||||
import warnings
|
||||
import weakref
|
||||
|
||||
import google.protobuf
|
||||
import numpy
|
||||
import setuptools
|
||||
|
||||
import paddle
|
||||
|
||||
NEED_SKIP_THIRD_PARTY_MODULES = {
|
||||
abc,
|
||||
collections,
|
||||
contextlib,
|
||||
copy,
|
||||
copyreg,
|
||||
dataclasses,
|
||||
enum,
|
||||
functools,
|
||||
google.protobuf,
|
||||
importlib,
|
||||
inspect,
|
||||
linecache,
|
||||
logging,
|
||||
multiprocessing,
|
||||
numpy,
|
||||
operator,
|
||||
os,
|
||||
posixpath,
|
||||
random,
|
||||
re,
|
||||
selectors,
|
||||
signal,
|
||||
tempfile,
|
||||
threading,
|
||||
tokenize,
|
||||
traceback,
|
||||
types,
|
||||
typing,
|
||||
unittest,
|
||||
weakref,
|
||||
_collections_abc,
|
||||
_weakrefset,
|
||||
codecs,
|
||||
uuid,
|
||||
setuptools,
|
||||
warnings,
|
||||
genericpath,
|
||||
}
|
||||
|
||||
if sys.version_info < (3, 11):
|
||||
import sre_compile
|
||||
import sre_parse
|
||||
|
||||
NEED_SKIP_THIRD_PARTY_MODULES.add(sre_compile)
|
||||
NEED_SKIP_THIRD_PARTY_MODULES.add(sre_parse)
|
||||
|
||||
if sys.version_info < (3, 12):
|
||||
import distutils
|
||||
|
||||
NEED_SKIP_THIRD_PARTY_MODULES.add(distutils)
|
||||
|
||||
|
||||
def _extend_skip_modules_if_exists(module_name: str):
|
||||
"""Extend skip modules set if the module exists."""
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
NEED_SKIP_THIRD_PARTY_MODULES.add(module)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
# Some modules may not be installed in the environment, so we try to import them.
|
||||
_extend_skip_modules_if_exists("coverage")
|
||||
_extend_skip_modules_if_exists("colorama")
|
||||
|
||||
|
||||
def _strip_init_py(s):
|
||||
return re.sub(r"__init__.py$", "", s)
|
||||
|
||||
|
||||
def _module_dir(m: types.ModuleType):
|
||||
module_file = getattr(m, "__file__", None)
|
||||
if module_file is None:
|
||||
return None
|
||||
return _strip_init_py(module_file)
|
||||
|
||||
|
||||
skip_file_names = {
|
||||
path
|
||||
for path in [_module_dir(m) for m in NEED_SKIP_THIRD_PARTY_MODULES]
|
||||
if path is not None
|
||||
}
|
||||
|
||||
|
||||
sot_path = os.path.dirname(__file__).rpartition(os.sep)[0] + os.sep
|
||||
paddle_path = sys.modules["paddle"].__file__.rpartition(os.sep)[0] + os.sep
|
||||
|
||||
skip_file_names.add(sot_path)
|
||||
skip_file_names.add(paddle_path)
|
||||
skip_file_names.add(
|
||||
"<frozen importlib",
|
||||
)
|
||||
skip_file_names.add("<__array_function__ internals>")
|
||||
|
||||
skip_file_name_re = re.compile(
|
||||
f"^({'|'.join(map(re.escape, skip_file_names))})"
|
||||
)
|
||||
|
||||
no_skip_code = {paddle.nn.Sequential.forward.__code__}
|
||||
|
||||
with_graph_codes = (
|
||||
paddle.nn.Layer.__call__.__code__,
|
||||
paddle.nn.Layer._dygraph_call_func.__code__,
|
||||
)
|
||||
|
||||
|
||||
def setup_skip_files():
|
||||
paddle.framework.core.eval_frame_skip_file_prefix(tuple(skip_file_names))
|
||||
paddle.framework.core.eval_frame_no_skip_codes(tuple(no_skip_code))
|
||||
paddle.framework.core.sot_setup_codes_with_graph(with_graph_codes)
|
||||
@@ -0,0 +1,21 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
from paddle.jit.profiler import (
|
||||
EventGuard as EventGuard,
|
||||
SotProfiler as SotProfiler,
|
||||
event_register as event_register,
|
||||
)
|
||||
|
||||
from .kernel_stats import SotStepProfilerGuard as SotStepProfilerGuard
|
||||
@@ -0,0 +1,236 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
|
||||
from paddle import profiler
|
||||
from paddle.base.core import tracer_event_type_to_string
|
||||
|
||||
from ..utils.envs import ENV_ENABLE_SOT_STEP_PROFILER
|
||||
|
||||
|
||||
class EventVisitor:
|
||||
def visit(self, event_node):
|
||||
event_type_name = tracer_event_type_to_string(event_node.type)
|
||||
visit_method_name = f"visit_{event_type_name}"
|
||||
if not hasattr(self, visit_method_name):
|
||||
self.generic_visit(event_node)
|
||||
return
|
||||
getattr(self, visit_method_name)(event_node)
|
||||
|
||||
def generic_visit(self, event_node):
|
||||
for child in event_node.children_node:
|
||||
self.visit(child)
|
||||
|
||||
def __call__(self, events):
|
||||
for event_node in events.values():
|
||||
self.visit(event_node)
|
||||
|
||||
|
||||
class KernelRunMode(Enum):
|
||||
Dygraph = 1
|
||||
Static = 2
|
||||
|
||||
|
||||
def safe_divide(a, b):
|
||||
# Avoid division by zero
|
||||
return a / b if b != 0 else 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class KernelInfo:
|
||||
name: str
|
||||
run_mode: KernelRunMode
|
||||
duration: float
|
||||
cuda_kernels: list[CudaKernelInfo]
|
||||
|
||||
|
||||
@dataclass
|
||||
class CudaKernelInfo:
|
||||
name: str
|
||||
duration: float
|
||||
|
||||
|
||||
# TODO(SigureMo): Split into multiple files by visitor type
|
||||
class KernelStatsVisitor(EventVisitor):
|
||||
SKIP_KERNEL_NAMES = {"full", "full_int_array", "shadow_feed"}
|
||||
KERNEL_NAME_REGEX = re.compile("(?P<kernel_name>.+) kernel launch")
|
||||
|
||||
def __init__(self):
|
||||
self.kernels = []
|
||||
|
||||
def get_kernel_name(self, event_name):
|
||||
if match_obj := self.KERNEL_NAME_REGEX.match(event_name):
|
||||
return match_obj.group("kernel_name")
|
||||
raise ValueError(f"Unexpected event name: {event_name}")
|
||||
|
||||
def calc_kernel_count(self, mode):
|
||||
return len(
|
||||
[
|
||||
kernel
|
||||
for kernel in self.kernels
|
||||
if (
|
||||
kernel.run_mode == mode
|
||||
and kernel.name not in KernelStatsVisitor.SKIP_KERNEL_NAMES
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def calc_kernel_duration(self, mode):
|
||||
return sum(
|
||||
[
|
||||
kernel.duration
|
||||
for kernel in self.kernels
|
||||
if (
|
||||
kernel.run_mode == mode
|
||||
and kernel.name not in KernelStatsVisitor.SKIP_KERNEL_NAMES
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def find_all_cuda_kernels(self, host_event):
|
||||
# TODO(SigureMo): Find a better way to find all CUDA kernels
|
||||
return [
|
||||
CudaKernelInfo(
|
||||
device_event.name, device_event.end_ns - device_event.start_ns
|
||||
)
|
||||
for runtime_event in host_event.runtime_node
|
||||
for device_event in runtime_event.device_node
|
||||
]
|
||||
|
||||
def visit_DygraphKernelLaunch(self, event_node):
|
||||
duration = event_node.end_ns - event_node.start_ns
|
||||
kernel_name = self.get_kernel_name(event_node.name)
|
||||
all_cuda_kernels = self.find_all_cuda_kernels(event_node)
|
||||
self.kernels.append(
|
||||
KernelInfo(
|
||||
kernel_name, KernelRunMode.Dygraph, duration, all_cuda_kernels
|
||||
)
|
||||
)
|
||||
self.generic_visit(event_node)
|
||||
|
||||
def visit_StaticKernelLaunch(self, event_node):
|
||||
duration = event_node.end_ns - event_node.start_ns
|
||||
kernel_name = self.get_kernel_name(event_node.name)
|
||||
all_cuda_kernels = self.find_all_cuda_kernels(event_node)
|
||||
self.kernels.append(
|
||||
KernelInfo(
|
||||
kernel_name, KernelRunMode.Static, duration, all_cuda_kernels
|
||||
)
|
||||
)
|
||||
self.generic_visit(event_node)
|
||||
|
||||
def summary(self) -> str:
|
||||
static_kernel_duration = self.calc_kernel_duration(KernelRunMode.Static)
|
||||
dygraph_kernel_duration = self.calc_kernel_duration(
|
||||
KernelRunMode.Dygraph
|
||||
)
|
||||
static_kernel_count = self.calc_kernel_count(KernelRunMode.Static)
|
||||
dygraph_kernel_count = self.calc_kernel_count(KernelRunMode.Dygraph)
|
||||
|
||||
percentage_static_kernel_count = safe_divide(
|
||||
static_kernel_count, static_kernel_count + dygraph_kernel_count
|
||||
)
|
||||
percentage_static_kernel_duration = safe_divide(
|
||||
static_kernel_duration,
|
||||
static_kernel_duration + dygraph_kernel_duration,
|
||||
)
|
||||
step_summary = ""
|
||||
step_summary += f"dygraph kernel count: {dygraph_kernel_count}\n"
|
||||
step_summary += f"static kernel count: {static_kernel_count}\n"
|
||||
step_summary += f"percentage static kernel count: {percentage_static_kernel_count:.2%}\n"
|
||||
|
||||
step_summary += f"dygraph kernel duration: {dygraph_kernel_duration / 1000000:.2f} ms\n"
|
||||
step_summary += f"static kernel duration: {static_kernel_duration / 1000000:.2f} ms\n"
|
||||
step_summary += f"percentage static kernel duration: {percentage_static_kernel_duration:.2%}\n"
|
||||
return step_summary
|
||||
|
||||
|
||||
class SotStepProfilerGuard:
|
||||
EXPORT_CHROME_TRACING_PATH = "./sot-chrome-tracing/"
|
||||
STEP_CNT = 0
|
||||
LAST_INFO_SUMMARY = None
|
||||
|
||||
def __init__(self, enable_kernel_stats=True, enable_chrome_tracing=False):
|
||||
self.enable_kernel_stats = enable_kernel_stats
|
||||
self.enable_chrome_tracing = enable_chrome_tracing
|
||||
self.started = False
|
||||
self.record_event = None
|
||||
self.summary = None
|
||||
|
||||
def _kernel_stats(self, prof) -> str:
|
||||
kernel_stats_visitor = KernelStatsVisitor()
|
||||
kernel_stats_visitor(prof.profiler_result.get_data())
|
||||
return kernel_stats_visitor.summary()
|
||||
|
||||
def collect_step_info_summary(self, prof) -> str:
|
||||
summary = ""
|
||||
if self.enable_kernel_stats:
|
||||
summary += self._kernel_stats(prof)
|
||||
if self.enable_chrome_tracing:
|
||||
# If you want to export chrome tracing, you can enable this flag
|
||||
# and view the tracing result in https://ui.perfetto.dev/#!/viewer
|
||||
profiler.export_chrome_tracing(
|
||||
SotStepProfilerGuard.EXPORT_CHROME_TRACING_PATH,
|
||||
f"step_{SotStepProfilerGuard.STEP_CNT:03d}",
|
||||
)(prof)
|
||||
return summary
|
||||
|
||||
def on_trace_ready(self, prof):
|
||||
summary = self.collect_step_info_summary(prof)
|
||||
if SotStepProfilerGuard.STEP_CNT == 0:
|
||||
coldstart_title = f"SOT step profiler info summary (ColdStart, step#{SotStepProfilerGuard.STEP_CNT}):"
|
||||
coldstart_report = f"{coldstart_title}\n{summary}"
|
||||
print(coldstart_report)
|
||||
self.summary = coldstart_report
|
||||
else:
|
||||
warmup_title = f"SOT step profiler info summary (Warmup, step#{SotStepProfilerGuard.STEP_CNT}):"
|
||||
warmup_report = f"{warmup_title}\n{summary}"
|
||||
if SotStepProfilerGuard.LAST_INFO_SUMMARY is None:
|
||||
atexit.register(
|
||||
lambda: print(SotStepProfilerGuard.LAST_INFO_SUMMARY)
|
||||
)
|
||||
SotStepProfilerGuard.LAST_INFO_SUMMARY = warmup_report
|
||||
self.summary = warmup_report
|
||||
|
||||
def start(self):
|
||||
if ENV_ENABLE_SOT_STEP_PROFILER.get():
|
||||
self.profiler = profiler.Profiler(
|
||||
targets=[
|
||||
profiler.ProfilerTarget.CPU,
|
||||
profiler.ProfilerTarget.GPU,
|
||||
],
|
||||
on_trace_ready=self.on_trace_ready,
|
||||
)
|
||||
self.profiler.start()
|
||||
self.started = True
|
||||
|
||||
def stop(self):
|
||||
if self.started:
|
||||
assert self.profiler is not None
|
||||
self.profiler.stop()
|
||||
self.profiler = None # Avoid to hold the profiler instance
|
||||
SotStepProfilerGuard.STEP_CNT += 1
|
||||
|
||||
def __enter__(self):
|
||||
self.start()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.stop()
|
||||
@@ -0,0 +1,68 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
from types import CodeType
|
||||
|
||||
T = TypeVar("T")
|
||||
P = ParamSpec("P")
|
||||
|
||||
NO_BREAKGRAPH_CODES: set[CodeType] = set()
|
||||
NO_FALLBACK_CODES: set[CodeType] = set()
|
||||
|
||||
|
||||
def assert_true(input: bool):
|
||||
assert input
|
||||
|
||||
|
||||
def print(*args, **kwargs):
|
||||
builtins.print("[Dygraph]", *args, **kwargs)
|
||||
|
||||
|
||||
def breakpoint():
|
||||
import paddle
|
||||
|
||||
old = paddle.framework.core.set_eval_frame(None)
|
||||
builtins.breakpoint() # noqa: T100
|
||||
paddle.framework.core.set_eval_frame(old)
|
||||
|
||||
|
||||
def check_no_breakgraph(fn: Callable[P, T]) -> Callable[P, T]:
|
||||
NO_BREAKGRAPH_CODES.add(fn.__code__)
|
||||
return fn
|
||||
|
||||
|
||||
def breakgraph():
|
||||
pass
|
||||
|
||||
|
||||
def check_no_fallback(fn: Callable[P, T]) -> Callable[P, T]:
|
||||
NO_FALLBACK_CODES.add(fn.__code__)
|
||||
return fn
|
||||
|
||||
|
||||
def fallback(recursive=False):
|
||||
pass
|
||||
|
||||
|
||||
def in_sot():
|
||||
return False
|
||||
@@ -0,0 +1,192 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..utils import log
|
||||
from .compile_cache import CompileSIRCache
|
||||
from .statement_ir import (
|
||||
ApiStatement,
|
||||
ASTStatement,
|
||||
CallStatement,
|
||||
LayerStatement,
|
||||
MethodStatement,
|
||||
ParametersHolder,
|
||||
StatementContext,
|
||||
StatementIR,
|
||||
StatementIRFactory,
|
||||
Symbol,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from paddle.static import InputSpec
|
||||
|
||||
|
||||
class StatementIRBuilder:
|
||||
"""
|
||||
A class to build a StatementIR.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
"""
|
||||
Reset the context.
|
||||
"""
|
||||
|
||||
self.statement_factory = StatementIRFactory()
|
||||
self._current_statement_ctxs = []
|
||||
self._current_sir: StatementIR = self.statement_factory.create()
|
||||
|
||||
@property
|
||||
def current_sir(self) -> StatementIR:
|
||||
"""
|
||||
Get the current SIR.
|
||||
"""
|
||||
return self._current_sir
|
||||
|
||||
def replace_current_sir(self, sir: StatementIR):
|
||||
"""
|
||||
Replace the current SIR with a new SIR.
|
||||
"""
|
||||
self._current_sir = sir
|
||||
self.statement_factory.update(sir)
|
||||
|
||||
def call_SIR(self, sirname, inputs, outputs, stacks):
|
||||
"""
|
||||
Call a SIR, which is a subgraph.
|
||||
"""
|
||||
|
||||
stmt = CallStatement(
|
||||
sirname, inputs, outputs, list(self._current_statement_ctxs), stacks
|
||||
)
|
||||
self.current_sir.add_statement(stmt)
|
||||
|
||||
def call_API(self, api, inputs, outputs, stacks):
|
||||
"""
|
||||
Call a paddle api.
|
||||
"""
|
||||
assert callable(api), "call_API must receive a paddle api."
|
||||
stmt = ApiStatement(
|
||||
api, inputs, outputs, list(self._current_statement_ctxs), stacks
|
||||
)
|
||||
self.current_sir.add_statement(stmt)
|
||||
|
||||
def call_METHOD(self, method_name, inputs, outputs, stacks):
|
||||
"""
|
||||
Call a method of a api. The API here can be python or Paddle
|
||||
"""
|
||||
assert isinstance(method_name, str), (
|
||||
"call_METHOD must method api name. string."
|
||||
)
|
||||
assert isinstance(inputs[0][0], Symbol), (
|
||||
"call_METHOD first argument must be Symbol Variable."
|
||||
)
|
||||
stmt = MethodStatement(
|
||||
method_name,
|
||||
inputs,
|
||||
outputs,
|
||||
list(self._current_statement_ctxs),
|
||||
stacks,
|
||||
)
|
||||
self.current_sir.add_statement(stmt)
|
||||
|
||||
def call_LAYER(self, layer, inputs, outputs, stacks):
|
||||
"""
|
||||
Call a layer of a api.
|
||||
"""
|
||||
stmt = LayerStatement(
|
||||
layer, inputs, outputs, list(self._current_statement_ctxs), stacks
|
||||
)
|
||||
self.current_sir.add_statement(stmt)
|
||||
|
||||
def call_AST(self, static_function, inputs, outputs, stacks):
|
||||
stmt = ASTStatement(
|
||||
static_function,
|
||||
inputs,
|
||||
outputs,
|
||||
list(self._current_statement_ctxs),
|
||||
stacks,
|
||||
)
|
||||
self.current_sir.add_statement(stmt)
|
||||
|
||||
def get_sir(self, name: str):
|
||||
"""
|
||||
Get a SIR from statement_factory.
|
||||
|
||||
Args:
|
||||
name (str): the name of SIR.
|
||||
|
||||
Returns:
|
||||
StatementIR: the SIR.
|
||||
"""
|
||||
return self.statement_factory[name]
|
||||
|
||||
@contextmanager
|
||||
def attach_statement_context_guard(self, ctx: StatementContext):
|
||||
"""
|
||||
Attach a statement context to the current SIR.
|
||||
"""
|
||||
self._current_statement_ctxs.append(ctx)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self._current_statement_ctxs.pop()
|
||||
|
||||
def finalize(self, ret_vals):
|
||||
current_sir: StatementIR = self.current_sir
|
||||
current_sir.inputs, current_sir.params = current_sir.analyse_inputs()
|
||||
current_sir.outputs = ret_vals
|
||||
log(2, "start subgraph compile and execution.\n")
|
||||
log(2, current_sir, "\n")
|
||||
return current_sir
|
||||
|
||||
def compile_do_nothing(self) -> Callable[..., Any]:
|
||||
"""
|
||||
Return a dummy function, which will return an empty list.
|
||||
|
||||
Args:
|
||||
ret_vals (list[Symbol]): the return values of the function.
|
||||
"""
|
||||
|
||||
class DummyFunc:
|
||||
def __call__(*args, **kwargs):
|
||||
return []
|
||||
|
||||
def graph_size(self):
|
||||
return 0
|
||||
|
||||
return DummyFunc()
|
||||
|
||||
def compile_fn(
|
||||
self,
|
||||
sir_name: str,
|
||||
parameters_holder: ParametersHolder,
|
||||
input_spec: tuple[InputSpec | None, ...],
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
start compile and return the python function, which must can be to_static without errors.
|
||||
"""
|
||||
static_func = CompileSIRCache()(
|
||||
self, sir_name, parameters_holder, input_spec, **kwargs
|
||||
)
|
||||
|
||||
return static_func
|
||||
@@ -0,0 +1,322 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
from paddle.jit.profiler import EventGuard, event_register
|
||||
|
||||
from ..infer_meta import convert_meta_to_input_spec
|
||||
from ..utils import (
|
||||
ENV_SOT_EXPORT,
|
||||
Cache,
|
||||
InfoCollector,
|
||||
NewSymbolHitRateInfo,
|
||||
Singleton,
|
||||
SIRToCodeMap,
|
||||
StepInfoManager,
|
||||
SubGraphInfo,
|
||||
SubGraphRelationInfo,
|
||||
log_do,
|
||||
)
|
||||
from .export import export
|
||||
from .interpreter import compile_sir
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle.static import InputSpec, Program
|
||||
|
||||
from .builder import StatementIRBuilder
|
||||
from .statement_ir import ParametersHolder
|
||||
|
||||
|
||||
def trace_back_frames():
|
||||
frame = inspect.currentframe()
|
||||
while frame.f_back is not None:
|
||||
frame = frame.f_back
|
||||
code = frame.f_code
|
||||
paddle.framework.core.sot_set_with_graph(code)
|
||||
|
||||
|
||||
def clear_eager_tensor_name(output_tensors):
|
||||
for output_tensor in output_tensors:
|
||||
output_tensor.name = ""
|
||||
|
||||
|
||||
def _is_builtin_op(op):
|
||||
dialect_name, opname = op.name().split(".")
|
||||
return dialect_name == "builtin"
|
||||
|
||||
|
||||
def _is_computation_op(op):
|
||||
return not _is_builtin_op(op) and op.name() not in ["pd_op.data"]
|
||||
|
||||
|
||||
class UniqueIdGenerator:
|
||||
def __init__(self):
|
||||
self._id = 0
|
||||
|
||||
def generate(self):
|
||||
self._id += 1
|
||||
return self._id
|
||||
|
||||
def __call__(self):
|
||||
return self.generate()
|
||||
|
||||
|
||||
class TensorIdAllocator(metaclass=Singleton):
|
||||
TENSOR_ID_ATTR = "__tensor_id__"
|
||||
|
||||
def __init__(self):
|
||||
self._id_generator = UniqueIdGenerator()
|
||||
|
||||
def allocate(self, tensor):
|
||||
if not hasattr(tensor, self.TENSOR_ID_ATTR):
|
||||
setattr(tensor, self.TENSOR_ID_ATTR, self._id_generator())
|
||||
return getattr(tensor, self.TENSOR_ID_ATTR)
|
||||
|
||||
|
||||
class FallbackWrapper:
|
||||
"""
|
||||
Used to store and call static graph methods generated by paddle.jit.to_static
|
||||
"""
|
||||
|
||||
def __init__(self, compiled_fn, SIR, is_training: bool):
|
||||
self.compiled_fn = compiled_fn
|
||||
self.partial_program = None
|
||||
self.concrete_program = None
|
||||
self.SIR = SIR # for debug
|
||||
self.is_training = is_training
|
||||
self.exported = False
|
||||
self.is_first_call = True
|
||||
|
||||
def graph_size(self):
|
||||
if self.partial_program is None:
|
||||
input_spec = convert_meta_to_input_spec(
|
||||
tuple(
|
||||
self.SIR.symbol_meta_map[symbol]
|
||||
for symbol in self.SIR.inputs
|
||||
)
|
||||
)
|
||||
(
|
||||
self.concrete_program,
|
||||
self.partial_program,
|
||||
) = self.compiled_fn.get_concrete_program(input_spec)
|
||||
self.partial_program.training = self.is_training
|
||||
global_block_ops = self.concrete_program.main_program.global_block().ops
|
||||
non_builtin_ops = list(filter(_is_computation_op, global_block_ops))
|
||||
return len(non_builtin_ops)
|
||||
|
||||
def collect_new_symbol_hit_rate(self, inputs, outputs):
|
||||
if not InfoCollector().need_collect(NewSymbolHitRateInfo):
|
||||
return
|
||||
input_tensor_ids = []
|
||||
output_tensor_ids = []
|
||||
assert len(inputs) == 1
|
||||
assert isinstance(inputs[0], tuple)
|
||||
for i, arg in enumerate(inputs[0]):
|
||||
assert isinstance(arg, paddle.Tensor), f"Expect Tensor, got {arg}"
|
||||
tensor_id = TensorIdAllocator().allocate(arg)
|
||||
input_tensor_ids.append(tensor_id)
|
||||
|
||||
for i, out in enumerate(outputs):
|
||||
assert isinstance(out, paddle.Tensor)
|
||||
tensor_id = TensorIdAllocator().allocate(out)
|
||||
output_tensor_ids.append(tensor_id)
|
||||
|
||||
InfoCollector().attach(
|
||||
NewSymbolHitRateInfo, input_tensor_ids, output_tensor_ids
|
||||
)
|
||||
|
||||
def collect_subgraph_relation(self, inputs, outputs, partial_program_layer):
|
||||
if not InfoCollector().need_collect(SubGraphRelationInfo):
|
||||
return
|
||||
input_shape_infos = []
|
||||
output_shape_infos = []
|
||||
forward_input_values = partial_program_layer.program.program_attr['fx']
|
||||
forward_output_values = partial_program_layer.program.program_attr['fo']
|
||||
assert len(inputs) == 1
|
||||
assert isinstance(inputs[0], tuple)
|
||||
assert len(inputs[0]) == len(forward_input_values)
|
||||
assert len(outputs) == len(forward_output_values)
|
||||
for i, arg in enumerate(inputs[0]):
|
||||
assert isinstance(arg, paddle.Tensor), f"Expect Tensor, got {arg}"
|
||||
tensor_id = TensorIdAllocator().allocate(arg)
|
||||
input_ir_shape = forward_input_values[i].shape
|
||||
input_real_shape = arg.shape
|
||||
input_shape_info = SubGraphRelationInfo.ConcreteShapeInfo(
|
||||
tensor_id, input_ir_shape, input_real_shape
|
||||
)
|
||||
input_shape_infos.append(input_shape_info)
|
||||
|
||||
for i, out in enumerate(outputs):
|
||||
assert isinstance(out, paddle.Tensor)
|
||||
tensor_id = TensorIdAllocator().allocate(out)
|
||||
output_ir_shape = forward_output_values[
|
||||
partial_program_layer._outputs.quick_index_map[i]
|
||||
].shape
|
||||
output_real_shape = out.shape
|
||||
output_shape_info = SubGraphRelationInfo.ConcreteShapeInfo(
|
||||
tensor_id, output_ir_shape, output_real_shape
|
||||
)
|
||||
output_shape_infos.append(output_shape_info)
|
||||
|
||||
InfoCollector().attach(
|
||||
SubGraphRelationInfo,
|
||||
self.SIR.name,
|
||||
input_shape_infos,
|
||||
output_shape_infos,
|
||||
self.is_first_call,
|
||||
self.graph_size(),
|
||||
)
|
||||
|
||||
def collect_subgraph_info(self, program: Program):
|
||||
if not InfoCollector().need_collect(SubGraphInfo):
|
||||
return
|
||||
|
||||
InfoCollector().attach(
|
||||
SubGraphInfo,
|
||||
str(program),
|
||||
self.graph_size(),
|
||||
self.SIR.name,
|
||||
)
|
||||
|
||||
def update_compile_time_info(self, SIR, partial_program_layer):
|
||||
if not self.is_first_call:
|
||||
return
|
||||
from ..opcode_translator.executor.executor_cache import (
|
||||
OpcodeExecutorCache,
|
||||
)
|
||||
|
||||
code = SIRToCodeMap().get(SIR)
|
||||
assert code is not None, f"Cannot find code for SIR: {SIR}"
|
||||
|
||||
OpcodeExecutorCache().compile_time_stats.setdefault(code, 0)
|
||||
OpcodeExecutorCache().compile_time_stats[code] += (
|
||||
partial_program_layer._compile_time_counter.get_total_time()
|
||||
)
|
||||
|
||||
@event_register(
|
||||
lambda self, *args, **kwargs: f"FallbackWrapper: {self.SIR.name}"
|
||||
)
|
||||
def __call__(self, *args, **kwargs):
|
||||
if StepInfoManager().need_back_trace:
|
||||
trace_back_frames()
|
||||
|
||||
log_do(
|
||||
2,
|
||||
lambda: print("[FallbackWrapper] start run SIR: \n", self.SIR),
|
||||
)
|
||||
log_do(
|
||||
4,
|
||||
lambda: print(
|
||||
self.compiled_fn.get_concrete_program(*args, **kwargs)[
|
||||
1
|
||||
].train_program
|
||||
),
|
||||
)
|
||||
if self.partial_program is None:
|
||||
with EventGuard("FallbackWrapper: get_concrete_program"):
|
||||
(
|
||||
self.concrete_program,
|
||||
self.partial_program,
|
||||
) = self.compiled_fn.get_concrete_program(*args, **kwargs)
|
||||
self.partial_program.training = self.is_training
|
||||
outputs = self.partial_program.sot_call(*args, **kwargs)
|
||||
|
||||
clear_eager_tensor_name(outputs)
|
||||
log_do(
|
||||
4,
|
||||
lambda: print("[CompileCache] run sir forward success."),
|
||||
)
|
||||
self.collect_new_symbol_hit_rate(args, outputs)
|
||||
self.collect_subgraph_relation(args, outputs, self.partial_program)
|
||||
self.collect_subgraph_info(self.concrete_program.main_program)
|
||||
self.update_compile_time_info(self.SIR, self.partial_program)
|
||||
if ENV_SOT_EXPORT.get() != "" and not self.exported:
|
||||
export(self.SIR, ENV_SOT_EXPORT.get())
|
||||
self.exported = True
|
||||
|
||||
self.is_first_call = False
|
||||
return outputs
|
||||
|
||||
|
||||
class CompileSIRCache(Cache, metaclass=Singleton):
|
||||
"""
|
||||
Cache the compiled function of SIR
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(weak=False)
|
||||
|
||||
def key_fn(
|
||||
self,
|
||||
builder: StatementIRBuilder,
|
||||
sir_name: str,
|
||||
parameters_holder: ParametersHolder,
|
||||
input_spec: tuple[InputSpec | None, ...],
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
generate a hash key for a SIR
|
||||
|
||||
Args:
|
||||
context: The context to compile
|
||||
sir_name: The name of the sir to compile
|
||||
build_strategy: The build strategy to compile
|
||||
|
||||
Returns:
|
||||
The hash key of the SIR
|
||||
"""
|
||||
sir = builder.get_sir(sir_name)
|
||||
# NOTE(dev): Is str(sir) a heavy operation ?
|
||||
hash_key = hash(
|
||||
(str(sir), *input_spec, id(parameters_holder), kwargs['training'])
|
||||
)
|
||||
return hash_key
|
||||
|
||||
def value_fn(
|
||||
self,
|
||||
builder: StatementIRBuilder,
|
||||
sir_name: str,
|
||||
parameters_holder: ParametersHolder,
|
||||
input_spec: tuple[InputSpec | None, ...],
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Generate static graph function
|
||||
|
||||
Args:
|
||||
context: The context to compile
|
||||
sir_name: The name of the sir to compile
|
||||
build_strategy: The build strategy to compile
|
||||
|
||||
Returns:
|
||||
The static graph function
|
||||
"""
|
||||
build_strategy = kwargs.get("build_strategy", None)
|
||||
backend = kwargs.get("backend", None)
|
||||
return FallbackWrapper(
|
||||
paddle.jit.to_static(
|
||||
compile_sir(builder, sir_name, parameters_holder),
|
||||
input_spec=[input_spec],
|
||||
build_strategy=build_strategy,
|
||||
backend=backend,
|
||||
full_graph=True,
|
||||
),
|
||||
builder.get_sir(sir_name),
|
||||
is_training=kwargs['training'],
|
||||
)
|
||||
@@ -0,0 +1,387 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
import os
|
||||
from itertools import chain
|
||||
|
||||
import paddle
|
||||
from paddle.utils import flatten
|
||||
|
||||
from ..utils import ConstTypes, ExportError, NameGenerator, get_api_fullname
|
||||
from .statement_ir import Symbol
|
||||
|
||||
|
||||
class PyStatement:
|
||||
tab = " " * 4
|
||||
|
||||
def __init__(self, *lines):
|
||||
self.sub_statement = []
|
||||
self.lines = lines
|
||||
|
||||
def get_lines(self, prefix=""):
|
||||
lines = [prefix + line for line in self.lines]
|
||||
for statement in self.sub_statement:
|
||||
lines.extend(statement.get_lines(self.tab + prefix))
|
||||
return lines
|
||||
|
||||
def add_sub(self, *lines):
|
||||
sub = PyStatement(*lines)
|
||||
self.sub_statement.append(sub)
|
||||
return sub
|
||||
|
||||
def __str__(self):
|
||||
return "\n".join(self.get_lines())
|
||||
|
||||
|
||||
class NameGener:
|
||||
def __init__(self, SIR):
|
||||
self.SIR = SIR
|
||||
self.name_map = {}
|
||||
self.param_name_generator = NameGenerator("self.parameter_")
|
||||
self.non_param_name_generator = NameGenerator("var_")
|
||||
|
||||
def __call__(self, var):
|
||||
return self.get_str(var)
|
||||
|
||||
def get_str(self, var):
|
||||
if isinstance(var, list):
|
||||
return self.get_list_str(var)
|
||||
elif isinstance(var, tuple):
|
||||
return self.get_tuple_str(var)
|
||||
elif isinstance(var, dict):
|
||||
return self.get_dict_str(var)
|
||||
elif isinstance(var, set):
|
||||
return self.get_set_str(var)
|
||||
else:
|
||||
return self.get_obj_str(var)
|
||||
|
||||
def get_list_str(self, list_):
|
||||
return "[{}]".format(", ".join(self.get_str(var) for var in list_))
|
||||
|
||||
def get_tuple_str(self, tuple_):
|
||||
return "({},)".format(", ".join(self.get_str(var) for var in tuple_))
|
||||
|
||||
def get_dict_str(self, dict_):
|
||||
return "{{{},}}".format(
|
||||
", ".join(
|
||||
f"{self.get_str(k)}: {self.get_str(v)}"
|
||||
for k, v in dict_.items()
|
||||
)
|
||||
)
|
||||
|
||||
def get_set_str(self, set_):
|
||||
return "{{{},}}".format(", ".join(self.get_str(var) for var in set_))
|
||||
|
||||
def get_obj_str(self, var):
|
||||
if isinstance(var, Symbol):
|
||||
if var not in self.name_map:
|
||||
self.register_symbol(var)
|
||||
return self.name_map[var]
|
||||
|
||||
elif isinstance(var, str):
|
||||
return f"'{var}'"
|
||||
else:
|
||||
return str(var)
|
||||
|
||||
def register_symbol(self, symbol):
|
||||
if symbol in self.SIR.param_symbol:
|
||||
name = self.param_name_generator.next()
|
||||
else:
|
||||
name = self.non_param_name_generator.next()
|
||||
self.name_map[symbol] = name
|
||||
|
||||
|
||||
class PyFileGen:
|
||||
def __init__(self, SIR):
|
||||
self.SIR = SIR
|
||||
self.roots = []
|
||||
|
||||
self.name_gener = NameGener(self.SIR)
|
||||
|
||||
self.SIR_sig = "||".join(
|
||||
f"{stmt.type}:{stmt.name}" for stmt in SIR.statements
|
||||
)
|
||||
|
||||
def new_root(self, *args):
|
||||
stmt = PyStatement(*args)
|
||||
self.roots.append(stmt)
|
||||
return stmt
|
||||
|
||||
def roots_to_string(self):
|
||||
lines = []
|
||||
for root in self.roots:
|
||||
lines.extend(root.get_lines())
|
||||
return "\n".join(lines)
|
||||
|
||||
def gen_py_codes(self):
|
||||
self.check_exportable()
|
||||
self.create_header()
|
||||
self.new_root("\n")
|
||||
self.create_layer()
|
||||
self.new_root("\n")
|
||||
self.create_inputs()
|
||||
self.new_root("\n")
|
||||
self.create_test()
|
||||
self.new_root("\n")
|
||||
self.create_tail()
|
||||
return self.roots_to_string()
|
||||
|
||||
def is_exportable_type(self, value):
|
||||
if (
|
||||
isinstance(value, (ConstTypes, Symbol, paddle.dtype))
|
||||
or value is Ellipsis # NOINT
|
||||
):
|
||||
return True
|
||||
if isinstance(value, slice):
|
||||
return (
|
||||
self.is_exportable_type(value.start)
|
||||
and self.is_exportable_type(value.stop)
|
||||
and self.is_exportable_type(value.step)
|
||||
)
|
||||
return False
|
||||
|
||||
def check_exportable(self):
|
||||
for stmt in self.SIR.statements:
|
||||
for inp in flatten(stmt.inputs):
|
||||
if not self.is_exportable_type(inp):
|
||||
raise ExportError(
|
||||
f"Not support create python file with input: {inp}"
|
||||
)
|
||||
|
||||
def create_header(self):
|
||||
self.new_root(
|
||||
f"# {self.SIR_sig}",
|
||||
"import paddle",
|
||||
"import unittest",
|
||||
"import numpy as np",
|
||||
)
|
||||
|
||||
def create_layer(self):
|
||||
layer_class = self.new_root("class LayerCase(paddle.nn.Layer):")
|
||||
|
||||
init_fn = layer_class.add_sub("def __init__(self):")
|
||||
init_fn.add_sub("super().__init__()")
|
||||
|
||||
for param in self.SIR.param_symbol:
|
||||
meta = self.SIR.symbol_meta_map[param].unwrap_unsafe()
|
||||
init_fn.add_sub(
|
||||
f"{self.name_gener(param)} = self.create_parameter(",
|
||||
f" shape={meta.shape},",
|
||||
f" dtype={meta.dtype},",
|
||||
")",
|
||||
)
|
||||
|
||||
for stmt in self.SIR.statements:
|
||||
if stmt.type == "layer":
|
||||
layer = stmt.layer()
|
||||
init_fn.add_sub(self.init_sub_layer(layer))
|
||||
|
||||
forward_definition = ["def forward(", " self,"]
|
||||
|
||||
for inp in self.SIR.inputs:
|
||||
if inp in self.SIR.non_param_symbol:
|
||||
meta = self.SIR.symbol_meta_map[inp]
|
||||
forward_definition.append(
|
||||
f" {self.name_gener(inp)}, # {meta}"
|
||||
)
|
||||
forward_definition.append("):")
|
||||
|
||||
forward_fn = layer_class.add_sub(*forward_definition)
|
||||
|
||||
for stmt in self.SIR.statements:
|
||||
forward_fn.add_sub(*self.create_stmt_line(stmt))
|
||||
|
||||
forward_fn.add_sub(
|
||||
"return {}".format(
|
||||
", ".join(self.name_gener(out) for out in self.SIR.outputs)
|
||||
)
|
||||
)
|
||||
|
||||
def create_inputs(self):
|
||||
create_paddle_inputs = self.new_root("def create_paddle_inputs():")
|
||||
self.new_root("\n")
|
||||
create_numpy_inputs = self.new_root("def create_numpy_inputs():")
|
||||
|
||||
paddle_inputs = ["inputs = ("]
|
||||
numpy_inputs = ["inputs = ("]
|
||||
|
||||
for inp in self.SIR.inputs:
|
||||
if inp in self.SIR.non_param_symbol:
|
||||
meta = self.SIR.symbol_meta_map[inp.name].unwrap_unsafe()
|
||||
shape_str = "[1]" if len(meta.shape) == 0 else str(meta.shape)
|
||||
if meta.dtype in (
|
||||
paddle.int8,
|
||||
paddle.int16,
|
||||
paddle.int32,
|
||||
paddle.int64,
|
||||
):
|
||||
paddle_inputs.append(
|
||||
f" paddle.randint(low=0, high=10, shape={shape_str}, dtype={meta.dtype}),"
|
||||
)
|
||||
numpy_inputs.append(
|
||||
" np.random.randint(low=0, high=10, size={}, dtype='{}'),".format(
|
||||
shape_str, str(meta.dtype).replace('paddle.', '')
|
||||
)
|
||||
)
|
||||
elif meta.dtype is paddle.bool:
|
||||
paddle_inputs.append(
|
||||
f" paddle.randint(low=0, high=2, shape={shape_str}, dtype=paddle.int32).cast(paddle.bool),"
|
||||
)
|
||||
numpy_inputs.append(
|
||||
f" np.random.randint(low=0, high=2, size={shape_str}, dtype='int').astype('bool'),"
|
||||
)
|
||||
else:
|
||||
paddle_inputs.append(
|
||||
f" paddle.rand(shape={shape_str}, dtype={meta.dtype}),"
|
||||
)
|
||||
numpy_inputs.append(
|
||||
" np.random.random(size={}).astype('{}'),".format(
|
||||
shape_str, str(meta.dtype).replace('paddle.', '')
|
||||
)
|
||||
)
|
||||
|
||||
paddle_inputs.append(")")
|
||||
paddle_inputs.append("return inputs")
|
||||
numpy_inputs.append(")")
|
||||
numpy_inputs.append("return inputs")
|
||||
|
||||
create_paddle_inputs.add_sub(*paddle_inputs)
|
||||
create_numpy_inputs.add_sub(*numpy_inputs)
|
||||
|
||||
def create_test(self):
|
||||
test_class = self.new_root("class TestLayer(unittest.TestCase):")
|
||||
|
||||
setup = test_class.add_sub("def setUp(self):")
|
||||
setup.add_sub("self.inputs = create_paddle_inputs()")
|
||||
setup.add_sub("self.net = LayerCase()")
|
||||
|
||||
train = test_class.add_sub(
|
||||
"def train(self, net, to_static, with_prim=False, with_cinn=False):"
|
||||
)
|
||||
train.add_sub(
|
||||
"if to_static:",
|
||||
" paddle.base.core._set_prim_all_enabled(with_prim)",
|
||||
" if with_cinn:",
|
||||
' assert with_prim, "with_cinn=True but with_prim=False is unsupported"',
|
||||
' net = paddle.jit.to_static(net, backend="CINN", full_graph=True)',
|
||||
" else:",
|
||||
" net = paddle.jit.to_static(net, backend=None, full_graph=True)",
|
||||
"paddle.seed(123)",
|
||||
"outs = net(*self.inputs)",
|
||||
"return outs",
|
||||
)
|
||||
|
||||
test_ast_cinn_static = test_class.add_sub(
|
||||
"def test_ast_prim_cinn(self):"
|
||||
)
|
||||
test_ast_cinn_static.add_sub(
|
||||
"st_out = self.train(self.net, to_static=True)",
|
||||
"cinn_out = self.train(self.net, to_static=True, with_prim=True, with_cinn=True)",
|
||||
"for st, cinn in zip(paddle.utils.flatten(st_out), paddle.utils.flatten(cinn_out)):",
|
||||
" np.testing.assert_allclose(st.numpy(), cinn.numpy(), atol=1e-8)",
|
||||
)
|
||||
|
||||
def create_tail(self):
|
||||
self.new_root(
|
||||
"if __name__ == '__main__':",
|
||||
" unittest.main()",
|
||||
)
|
||||
|
||||
def init_sub_layer(self, layer, layer_name):
|
||||
# TODO @wuzhanfei need more efficient way to create a sub layer
|
||||
# now, we just close call_Layer behavior
|
||||
raise ExportError("Not support create sub layer now.")
|
||||
|
||||
def create_input_string(self, args, kwargs):
|
||||
return ", ".join(
|
||||
chain(
|
||||
(self.name_gener(arg) for arg in args),
|
||||
(f"{k}={self.name_gener(v)}" for k, v in kwargs.items()),
|
||||
)
|
||||
)
|
||||
|
||||
def create_unpack_output_string(self, outputs):
|
||||
path = ["out"]
|
||||
result = []
|
||||
|
||||
def search(outputs, path, result):
|
||||
if isinstance(outputs, (list, tuple)):
|
||||
search_sequence(outputs, path, result)
|
||||
elif isinstance(outputs, dict):
|
||||
search_dict(outputs, path, result)
|
||||
elif isinstance(outputs, Symbol):
|
||||
result.append(self.name_gener(outputs) + " = " + "".join(path))
|
||||
|
||||
def search_sequence(outputs, path, result):
|
||||
for idx, out in enumerate(outputs):
|
||||
path.append(f"[{idx}]")
|
||||
search(out, path, result)
|
||||
path.pop()
|
||||
|
||||
def search_dict(outputs, path, result):
|
||||
for k, out in outputs.items():
|
||||
path.append(f"[{k}]")
|
||||
search(out, path, result)
|
||||
path.pop()
|
||||
|
||||
search(outputs, path, result)
|
||||
return result
|
||||
|
||||
def create_stmt_line(self, stmt):
|
||||
return getattr(self, "create_" + stmt.type + "_stmt")(stmt)
|
||||
|
||||
def create_api_stmt(self, stmt):
|
||||
args, kwargs = stmt.inputs
|
||||
input_str = self.create_input_string(args, kwargs)
|
||||
api = stmt.api
|
||||
api_str = get_api_fullname(api)
|
||||
if api_str is None:
|
||||
raise ExportError(f"Can not find module of {api}")
|
||||
if isinstance(stmt.outputs, Symbol):
|
||||
return [f"{self.name_gener(stmt.outputs)} = {api_str}({input_str})"]
|
||||
else:
|
||||
compute_code = f"out = {api_str}({input_str})"
|
||||
unpack_codes = self.create_unpack_output_string(stmt.outputs)
|
||||
return [compute_code, *unpack_codes]
|
||||
|
||||
def create_method_stmt(self, stmt):
|
||||
args, kwargs = stmt.inputs
|
||||
input_str = self.create_input_string(args[1:], kwargs)
|
||||
method_str = self.name_gener(args[0]) + "." + stmt.method
|
||||
if isinstance(stmt.outputs, Symbol):
|
||||
return [
|
||||
f"{self.name_gener(stmt.outputs)} = {method_str}({input_str})"
|
||||
]
|
||||
else:
|
||||
compute_code = f"out = {method_str}({input_str})"
|
||||
unpack_codes = self.create_unpack_output_string(stmt.outputs)
|
||||
return [compute_code, *unpack_codes]
|
||||
|
||||
|
||||
def export(SIR, path):
|
||||
try:
|
||||
pygen = PyFileGen(SIR)
|
||||
string = pygen.gen_py_codes()
|
||||
except ExportError as e:
|
||||
print(f"[SOT] Export {SIR.name} Failed:", e)
|
||||
return
|
||||
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
|
||||
with open(os.path.join(path, f"{SIR.name}.py"), "w") as f:
|
||||
f.write(string)
|
||||
print(
|
||||
f"[SOT] Export {SIR.name} Success with size {len(SIR.statements)}"
|
||||
)
|
||||
@@ -0,0 +1,220 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
from paddle.jit.dy2static.utils import compose_guards
|
||||
from paddle.utils import to_sequence
|
||||
|
||||
from ..utils import (
|
||||
InnerError,
|
||||
log_do,
|
||||
map_if,
|
||||
map_if_extend,
|
||||
)
|
||||
from .statement_ir import (
|
||||
ParametersHolder,
|
||||
StatementContext,
|
||||
StatementContextRegistry,
|
||||
Symbol,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .builder import StatementIRBuilder
|
||||
from .statement_ir import Statement, StatementIR
|
||||
|
||||
|
||||
def replace_symbol(
|
||||
values: list[Symbol] | list[object], state: dict[str, Symbol]
|
||||
):
|
||||
"""
|
||||
Replaces Symbol objects with their corresponding values.
|
||||
|
||||
Args:
|
||||
values: A list of values that may contain Symbol objects.
|
||||
state: A dict mapping Symbol names to their corresponding values.
|
||||
|
||||
Returns:
|
||||
A new list with Symbol objects replaced by their corresponding values in the state dict.
|
||||
"""
|
||||
# deal with list / map etc.
|
||||
values = map_if_extend(
|
||||
values,
|
||||
pred=lambda x: isinstance(x, Symbol),
|
||||
true_fn=lambda x: state[x.name],
|
||||
false_fn=lambda x: x,
|
||||
)
|
||||
return values
|
||||
|
||||
|
||||
def _append_opstack_between(start, end, stack):
|
||||
# The range is [start, end)
|
||||
for op in for_each_ops_between(start, end):
|
||||
op.callstack = stack
|
||||
|
||||
|
||||
def for_each_ops_between(start, end):
|
||||
# [start, end)
|
||||
program = paddle.static.default_main_program()
|
||||
ops = program.global_block().ops[start:end]
|
||||
yield from ops
|
||||
|
||||
|
||||
def opnum_in_program():
|
||||
program = paddle.static.default_main_program()
|
||||
return len(program.global_block().ops)
|
||||
|
||||
|
||||
class Interpreter:
|
||||
"""
|
||||
Interpreter is used to interpret and execute SIR.
|
||||
"""
|
||||
|
||||
def __init__(self, builder: StatementIRBuilder):
|
||||
self._builder = builder
|
||||
|
||||
def get_sir(self, name: str) -> StatementIR:
|
||||
"""
|
||||
Returns the StatementIR object by given name.
|
||||
|
||||
Args:
|
||||
name: The name of the StatementIR.
|
||||
|
||||
Returns:
|
||||
The StatementIR object with the given name.
|
||||
"""
|
||||
return self._builder.get_sir(name)
|
||||
|
||||
def run_sir(self, name: str, state: dict[str, Symbol]):
|
||||
"""
|
||||
Runs the StatementIR with the given name using the provided state.
|
||||
|
||||
Args:
|
||||
name: The name of the given StatementIR to run.
|
||||
state: A dict mapping Symbol names to their corresponding values.
|
||||
|
||||
Returns:
|
||||
A list of the Symbol of the StatementIR after execution.
|
||||
"""
|
||||
|
||||
def _set(v, s):
|
||||
state[s.name] = v
|
||||
|
||||
SIR = self.get_sir(name)
|
||||
for stmt in SIR.statements:
|
||||
stmt: Statement
|
||||
before_stmt_opnum = opnum_in_program()
|
||||
inputs = replace_symbol(stmt.inputs, state)
|
||||
|
||||
with create_context_guard(stmt.contexts)():
|
||||
outs = getattr(self, stmt.type)(stmt, inputs)
|
||||
|
||||
if len(to_sequence(outs)) != len(to_sequence(stmt.outputs)):
|
||||
raise InnerError("Number output mismatch, some error happen.")
|
||||
|
||||
log_do(
|
||||
3,
|
||||
lambda: _append_opstack_between(
|
||||
before_stmt_opnum, opnum_in_program() + 1, stmt.stmt_stack
|
||||
),
|
||||
)
|
||||
|
||||
map_if(
|
||||
outs,
|
||||
stmt.outputs,
|
||||
pred=lambda v, s: isinstance(s, Symbol),
|
||||
true_fn=lambda v, s: _set(v, s),
|
||||
false_fn=lambda v, s: None,
|
||||
)
|
||||
# fetch outputs
|
||||
return replace_symbol(SIR.outputs, state)
|
||||
|
||||
def api(self, stmt, inputs):
|
||||
args, kwargs = inputs
|
||||
return stmt.api(*args, **kwargs)
|
||||
|
||||
def method(self, stmt, inputs):
|
||||
args, kwargs = inputs
|
||||
var = args[0]
|
||||
return getattr(var, stmt.method)(*args[1:], **kwargs)
|
||||
|
||||
def layer(self, stmt, inputs):
|
||||
args, kwargs = inputs
|
||||
layer = stmt.layer()
|
||||
assert layer is not None, "SIR bound layer is None."
|
||||
return layer(*args, **kwargs)
|
||||
|
||||
def AST(self, stmt, inputs):
|
||||
args, kwargs = inputs
|
||||
return stmt.converted_func(*args, **kwargs)
|
||||
|
||||
|
||||
def compile_sir(
|
||||
builder: StatementIRBuilder, name: str, parameters_holder: ParametersHolder
|
||||
):
|
||||
"""
|
||||
Compile a SIR to a new function
|
||||
|
||||
Args:
|
||||
context: The context to compile
|
||||
name: The name of the sir to compile
|
||||
|
||||
"""
|
||||
|
||||
@paddle.jit.not_to_static
|
||||
def wrapper(args):
|
||||
"""
|
||||
This function will be decorated by paddle.to_static.
|
||||
so the args is variables, not eager tensors.
|
||||
"""
|
||||
interpreter = Interpreter(builder)
|
||||
SIR = interpreter.get_sir(name)
|
||||
state = prepare_state(SIR, args, parameters_holder)
|
||||
return interpreter.run_sir(name, state)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def prepare_state(
|
||||
SIR: StatementIR, inputs, parameters_holder: ParametersHolder
|
||||
):
|
||||
state = {}
|
||||
# bind inputs
|
||||
assert len(SIR.inputs) == len(inputs), "Inputs length mismatch."
|
||||
for sir_inp, inp in zip(SIR.inputs, inputs):
|
||||
state[sir_inp.name] = inp
|
||||
|
||||
for sir_param in SIR.params:
|
||||
state[sir_param.name] = paddle.base.dygraph.base._convert_into_variable(
|
||||
parameters_holder.get(sir_param.name)
|
||||
)
|
||||
|
||||
return state
|
||||
|
||||
|
||||
def create_context_guard(contexts: list[StatementContext]):
|
||||
guards = list(
|
||||
map(
|
||||
lambda ctx: (
|
||||
lambda: StatementContextRegistry.get_context_guard(type(ctx))(
|
||||
ctx
|
||||
)
|
||||
),
|
||||
contexts,
|
||||
)
|
||||
)
|
||||
return compose_guards(*guards)
|
||||
@@ -0,0 +1,419 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import weakref
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
from weakref import WeakValueDictionary
|
||||
|
||||
import paddle
|
||||
from paddle.jit.dy2static.utils import parameters_persistent_mode_is_enabled
|
||||
from paddle.jit.utils import OrderedSet
|
||||
from paddle.utils import flatten, map_structure
|
||||
|
||||
from ..utils import (
|
||||
InnerError,
|
||||
NameGenerator,
|
||||
Singleton,
|
||||
flatten_extend,
|
||||
get_api_fullname,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
from contextlib import AbstractContextManager
|
||||
|
||||
|
||||
_StatementContextT = TypeVar("_StatementContextT", bound="StatementContext")
|
||||
|
||||
|
||||
class ParametersHolder:
|
||||
def __init__(self):
|
||||
self._params = WeakValueDictionary[
|
||||
str, paddle.base.framework.EagerParamBase
|
||||
]()
|
||||
|
||||
def set(self, name, param):
|
||||
self._params[name] = param
|
||||
|
||||
def get(self, name):
|
||||
if (param := self._params.get(name)) is None:
|
||||
raise InnerError(
|
||||
f"Parameter '{name}' not found in ParametersHolder."
|
||||
)
|
||||
return param
|
||||
|
||||
def copy(self):
|
||||
new_holder = ParametersHolder()
|
||||
new_holder._params = self._params.copy()
|
||||
return new_holder
|
||||
|
||||
|
||||
class Reference: # to unify weak_ref and strong_ref
|
||||
def __init__(self, value, is_weak):
|
||||
self.is_weak = is_weak
|
||||
if is_weak is True:
|
||||
self.ref = weakref.ref(value)
|
||||
else:
|
||||
self.ref = value
|
||||
|
||||
def __call__(self):
|
||||
if self.is_weak is True:
|
||||
return self.ref()
|
||||
else:
|
||||
return self.ref
|
||||
|
||||
|
||||
class Symbol:
|
||||
"""
|
||||
Symbol is used to distinguish a string and a `math variable`.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
def __str__(self):
|
||||
return self.name
|
||||
|
||||
def __repr__(self):
|
||||
return str(self)
|
||||
|
||||
def __eq__(self, other):
|
||||
if isinstance(other, str):
|
||||
return self.name == other
|
||||
return self.name == other.name
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self.name)
|
||||
|
||||
def __deepcopy__(self, memo=None):
|
||||
return Symbol(self.name)
|
||||
|
||||
|
||||
class StatementContext: ...
|
||||
|
||||
|
||||
class StatementContextRegistry:
|
||||
_ctx_map: dict[
|
||||
type[Any],
|
||||
Callable[[Any], AbstractContextManager[None]],
|
||||
] = {}
|
||||
|
||||
@classmethod
|
||||
def register_context_guard(
|
||||
cls,
|
||||
ctx_cls: type[_StatementContextT],
|
||||
handler: Callable[[_StatementContextT], AbstractContextManager[None]],
|
||||
):
|
||||
"""
|
||||
Register a context handler for the given context.
|
||||
"""
|
||||
if ctx_cls in cls._ctx_map:
|
||||
raise ValueError(f"Context {ctx_cls} is already registered.")
|
||||
cls._ctx_map[ctx_cls] = handler
|
||||
|
||||
@classmethod
|
||||
def register_context(
|
||||
cls,
|
||||
handler: Callable[[_StatementContextT], AbstractContextManager[None]],
|
||||
):
|
||||
def decorator(ctx_cls: type[_StatementContextT]):
|
||||
cls.register_context_guard(ctx_cls, handler)
|
||||
return ctx_cls
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def get_context_guard(
|
||||
cls,
|
||||
ctx_cls: type[_StatementContextT],
|
||||
) -> Callable[[_StatementContextT], AbstractContextManager[None]]:
|
||||
"""
|
||||
Get the context handler for the given context.
|
||||
"""
|
||||
if ctx_cls not in cls._ctx_map:
|
||||
raise ValueError(f"Context {ctx_cls} is not registered.")
|
||||
return cls._ctx_map[ctx_cls]
|
||||
|
||||
|
||||
class Statement:
|
||||
"""
|
||||
Statement is used to represent a sentence of code for building the neural network model,
|
||||
which has four types: "call", "api", "method", and "layer".
|
||||
|
||||
Note:
|
||||
Statement temporarily does not support control flow.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
type: str,
|
||||
name: str,
|
||||
inputs: list[Symbol],
|
||||
outputs: list[Symbol],
|
||||
contexts: list[StatementContext],
|
||||
stacks: list[str],
|
||||
):
|
||||
assert type in ["call", "api", "method", "layer", "AST"]
|
||||
self.name = name
|
||||
self.inputs = inputs # (list of Symbols, dict of Symbols)
|
||||
self.outputs = outputs # list of Symbol | PythonObj
|
||||
self.contexts = contexts # list of StatementContext
|
||||
self.stmt_stack = (
|
||||
stacks # a list of string to record the source code callstack.
|
||||
)
|
||||
self.type = type
|
||||
|
||||
def __str__(self):
|
||||
return "{} || {} = {} ({}) ".format(
|
||||
self.type + " " * (10 - len(self.type)),
|
||||
self.to_string(self.outputs),
|
||||
self.name,
|
||||
self.to_string(self.inputs),
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return self.__str__()
|
||||
|
||||
@staticmethod
|
||||
def to_string(inps):
|
||||
return ", ".join(repr(x) for x in flatten(inps))
|
||||
|
||||
|
||||
class CallStatement(Statement):
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
inputs: list[Symbol],
|
||||
outputs: list[Symbol],
|
||||
contexts: list[StatementContext],
|
||||
stacks: list[str],
|
||||
):
|
||||
super().__init__("call", name, inputs, outputs, contexts, stacks)
|
||||
self.sir_name = name
|
||||
|
||||
|
||||
class ApiStatement(Statement):
|
||||
def __init__(
|
||||
self,
|
||||
api: Callable,
|
||||
inputs: list[Symbol],
|
||||
outputs: list[Symbol],
|
||||
contexts: list[StatementContext],
|
||||
stacks: list[str],
|
||||
):
|
||||
fullname = get_api_fullname(api)
|
||||
if fullname is None:
|
||||
fullname = "paddle." + api.__name__
|
||||
super().__init__("api", fullname, inputs, outputs, contexts, stacks)
|
||||
self.api = api
|
||||
|
||||
|
||||
class MethodStatement(Statement):
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
inputs: list[Symbol],
|
||||
outputs: list[Symbol],
|
||||
contexts: list[StatementContext],
|
||||
stacks: list[str],
|
||||
):
|
||||
super().__init__("method", name, inputs, outputs, contexts, stacks)
|
||||
self.method = name
|
||||
|
||||
|
||||
class LayerStatement(Statement):
|
||||
def __init__(
|
||||
self,
|
||||
layer: Reference, # Reference of paddle.nn.Layer
|
||||
inputs: list[Symbol],
|
||||
outputs: list[Symbol],
|
||||
contexts: list[StatementContext],
|
||||
stacks: list[str],
|
||||
):
|
||||
if isinstance(layer, Reference):
|
||||
name = layer().__class__.__name__
|
||||
else:
|
||||
name = layer.__class__.__name__
|
||||
super().__init__(
|
||||
"layer",
|
||||
name,
|
||||
inputs,
|
||||
outputs,
|
||||
contexts,
|
||||
stacks,
|
||||
)
|
||||
self.layer = layer
|
||||
|
||||
|
||||
class ASTStatement(Statement):
|
||||
def __init__(
|
||||
self,
|
||||
static_function,
|
||||
inputs: list[Symbol],
|
||||
outputs: list[Symbol],
|
||||
contexts: list[StatementContext],
|
||||
stacks: list[str],
|
||||
):
|
||||
# this dygraph_function always has attr __code__, which is checked before
|
||||
dygraph_func = static_function.dygraph_function
|
||||
super().__init__(
|
||||
"AST",
|
||||
dygraph_func.__code__.co_name,
|
||||
inputs,
|
||||
outputs,
|
||||
contexts,
|
||||
stacks,
|
||||
)
|
||||
converted_func = paddle.jit.dy2static.convert_to_static(dygraph_func)
|
||||
func_self = getattr(dygraph_func, '__self__', None)
|
||||
if func_self is not None:
|
||||
converted_func = functools.partial(converted_func, func_self)
|
||||
self.converted_func = converted_func
|
||||
|
||||
|
||||
class StatementIR:
|
||||
"""
|
||||
StatementIR is the carrier that records the code for building the neural network model.It is
|
||||
a representation of a purely computational structure, and does not care about specific values.
|
||||
The function converted from StatementIR can ensure that it can be turned into a static state.
|
||||
In this way, we can reuse the original `to_static` function to realize the execution of the static graph.
|
||||
|
||||
Note:
|
||||
Don't create by yourself, just use the StatementIRFactory.create()
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
self.inputs: list[Symbol] = []
|
||||
self.params: list[Symbol] = []
|
||||
self.outputs: list[Symbol] = []
|
||||
self.statements: list[Statement] = []
|
||||
|
||||
self.symbol_meta_map = {}
|
||||
self.param_symbol = set()
|
||||
self.non_param_symbol = set()
|
||||
|
||||
@property
|
||||
def input_with_params(self):
|
||||
return self.inputs + self.params
|
||||
|
||||
def __len__(self):
|
||||
return len(self.statements)
|
||||
|
||||
def __deepcopy__(self, memo=None):
|
||||
new_sir = StatementIR(self.name)
|
||||
new_sir.inputs = list(self.inputs)
|
||||
new_sir.params = list(self.params)
|
||||
new_sir.outputs = list(self.outputs)
|
||||
new_sir.statements = list(self.statements)
|
||||
new_sir.symbol_meta_map = dict(self.symbol_meta_map.items())
|
||||
new_sir.param_symbol = set(self.param_symbol)
|
||||
new_sir.non_param_symbol = set(self.non_param_symbol)
|
||||
return new_sir
|
||||
|
||||
def set_parameter_info(self, params, non_params):
|
||||
self.param_symbol.update(params)
|
||||
self.non_param_symbol.update(non_params)
|
||||
|
||||
def set_symbol_meta_map(self, meta_map):
|
||||
# if the meta of a input symbol inplace changed, we should get the origin meta as input of SIR
|
||||
meta_map.update(self.symbol_meta_map)
|
||||
self.symbol_meta_map = meta_map
|
||||
|
||||
def add_input(self, input):
|
||||
self.inputs.append(input)
|
||||
|
||||
def add_output(self, output):
|
||||
self.outputs.append(output)
|
||||
|
||||
def add_statement(self, statement):
|
||||
assert isinstance(statement, Statement)
|
||||
self.statements.append(statement)
|
||||
|
||||
def analyse_inputs(self):
|
||||
used_symbols = OrderedSet()
|
||||
generated_symbols = OrderedSet()
|
||||
for stmt in self.statements:
|
||||
for inp in flatten_extend(stmt.inputs):
|
||||
if isinstance(inp, Symbol) and inp not in generated_symbols:
|
||||
used_symbols.add(inp)
|
||||
for out in flatten_extend(stmt.outputs):
|
||||
if isinstance(out, Symbol):
|
||||
generated_symbols.add(out)
|
||||
|
||||
used_symbols = sorted(used_symbols, key=lambda x: x.name)
|
||||
if not parameters_persistent_mode_is_enabled():
|
||||
return used_symbols, []
|
||||
input_symbols = [
|
||||
symbol for symbol in used_symbols if symbol not in self.param_symbol
|
||||
]
|
||||
param_symbols = [
|
||||
symbol for symbol in used_symbols if symbol in self.param_symbol
|
||||
]
|
||||
return input_symbols, param_symbols
|
||||
|
||||
def __str__(self):
|
||||
strs = []
|
||||
strs.append(f"StatementIR: {self.name}")
|
||||
strs.append(f" inputs: {map_structure(lambda x: x.name, self.inputs)}")
|
||||
strs.append(f" params: {map_structure(lambda x: x.name, self.params)}")
|
||||
strs.append(
|
||||
f" outputs: {map_structure(lambda x: x.name, self.outputs)}"
|
||||
)
|
||||
strs.append(" statements: ")
|
||||
for stmt in self.statements:
|
||||
strs.append(f" {stmt}")
|
||||
return "\n".join(strs)
|
||||
|
||||
def __repr__(self):
|
||||
return self.__str__()
|
||||
|
||||
|
||||
class StatementIRFactory(metaclass=Singleton):
|
||||
"""
|
||||
It is used to create a StatementIR.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.cache = {}
|
||||
self.name_generator = NameGenerator("SIR_")
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.cache[key]
|
||||
|
||||
def create(self, input_name=None):
|
||||
if input_name:
|
||||
name = input_name
|
||||
else:
|
||||
name = self.name_generator.next()
|
||||
|
||||
sir = StatementIR(name)
|
||||
self.cache[name] = sir
|
||||
return sir
|
||||
|
||||
def update(self, stmt_ir):
|
||||
name = stmt_ir.name
|
||||
self.cache[name] = stmt_ir
|
||||
|
||||
def clear(self):
|
||||
want_clear = [
|
||||
key
|
||||
for key in self.cache.keys()
|
||||
if self.name_generator.match_name(key)
|
||||
]
|
||||
for key in want_clear:
|
||||
del self.cache[key]
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) 2025 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.
|
||||
@@ -0,0 +1,262 @@
|
||||
# Copyright (c) 2025 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
|
||||
from ..utils.exceptions import InnerError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..opcode_translator.executor.guard import StringifiedExpression
|
||||
|
||||
|
||||
class ConstraintNode:
|
||||
def __init__(self, inputs: list[ConstraintNode]):
|
||||
self.inputs = inputs
|
||||
|
||||
def create_guard_expr(
|
||||
self, extern_vars: dict[str, StringifiedExpression]
|
||||
) -> StringifiedExpression:
|
||||
raise NotImplementedError
|
||||
|
||||
def create_guard_node(
|
||||
self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
|
||||
) -> paddle.framework.core.ExprNodeBase:
|
||||
raise NotImplementedError(
|
||||
f"{self.__class__.__name__}.create_guard_node is not implemented"
|
||||
)
|
||||
|
||||
|
||||
class LeafConstraintNode(ConstraintNode):
|
||||
def __init__(self):
|
||||
super().__init__([])
|
||||
|
||||
|
||||
class UnaryConstraintNode(ConstraintNode):
|
||||
READABLE_SYMBOL: str
|
||||
|
||||
def __init__(self, input: ConstraintNode):
|
||||
super().__init__([input])
|
||||
self.input = input
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}({self.input})"
|
||||
|
||||
def create_guard_expr(
|
||||
self, extern_vars: dict[str, StringifiedExpression]
|
||||
) -> StringifiedExpression:
|
||||
from ..opcode_translator.executor.guard import (
|
||||
StringifiedExpression,
|
||||
union_free_vars,
|
||||
)
|
||||
|
||||
input = self.input.create_guard_expr(extern_vars)
|
||||
return StringifiedExpression(
|
||||
f"{self.READABLE_SYMBOL}({{}})",
|
||||
[input],
|
||||
union_free_vars(input.free_vars),
|
||||
)
|
||||
|
||||
def create_guard_node(
|
||||
self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
|
||||
) -> paddle.framework.core.ExprNodeBase:
|
||||
input = self.input.create_guard_node(extern_vars)
|
||||
return paddle.framework.core.UnaryExprNode(input, self.READABLE_SYMBOL)
|
||||
|
||||
|
||||
class BinaryConstraintNode(ConstraintNode):
|
||||
READABLE_SYMBOL: str
|
||||
|
||||
def __init__(self, lhs: ConstraintNode, rhs: ConstraintNode):
|
||||
super().__init__([lhs, rhs])
|
||||
self.lhs = lhs
|
||||
self.rhs = rhs
|
||||
|
||||
def create_guard_expr(
|
||||
self, extern_vars: dict[str, StringifiedExpression]
|
||||
) -> StringifiedExpression:
|
||||
from ..opcode_translator.executor.guard import (
|
||||
StringifiedExpression,
|
||||
union_free_vars,
|
||||
)
|
||||
|
||||
lhs = self.lhs.create_guard_expr(extern_vars)
|
||||
rhs = self.rhs.create_guard_expr(extern_vars)
|
||||
return StringifiedExpression(
|
||||
f"({{}} {self.READABLE_SYMBOL} {{}})",
|
||||
[lhs, rhs],
|
||||
union_free_vars(lhs.free_vars, rhs.free_vars),
|
||||
)
|
||||
|
||||
def create_guard_node(
|
||||
self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
|
||||
) -> paddle.framework.core.ExprNodeBase:
|
||||
lhs = self.lhs.create_guard_node(extern_vars)
|
||||
rhs = self.rhs.create_guard_node(extern_vars)
|
||||
return paddle.framework.core.BinaryExprNode(
|
||||
lhs, rhs, self.READABLE_SYMBOL
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.__class__.__name__}({self.lhs}, {self.rhs})"
|
||||
|
||||
|
||||
class ConstantConstraintNode(LeafConstraintNode):
|
||||
def __init__(self, value):
|
||||
super().__init__()
|
||||
self.value = value
|
||||
|
||||
def create_guard_expr(
|
||||
self, extern_vars: dict[str, StringifiedExpression]
|
||||
) -> StringifiedExpression:
|
||||
from ..opcode_translator.executor.guard import (
|
||||
StringifiedExpression,
|
||||
)
|
||||
|
||||
return StringifiedExpression(f"{self.value!r}", [], {})
|
||||
|
||||
def create_guard_node(
|
||||
self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
|
||||
) -> paddle.framework.core.ExprNodeBase:
|
||||
return paddle.framework.core.ConstantExprNode(self.value)
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.__class__.__name__}({self.value})"
|
||||
|
||||
|
||||
class SymbolicConstraintNode(LeafConstraintNode):
|
||||
def __init__(self, name: str):
|
||||
super().__init__()
|
||||
self.name = name
|
||||
|
||||
def create_guard_expr(
|
||||
self, extern_vars: dict[str, StringifiedExpression]
|
||||
) -> StringifiedExpression:
|
||||
from ..opcode_translator.executor.guard import (
|
||||
StringifiedExpression,
|
||||
union_free_vars,
|
||||
)
|
||||
|
||||
if self.name not in extern_vars:
|
||||
raise InnerError(
|
||||
f"Symbolic variable {self.name} not found in extern_vars."
|
||||
)
|
||||
return StringifiedExpression(
|
||||
"{}",
|
||||
[extern_vars[self.name]],
|
||||
union_free_vars(extern_vars[self.name].free_vars),
|
||||
)
|
||||
|
||||
def create_guard_node(
|
||||
self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
|
||||
) -> paddle.framework.core.ExprNodeBase:
|
||||
if self.name not in extern_vars:
|
||||
raise InnerError(
|
||||
f"Symbolic variable {self.name} not found in extern_vars."
|
||||
)
|
||||
return extern_vars[self.name]
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.__class__.__name__}({self.name})"
|
||||
|
||||
|
||||
class NegativeConstraintNode(UnaryConstraintNode):
|
||||
READABLE_SYMBOL = "-"
|
||||
|
||||
|
||||
class BitwiseNotConstraintNode(UnaryConstraintNode):
|
||||
READABLE_SYMBOL = "~"
|
||||
|
||||
|
||||
class AddConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "+"
|
||||
|
||||
|
||||
class SubConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "-"
|
||||
|
||||
|
||||
class MulConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "*"
|
||||
|
||||
|
||||
class TrueDivConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "/"
|
||||
|
||||
|
||||
class FloorDivConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "//"
|
||||
|
||||
|
||||
class ModConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "%"
|
||||
|
||||
|
||||
class PowConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "**"
|
||||
|
||||
|
||||
class BitwiseLShiftConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "<<"
|
||||
|
||||
|
||||
class BitwiseRShiftConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = ">>"
|
||||
|
||||
|
||||
class BitwiseAndConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "&"
|
||||
|
||||
|
||||
class BitwiseOrConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "|"
|
||||
|
||||
|
||||
class BitwiseXorConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "^"
|
||||
|
||||
|
||||
class LogicalToBoolConstraintNode(UnaryConstraintNode):
|
||||
READABLE_SYMBOL = "bool"
|
||||
|
||||
|
||||
class LogicalNotConstraintNode(UnaryConstraintNode):
|
||||
READABLE_SYMBOL = "not"
|
||||
|
||||
|
||||
class EqualConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "=="
|
||||
|
||||
|
||||
class NotEqualConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "!="
|
||||
|
||||
|
||||
class LessThanConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "<"
|
||||
|
||||
|
||||
class LessEqualConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = "<="
|
||||
|
||||
|
||||
class GreaterThanConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = ">"
|
||||
|
||||
|
||||
class GreaterEqualConstraintNode(BinaryConstraintNode):
|
||||
READABLE_SYMBOL = ">="
|
||||
@@ -0,0 +1,104 @@
|
||||
# Copyright (c) 2025 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..utils.magic_methods import BinaryOp, UnaryOp
|
||||
|
||||
|
||||
def symbolic_to_bool(x):
|
||||
# Unified api for python number and paddle Tensor
|
||||
return x != 0
|
||||
|
||||
|
||||
def symbolic_not(x):
|
||||
return x == 0
|
||||
|
||||
|
||||
def symbolic_truediv(x, y):
|
||||
# NOTE(SigureMo): In Paddle, the truediv maybe has precision issue.
|
||||
# For example, paddle.tensor(168) / 7, in Python it should be 24.0,
|
||||
# but in Paddle it will construct a Scale OP, which will calculate
|
||||
# as 168 * (1 / 7) = 24.00000191, which may cause some unexpected
|
||||
# bugs. So we cast the tensor and scalar both to float64 to avoid
|
||||
# this issue.
|
||||
is_need_cast_tensor = lambda v: (
|
||||
isinstance(v, paddle.pir.Value) and v.dtype is not paddle.float64
|
||||
)
|
||||
cast_tensor_if_needed = lambda v: (
|
||||
v.cast(paddle.float64) if is_need_cast_tensor(v) else v
|
||||
)
|
||||
cast_scalar_if_needed = lambda v: (
|
||||
paddle.full([], v, dtype=paddle.float64)
|
||||
if isinstance(v, (int, float))
|
||||
else v
|
||||
)
|
||||
cast_if_needed = lambda v: cast_tensor_if_needed(cast_scalar_if_needed(v))
|
||||
has_tensor_need_cast = is_need_cast_tensor(x) or is_need_cast_tensor(y)
|
||||
if not has_tensor_need_cast:
|
||||
return operator.truediv(x, y)
|
||||
x = cast_if_needed(x)
|
||||
y = cast_if_needed(y)
|
||||
return operator.truediv(x, y)
|
||||
|
||||
|
||||
# All symbolic operations need unified for python number and paddle Tensor
|
||||
SYMBOLIC_UNARY_MATH_OPS: list[UnaryOp] = [
|
||||
# Basic
|
||||
operator.neg,
|
||||
# Bitwise
|
||||
operator.invert,
|
||||
]
|
||||
SYMBOLIC_BINARY_MATH_OPS: list[BinaryOp] = [
|
||||
# Basic
|
||||
operator.add,
|
||||
operator.sub,
|
||||
operator.mul,
|
||||
symbolic_truediv,
|
||||
operator.floordiv,
|
||||
operator.pow,
|
||||
operator.mod,
|
||||
# Bitwise
|
||||
operator.lshift,
|
||||
operator.rshift,
|
||||
operator.and_,
|
||||
operator.or_,
|
||||
operator.xor,
|
||||
]
|
||||
SYMBOLIC_UNARY_LOGICAL_OPS: list[UnaryOp] = [
|
||||
symbolic_to_bool,
|
||||
symbolic_not,
|
||||
]
|
||||
SYMBOLIC_BINARY_LOGICAL_OPS: list[BinaryOp] = [
|
||||
operator.eq,
|
||||
operator.ne,
|
||||
operator.lt,
|
||||
operator.le,
|
||||
operator.gt,
|
||||
operator.ge,
|
||||
]
|
||||
SYMBOLIC_MATH_OPS = SYMBOLIC_UNARY_MATH_OPS + SYMBOLIC_BINARY_MATH_OPS
|
||||
SYMBOLIC_MATH_OPS = SYMBOLIC_UNARY_MATH_OPS + SYMBOLIC_BINARY_MATH_OPS
|
||||
SYMBOLIC_UNARY_OPS: list[UnaryOp] = (
|
||||
SYMBOLIC_UNARY_MATH_OPS + SYMBOLIC_UNARY_LOGICAL_OPS
|
||||
)
|
||||
SYMBOLIC_BINARY_OPS: list[BinaryOp] = (
|
||||
SYMBOLIC_BINARY_MATH_OPS + SYMBOLIC_BINARY_LOGICAL_OPS
|
||||
)
|
||||
@@ -0,0 +1,57 @@
|
||||
# Copyright (c) 2025 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
_T = TypeVar("_T", "int", "float", "bool")
|
||||
|
||||
|
||||
class SymbolicValue(Generic[_T]):
|
||||
example_value: _T | None
|
||||
|
||||
def __init__(self, example_value=None):
|
||||
self.example_value = example_value
|
||||
|
||||
def __repr__(self) -> str:
|
||||
if self.is_backed():
|
||||
return f"{self.__class__.__name__}({self.example_value})"
|
||||
return f"{self.__class__.__name__}()"
|
||||
|
||||
def get_static_type(self) -> type[_T]:
|
||||
raise NotImplementedError("get_py_type is not implemented.")
|
||||
|
||||
def is_backed(self):
|
||||
return self.example_value is not None
|
||||
|
||||
def get_example_value(self) -> _T:
|
||||
if self.example_value is None:
|
||||
raise ValueError(f"{self} is not backed by a value.")
|
||||
return self.example_value
|
||||
|
||||
|
||||
class SymbolicBool(SymbolicValue):
|
||||
def get_static_type(self) -> type[bool]:
|
||||
return bool
|
||||
|
||||
|
||||
class SymbolicInt(SymbolicValue):
|
||||
def get_static_type(self) -> type[int]:
|
||||
return int
|
||||
|
||||
|
||||
class SymbolicFloat(SymbolicValue):
|
||||
def get_static_type(self) -> type[float]:
|
||||
return float
|
||||
@@ -0,0 +1,118 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
import paddle
|
||||
|
||||
from .opcode_translator import eval_frame_callback
|
||||
from .profiler import SotStepProfilerGuard
|
||||
from .utils import (
|
||||
InfoCollector,
|
||||
StepInfoManager,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
|
||||
|
||||
def symbolic_translate(fn: Callable[P, R], **kwargs) -> Callable[P, R]:
|
||||
"""
|
||||
This function is the entry point of PaddleSOT. It sets eval_frame_callback before input
|
||||
function to achieve Opcode-level translation. The translation process depends on the
|
||||
simulation execution, in which information will be collected, especially the network
|
||||
code. After the simulation execution is completed, the network code will be compiled
|
||||
into a static graph Program to improve performance.
|
||||
|
||||
Args:
|
||||
fn: The input function.
|
||||
|
||||
Returns:
|
||||
Callable, The wrapped function.
|
||||
|
||||
Examples:
|
||||
>>> # doctest: +SKIP("Could not get source code of function foo."")
|
||||
>>> import paddle
|
||||
>>> import numpy as np
|
||||
>>> from sot.translate import symbolic_translate
|
||||
>>> def foo(cond: paddle.Tensor, x: paddle.Tensor):
|
||||
... x += 1
|
||||
... if cond:
|
||||
... x += 1
|
||||
... else:
|
||||
... x -= 1
|
||||
... return x
|
||||
>>> symbolic_translate_foo = symbolic_translate(foo)
|
||||
>>> # For the true branch, the output is 2.
|
||||
>>> cond = paddle.to_tensor(True)
|
||||
>>> x = paddle.to_tensor(0)
|
||||
>>> dygraph_out = foo(cond, x)
|
||||
>>> symbolic_translate_out = symbolic_translate_foo(cond, x)
|
||||
>>> dygraph_out
|
||||
Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
2)
|
||||
>>> symbolic_translate_out
|
||||
Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
2)
|
||||
>>> np.testing.assert_allclose(dygraph_out.numpy(), symbolic_translate_out.numpy())
|
||||
>>> # For the false branch, the output is 0.
|
||||
>>> cond = paddle.to_tensor(False)
|
||||
>>> dygraph_out = foo(cond, x)
|
||||
>>> symbolic_translate_out = symbolic_translate_foo(cond, x)
|
||||
>>> dygraph_out
|
||||
Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
0)
|
||||
>>> symbolic_translate_out
|
||||
Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
0)
|
||||
>>> np.testing.assert_allclose(dygraph_out.numpy(), symbolic_translate_out.numpy())
|
||||
|
||||
"""
|
||||
|
||||
if not paddle.framework.use_pir_api():
|
||||
raise RuntimeError(
|
||||
"SOT is only supported when running in PIR mode. Please set the environment variable "
|
||||
"FLAGS_enable_pir_api=1 to enable it."
|
||||
)
|
||||
|
||||
kwargs.setdefault('training', True)
|
||||
|
||||
def callback(frame):
|
||||
return eval_frame_callback(frame, **kwargs)
|
||||
|
||||
def impl(*args: P.args, **kwargs: P.kwargs) -> R:
|
||||
assert hasattr(fn, "__code__"), (
|
||||
"Target function doesn't have code for simulating."
|
||||
)
|
||||
with StepInfoManager().step_guard(fn.__code__), SotStepProfilerGuard():
|
||||
InfoCollector().clear_step_info()
|
||||
paddle.framework.core.set_eval_frame(callback)
|
||||
try:
|
||||
outs = fn(*args, **kwargs)
|
||||
except Exception as e:
|
||||
raise e
|
||||
finally:
|
||||
paddle.framework.core.set_eval_frame(None)
|
||||
|
||||
InfoCollector().print_step_report()
|
||||
return outs
|
||||
|
||||
return impl
|
||||
@@ -0,0 +1,128 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from .call_ast_utils import get_static_function, try_ast_func # noqa: F401
|
||||
from .envs import ( # noqa: F401
|
||||
ENV_MIN_GRAPH_SIZE,
|
||||
ENV_SOT_ALLOW_DYNAMIC_SHAPE,
|
||||
ENV_SOT_CE_DEBUG_MODE,
|
||||
ENV_SOT_COLLECT_INFO,
|
||||
ENV_SOT_ENABLE_0_SIZE_FALLBACK,
|
||||
ENV_SOT_ENABLE_COMPILE_TIME_LIMIT,
|
||||
ENV_SOT_ENABLE_FASTER_GUARD,
|
||||
ENV_SOT_ENABLE_GUARD_TREE,
|
||||
ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
|
||||
ENV_SOT_EXPORT,
|
||||
ENV_SOT_FORCE_FALLBACK_SIR_IDS,
|
||||
ENV_SOT_LOG_LEVEL,
|
||||
ENV_SOT_SERIALIZE_INFO,
|
||||
ENV_SOT_TRACE_NUMPY,
|
||||
ENV_SOT_UNSAFE_CACHE_FASTPATH,
|
||||
ENV_SOT_WITH_CONTROL_FLOW,
|
||||
ENV_STRICT_MODE,
|
||||
PEP508LikeEnvironmentVariable,
|
||||
allow_dynamic_shape_guard,
|
||||
enable_0_size_fallback_guard,
|
||||
export_guard,
|
||||
faster_guard_guard,
|
||||
guard_tree_guard,
|
||||
min_graph_size_guard,
|
||||
sot_step_profiler_guard,
|
||||
specialized_dim_numbers_guard,
|
||||
strict_mode_guard,
|
||||
with_control_flow_guard,
|
||||
)
|
||||
from .exceptions import ( # noqa: F401
|
||||
BreakGraphError,
|
||||
BreakGraphReasonBase,
|
||||
BuiltinFunctionBreak,
|
||||
ConditionalFallbackError,
|
||||
DataDependencyControlFlowBreak,
|
||||
DataDependencyDynamicShapeBreak,
|
||||
DataDependencyOperationBreak,
|
||||
ExportError,
|
||||
FallbackError,
|
||||
InnerError,
|
||||
PsdbBreakReason,
|
||||
SotCapturedException,
|
||||
SotCapturedExceptionFactory,
|
||||
SotErrorBase,
|
||||
UnsupportedIteratorBreak,
|
||||
UnsupportedOperationBreak,
|
||||
inner_error_default_handler,
|
||||
)
|
||||
from .info_collector import ( # noqa: F401
|
||||
BreakGraphReasonInfo,
|
||||
CompileCountInfo,
|
||||
InfoCollector,
|
||||
NewSymbolHitRateInfo,
|
||||
SubGraphInfo,
|
||||
SubGraphRelationInfo,
|
||||
)
|
||||
from .magic_methods import magic_method_builtin_dispatch # noqa: F401
|
||||
from .numpy_utils import ( # noqa: F401
|
||||
NUMPY_API_SUPPORTED_DICT,
|
||||
)
|
||||
from .paddle_api_config import ( # noqa: F401
|
||||
get_tensor_methods,
|
||||
is_break_graph_tensor_methods,
|
||||
is_directly_run_api,
|
||||
is_inplace_api,
|
||||
is_not_supported_paddle_layer,
|
||||
)
|
||||
from .utils import ( # noqa: F401
|
||||
Cache,
|
||||
ConstTypes,
|
||||
NameGenerator,
|
||||
ResumeFnNameFactory,
|
||||
Singleton,
|
||||
SIRToCodeMap,
|
||||
SotUndefinedVar,
|
||||
StepInfoManager,
|
||||
already_unified_in_dynamic_and_static_graph,
|
||||
count_if,
|
||||
current_symbol_registry,
|
||||
do_until_stop_iteration,
|
||||
execute_time,
|
||||
flatten,
|
||||
flatten_extend,
|
||||
get_api_fullname,
|
||||
get_min_non_specialized_number,
|
||||
get_numpy_ufuncs,
|
||||
get_obj_stable_repr,
|
||||
get_unbound_method,
|
||||
hashable,
|
||||
in_paddle_module,
|
||||
is_break_graph_api,
|
||||
is_builtin_fn,
|
||||
is_comprehensive_name,
|
||||
is_namedtuple_class,
|
||||
is_paddle_api,
|
||||
is_strict_mode,
|
||||
list_contain_by_id,
|
||||
list_find_index_by_id,
|
||||
log,
|
||||
log_do,
|
||||
log_enabled,
|
||||
log_format,
|
||||
log_once,
|
||||
map_if,
|
||||
map_if_extend,
|
||||
meta_str,
|
||||
need_capture_control_flow,
|
||||
no_eval_frame,
|
||||
printable,
|
||||
switch_symbol_registry,
|
||||
update_list_inplace,
|
||||
)
|
||||
@@ -0,0 +1,95 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
import inspect
|
||||
import types
|
||||
|
||||
import paddle
|
||||
|
||||
from .envs import ENV_SOT_WITH_CONTROL_FLOW
|
||||
from .exceptions import InnerError
|
||||
from .utils import Singleton
|
||||
|
||||
try_ast_codes = set()
|
||||
|
||||
|
||||
def try_ast_func(func):
|
||||
def _is_wrapped(f):
|
||||
return hasattr(f, '__wrapped__')
|
||||
|
||||
unwrapped_f = func
|
||||
if hasattr(unwrapped_f, "__code__"):
|
||||
try_ast_codes.add(func.__code__)
|
||||
|
||||
while _is_wrapped(unwrapped_f):
|
||||
unwrapped_f = unwrapped_f.__wrapped__
|
||||
if hasattr(unwrapped_f, "__code__"):
|
||||
try_ast_codes.add(func.__code__)
|
||||
|
||||
return func
|
||||
|
||||
|
||||
class StaticFunctionManager(metaclass=Singleton):
|
||||
def __init__(self):
|
||||
self.code_map = {}
|
||||
|
||||
def ast_transform_with_frame(self, frame):
|
||||
code = frame.f_code
|
||||
if code not in try_ast_codes:
|
||||
return None
|
||||
if code not in self.code_map:
|
||||
if code.co_name.startswith("#") or code.co_name.startswith("$"):
|
||||
self.code_map[code] = None
|
||||
elif len(code.co_cellvars) + len(code.co_freevars) != 0:
|
||||
self.code_map[code] = None
|
||||
else:
|
||||
function = types.FunctionType(
|
||||
code,
|
||||
frame.f_globals,
|
||||
code.co_name,
|
||||
(),
|
||||
(),
|
||||
)
|
||||
function = paddle.jit.to_static(function, full_graph=True)
|
||||
self.code_map[code] = function
|
||||
|
||||
return self.code_map[code]
|
||||
|
||||
def ast_transform_with_callable(self, fn):
|
||||
if not inspect.isfunction(fn) or not hasattr(fn, "__code__"):
|
||||
return None
|
||||
|
||||
code = fn.__code__
|
||||
if code not in try_ast_codes:
|
||||
return None
|
||||
if code not in self.code_map:
|
||||
if code.co_name.startswith("#") or code.co_name.startswith("$"):
|
||||
self.code_map[code] = None
|
||||
elif len(code.co_cellvars) + len(code.co_freevars) != 0:
|
||||
self.code_map[code] = None
|
||||
else:
|
||||
self.code_map[code] = paddle.jit.to_static(fn, full_graph=True)
|
||||
|
||||
return self.code_map[code]
|
||||
|
||||
|
||||
def get_static_function(obj, type_):
|
||||
if ENV_SOT_WITH_CONTROL_FLOW.get():
|
||||
if type_ == "eval_frame":
|
||||
return StaticFunctionManager().ast_transform_with_frame(obj)
|
||||
elif type_ == "inline_call":
|
||||
return StaticFunctionManager().ast_transform_with_callable(obj)
|
||||
else:
|
||||
raise InnerError(f"Can not get static function with type {type_}.")
|
||||
return None
|
||||
@@ -0,0 +1,238 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
|
||||
from paddle.utils.environments import (
|
||||
BooleanEnvironmentVariable,
|
||||
EnvironmentVariable,
|
||||
EnvironmentVariableGuard,
|
||||
IntegerEnvironmentVariable,
|
||||
StringEnvironmentVariable,
|
||||
)
|
||||
|
||||
|
||||
class PEP508LikeEnvironmentVariable(EnvironmentVariable[dict[str, list[str]]]):
|
||||
"""
|
||||
Environment variable parser following PEP 508 extras specification syntax.
|
||||
https://peps.python.org/pep-0508/
|
||||
|
||||
Processes strings using PEP 508-style bracket notation for optional components:
|
||||
"feat1[opt1,opt2], feat2[opt3,opt4]" -> {'feat1': ['opt1', 'opt2'], 'feat2': ['opt3', 'opt4']}
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, default: dict[str, list[str]]):
|
||||
super().__init__(name, default)
|
||||
assert isinstance(default, dict), "default must be a dict"
|
||||
|
||||
def parse_from_string(self) -> dict[str, list[str]]:
|
||||
env_var = os.getenv(self.name)
|
||||
if env_var is None or env_var == "":
|
||||
return self.default
|
||||
items = self.split_by_unbracketed_commas(env_var)
|
||||
ret = {}
|
||||
for item in items:
|
||||
ret.update(self.parse_parameterized_key(item))
|
||||
return ret
|
||||
|
||||
def convert_to_string(self, value: dict[str, list[str]]) -> str:
|
||||
assert isinstance(value, dict), "The input must be a dict"
|
||||
assert all(isinstance(x, str) for x in value.keys()), (
|
||||
"Keys must be a string"
|
||||
)
|
||||
assert all(isinstance(x, list) for x in value.values()), (
|
||||
"Values must be a list"
|
||||
)
|
||||
|
||||
env_list = []
|
||||
for k, v in value.items():
|
||||
env_list.append(f"{k}" + (f"[{','.join(v)}]" if len(v) else ""))
|
||||
|
||||
return ",".join(env_list)
|
||||
|
||||
@staticmethod
|
||||
def split_by_unbracketed_commas(input_str: str) -> list[str]:
|
||||
"""Split string by commas that are not enclosed in square brackets"""
|
||||
# "feat1[opt1,opt2], feat2[opt3], feat3" -> ["feat1[opt1,opt2]", "feat2[opt3]", "feat3"]
|
||||
bracket_depth = 0
|
||||
split_parts = []
|
||||
_start = 0
|
||||
|
||||
for _current, char in enumerate(input_str):
|
||||
if char == "[":
|
||||
bracket_depth += 1
|
||||
elif char == "]":
|
||||
bracket_depth = max(
|
||||
0, bracket_depth - 1
|
||||
) # Prevent negative depth
|
||||
|
||||
if char == "," and bracket_depth == 0:
|
||||
split_parts.append(input_str[_start:_current].strip())
|
||||
_start = _current + 1 # Skip comma
|
||||
|
||||
# Add remaining content after last comma
|
||||
if remaining := input_str[_start:].strip():
|
||||
split_parts.append(remaining)
|
||||
|
||||
return split_parts
|
||||
|
||||
@staticmethod
|
||||
def parse_parameterized_key(input_str: str) -> dict[str, list[str]]:
|
||||
"""Parse key with parameters in brackets into a dictionary."""
|
||||
|
||||
start_bracket = input_str.find("[")
|
||||
end_bracket = input_str.rfind("]")
|
||||
|
||||
if start_bracket == -1 or end_bracket == -1:
|
||||
return {input_str: []}
|
||||
|
||||
parameter_key = input_str[:start_bracket].strip()
|
||||
|
||||
# Extract and clean parameters
|
||||
parameters_str = input_str[start_bracket + 1 : end_bracket]
|
||||
parameter_values = [
|
||||
v.strip() for v in parameters_str.split(",") if v.strip()
|
||||
]
|
||||
|
||||
return {parameter_key: parameter_values}
|
||||
|
||||
|
||||
ENV_MIN_GRAPH_SIZE = IntegerEnvironmentVariable("MIN_GRAPH_SIZE", 10)
|
||||
ENV_SOT_LOG_LEVEL = IntegerEnvironmentVariable("SOT_LOG_LEVEL", 0)
|
||||
ENV_STRICT_MODE = BooleanEnvironmentVariable("STRICT_MODE", False)
|
||||
ENV_SOT_WITH_CONTROL_FLOW = BooleanEnvironmentVariable(
|
||||
"SOT_WITH_CONTROL_FLOW", True
|
||||
)
|
||||
ENV_SOT_EXPORT = StringEnvironmentVariable("SOT_EXPORT", "")
|
||||
ENV_SOT_ALLOW_DYNAMIC_SHAPE = BooleanEnvironmentVariable(
|
||||
"SOT_ALLOW_DYNAMIC_SHAPE",
|
||||
# Enable SOT dynamic shape as default in PIR mode
|
||||
True,
|
||||
)
|
||||
ENV_SOT_ENABLE_FASTER_GUARD = BooleanEnvironmentVariable(
|
||||
"SOT_ENABLE_FASTER_GUARD",
|
||||
False,
|
||||
)
|
||||
ENV_SOT_ENABLE_STRICT_GUARD_CHECK = BooleanEnvironmentVariable(
|
||||
"SOT_ENABLE_STRICT_GUARD_CHECK",
|
||||
False,
|
||||
)
|
||||
ENV_SOT_ENABLE_GUARD_TREE = BooleanEnvironmentVariable(
|
||||
"SOT_ENABLE_GUARD_TREE",
|
||||
False,
|
||||
)
|
||||
ENV_ENABLE_SOT_STEP_PROFILER = BooleanEnvironmentVariable(
|
||||
"ENABLE_SOT_STEP_PROFILER", False
|
||||
)
|
||||
ENV_SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE = BooleanEnvironmentVariable(
|
||||
"SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE", False
|
||||
)
|
||||
ENV_SOT_COLLECT_INFO = PEP508LikeEnvironmentVariable("SOT_COLLECT_INFO", {})
|
||||
ENV_SOT_SERIALIZE_INFO = BooleanEnvironmentVariable("SOT_SERIALIZE_INFO", False)
|
||||
ENV_SOT_CE_DEBUG_MODE = BooleanEnvironmentVariable("SOT_CE_DEBUG_MODE", False)
|
||||
ENV_SOT_FORCE_FALLBACK_SIR_IDS = StringEnvironmentVariable(
|
||||
"SOT_FORCE_FALLBACK_SIR_IDS", ""
|
||||
)
|
||||
ENV_SOT_TRACE_NUMPY = BooleanEnvironmentVariable("ENV_SOT_TRACE_NUMPY", True)
|
||||
ENV_SOT_UNSAFE_CACHE_FASTPATH = BooleanEnvironmentVariable(
|
||||
"SOT_UNSAFE_CACHE_FASTPATH", False
|
||||
)
|
||||
ENV_SOT_ENABLE_0_SIZE_FALLBACK = BooleanEnvironmentVariable(
|
||||
"SOT_ENABLE_0_SIZE_FALLBACK", True
|
||||
)
|
||||
ENV_SOT_SPECIALIZED_DIM_NUMBERS = StringEnvironmentVariable(
|
||||
"SOT_SPECIALIZED_DIM_NUMBERS", "0"
|
||||
)
|
||||
ENV_SOT_ENABLE_COMPILE_TIME_LIMIT = BooleanEnvironmentVariable(
|
||||
"SOT_ENABLE_COMPILE_TIME_LIMIT", True
|
||||
)
|
||||
|
||||
|
||||
def update_ce_flags():
|
||||
if not ENV_SOT_CE_DEBUG_MODE.get():
|
||||
return
|
||||
# Enable information collection flags to facilitate debugging and analysis
|
||||
|
||||
collected_info_item: dict[str, list[str]] = ENV_SOT_COLLECT_INFO.get()
|
||||
collected_info_item.setdefault("breakgraph_reason", [])
|
||||
collected_info_item.setdefault("subgraph_info", [])
|
||||
|
||||
ENV_SOT_COLLECT_INFO.set(collected_info_item)
|
||||
ENV_SOT_SERIALIZE_INFO.set(True)
|
||||
|
||||
|
||||
update_ce_flags()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def strict_mode_guard(value: bool):
|
||||
with EnvironmentVariableGuard(ENV_STRICT_MODE, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def min_graph_size_guard(value: int):
|
||||
with EnvironmentVariableGuard(ENV_MIN_GRAPH_SIZE, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def with_control_flow_guard(value: bool):
|
||||
with EnvironmentVariableGuard(ENV_SOT_WITH_CONTROL_FLOW, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def export_guard(value: str):
|
||||
with EnvironmentVariableGuard(ENV_SOT_EXPORT, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def allow_dynamic_shape_guard(value: bool):
|
||||
with EnvironmentVariableGuard(ENV_SOT_ALLOW_DYNAMIC_SHAPE, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def faster_guard_guard(value: bool):
|
||||
with EnvironmentVariableGuard(ENV_SOT_ENABLE_FASTER_GUARD, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def guard_tree_guard(value: bool):
|
||||
with EnvironmentVariableGuard(ENV_SOT_ENABLE_GUARD_TREE, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def sot_step_profiler_guard(value: bool):
|
||||
with EnvironmentVariableGuard(ENV_ENABLE_SOT_STEP_PROFILER, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def specialized_dim_numbers_guard(value: str):
|
||||
with EnvironmentVariableGuard(ENV_SOT_SPECIALIZED_DIM_NUMBERS, value):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def enable_0_size_fallback_guard(value: bool):
|
||||
with EnvironmentVariableGuard(ENV_SOT_ENABLE_0_SIZE_FALLBACK, value):
|
||||
yield
|
||||
@@ -0,0 +1,488 @@
|
||||
# Copyright (c) 2023 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.
|
||||
from __future__ import annotations
|
||||
|
||||
import traceback
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .info_collector import BreakGraphReasonInfo
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..opcode_translator.executor.variables.base import VariableBase
|
||||
|
||||
|
||||
class BreakGraphReasonBase:
|
||||
"""Base class for representing reasons why graph execution was interrupted.
|
||||
|
||||
Attributes:
|
||||
reason_str (str): Description of the break reason
|
||||
file_path (str): Path to the file where break occurred
|
||||
line_number (int): Line number where break occurred
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reason_str,
|
||||
file_path="",
|
||||
line_number=-1,
|
||||
):
|
||||
self.reason_str = reason_str
|
||||
self.file_path = file_path
|
||||
self.line_number = line_number
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.reason_str}"
|
||||
|
||||
|
||||
class DataDependencyBreak(BreakGraphReasonBase):
|
||||
pass
|
||||
|
||||
|
||||
class DataDependencyControlFlowBreak(DataDependencyBreak):
|
||||
"""Break reason for control flow execution."""
|
||||
|
||||
def __init__(self, reason_str=None, file_path="", line_number=-1):
|
||||
if reason_str is None:
|
||||
reason_str = "OpcodeInlineExecutor want break graph when simulate control flow."
|
||||
|
||||
super().__init__(
|
||||
reason_str,
|
||||
file_path,
|
||||
line_number,
|
||||
)
|
||||
|
||||
|
||||
class DataDependencyDynamicShapeBreak(DataDependencyBreak):
|
||||
pass
|
||||
|
||||
|
||||
class DataDependencyOperationBreak(DataDependencyBreak):
|
||||
pass
|
||||
|
||||
|
||||
class UnsupportedOperationBreak(BreakGraphReasonBase):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
left_type=None,
|
||||
right_type=None,
|
||||
operator=None,
|
||||
reason_str=None,
|
||||
file_path="",
|
||||
line_number=-1,
|
||||
):
|
||||
if reason_str is None:
|
||||
reason_str = f"Unsupported operator '{operator}' between {left_type} and {right_type}"
|
||||
super().__init__(reason_str, file_path, line_number)
|
||||
|
||||
|
||||
class UnsupportedPaddleAPIBreak(UnsupportedOperationBreak):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
fn_name=None,
|
||||
reason_str=None,
|
||||
file_path="",
|
||||
line_number=-1,
|
||||
):
|
||||
if reason_str is None:
|
||||
reason_str = f"Not support Paddlepaddle API: {fn_name}"
|
||||
|
||||
super().__init__(
|
||||
reason_str=reason_str,
|
||||
file_path=file_path,
|
||||
line_number=line_number,
|
||||
)
|
||||
|
||||
|
||||
class UnsupportedNumPyAPIBreak(UnsupportedOperationBreak):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
fn_name=None,
|
||||
reason_str=None,
|
||||
file_path="",
|
||||
line_number=-1,
|
||||
):
|
||||
if reason_str is None:
|
||||
reason_str = f"Not support NumPy API: {fn_name}"
|
||||
|
||||
super().__init__(
|
||||
reason_str=reason_str,
|
||||
file_path=file_path,
|
||||
line_number=line_number,
|
||||
)
|
||||
|
||||
|
||||
class UnsupportedRandomAPIBreak(UnsupportedOperationBreak):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
fn_name=None,
|
||||
reason_str=None,
|
||||
file_path="",
|
||||
line_number=-1,
|
||||
):
|
||||
if reason_str is None:
|
||||
reason_str = f"Random function {fn_name} is not supported."
|
||||
|
||||
super().__init__(
|
||||
reason_str=reason_str,
|
||||
file_path=file_path,
|
||||
line_number=line_number,
|
||||
)
|
||||
|
||||
|
||||
class ForceBreak(UnsupportedOperationBreak):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
reason_str=None,
|
||||
file_path="",
|
||||
line_number=-1,
|
||||
):
|
||||
if reason_str is None:
|
||||
reason_str = "Force break graph execution"
|
||||
|
||||
super().__init__(
|
||||
reason_str=reason_str,
|
||||
file_path=file_path,
|
||||
line_number=line_number,
|
||||
)
|
||||
|
||||
|
||||
class BuiltinFunctionBreak(UnsupportedOperationBreak):
|
||||
"""Break reason for unsupported built-in function calls.
|
||||
|
||||
Args:
|
||||
fn_name (str): Name of the builtin function
|
||||
arg_types (list): Types of the arguments passed to the function
|
||||
file_path (str): Path to the file where break occurred
|
||||
line_number (int): Line number where break occurred
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
fn_name=None,
|
||||
arg_types=None,
|
||||
reason_str=None,
|
||||
file_path="",
|
||||
line_number=-1,
|
||||
):
|
||||
if reason_str is None:
|
||||
reason_str = f"Not support builtin function: {fn_name} with args: Args({arg_types})"
|
||||
|
||||
super().__init__(
|
||||
reason_str=reason_str,
|
||||
file_path=file_path,
|
||||
line_number=line_number,
|
||||
)
|
||||
|
||||
|
||||
class SideEffectBreak(BreakGraphReasonBase):
|
||||
pass
|
||||
|
||||
|
||||
class UnsupportedIteratorBreak(SideEffectBreak):
|
||||
pass
|
||||
|
||||
|
||||
class InlineCallBreak(BreakGraphReasonBase):
|
||||
pass
|
||||
|
||||
|
||||
class FallbackInlineCallBreak(InlineCallBreak):
|
||||
pass
|
||||
|
||||
|
||||
class BreakGraphInlineCallBreak(InlineCallBreak):
|
||||
pass
|
||||
|
||||
|
||||
class OtherInlineCallBreak(InlineCallBreak):
|
||||
pass
|
||||
|
||||
|
||||
class DygraphInconsistentWithStaticBreak(BreakGraphReasonBase):
|
||||
pass
|
||||
|
||||
|
||||
class PsdbBreakReason(BreakGraphReasonBase):
|
||||
pass
|
||||
|
||||
|
||||
class InferMetaBreak(BreakGraphReasonBase):
|
||||
"""Break reason during meta information inference phase."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class NullMetaBreak(BreakGraphReasonBase):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
file_path="",
|
||||
line_number=-1,
|
||||
):
|
||||
super().__init__(
|
||||
"Access attribute from null meta", file_path, line_number
|
||||
)
|
||||
|
||||
|
||||
class SotErrorBase(Exception):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
from ..opcode_translator.breakpoint import BreakpointManager
|
||||
|
||||
BreakpointManager().on_event(f"{self.__class__.__name__}")
|
||||
|
||||
def print(self):
|
||||
lines = traceback.format_tb(self.__traceback__)
|
||||
print("".join(lines))
|
||||
|
||||
|
||||
class InnerError(SotErrorBase):
|
||||
pass
|
||||
|
||||
|
||||
class HasNoAttributeError(InnerError):
|
||||
pass
|
||||
|
||||
|
||||
class FallbackError(SotErrorBase):
|
||||
def __init__(self, msg, disable_eval_frame=False):
|
||||
super().__init__(msg)
|
||||
self.disable_eval_frame = disable_eval_frame
|
||||
|
||||
|
||||
class ConditionalFallbackError(FallbackError): ...
|
||||
|
||||
|
||||
# raise in inline function call strategy.
|
||||
class BreakGraphError(SotErrorBase):
|
||||
def __init__(self, reason: BreakGraphReasonBase = None):
|
||||
super().__init__(str(reason))
|
||||
|
||||
if not isinstance(reason, BreakGraphReasonBase):
|
||||
raise ValueError(
|
||||
"reason must be a subclass of BreakGraphReasonBase"
|
||||
)
|
||||
|
||||
self.reason = reason
|
||||
BreakGraphReasonInfo.collect_break_graph_reason(reason)
|
||||
|
||||
|
||||
def inner_error_default_handler(func, message_fn):
|
||||
"""Wrap function and an error handling function and throw an InnerError."""
|
||||
|
||||
def impl(*args, **kwargs):
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except SotErrorBase as e:
|
||||
raise e
|
||||
except Exception as e:
|
||||
message = message_fn(*args, **kwargs)
|
||||
origin_exception_message = "\n".join(
|
||||
traceback.format_exception(type(e), e, e.__traceback__)
|
||||
)
|
||||
raise InnerError(
|
||||
f"{message}\nOrigin Exception is: \n {origin_exception_message}"
|
||||
) from e
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
class ExportError(SotErrorBase):
|
||||
pass
|
||||
|
||||
|
||||
class SotExtraInfo:
|
||||
SOT_EXTRA_INFO_ATTR_NAME = "__SOT_EXTRA_INFO__"
|
||||
|
||||
def __init__(self, *, need_breakgraph: bool = False):
|
||||
self.need_breakgraph = need_breakgraph
|
||||
|
||||
def set_need_breakgraph(self, need_breakgraph: bool):
|
||||
self.need_breakgraph = need_breakgraph
|
||||
|
||||
def attach(self, err: BaseException):
|
||||
setattr(err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, self)
|
||||
|
||||
@staticmethod
|
||||
def default() -> SotExtraInfo:
|
||||
return SotExtraInfo()
|
||||
|
||||
@staticmethod
|
||||
def from_exception(err: BaseException) -> SotExtraInfo:
|
||||
info = getattr(
|
||||
err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, SotExtraInfo.default()
|
||||
)
|
||||
setattr(err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, info)
|
||||
return info
|
||||
|
||||
|
||||
class SotCapturedException(SotErrorBase):
|
||||
# Represents an exception encountered during bytecode execution simulation.
|
||||
# This exception is used by SOT to handle Python exceptions by mapping them to
|
||||
# SotCapturedException for consistent exception handling in the simulation process.
|
||||
...
|
||||
|
||||
|
||||
class SotCapturedLookupError(SotCapturedException): ...
|
||||
|
||||
|
||||
class SotCapturedIndexError(SotCapturedLookupError): ...
|
||||
|
||||
|
||||
class SotCapturedKeyError(SotCapturedLookupError): ...
|
||||
|
||||
|
||||
class SotCapturedArithmeticError(SotCapturedException): ...
|
||||
|
||||
|
||||
class SotCapturedFloatingPointError(SotCapturedArithmeticError): ...
|
||||
|
||||
|
||||
class SotCapturedOverflowError(SotCapturedArithmeticError): ...
|
||||
|
||||
|
||||
class SotCapturedZeroDivisionError(SotCapturedArithmeticError): ...
|
||||
|
||||
|
||||
class SotCapturedImportError(SotCapturedException): ...
|
||||
|
||||
|
||||
class SotCapturedModuleNotFoundError(SotCapturedImportError): ...
|
||||
|
||||
|
||||
class SotCapturedRuntimeError(SotCapturedException): ...
|
||||
|
||||
|
||||
class SotCapturedNotImplementedError(SotCapturedRuntimeError): ...
|
||||
|
||||
|
||||
class SotCapturedRecursionError(SotCapturedRuntimeError): ...
|
||||
|
||||
|
||||
class SotCapturedNameError(SotCapturedException): ...
|
||||
|
||||
|
||||
class SotCapturedUnboundLocalError(SotCapturedNameError): ...
|
||||
|
||||
|
||||
class SotCapturedSyntaxError(SotCapturedException): ...
|
||||
|
||||
|
||||
class SotCapturedIndentationError(SotCapturedSyntaxError): ...
|
||||
|
||||
|
||||
class SotCapturedTabError(SotCapturedIndentationError): ...
|
||||
|
||||
|
||||
class SotCapturedOSError(SotCapturedException): ...
|
||||
|
||||
|
||||
class SotCapturedFileExistsError(SotCapturedOSError): ...
|
||||
|
||||
|
||||
class SotCapturedFileNotFoundError(SotCapturedOSError): ...
|
||||
|
||||
|
||||
class SotCapturedIsADirectoryError(SotCapturedOSError): ...
|
||||
|
||||
|
||||
class SotCapturedNotADirectoryError(SotCapturedOSError): ...
|
||||
|
||||
|
||||
class SotCapturedPermissionError(SotCapturedOSError): ...
|
||||
|
||||
|
||||
class SotCapturedTimeoutError(SotCapturedOSError): ...
|
||||
|
||||
|
||||
class SotCapturedStopIteration(SotCapturedOSError): ...
|
||||
|
||||
|
||||
class SotCapturedExceptionFactory:
|
||||
# This dictionary maps common built-in Python Exception types to their corresponding SotCapturedException
|
||||
# types, preserving the original exception hierarchy for proper inheritance behavior.
|
||||
# Reference: https://docs.python.org/3/library/exceptions.html#exception-hierarchy
|
||||
MAPPING = {
|
||||
Exception: SotCapturedException,
|
||||
LookupError: SotCapturedLookupError,
|
||||
IndexError: SotCapturedIndexError,
|
||||
KeyError: SotCapturedKeyError,
|
||||
ArithmeticError: SotCapturedArithmeticError,
|
||||
FloatingPointError: SotCapturedFloatingPointError,
|
||||
OverflowError: SotCapturedOverflowError,
|
||||
ZeroDivisionError: SotCapturedZeroDivisionError,
|
||||
ImportError: SotCapturedImportError,
|
||||
ModuleNotFoundError: SotCapturedModuleNotFoundError,
|
||||
RuntimeError: SotCapturedRuntimeError,
|
||||
NotImplementedError: SotCapturedNotImplementedError,
|
||||
NameError: SotCapturedNameError,
|
||||
UnboundLocalError: SotCapturedUnboundLocalError,
|
||||
SyntaxError: SotCapturedSyntaxError,
|
||||
IndentationError: SotCapturedIndentationError,
|
||||
TabError: SotCapturedTabError,
|
||||
OSError: SotCapturedOSError,
|
||||
FileExistsError: SotCapturedFileExistsError,
|
||||
FileNotFoundError: SotCapturedFileNotFoundError,
|
||||
IsADirectoryError: SotCapturedIsADirectoryError,
|
||||
NotADirectoryError: SotCapturedNotADirectoryError,
|
||||
PermissionError: SotCapturedPermissionError,
|
||||
TimeoutError: SotCapturedTimeoutError,
|
||||
StopIteration: SotCapturedStopIteration,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get(
|
||||
cls,
|
||||
exc_type: type[Exception],
|
||||
) -> type[SotCapturedException]:
|
||||
if isinstance(exc_type, type) and issubclass(
|
||||
exc_type, SotCapturedException
|
||||
):
|
||||
return exc_type
|
||||
|
||||
if exc_type not in cls.MAPPING:
|
||||
name = getattr(exc_type, "__name__", str(exc_type))
|
||||
cls.MAPPING[exc_type] = type(
|
||||
f"SotCaptured{name}", (SotCapturedException,), {}
|
||||
)
|
||||
return cls.MAPPING[exc_type]
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
origin_exc: Exception,
|
||||
tracked_args: list[VariableBase] | None = None,
|
||||
) -> SotCapturedException:
|
||||
# transform an Exception to SotCapturedException
|
||||
exc_type = origin_exc.__class__
|
||||
|
||||
new_exc_type = cls.get(exc_type)
|
||||
new_exc = new_exc_type(*origin_exc.args)
|
||||
new_exc.__cause__ = origin_exc.__cause__
|
||||
new_exc.__context__ = origin_exc.__context__
|
||||
new_exc.__suppress_context__ = origin_exc.__suppress_context__
|
||||
new_exc.__traceback__ = origin_exc.__traceback__
|
||||
|
||||
# Propagating Exception Parameters through SotCapturedException
|
||||
if tracked_args is None:
|
||||
tracked_args = []
|
||||
new_exc.tracked_args = tracked_args
|
||||
|
||||
return new_exc
|
||||
@@ -0,0 +1,468 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import base64
|
||||
import json
|
||||
import sys
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
from .envs import ENV_SOT_COLLECT_INFO, ENV_SOT_SERIALIZE_INFO
|
||||
from .utils import Singleton
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import types
|
||||
|
||||
from .exceptions import BreakGraphReasonBase
|
||||
|
||||
PREFIX = "<sot>"
|
||||
SUFFIX = "</sot>"
|
||||
ENCODING = "utf-8"
|
||||
|
||||
|
||||
def try_import_graphviz():
|
||||
try:
|
||||
import graphviz
|
||||
|
||||
return graphviz
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
class InfoType(Enum):
|
||||
STEP_INFO = 0
|
||||
E2E_INFO = 1
|
||||
|
||||
|
||||
class InfoCollector(metaclass=Singleton):
|
||||
def __init__(self):
|
||||
self._step_info: dict[str, list[InfoBase]] = {}
|
||||
self._e2e_info: dict[str, list[InfoBase]] = {}
|
||||
|
||||
def get_info_dict(self, info_type: InfoType) -> dict[str, list[InfoBase]]:
|
||||
if info_type == InfoType.STEP_INFO:
|
||||
return self._step_info
|
||||
else:
|
||||
return self._e2e_info
|
||||
|
||||
def attach(self, cls: type[InfoBase], *args, **kwargs) -> None:
|
||||
if self.need_collect(cls):
|
||||
info = cls(*args, **kwargs)
|
||||
self.register(info)
|
||||
|
||||
def register(self, info: InfoBase) -> None:
|
||||
info_class_name = info.__class__.__name__
|
||||
info_type = info.TYPE
|
||||
info_dict = self.get_info_dict(info_type)
|
||||
info_dict.setdefault(info_class_name, [])
|
||||
info_dict[info_class_name].append(info)
|
||||
|
||||
def need_collect(self, cls: type[InfoBase]) -> bool:
|
||||
return cls.SHORT_NAME in ENV_SOT_COLLECT_INFO.get()
|
||||
|
||||
def clear_step_info(self):
|
||||
self._step_info.clear()
|
||||
|
||||
def clear_e2e_info(self):
|
||||
self._e2e_info.clear()
|
||||
|
||||
def clear(self):
|
||||
self.clear_step_info()
|
||||
self.clear_e2e_info()
|
||||
|
||||
def print_step_report(self):
|
||||
self.print_report(InfoType.STEP_INFO)
|
||||
|
||||
def print_e2e_info_atexit(self) -> None:
|
||||
def atexit_hook():
|
||||
self.print_report(InfoType.E2E_INFO)
|
||||
sys.stdout.flush()
|
||||
self.clear()
|
||||
|
||||
atexit.register(atexit_hook)
|
||||
|
||||
def print_report(self, info_type: InfoType) -> None:
|
||||
if info_dict := self.get_info_dict(info_type):
|
||||
print(self.generate_report(info_dict))
|
||||
|
||||
def generate_report(self, info_dict: dict[str, list[InfoBase]]) -> str:
|
||||
report = ""
|
||||
for info_class_name, info_list in info_dict.items():
|
||||
cls = info_list[0].__class__
|
||||
report += f"{info_class_name} ({cls.SHORT_NAME}):\n"
|
||||
if ENV_SOT_SERIALIZE_INFO.get():
|
||||
report += cls.json_report(info_list)
|
||||
else:
|
||||
report += cls.summary(info_list)
|
||||
report += "\n"
|
||||
return report
|
||||
|
||||
|
||||
InfoCollector().print_e2e_info_atexit()
|
||||
|
||||
|
||||
class InfoBase(ABC):
|
||||
SHORT_NAME: ClassVar[str]
|
||||
TYPE: ClassVar[InfoType]
|
||||
|
||||
def __init__(self): ...
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def summary(cls, history: list[Self]) -> str: ...
|
||||
|
||||
@classmethod
|
||||
def serialize(cls, obj: dict[str:Any]) -> str:
|
||||
json_data = json.dumps(obj)
|
||||
b64_bytes = base64.b64encode(json_data.encode(ENCODING))
|
||||
|
||||
return b64_bytes.decode(ENCODING)
|
||||
|
||||
@classmethod
|
||||
def deserialize(cls, data: bytes | str) -> dict:
|
||||
if isinstance(data, str):
|
||||
data = data.encode(ENCODING)
|
||||
json_str = base64.b64decode(data).decode(ENCODING)
|
||||
|
||||
return json.loads(json_str)
|
||||
|
||||
|
||||
class NewSymbolHitRateInfo(InfoBase):
|
||||
SHORT_NAME = "new_symbol_hit_rate"
|
||||
TYPE = InfoType.STEP_INFO
|
||||
|
||||
def __init__(
|
||||
self, input_tensor_ids: list[int], output_tensor_ids: list[int]
|
||||
):
|
||||
super().__init__()
|
||||
self.input_tensor_ids = input_tensor_ids
|
||||
self.output_tensor_ids = output_tensor_ids
|
||||
|
||||
@classmethod
|
||||
def summary(cls, history: list[Self]) -> str:
|
||||
if len(history) == 0:
|
||||
return f"No {cls.SHORT_NAME} info"
|
||||
if len(history) == 1:
|
||||
return "Only one subgraph is generated"
|
||||
known_tensor_ids = set()
|
||||
hit_count = 0
|
||||
all_count = sum([len(info.input_tensor_ids) for info in history[1:]])
|
||||
for i, info in enumerate(history):
|
||||
for tensor_id in info.input_tensor_ids:
|
||||
# Skip the first graph
|
||||
if i == 0:
|
||||
continue
|
||||
if tensor_id in known_tensor_ids:
|
||||
hit_count += 1
|
||||
for tensor_id in info.output_tensor_ids:
|
||||
known_tensor_ids.add(tensor_id)
|
||||
summary = f"All tensor count: {all_count}, hit count: {hit_count}\n"
|
||||
summary += f"Hit rate: {hit_count / all_count:.2f}"
|
||||
return summary
|
||||
|
||||
@classmethod
|
||||
def json_report(cls, history: list[Self]) -> str:
|
||||
# TODO: need to support serialize the output
|
||||
return cls.summary(history)
|
||||
|
||||
|
||||
class SubGraphRelationInfo(InfoBase):
|
||||
SHORT_NAME = "subgraph_relation"
|
||||
TYPE = InfoType.STEP_INFO
|
||||
STEP_UNIQUE_ID = 0
|
||||
|
||||
class ConcreteShapeInfo(NamedTuple):
|
||||
id: int
|
||||
ir_shape: list[int]
|
||||
real_shape: list[int]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
subgraph_name: str,
|
||||
input_shape_infos: list[SubGraphRelationInfo.ConcreteShapeInfo],
|
||||
output_shape_infos: list[SubGraphRelationInfo.ConcreteShapeInfo],
|
||||
is_first_call: bool,
|
||||
graph_size: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.subgraph_name = subgraph_name
|
||||
self.input_shape_infos = input_shape_infos
|
||||
self.output_shape_infos = output_shape_infos
|
||||
self.is_first_call = is_first_call
|
||||
self.graph_size = graph_size
|
||||
|
||||
@classmethod
|
||||
def summary(cls, history: list[Self]) -> str:
|
||||
# TODO: attach input shape (with dynamic shape info)
|
||||
cls.STEP_UNIQUE_ID += 1
|
||||
if len(history) == 0:
|
||||
return f"No {cls.SHORT_NAME} info"
|
||||
if all(not subgraph_info.is_first_call for subgraph_info in history):
|
||||
return "All subgraph are not the first call"
|
||||
graphviz = try_import_graphviz()
|
||||
if graphviz is None:
|
||||
return "Please install graphviz to show the subgraph relation"
|
||||
dot = graphviz.Digraph()
|
||||
shape_infos = [
|
||||
shape_info
|
||||
for info in history
|
||||
for shape_info in info.input_shape_infos + info.output_shape_infos
|
||||
]
|
||||
|
||||
def to_tensor_node_name(
|
||||
shape_info: SubGraphRelationInfo.ConcreteShapeInfo,
|
||||
):
|
||||
return f"tensor_{shape_info.id}"
|
||||
|
||||
visited_shape = set()
|
||||
for shape_info in shape_infos:
|
||||
if shape_info.id in visited_shape:
|
||||
continue
|
||||
visited_shape.add(shape_info.id)
|
||||
dot.node(
|
||||
to_tensor_node_name(shape_info),
|
||||
f"Tensor {shape_info.id} shape={shape_info.real_shape}",
|
||||
shape="rect",
|
||||
)
|
||||
for i, info in enumerate(history):
|
||||
subgraph_id = f"subgraph_{i}"
|
||||
dot.node(
|
||||
subgraph_id,
|
||||
f"Subgraph {i} ({info.subgraph_name}, size={info.graph_size})",
|
||||
shape="oval",
|
||||
fillcolor="cyan" if info.is_first_call else None,
|
||||
style="filled" if info.is_first_call else None,
|
||||
)
|
||||
for shape_info in info.input_shape_infos:
|
||||
dot.edge(
|
||||
to_tensor_node_name(shape_info),
|
||||
subgraph_id,
|
||||
label=str(shape_info.ir_shape),
|
||||
)
|
||||
for shape_info in info.output_shape_infos:
|
||||
dot.edge(
|
||||
subgraph_id,
|
||||
to_tensor_node_name(shape_info),
|
||||
label=str(shape_info.ir_shape),
|
||||
)
|
||||
|
||||
directory = Path(".") / "subgraph_relation"
|
||||
directory.mkdir(exist_ok=True, parents=True)
|
||||
filename = f"subgraph_relation_{cls.STEP_UNIQUE_ID}"
|
||||
dot.render(directory / filename, format="svg", cleanup=True)
|
||||
return f"Please check {directory / filename}.svg for subgraph relation"
|
||||
|
||||
@classmethod
|
||||
def json_report(cls, history: list[Self]) -> str:
|
||||
# TODO: need to support serialize the output
|
||||
return cls.summary(history)
|
||||
|
||||
|
||||
class CompileCountInfo(InfoBase):
|
||||
SHORT_NAME = "compile_count"
|
||||
TYPE = InfoType.E2E_INFO
|
||||
|
||||
def __init__(self, code: types.CodeType):
|
||||
super().__init__()
|
||||
self.code = code
|
||||
|
||||
@classmethod
|
||||
def summary(cls, history: list[Self]) -> str:
|
||||
if len(history) == 0:
|
||||
return f"No {cls.SHORT_NAME} info"
|
||||
code_count = {}
|
||||
for info in history:
|
||||
code_count[info.code] = code_count.get(info.code, 0) + 1
|
||||
summary_lines = []
|
||||
for code, count in sorted(
|
||||
code_count.items(), key=lambda x: x[1], reverse=True
|
||||
):
|
||||
filename, lineno = code.co_filename, code.co_firstlineno
|
||||
summary_lines.append(
|
||||
f" {code.co_name} ({filename}:{lineno}): {count}"
|
||||
)
|
||||
summary = "\n".join(summary_lines)
|
||||
return summary
|
||||
|
||||
@classmethod
|
||||
def json_report(cls, history: list[Self]) -> str:
|
||||
# TODO: need to support serialize the output
|
||||
return cls.summary(history)
|
||||
|
||||
|
||||
class BreakGraphReasonInfo(InfoBase):
|
||||
SHORT_NAME = "breakgraph_reason"
|
||||
TYPE = InfoType.E2E_INFO
|
||||
|
||||
def __init__(self, reason: BreakGraphReasonBase):
|
||||
super().__init__()
|
||||
self.reason = reason
|
||||
|
||||
@classmethod
|
||||
def classify(cls, history: list[Self]) -> str:
|
||||
reasons_dict = {}
|
||||
|
||||
for info in history:
|
||||
name = info.reason.__class__.__name__
|
||||
if name not in reasons_dict:
|
||||
reasons_dict[name] = []
|
||||
reasons_dict[name].append(str(info.reason))
|
||||
|
||||
sorted_reasons = list(reasons_dict.items())
|
||||
sorted_reasons.sort(key=lambda x: len(x[1]), reverse=True)
|
||||
|
||||
return reasons_dict, sorted_reasons
|
||||
|
||||
@classmethod
|
||||
def summary(cls, history: list[Self]) -> str:
|
||||
reason_dict, reason_list = cls.classify(history)
|
||||
|
||||
return "\n".join(
|
||||
[
|
||||
f"{name} ({len(reasons)}):\n\t" + "\n\t".join(reasons)
|
||||
for name, reasons in reason_list
|
||||
]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def json_report(cls, history: list[Self]) -> str:
|
||||
reason_dict, sorted_reasons = cls.classify(history)
|
||||
reason_dict["count"] = {k: len(v) for k, v in sorted_reasons}
|
||||
serialized = cls.serialize({cls.SHORT_NAME: reason_dict})
|
||||
|
||||
return f"{PREFIX}{serialized}{SUFFIX}"
|
||||
|
||||
@classmethod
|
||||
def restore_from_string(cls, serialized: str) -> list[Self]:
|
||||
# This method is the inverse of json_report
|
||||
|
||||
from paddle.jit.sot.utils import exceptions
|
||||
|
||||
history = []
|
||||
obj = cls.deserialize(serialized)[cls.SHORT_NAME]
|
||||
obj.pop("count")
|
||||
|
||||
for classname in obj:
|
||||
ReasonClass = getattr(exceptions, classname, None)
|
||||
for reason in obj[classname]:
|
||||
history.append(cls(ReasonClass(reason_str=reason)))
|
||||
|
||||
return history
|
||||
|
||||
@staticmethod
|
||||
def collect_break_graph_reason(reason: BreakGraphReasonBase):
|
||||
if not InfoCollector().need_collect(BreakGraphReasonInfo):
|
||||
return
|
||||
|
||||
InfoCollector().attach(BreakGraphReasonInfo, reason)
|
||||
|
||||
|
||||
class SubGraphInfo(InfoBase):
|
||||
SHORT_NAME = "subgraph_info"
|
||||
TYPE = InfoType.STEP_INFO
|
||||
|
||||
def __init__(self, graph: str, op_num: int, sir_name: str):
|
||||
# NOTE: All data should be serializable
|
||||
super().__init__()
|
||||
self.graph = graph
|
||||
self.op_num = op_num
|
||||
self.sir_name = sir_name
|
||||
|
||||
def __str__(self):
|
||||
return (
|
||||
f"[SIR Name] {self.sir_name} [OpNum] {self.op_num}\n{self.graph}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def summary(cls, history: list[Self]) -> str:
|
||||
num_of_subgraph = len(history)
|
||||
sum_of_op_num = sum(item.op_num for item in history)
|
||||
|
||||
need_details = "details" in ENV_SOT_COLLECT_INFO.get().get(
|
||||
cls.SHORT_NAME, []
|
||||
)
|
||||
|
||||
details = ""
|
||||
if need_details:
|
||||
details = "\n".join(
|
||||
[
|
||||
f"[SubGraphIdx] {idx} {info}"
|
||||
for idx, info in enumerate(map(str, history))
|
||||
]
|
||||
)
|
||||
|
||||
summary = f"[Number of subgraph] {num_of_subgraph} [Sum of opnum] {sum_of_op_num}"
|
||||
|
||||
return f"{summary}\n{details}"
|
||||
|
||||
@classmethod
|
||||
def json_report(cls, history: list[Self]) -> str:
|
||||
need_details = "details" in ENV_SOT_COLLECT_INFO.get().get(
|
||||
cls.SHORT_NAME, []
|
||||
)
|
||||
|
||||
aggregated_info_list = []
|
||||
for idx, record in enumerate(history):
|
||||
entry_data = {}
|
||||
|
||||
entry_data["SIR_name"] = record.sir_name
|
||||
entry_data["OpNum"] = record.op_num
|
||||
entry_data["Graph"] = ""
|
||||
if need_details:
|
||||
entry_data["Graph"] = str(record.graph)
|
||||
aggregated_info_list.append(entry_data)
|
||||
|
||||
serialized = cls.serialize({cls.SHORT_NAME: aggregated_info_list})
|
||||
|
||||
return f"{PREFIX}{serialized}{SUFFIX}"
|
||||
|
||||
@classmethod
|
||||
def restore_from_string(cls, serialized: str) -> list[Self]:
|
||||
# This method is the inverse of json_report
|
||||
|
||||
history = []
|
||||
obj = cls.deserialize(serialized)[cls.SHORT_NAME]
|
||||
|
||||
for entry in obj:
|
||||
history.append(
|
||||
SubGraphInfo(
|
||||
graph=entry["Graph"],
|
||||
op_num=entry["OpNum"],
|
||||
sir_name=entry["SIR_name"],
|
||||
)
|
||||
)
|
||||
|
||||
return history
|
||||
|
||||
def __eq__(self, other):
|
||||
need_graph_equal = "details" in ENV_SOT_COLLECT_INFO.get().get(
|
||||
self.SHORT_NAME, []
|
||||
)
|
||||
|
||||
graph_equal_or_not = True
|
||||
if need_graph_equal:
|
||||
graph_equal_or_not = self.graph == other.graph
|
||||
|
||||
return (
|
||||
graph_equal_or_not
|
||||
and self.op_num == other.op_num
|
||||
and self.sir_name == other.sir_name
|
||||
)
|
||||
@@ -0,0 +1,155 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .utils import hashable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
BinaryOp = Callable[[Any, Any], Any]
|
||||
UnaryOp = Callable[[Any], Any]
|
||||
|
||||
|
||||
INPLACE_BINARY_OPS_TO_MAGIC_NAMES: dict[BinaryOp, tuple[str, BinaryOp]] = {
|
||||
# inplace op fn: (magic name, non-inplace op fn)
|
||||
operator.iadd: ("__iadd__", operator.add),
|
||||
operator.iand: ("__iand__", operator.and_),
|
||||
operator.iconcat: ("__iconcat__", operator.concat),
|
||||
operator.ifloordiv: ("__ifloordiv__", operator.floordiv),
|
||||
operator.ilshift: ("__ilshift__", operator.lshift),
|
||||
operator.imatmul: ("__imatmul__", operator.matmul),
|
||||
operator.imod: ("__imod__", operator.mod),
|
||||
operator.imul: ("__imul__", operator.mul),
|
||||
operator.ior: ("__ior__", operator.or_),
|
||||
operator.ipow: ("__ipow__", operator.pow),
|
||||
operator.irshift: ("__irshift__", operator.rshift),
|
||||
operator.isub: ("__isub__", operator.sub),
|
||||
operator.itruediv: ("__itruediv__", operator.truediv),
|
||||
operator.ixor: ("__ixor__", operator.xor),
|
||||
}
|
||||
|
||||
NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES: dict[
|
||||
BinaryOp, tuple[str, str | None]
|
||||
] = {
|
||||
# op fn: (magic name, reverse magic name)
|
||||
operator.add: ("__add__", "__radd__"),
|
||||
operator.and_: ("__and__", "__rand__"),
|
||||
operator.contains: ("__contains__", None),
|
||||
operator.delitem: ("__delitem__", None),
|
||||
operator.eq: ("__eq__", "__eq__"),
|
||||
operator.floordiv: ("__floordiv__", "__rfloordiv__"),
|
||||
operator.ge: ("__ge__", "__le__"),
|
||||
operator.getitem: ("__getitem__", None),
|
||||
operator.gt: ("__gt__", "__lt__"),
|
||||
operator.le: ("__le__", "__ge__"),
|
||||
operator.lshift: ("__lshift__", "__rlshift__"),
|
||||
operator.lt: ("__lt__", "__gt__"),
|
||||
operator.matmul: ("__matmul__", "__rmatmul__"),
|
||||
operator.mod: ("__mod__", "__rmod__"),
|
||||
operator.mul: ("__mul__", "__rmul__"),
|
||||
operator.ne: ("__ne__", "__ne__"),
|
||||
operator.or_: ("__or__", "__ror__"),
|
||||
operator.pow: ("__pow__", "__rpow__"),
|
||||
operator.rshift: ("__rshift__", "__rrshift__"),
|
||||
operator.sub: ("__sub__", "__rsub__"),
|
||||
operator.truediv: ("__truediv__", "__rtruediv__"),
|
||||
operator.xor: ("__xor__", "__rxor__"),
|
||||
}
|
||||
|
||||
UNARY_OPS_TO_MAGIC_NAMES: dict[UnaryOp, str] = {
|
||||
operator.neg: "__neg__",
|
||||
operator.invert: "__invert__",
|
||||
operator.pos: "__pos__",
|
||||
operator.abs: "__abs__",
|
||||
operator.index: "__index__",
|
||||
operator.inv: "__inv__",
|
||||
operator.truth: "__bool__",
|
||||
bool: "__bool__",
|
||||
abs: "__abs__",
|
||||
float: "__float__",
|
||||
len: "__len__",
|
||||
int: "__int__",
|
||||
complex: "__complex__",
|
||||
}
|
||||
# TODO(SigureMo): support any, all, sum
|
||||
|
||||
|
||||
INPLACE_BINARY_OPS = set(INPLACE_BINARY_OPS_TO_MAGIC_NAMES.keys())
|
||||
NON_INPLACE_BINARY_OPS = set(NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES.keys())
|
||||
BINARY_OPS = INPLACE_BINARY_OPS | NON_INPLACE_BINARY_OPS
|
||||
UNARY_OPS = set(UNARY_OPS_TO_MAGIC_NAMES.keys())
|
||||
|
||||
|
||||
# NOTE: Both operator.pow and operator.ipow should be considered for inclusion in this list,
|
||||
# as they raise ZeroDivisionError when evaluating 0^n where n < 0 (division by zero).
|
||||
NEED_GUARD_ZERO_DIVISION_ERROR_OPS: list[BinaryOp] = [
|
||||
operator.floordiv,
|
||||
operator.truediv,
|
||||
operator.mod,
|
||||
operator.ifloordiv,
|
||||
operator.itruediv,
|
||||
operator.imod,
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class MagicMethod:
|
||||
name: str
|
||||
is_inplace: bool = False
|
||||
is_reverse: bool = False
|
||||
|
||||
|
||||
def magic_method_builtin_dispatch(fn: BinaryOp | UnaryOp) -> list[MagicMethod]:
|
||||
if not hashable(fn):
|
||||
return []
|
||||
if fn in INPLACE_BINARY_OPS:
|
||||
inplace_magic_name, non_inplace_op = INPLACE_BINARY_OPS_TO_MAGIC_NAMES[
|
||||
fn
|
||||
]
|
||||
return [
|
||||
MagicMethod(inplace_magic_name, is_inplace=True),
|
||||
*magic_method_builtin_dispatch(non_inplace_op),
|
||||
]
|
||||
elif fn in NON_INPLACE_BINARY_OPS:
|
||||
magic_name, reverse_magic_name = NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES[
|
||||
fn
|
||||
]
|
||||
magic_methods = [MagicMethod(magic_name)]
|
||||
if reverse_magic_name is not None:
|
||||
magic_methods.append(
|
||||
MagicMethod(reverse_magic_name, is_reverse=True)
|
||||
)
|
||||
return magic_methods
|
||||
elif fn in UNARY_OPS:
|
||||
magic_name = UNARY_OPS_TO_MAGIC_NAMES[fn]
|
||||
return [MagicMethod(magic_name)]
|
||||
return []
|
||||
|
||||
|
||||
def non_inplace_op_to_inplace_op(
|
||||
fn: BinaryOp,
|
||||
) -> BinaryOp | None:
|
||||
for inplace_op, (
|
||||
_,
|
||||
non_inplace_op,
|
||||
) in INPLACE_BINARY_OPS_TO_MAGIC_NAMES.items():
|
||||
if fn is non_inplace_op:
|
||||
return inplace_op
|
||||
return None
|
||||
@@ -0,0 +1,25 @@
|
||||
# Copyright (c) 2025 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.
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
|
||||
NUMPY_API_SUPPORTED_DICT = {
|
||||
np.add: paddle.add,
|
||||
np.subtract: paddle.subtract,
|
||||
np.multiply: paddle.multiply,
|
||||
np.divide: paddle.divide,
|
||||
np.equal: paddle.equal,
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
import inspect
|
||||
|
||||
import paddle
|
||||
|
||||
|
||||
def is_inplace_api(func):
|
||||
inplace_apis = {paddle.static.setitem}
|
||||
return func in inplace_apis
|
||||
|
||||
|
||||
def get_tensor_methods():
|
||||
return [
|
||||
member_name
|
||||
for member_name, member in inspect.getmembers(paddle.pir.Value)
|
||||
if inspect.isfunction(member) or inspect.ismethoddescriptor(member)
|
||||
]
|
||||
|
||||
|
||||
def get_paddle_api():
|
||||
modules = [
|
||||
paddle,
|
||||
paddle.nn.functional,
|
||||
paddle.nn.quant,
|
||||
paddle.incubate.nn.functional,
|
||||
paddle.linalg,
|
||||
paddle.signal,
|
||||
paddle.fft,
|
||||
paddle.vision.ops,
|
||||
paddle.metric,
|
||||
paddle.geometric,
|
||||
]
|
||||
distributed_apis = [
|
||||
paddle.distributed.all_reduce,
|
||||
paddle.distributed.shard_tensor,
|
||||
paddle.distributed.reshard,
|
||||
paddle.distributed.all_gather,
|
||||
paddle.distributed.alltoall,
|
||||
paddle.distributed.barrier,
|
||||
paddle.distributed.recv,
|
||||
paddle.distributed.send,
|
||||
paddle.distributed.broadcast,
|
||||
paddle.distributed.unshard_dtensor,
|
||||
paddle.distributed.auto_parallel.api.dtensor_to_local,
|
||||
paddle.distributed.auto_parallel.api.dtensor_from_local,
|
||||
paddle.distributed.auto_parallel.api.moe_global_mesh_tensor,
|
||||
paddle.distributed.auto_parallel.api.moe_sub_mesh_tensors,
|
||||
]
|
||||
special_paddle_apis = [
|
||||
paddle.tensor.fill_constant,
|
||||
paddle.tensor.top_p_sampling,
|
||||
]
|
||||
non_operator_related_apis = [
|
||||
paddle.in_dynamic_mode,
|
||||
paddle.save,
|
||||
paddle.load,
|
||||
paddle.get_cuda_rng_state,
|
||||
paddle.set_rng_state,
|
||||
paddle.set_cuda_rng_state,
|
||||
paddle.get_rng_state,
|
||||
paddle.set_default_dtype,
|
||||
paddle.check_shape,
|
||||
paddle.summary,
|
||||
paddle.finfo,
|
||||
paddle.iinfo,
|
||||
paddle.enable_static,
|
||||
paddle.disable_static,
|
||||
paddle.is_grad_enabled,
|
||||
]
|
||||
# TODO: users should not call static_apis, but we need to use, so add static_apis here temporary
|
||||
static_apis = [paddle.static.setitem, paddle.static.accuracy]
|
||||
paddle_api_list = []
|
||||
for module in modules:
|
||||
for fn_name in getattr(module, "__all__", []):
|
||||
fn = getattr(module, fn_name)
|
||||
if inspect.isfunction(fn):
|
||||
paddle_api_list.append(fn)
|
||||
return list(
|
||||
set(special_paddle_apis)
|
||||
| set(distributed_apis)
|
||||
| set(static_apis)
|
||||
| set(paddle_api_list) - set(non_operator_related_apis)
|
||||
)
|
||||
|
||||
|
||||
paddle_api_list = get_paddle_api()
|
||||
|
||||
# TODO(Aurelius84): It seems that we use it to judge 'in_paddle_module()'.
|
||||
# Bug what does 'is_paddle_module' really means? Is all paddle.xx sub module
|
||||
# considered as paddle module？
|
||||
paddle_api_module_prefix = {
|
||||
"paddle.nn.functional",
|
||||
}
|
||||
|
||||
break_graph_functions = set()
|
||||
break_graph_layer_classes = set()
|
||||
break_graph_tensor_method = {
|
||||
'register_hook',
|
||||
'numpy',
|
||||
'clear_gradient',
|
||||
'tolist',
|
||||
'item',
|
||||
# TODO: Browse all possible functions and make prior judgments.
|
||||
}
|
||||
|
||||
not_supported_paddle_layer = {paddle.nn.RNN}
|
||||
|
||||
|
||||
def is_not_supported_paddle_layer(layer_class):
|
||||
return layer_class in not_supported_paddle_layer
|
||||
|
||||
|
||||
def is_break_graph_tensor_methods(method_name):
|
||||
return method_name in break_graph_tensor_method
|
||||
|
||||
|
||||
def add_break_graph_function(fn):
|
||||
break_graph_functions.add(fn)
|
||||
|
||||
|
||||
def add_break_graph_layer_class(layer_class: type[paddle.nn.Layer]):
|
||||
break_graph_layer_classes.add(layer_class)
|
||||
|
||||
|
||||
def is_directly_run_api(api):
|
||||
from .utils import hashable
|
||||
|
||||
if not hashable(api):
|
||||
return False
|
||||
NATIVE_CODE_PURE_FUNCTIONS = {
|
||||
paddle.base.libpaddle.is_compiled_with_avx,
|
||||
paddle.base.libpaddle.is_compiled_with_cuda,
|
||||
paddle.base.libpaddle.is_compiled_with_cudnn_frontend,
|
||||
paddle.base.libpaddle.is_compiled_with_rocm,
|
||||
paddle.base.libpaddle.is_compiled_with_custom_device,
|
||||
paddle.base.libpaddle.is_compiled_with_ipu,
|
||||
paddle.base.libpaddle.is_compiled_with_xpu,
|
||||
paddle.base.libpaddle.is_compiled_with_mkldnn,
|
||||
paddle.base.libpaddle.is_compiled_with_onednn,
|
||||
paddle.base.libpaddle.is_compiled_with_nccl,
|
||||
paddle.base.libpaddle.is_compiled_with_mpi,
|
||||
paddle.base.libpaddle.is_compiled_with_mpi_aware,
|
||||
paddle.base.libpaddle.is_compiled_with_cinn,
|
||||
paddle.base.libpaddle.is_compiled_with_distribute,
|
||||
paddle.base.libpaddle.is_compiled_with_brpc,
|
||||
paddle.base.libpaddle.is_compiled_with_dist,
|
||||
paddle.base.libpaddle.is_compiled_with_flagcx,
|
||||
}
|
||||
|
||||
if hasattr(paddle.base.libpaddle, "get_device_properties"):
|
||||
NATIVE_CODE_PURE_FUNCTIONS.add(
|
||||
paddle.base.libpaddle.get_device_properties
|
||||
)
|
||||
|
||||
return api in NATIVE_CODE_PURE_FUNCTIONS
|
||||
@@ -0,0 +1,541 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import copy
|
||||
import inspect
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
import weakref
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import is_dataclass
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
from weakref import WeakValueDictionary
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.jit.dy2static.utils import (
|
||||
TransformOptions,
|
||||
dataclass_as_dict,
|
||||
dataclass_from_dict,
|
||||
)
|
||||
from paddle.utils import flatten, map_structure
|
||||
|
||||
from .envs import (
|
||||
ENV_SOT_LOG_LEVEL,
|
||||
ENV_SOT_SPECIALIZED_DIM_NUMBERS,
|
||||
ENV_STRICT_MODE,
|
||||
)
|
||||
from .paddle_api_config import (
|
||||
break_graph_functions,
|
||||
paddle_api_list,
|
||||
paddle_api_module_prefix,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from paddle._typing import NestedStructure
|
||||
|
||||
T = TypeVar("T")
|
||||
T1 = TypeVar("T1")
|
||||
T2 = TypeVar("T2")
|
||||
T3 = TypeVar("T3")
|
||||
ConstTypes = (int, float, str, bool, type(None), bytes)
|
||||
|
||||
|
||||
class Singleton(type):
|
||||
_instances: dict[Any, Any] = {}
|
||||
|
||||
def __call__(cls, *args: Any, **kwargs: Any):
|
||||
if cls not in cls._instances:
|
||||
cls._instances[cls] = super().__call__(*args, **kwargs)
|
||||
return cls._instances[cls]
|
||||
|
||||
|
||||
class NameGenerator:
|
||||
def __init__(self, prefix):
|
||||
self.counter = 0
|
||||
self.prefix = prefix
|
||||
|
||||
def next(self):
|
||||
name = self.prefix + str(self.counter)
|
||||
self.counter += 1
|
||||
return name
|
||||
|
||||
def match_name(self, name: str) -> bool:
|
||||
return name.startswith(self.prefix)
|
||||
|
||||
|
||||
class SymbolRegistry:
|
||||
def __init__(self):
|
||||
self.symbol_generator = NameGenerator(prefix="___t_")
|
||||
self.tmp_names_record = OrderedDict()
|
||||
self.declared_symbols: set[str] = set()
|
||||
self.symbol_table = {}
|
||||
|
||||
def next_symbol(self) -> str:
|
||||
return self.symbol_generator.next()
|
||||
|
||||
def request_symbol(self, expr: str) -> str:
|
||||
if expr in self.symbol_table:
|
||||
return self.symbol_table[expr]
|
||||
symbol = self.next_symbol()
|
||||
self.symbol_table[expr] = symbol
|
||||
return symbol
|
||||
|
||||
def gen_expr(self, expr: str, gen_expr_fn):
|
||||
symbol = self.symbol_table[expr]
|
||||
if symbol in self.declared_symbols:
|
||||
return symbol
|
||||
self.declared_symbols.add(symbol)
|
||||
return f"({symbol} := ({gen_expr_fn()}))"
|
||||
|
||||
|
||||
_symbol_registry = SymbolRegistry()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def switch_symbol_registry():
|
||||
global _symbol_registry
|
||||
original_registry = _symbol_registry
|
||||
_symbol_registry = SymbolRegistry()
|
||||
yield
|
||||
_symbol_registry = original_registry
|
||||
|
||||
|
||||
def current_symbol_registry():
|
||||
global _symbol_registry
|
||||
return _symbol_registry
|
||||
|
||||
|
||||
class ResumeFnNameFactory(metaclass=Singleton):
|
||||
def __init__(self) -> None:
|
||||
self.gen = NameGenerator('resume_')
|
||||
|
||||
def next(self):
|
||||
name = self.gen.next()
|
||||
return name
|
||||
|
||||
|
||||
class SIRToCodeMap(metaclass=Singleton):
|
||||
def __init__(self):
|
||||
self._map = {}
|
||||
|
||||
def register(self, sir, code):
|
||||
self._map[sir.name] = code
|
||||
|
||||
def get(self, sir):
|
||||
return self._map.get(sir.name)
|
||||
|
||||
|
||||
def log(level, *args):
|
||||
cur_level = ENV_SOT_LOG_LEVEL.get()
|
||||
if level <= cur_level:
|
||||
print(*args, end="", flush=True)
|
||||
|
||||
|
||||
def log_do(level, fn):
|
||||
cur_level = ENV_SOT_LOG_LEVEL.get()
|
||||
if level <= cur_level:
|
||||
fn()
|
||||
|
||||
|
||||
def log_format(level, str, *args):
|
||||
cur_level = ENV_SOT_LOG_LEVEL.get()
|
||||
if level <= cur_level:
|
||||
print(str.format(*args), end="", flush=True)
|
||||
|
||||
|
||||
def log_enabled(level):
|
||||
return level <= ENV_SOT_LOG_LEVEL.get()
|
||||
|
||||
|
||||
@lru_cache
|
||||
def log_once(msg):
|
||||
print(msg, flush=True)
|
||||
|
||||
|
||||
def no_eval_frame(func):
|
||||
def no_eval_frame_func(*args, **kwargs):
|
||||
old_cb = paddle.framework.core.set_eval_frame(None)
|
||||
try:
|
||||
retval = func(*args, **kwargs)
|
||||
except:
|
||||
raise
|
||||
finally:
|
||||
paddle.framework.core.set_eval_frame(old_cb)
|
||||
return retval
|
||||
|
||||
return no_eval_frame_func
|
||||
|
||||
|
||||
def is_comprehensive_name(name):
|
||||
return name in ["<listcomp>", "<dictcomp>", "<setcomp>", "<genexpr>"]
|
||||
|
||||
|
||||
def is_paddle_api(func):
|
||||
if isinstance(func, paddle.nn.Layer): # ignore all the classes
|
||||
return False
|
||||
if hasattr(func, "__self__"): # ignore all the methods
|
||||
return False
|
||||
if inspect.isclass(
|
||||
func
|
||||
): # paddle.Tensor should not be wrapped, but how about other situations?
|
||||
return False
|
||||
return in_paddle_module(func) or func in paddle_api_list
|
||||
|
||||
|
||||
def already_unified_in_dynamic_and_static_graph(fn):
|
||||
if is_paddle_api(fn):
|
||||
return True
|
||||
return not TransformOptions.check_fn_need_transform(
|
||||
fn, TransformOptions.ToStaticMode.SOT
|
||||
)
|
||||
|
||||
|
||||
def need_capture_control_flow(fn):
|
||||
return TransformOptions.check_fn_need_capture_control_flow(fn)
|
||||
|
||||
|
||||
def is_builtin_fn(fn):
|
||||
special_builtin_fns = [weakref.ref]
|
||||
if fn in special_builtin_fns:
|
||||
return True
|
||||
if isinstance(fn, types.BuiltinFunctionType):
|
||||
return True
|
||||
for member_name, member in inspect.getmembers(builtins):
|
||||
if member is fn and isinstance(member, type):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def in_paddle_module(func):
|
||||
if hasattr(func, "__module__"):
|
||||
module_str = func.__module__
|
||||
if module_str is None:
|
||||
return False
|
||||
log(5, "find paddle function with __module__: ", module_str, "\n")
|
||||
if hasattr(func, "__name__"):
|
||||
log(
|
||||
5, " with __name__ : ", func.__name__, "\n"
|
||||
)
|
||||
log(5, " with results : ")
|
||||
for prefix in paddle_api_module_prefix:
|
||||
if module_str.startswith(prefix):
|
||||
log(5, " True\n")
|
||||
return True
|
||||
log(5, " False\n")
|
||||
return False
|
||||
|
||||
|
||||
def is_break_graph_api(func):
|
||||
return func in break_graph_functions
|
||||
|
||||
|
||||
def is_namedtuple_class(cls):
|
||||
if not inspect.isclass(cls):
|
||||
return False
|
||||
if not issubclass(cls, tuple):
|
||||
return False
|
||||
# The signature created by nametuple function
|
||||
namedtuple_attrs = {"_make", "_asdict", "_fields", "_replace"}
|
||||
cls_attrs = set(dir(cls))
|
||||
return namedtuple_attrs.issubset(cls_attrs)
|
||||
|
||||
|
||||
def map_if(
|
||||
*structures: NestedStructure[T1],
|
||||
pred: Callable[[T1], bool],
|
||||
true_fn: Callable[[T1], T2],
|
||||
false_fn: Callable[[T1], T3],
|
||||
) -> NestedStructure[T2 | T3]:
|
||||
def replace(*args):
|
||||
if pred(*args):
|
||||
return true_fn(*args)
|
||||
return false_fn(*args)
|
||||
|
||||
return map_structure(replace, *structures)
|
||||
|
||||
|
||||
def flatten_extend(structure):
|
||||
for item in flatten(structure):
|
||||
if isinstance(item, slice):
|
||||
yield item.start
|
||||
yield item.stop
|
||||
yield item.step
|
||||
else:
|
||||
yield item
|
||||
|
||||
|
||||
def map_if_extend(structure, pred, true_fn, false_fn):
|
||||
"""support extended structures like slice and SliceVariable"""
|
||||
|
||||
def wrapped_pred(x):
|
||||
if isinstance(x, slice):
|
||||
return True
|
||||
if is_dataclass(x) and not isinstance(x, type):
|
||||
return True
|
||||
return pred(x)
|
||||
|
||||
def wrapped_true_fn(x):
|
||||
if isinstance(x, (slice)):
|
||||
l = [x.start, x.stop, x.step]
|
||||
l = map_if_extend(l, pred, true_fn, false_fn)
|
||||
return slice(*l)
|
||||
|
||||
if is_dataclass(x) and not isinstance(x, type):
|
||||
dt_dict = dataclass_as_dict(x)
|
||||
dt_dict = map_if_extend(dt_dict, pred, true_fn, false_fn)
|
||||
return dataclass_from_dict(type(x), dt_dict)
|
||||
|
||||
return true_fn(x)
|
||||
|
||||
return map_if(
|
||||
structure, pred=wrapped_pred, true_fn=wrapped_true_fn, false_fn=false_fn
|
||||
)
|
||||
|
||||
|
||||
def count_if(*structures, pred):
|
||||
def is_true(*args):
|
||||
if pred(*args):
|
||||
return 1
|
||||
return 0
|
||||
|
||||
return sum(flatten(map_structure(is_true, *structures)))
|
||||
|
||||
|
||||
class Cache:
|
||||
def __init__(self, weak=False, copy=False):
|
||||
if not weak:
|
||||
self.cache = {}
|
||||
else:
|
||||
self.cache = WeakValueDictionary()
|
||||
self.hit_num = 0
|
||||
self.copy = copy
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
cache_key = self.key_fn(*args, **kwargs)
|
||||
if not hashable(cache_key):
|
||||
return self.value_fn(*args, **kwargs)
|
||||
if cache_key in self.cache:
|
||||
log(5, "cache hit: ", cache_key, "\n")
|
||||
self.hit_num += 1
|
||||
cache_item = self.cache[cache_key]
|
||||
if self.copy:
|
||||
cache_item = copy.deepcopy(cache_item)
|
||||
return cache_item
|
||||
value = self.value_fn(*args, **kwargs)
|
||||
self.cache[cache_key] = value
|
||||
return value
|
||||
|
||||
def clear(self):
|
||||
self.cache.clear()
|
||||
self.hit_num = 0
|
||||
|
||||
def key_fn(self, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
def value_fn(self, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def execute_time(func):
|
||||
def wrapper(*args, **kwargs):
|
||||
start_time = time.time()
|
||||
result = func(*args, **kwargs)
|
||||
end_time = time.time()
|
||||
execution_time = end_time - start_time
|
||||
print("Execute time:", execution_time)
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def meta_str(shape, dtype, stop_gradient):
|
||||
return f"(shape: {shape}, dtype: {dtype}, stop_gradient: {stop_gradient})"
|
||||
|
||||
|
||||
def is_strict_mode():
|
||||
return ENV_STRICT_MODE.get()
|
||||
|
||||
|
||||
def list_find_index_by_id(li: list[Any], item: Any) -> int:
|
||||
return [id(it) for it in li].index(id(item))
|
||||
|
||||
|
||||
def list_contain_by_id(li: list[Any], item: Any) -> int:
|
||||
return id(item) in [id(it) for it in li]
|
||||
|
||||
|
||||
def get_unbound_method(obj, name):
|
||||
# TODO(dev): Consider the case of patching methods to instances
|
||||
return getattr(obj.__class__, name)
|
||||
|
||||
|
||||
class SotUndefinedVar(metaclass=Singleton):
|
||||
pass
|
||||
|
||||
|
||||
def hashable(obj):
|
||||
try:
|
||||
hash(obj)
|
||||
return True
|
||||
except TypeError as e:
|
||||
return False
|
||||
|
||||
|
||||
def printable(obj):
|
||||
try:
|
||||
str(obj)
|
||||
return True
|
||||
except Exception as e:
|
||||
return False
|
||||
|
||||
|
||||
class StepInfo:
|
||||
BACK_TRACE_STEPS = 20
|
||||
|
||||
def __init__(self):
|
||||
self.step_count = -1
|
||||
|
||||
def need_back_trace(self):
|
||||
return self.step_count < self.BACK_TRACE_STEPS
|
||||
|
||||
|
||||
class StepInfoManager(metaclass=Singleton):
|
||||
def __init__(self):
|
||||
self.step_record = {}
|
||||
self.current_code = None
|
||||
self.current_step_info = None
|
||||
|
||||
@contextmanager
|
||||
def step_guard(self, code):
|
||||
try:
|
||||
old_code = self.current_code
|
||||
old_info = self.current_step_info
|
||||
|
||||
self.current_code = code
|
||||
if code not in self.step_record:
|
||||
self.step_record[code] = StepInfo()
|
||||
self.current_step_info = self.step_record[code]
|
||||
|
||||
self.current_step_info.step_count += 1
|
||||
yield
|
||||
finally:
|
||||
self.current_code = old_code
|
||||
self.current_step_info = old_info
|
||||
|
||||
@property
|
||||
def need_back_trace(self):
|
||||
return (
|
||||
self.current_step_info is not None
|
||||
and self.current_step_info.need_back_trace()
|
||||
)
|
||||
|
||||
@property
|
||||
def current_step(self):
|
||||
return self.current_step_info.step_count
|
||||
|
||||
def clear(self):
|
||||
self.step_record.clear()
|
||||
self.current_code = None
|
||||
self.current_step = -1
|
||||
|
||||
|
||||
def get_api_fullname(api):
|
||||
api_name = api.__name__
|
||||
module_str = api.__module__
|
||||
while len(module_str) > 0:
|
||||
if module_str not in sys.modules:
|
||||
return api_name
|
||||
module = sys.modules[module_str]
|
||||
if hasattr(module, api_name):
|
||||
return module_str + "." + api_name
|
||||
module_str = module_str.rpartition(".")[0]
|
||||
return None
|
||||
|
||||
|
||||
def get_numpy_ufuncs():
|
||||
ufuncs = [
|
||||
ufunc
|
||||
for _, ufunc in inspect.getmembers(
|
||||
np, lambda member: isinstance(member, np.ufunc)
|
||||
)
|
||||
]
|
||||
unary_ufuncs = filter(lambda ufunc: ufunc.nin == 1, ufuncs)
|
||||
binary_ufuncs = filter(lambda ufunc: ufunc.nin == 2, ufuncs)
|
||||
return list(unary_ufuncs), list(binary_ufuncs)
|
||||
|
||||
|
||||
def do_until_stop_iteration(fn: Callable[[], T]) -> list[T]:
|
||||
from paddle.jit.sot.utils.exceptions import SotCapturedStopIteration
|
||||
|
||||
res = []
|
||||
while True:
|
||||
try:
|
||||
res.append(fn())
|
||||
except SotCapturedStopIteration:
|
||||
break
|
||||
return res
|
||||
|
||||
|
||||
def update_list_inplace(
|
||||
original_list: list[T], new_contents: list[T]
|
||||
) -> list[T]:
|
||||
original_list.clear()
|
||||
original_list.extend(new_contents)
|
||||
return original_list
|
||||
|
||||
|
||||
def get_obj_stable_repr(obj) -> str:
|
||||
if hasattr(obj, '__qualname__'):
|
||||
return obj.__qualname__
|
||||
if hasattr(obj, '__name__'):
|
||||
return obj.__name__
|
||||
|
||||
class_name = obj.__class__.__name__
|
||||
|
||||
# If module is available and not __main__, include it
|
||||
if hasattr(obj, "__class__") and hasattr(obj.__class__, "__module__"):
|
||||
module = obj.__class__.__module__
|
||||
if module not in ("__main__", "builtins"):
|
||||
return f"{module}.{class_name}()"
|
||||
|
||||
return f"{class_name}()"
|
||||
|
||||
|
||||
def get_min_non_specialized_number() -> int:
|
||||
specialized_dim_numbers_raw_str = (
|
||||
ENV_SOT_SPECIALIZED_DIM_NUMBERS.get().lower()
|
||||
)
|
||||
assert specialized_dim_numbers_raw_str in [
|
||||
"no",
|
||||
"0",
|
||||
"01",
|
||||
], f"Unsupported specialized_dim_numbers: {specialized_dim_numbers_raw_str}"
|
||||
to_min_non_specialized_number = {
|
||||
# specialized numbers, minimum non-specialized number
|
||||
"no": 0,
|
||||
"0": 1,
|
||||
"01": 2,
|
||||
}
|
||||
return to_min_non_specialized_number[specialized_dim_numbers_raw_str]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,280 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Generic, TypeVar
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable, Iterator
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class OrderedSet(Generic[T]):
|
||||
"""
|
||||
A set that preserves the order of insertion.
|
||||
"""
|
||||
|
||||
_data: dict[T, None]
|
||||
|
||||
def __init__(self, items: Iterable[T] | None = None):
|
||||
"""
|
||||
Examples:
|
||||
>>> s = OrderedSet([1, 2, 3])
|
||||
>>> s
|
||||
OrderedSet(1, 2, 3)
|
||||
>>> s = OrderedSet()
|
||||
>>> s
|
||||
OrderedSet()
|
||||
"""
|
||||
self._data = dict.fromkeys(items) if items is not None else {}
|
||||
|
||||
def __iter__(self) -> Iterator[T]:
|
||||
"""
|
||||
Examples:
|
||||
>>> s = OrderedSet([1, 2, 3])
|
||||
>>> for item in s:
|
||||
... print(item)
|
||||
1
|
||||
2
|
||||
3
|
||||
"""
|
||||
return iter(self._data)
|
||||
|
||||
def __or__(self, other: OrderedSet[T]) -> OrderedSet[T]:
|
||||
"""
|
||||
Union two sets.
|
||||
|
||||
Args:
|
||||
other: Another set to be unioned.
|
||||
|
||||
Returns:
|
||||
The union of two sets.
|
||||
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([2, 3, 4])
|
||||
>>> s1 | s2
|
||||
OrderedSet(1, 2, 3, 4)
|
||||
"""
|
||||
return OrderedSet(list(self) + list(other))
|
||||
|
||||
def __ior__(self, other: OrderedSet[T]):
|
||||
"""
|
||||
Union two sets in place.
|
||||
|
||||
Args:
|
||||
other: Another set to be unioned.
|
||||
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([2, 3, 4])
|
||||
>>> s1 |= s2
|
||||
>>> s1
|
||||
OrderedSet(1, 2, 3, 4)
|
||||
"""
|
||||
self._data.update(dict.fromkeys(other))
|
||||
return self
|
||||
|
||||
def __and__(self, other: OrderedSet[T]) -> OrderedSet[T]:
|
||||
"""
|
||||
Intersect two sets.
|
||||
|
||||
Args:
|
||||
other: Another set to be intersected.
|
||||
|
||||
Returns:
|
||||
The intersection of two sets.
|
||||
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([2, 3, 4])
|
||||
>>> s1 & s2
|
||||
OrderedSet(2, 3)
|
||||
"""
|
||||
return OrderedSet([item for item in self if item in other])
|
||||
|
||||
def __iand__(self, other: OrderedSet[T]):
|
||||
"""
|
||||
Intersect two sets in place.
|
||||
|
||||
Args:
|
||||
other: Another set to be intersected.
|
||||
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([2, 3, 4])
|
||||
>>> s1 &= s2
|
||||
>>> s1
|
||||
OrderedSet(2, 3)
|
||||
"""
|
||||
self._data = {item: None for item in self if item in other}
|
||||
return self
|
||||
|
||||
def __sub__(self, other: OrderedSet[T]) -> OrderedSet[T]:
|
||||
"""
|
||||
Subtract two sets.
|
||||
|
||||
Args:
|
||||
other: Another set to be subtracted.
|
||||
|
||||
Returns:
|
||||
The subtraction of two sets.
|
||||
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([2, 3, 4])
|
||||
>>> s1 - s2
|
||||
OrderedSet(1)
|
||||
"""
|
||||
return OrderedSet([item for item in self if item not in other])
|
||||
|
||||
def __isub__(self, other: OrderedSet[T]):
|
||||
"""
|
||||
Subtract two sets in place.
|
||||
|
||||
Args:
|
||||
other: Another set to be subtracted.
|
||||
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([2, 3, 4])
|
||||
>>> s1 -= s2
|
||||
>>> s1
|
||||
OrderedSet(1)
|
||||
"""
|
||||
self._data = {item: None for item in self if item not in other}
|
||||
return self
|
||||
|
||||
def __xor__(self, other: OrderedSet[T]) -> OrderedSet[T]:
|
||||
"""
|
||||
Symmetric difference of two sets.
|
||||
|
||||
Args:
|
||||
other: Another set to be xor'ed.
|
||||
|
||||
Returns:
|
||||
The symmetric difference of two sets.
|
||||
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([2, 3, 4])
|
||||
>>> s1 ^ s2
|
||||
OrderedSet(1, 4)
|
||||
"""
|
||||
return OrderedSet(
|
||||
[item for item in self if item not in other]
|
||||
) | OrderedSet([item for item in other if item not in self])
|
||||
|
||||
def __ixor__(self, other: OrderedSet[T]):
|
||||
"""
|
||||
Symmetric difference of two sets in place.
|
||||
|
||||
Args:
|
||||
other: Another set to be xor'ed.
|
||||
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([2, 3, 4])
|
||||
>>> s1 ^= s2
|
||||
>>> s1
|
||||
OrderedSet(1, 4)
|
||||
"""
|
||||
self._data = {
|
||||
**{item: None for item in self if item not in other},
|
||||
**{item: None for item in other if item not in self},
|
||||
}
|
||||
return self
|
||||
|
||||
def add(self, item: T):
|
||||
"""
|
||||
Add an item to the set.
|
||||
|
||||
Args:
|
||||
item: The item to be added.
|
||||
|
||||
Examples:
|
||||
>>> s = OrderedSet([1, 2, 3])
|
||||
>>> s.add(4)
|
||||
>>> s
|
||||
OrderedSet(1, 2, 3, 4)
|
||||
"""
|
||||
self._data.setdefault(item)
|
||||
|
||||
def remove(self, item: T):
|
||||
"""
|
||||
Remove an item from the set.
|
||||
|
||||
Args:
|
||||
item: The item to be removed.
|
||||
|
||||
Examples:
|
||||
>>> s = OrderedSet([1, 2, 3])
|
||||
>>> s.remove(2)
|
||||
>>> s
|
||||
OrderedSet(1, 3)
|
||||
"""
|
||||
del self._data[item]
|
||||
|
||||
def __contains__(self, item: T) -> bool:
|
||||
"""
|
||||
Examples:
|
||||
>>> s = OrderedSet([1, 2, 3])
|
||||
>>> 1 in s
|
||||
True
|
||||
>>> 4 in s
|
||||
False
|
||||
"""
|
||||
return item in self._data
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""
|
||||
Examples:
|
||||
>>> s = OrderedSet([1, 2, 3])
|
||||
>>> len(s)
|
||||
3
|
||||
"""
|
||||
return len(self._data)
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
"""
|
||||
Examples:
|
||||
>>> s = OrderedSet([1, 2, 3])
|
||||
>>> bool(s)
|
||||
True
|
||||
>>> s = OrderedSet()
|
||||
>>> bool(s)
|
||||
False
|
||||
"""
|
||||
return bool(self._data)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
"""
|
||||
Examples:
|
||||
>>> s1 = OrderedSet([1, 2, 3])
|
||||
>>> s2 = OrderedSet([1, 2, 3])
|
||||
>>> s1 == s2
|
||||
True
|
||||
>>> s3 = OrderedSet([3, 2, 1])
|
||||
>>> s1 == s3
|
||||
False
|
||||
"""
|
||||
if not isinstance(other, OrderedSet):
|
||||
return NotImplemented
|
||||
return list(self) == list(other)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
data_repr = ", ".join(map(repr, self._data))
|
||||
return f"OrderedSet({data_repr})"
|
||||
Reference in New Issue
Block a user