wehub-resource-sync
448 lines
14 KiB
Python
448 lines
14 KiB
Python
# 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
|