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