chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
Reference in New Issue
Block a user