def make_control_flow_handlers(self, cont_n, status_n, expected_return,
has_cont, has_break):
'''
Create the statements in charge of gathering control flow information
for the static_if result, and executes the expected control flow
instruction
'''
if expected_return:
assign = cont_ass = [ast.Assign(
[ast.Tuple(expected_return, ast.Store())],
ast.Name(cont_n, ast.Load(), None, None), None)]
else:
assign = cont_ass = []
if has_cont:
cmpr = ast.Compare(ast.Name(status_n, ast.Load(), None, None),
[ast.Eq()], [ast.Constant(LOOP_CONT, None)])
cont_ass = [ast.If(cmpr,
deepcopy(assign) + [ast.Continue()],
cont_ass)]
if has_break:
cmpr = ast.Compare(ast.Name(status_n, ast.Load(), None, None),
[ast.Eq()], [ast.Constant(LOOP_BREAK, None)])
cont_ass = [ast.If(cmpr,
deepcopy(assign) + [ast.Break()],
cont_ass)]
return cont_ass
def visit_For(self, node):
modified_node = self.generic_visit(node)
continue_flags = self.for_continue_stack.pop()
for flag in continue_flags:
node.body.insert(
0,
gast.Assign(targets=[
gast.Name(id=flag, ctx=gast.Store(), annotation=None)
],
value=gast.NameConstant(value=False)))
breaked_flags = self.for_breaked_stack.pop()
bool_values = []
for flag in breaked_flags:
node.body.insert(
0,
gast.Assign(targets=[
gast.Name(id=flag, ctx=gast.Store(), annotation=None)
],
value=gast.NameConstant(value=False)))
bool_values.append(
gast.Name(id=flag, ctx=gast.Load(), annotation=None))
if len(bool_values) > 0:
if len(bool_values) == 1:
cond = bool_values[0]
elif len(bool_values) > 1:
cond = gast.BoolOp(op=gast.Or(), values=bool_values)
if isinstance(modified_node, gast.For):
modified_node.body.append(
gast.Assign(targets=[
gast.Name(id=self.keepgoing_flag,
ctx=gast.Store(),
annotation=None)
],
value=gast.UnaryOp(op=gast.Not(),
operand=cond)))
modified_node.body.append(
gast.If(test=cond, body=[gast.Break()], orelse=[]))
elif isinstance(modified_node, gast.If):
if isinstance(modified_node.body[0], gast.For):
modified_node.body[0].body.append(
gast.Assign(targets=[
gast.Name(id=self.keepgoing_flag,
ctx=gast.Store(),
annotation=None)
],
value=gast.UnaryOp(op=gast.Not(),
operand=cond)))
modified_node.body[0].body.append(
gast.If(test=cond, body=[gast.Break()], orelse=[]))
return modified_node
def generic_visit(self, node):
if isinstance(node, gast.stmt):
if self.stack_has_flags(
self.for_continued_stack) or self.stack_has_flags(
self.for_breaked_stack) or self.stack_has_flags(
self.func_returned_stack):
bool_values = []
if self.stack_has_flags(self.for_continued_stack):
continued_id = len(self.for_continued_stack)
bool_values.append(
gast.UnaryOp(op=gast.Not(),
operand=gast.Name(id=self.continued_flag +
str(continued_id),
ctx=gast.Load(),
annotation=None,
type_comment=None)))
if self.stack_has_flags(self.for_breaked_stack):
breaked_id = len(self.for_breaked_stack)
bool_values.append(
gast.UnaryOp(op=gast.Not(),
operand=gast.Name(id=self.breaked_flag +
str(breaked_id),
ctx=gast.Load(),
annotation=None,
type_comment=None)))
if self.stack_has_flags(self.func_returned_stack):
returned_id = len(self.func_returned_stack)
bool_values.append(
gast.UnaryOp(op=gast.Not(),
operand=gast.Name(id=self.returned_flag +
str(returned_id),
ctx=gast.Load(),
annotation=None,
type_comment=None)))
if isinstance(node, gast.For):
self.for_continued_stack.append(False)
self.for_breaked_stack.append(False)
elif isinstance(node, gast.FunctionDef):
self.func_returned_stack.append(False)
modified_node = super().generic_visit(node)
if len(bool_values) == 1:
cond = bool_values[0]
else:
cond = gast.BoolOp(op=gast.And(), values=bool_values)
replacement = gast.If(test=cond,
body=[modified_node],
orelse=[])
ret = gast.copy_location(replacement, node)
else:
if isinstance(node, gast.For):
self.for_continued_stack.append(False)
self.for_breaked_stack.append(False)
elif isinstance(node, gast.FunctionDef):
self.func_returned_stack.append(False)
ret = super().generic_visit(node)
else:
ret = super().generic_visit(node)
return ret
def visit_For(self, node):
modified_node = self.generic_visit(node)
continued_id = len(self.for_continued_stack)
continued_flags = self.for_continued_stack.pop()
if continued_flags:
node.body.insert(
0,
gast.Assign(targets=[
gast.Name(id=self.continued_flag + str(continued_id),
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=gast.Constant(value=False, kind=None)))
breaked_id = len(self.for_breaked_stack)
breaked_flags = self.for_breaked_stack.pop()
bool_values = []
if breaked_flags:
node.body.insert(
0,
gast.Assign(targets=[
gast.Name(id=self.breaked_flag + str(breaked_id),
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=gast.Constant(value=False, kind=None)))
bool_values.append(
gast.Name(id=self.breaked_flag + str(breaked_id),
ctx=gast.Load(),
annotation=None,
type_comment=None))
if len(self.func_returned_stack) > 0:
returned_id = len(self.func_returned_stack)
returned_flags = self.func_returned_stack[-1]
if returned_flags:
bool_values.append(
gast.Name(id=self.returned_flag + str(returned_id),
ctx=gast.Load(),
annotation=None,
type_comment=None))
if len(bool_values) > 0:
if len(bool_values) == 1:
cond = bool_values[0]
elif len(bool_values) > 1:
cond = gast.BoolOp(op=gast.Or(), values=bool_values)
node.body.append(
gast.Assign(targets=[
gast.Name(id=self.keepgoing_flag,
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=gast.UnaryOp(op=gast.Not(), operand=cond)))
node.body.append(gast.If(test=cond, body=[gast.Break()],
orelse=[]))
return modified_node
def visit_If(self, node):
if self.isompdirective(node.test):
self.visit(ast.Expr(node.test))
return self.visit(
ast.If(ast.Constant(1, None), node.body, node.orelse))
else:
return self.attach_data(node)
def attach_data(self, node):
'''Generic method called for visit_XXXX() with XXXX in
GatherOMPData.statements list
'''
if self.current:
for curr in self.current:
md = OMPDirective(curr)
metadata.add(node, md)
self.current = list()
# add a Pass to hold some directives
for field_name, field in ast.iter_fields(node):
if field_name in GatherOMPData.statement_lists:
if(field and
isinstance(field[-1], ast.Expr) and
self.isompdirective(field[-1].value)):
field.append(ast.Pass())
self.generic_visit(node)
# add an If to hold scoping OpenMP directives
directives = metadata.get(node, OMPDirective)
field_names = {n for n, _ in ast.iter_fields(node)}
has_no_scope = field_names.isdisjoint(GatherOMPData.statement_lists)
if directives and has_no_scope:
# some directives create a scope, but the holding stmt may not
# artificially create one here if needed
sdirective = ''.join(d.s for d in directives)
scoping = ('parallel', 'task', 'section')
if any(s in sdirective for s in scoping):
metadata.clear(node, OMPDirective)
node = ast.If(ast.Num(1), [node], [])
for directive in directives:
metadata.add(node, directive)
return node
def test_unparse(self):
node = gast.If(test=gast.Constant(1, kind=None),
body=[
gast.Assign(targets=[
gast.Name('a',
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=gast.Name('b',
ctx=gast.Load(),
annotation=None,
type_comment=None))
],
orelse=[
gast.Assign(targets=[
gast.Name('a',
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=gast.Constant('c', kind=None))
])
source = parser.unparse(node, indentation=' ')
self.assertEqual(
textwrap.dedent("""
# coding=utf-8
if 1:
a = b
else:
a = 'c'
""").strip(), source.strip())
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
if isinstance(parent_node_of_return, gast.If):
# Prepend control flow boolean nodes such as '__return@1 = False'
node_str = "{} = paddle.jit.dy2static.create_bool_as_type({}, False)".format(
return_name,
ast_to_source_code(parent_node_of_return.test).strip())
assign_false_node = gast.parse(node_str).body[0]
stmt_list[i:i] = [assign_false_node]
return True
def visit_If(self, node):
self.generic_visit(node)
try:
if ast.literal_eval(node.test):
if not metadata.get(node, OMPDirective):
self.update = True
return node.body
else:
if not metadata.get(node, OMPDirective):
self.update = True
return node.orelse
except ValueError:
# not a constant expression
pass
have_body = any(not isinstance(x, ast.Pass) for x in node.body)
have_else = any(not isinstance(x, ast.Pass) for x in node.orelse)
# If the "body" is empty but "else content" is useful, switch branches
# and remove else content
if not have_body and have_else:
test = ast.UnaryOp(op=ast.Not(), operand=node.test)
self.update = True
return ast.If(test=test, body=node.orelse, orelse=list())
# if neither "if" and "else" are useful, keep test if it is not pure
elif not have_body:
self.update = True
if node.test in self.pure_expressions:
return ast.Pass()
else:
node = ast.Expr(value=node.test)
self.generic_visit(node)
return node
def _process_body_item(self, node):
if isinstance(node, gast.Assign) and (node.value.id == 'y'):
if_node = gast.If(
gast.Name(
'x', ctx=gast.Load(), annotation=None, type_comment=None),
[node], [])
return if_node, if_node.body
return node, None
def visit_Compare(self, node):
node = self.generic_visit(node)
if len(node.ops) > 1:
# in case we have more than one compare operator
# we generate an auxiliary function
# that lazily evaluates the needed parameters
imported_ids = self.passmanager.gather(ImportedIds, node, self.ctx)
imported_ids = sorted(imported_ids)
binded_args = [ast.Name(i, ast.Load(), None) for i in imported_ids]
# name of the new function
forged_name = "{0}_compare{1}".format(self.prefix,
len(self.compare_functions))
# call site
call = ast.Call(ast.Name(forged_name, ast.Load(), None),
binded_args, [])
# new function
arg_names = [ast.Name(i, ast.Param(), None) for i in imported_ids]
args = ast.arguments(arg_names, None, [], [], None, [])
body = [] # iteratively fill the body (yeah, feel your body!)
if is_trivially_copied(node.left):
prev_holder = node.left
else:
body.append(
ast.Assign([ast.Name('$0', ast.Store(), None)], node.left))
prev_holder = ast.Name('$0', ast.Load(), None)
for i, exp in enumerate(node.comparators):
if is_trivially_copied(exp):
holder = exp
else:
body.append(
ast.Assign(
[ast.Name('${}'.format(i + 1), ast.Store(), None)],
exp))
holder = ast.Name('${}'.format(i + 1), ast.Load(), None)
cond = ast.Compare(prev_holder, [node.ops[i]], [holder])
body.append(
ast.If(
cond, [ast.Pass()],
[ast.Return(path_to_attr(('__builtin__', 'False')))]))
prev_holder = holder
body.append(ast.Return(path_to_attr(('__builtin__', 'True'))))
forged_fdef = ast.FunctionDef(forged_name, args, body, [], None)
self.compare_functions.append(forged_fdef)
return call
else:
return node
def visit_If(self, node):
self.generic_visit(node)
body_depth = max(n._statement_depth for n in node.body)
orelse_depth = max(n._statement_depth for n in node.orelse)
if orelse_depth >= body_depth:
new_node = gast.If(test=self.invert(node.test),
body=node.orelse,
orelse=node.body)
# ensure newly created node has a depth annotation too.
new_node._statement_depth = node._statement_depth
return new_node
else:
return node
def visit_If(self, node):
node = self.generic_visit(node)
node = self.add_mask(node, node.test)
nodes = [node]
if len(node.orelse) > 0:
test_inverse = gast.Call(
gast.Attribute(node.test, gast.Name('eq', gast.Load(), None),
gast.Load()), [gast.Num(0)], [])
else_node = gast.If(any_active(test_inverse), node.orelse, [])
node.orelse = []
self.add_mask(else_node, test_inverse)
nodes.append(else_node)
node.test = any_active(node.test)
return nodes
def generate_If(self):
"""Generate an If node."""
test = self.generate_Compare()
# Generate true branch statements
body = self.sample_node_list(low=1,
high=N_CONTROLFLOW_STATEMENTS // 2,
generator=self.generate_statement)
# Generate false branch statements
orelse = self.sample_node_list(low=1,
high=N_CONTROLFLOW_STATEMENTS // 2,
generator=self.generate_statement)
node = gast.If(test, body, orelse)
return node
def wrap_in_ifs(node, ifs):
"""
Wrap comprehension content in all possibles if clauses.
Examples
--------
>> [i for i in range(2) if i < 3 if 0 < i]
Becomes
>> for i in range(2):
>> if i < 3:
>> if 0 < i:
>> ... the code from `node` ...
Note the nested ifs clauses.
"""
return reduce(lambda n, if_: ast.If(if_, [n], []), ifs, node)
def _replace_after_node_to_if_in_stmt_list(self, stmt_list, node,
return_name):
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]
return True
def test_ast_to_source(self):
node = gast.If(
test=gast.Num(1),
body=[
gast.Assign(targets=[gast.Name('a', gast.Store(), None)],
value=gast.Name('b', gast.Load(), None))
],
orelse=[
gast.Assign(targets=[gast.Name('a', gast.Store(), None)],
value=gast.Str('c'))
])
source = compiler.ast_to_source(node, indentation=' ')
self.assertEqual(
textwrap.dedent("""
if 1:
a = b
else:
a = 'c'
""").strip(), source.strip())
def generic_visit(self, node):
if isinstance(node, gast.stmt):
if (len(self.for_continue_stack) > 0
and len(self.for_continue_stack[-1]) > 0) or (
len(self.for_breaked_stack) > 0
and len(self.for_breaked_stack[-1]) > 0):
bool_values = []
if (len(self.for_continue_stack) > 0
and len(self.for_continue_stack[-1]) > 0):
for flag in self.for_continue_stack[-1]:
bool_values.append(
gast.UnaryOp(op=gast.Not(),
operand=gast.Name(id=flag,
ctx=gast.Load(),
annotation=None)))
if (len(self.for_breaked_stack) > 0
and len(self.for_breaked_stack[-1]) > 0):
for flag in self.for_breaked_stack[-1]:
bool_values.append(
gast.UnaryOp(op=gast.Not(),
operand=gast.Name(id=flag,
ctx=gast.Load(),
annotation=None)))
if isinstance(node, gast.For):
self.for_continue_stack.append([])
self.for_breaked_stack.append([])
node = super().generic_visit(node)
if len(bool_values) == 1:
cond = bool_values[0]
else:
cond = gast.BoolOp(op=gast.And(), values=bool_values)
replacement = gast.If(test=cond, body=[node], orelse=[])
ret = gast.copy_location(replacement, node)
else:
if isinstance(node, gast.For):
self.for_continue_stack.append([])
self.for_breaked_stack.append([])
ret = super().generic_visit(node)
else:
ret = super().generic_visit(node)
return ret
def make_fake(stmts):
return ast.If(ast.Constant(0, None), stmts, [])
def visit_If(self, node):
self.generic_visit(node)
if node.test not in self.static_expressions:
return node
imported_ids = self.passmanager.gather(ImportedIds, node, self.ctx)
assigned_ids_left = set(
self.passmanager.gather(IsAssigned, self.make_fake(node.body),
self.ctx).keys())
assigned_ids_right = set(
self.passmanager.gather(IsAssigned, self.make_fake(node.orelse),
self.ctx).keys())
assigned_ids_both = assigned_ids_left.union(assigned_ids_right)
imported_ids.update(i for i in assigned_ids_left
if i not in assigned_ids_right)
imported_ids.update(i for i in assigned_ids_right
if i not in assigned_ids_left)
imported_ids = sorted(imported_ids)
assigned_ids = sorted(assigned_ids_both)
true_has_return = self.passmanager.gather(HasReturn,
self.make_fake(node.body),
self.ctx)
false_has_return = self.passmanager.gather(HasReturn,
self.make_fake(node.orelse),
self.ctx)
has_return = true_has_return or false_has_return
func_true = outline(self.true_name(), imported_ids, assigned_ids,
node.body, has_return)
func_false = outline(self.false_name(), imported_ids, assigned_ids,
node.orelse, has_return)
self.new_functions.extend((func_true, func_false))
actual_call = self.make_dispatcher(node.test, func_true, func_false,
imported_ids)
expected_return = [
ast.Name(ii, ast.Load(), None) for ii in assigned_ids
]
if has_return:
n = len(self.new_functions)
fast_return = [
ast.Name("$status{}".format(n), ast.Load(), None),
ast.Name("$return{}".format(n), ast.Load(), None),
ast.Name("$cont{}".format(n), ast.Load(), None)
]
if expected_return:
cont_ass = [
ast.Assign([ast.Tuple(expected_return, ast.Store())],
ast.Name("$cont{}".format(n), ast.Load(), None))
]
else:
cont_ass = []
return [
ast.Assign([ast.Tuple(fast_return, ast.Store())], actual_call),
ast.If(ast.Name("$status{}".format(n), ast.Load(), None), [
ast.Return(
ast.Name("$return{}".format(n), ast.Load(), None))
], cont_ass)
]
elif expected_return:
return ast.Assign([ast.Tuple(expected_return, ast.Store())],
actual_call)
else:
return ast.Expr(actual_call)
def build(test, body, orelse):
return gast.If(test=test, body=body, orelse=orelse)
def visit_If(self, node):
if node.test not in self.static_expressions:
return self.generic_visit(node)
imported_ids = self.gather(ImportedIds, node)
assigned_ids_left = self.escaping_ids(node, node.body)
assigned_ids_right = self.escaping_ids(node, node.orelse)
assigned_ids_both = assigned_ids_left.union(assigned_ids_right)
imported_ids.update(i for i in assigned_ids_left
if i not in assigned_ids_right)
imported_ids.update(i for i in assigned_ids_right
if i not in assigned_ids_left)
imported_ids = sorted(imported_ids)
assigned_ids = sorted(assigned_ids_both)
fbody = self.make_fake(node.body)
true_has_return = self.gather(HasReturn, fbody)
true_has_break = self.gather(HasBreak, fbody)
true_has_cont = self.gather(HasContinue, fbody)
felse = self.make_fake(node.orelse)
false_has_return = self.gather(HasReturn, felse)
false_has_break = self.gather(HasBreak, felse)
false_has_cont = self.gather(HasContinue, felse)
has_return = true_has_return or false_has_return
has_break = true_has_break or false_has_break
has_cont = true_has_cont or false_has_cont
self.generic_visit(node)
func_true = outline(self.true_name(), imported_ids, assigned_ids,
node.body, has_return, has_break, has_cont)
func_false = outline(self.false_name(), imported_ids, assigned_ids,
node.orelse, has_return, has_break, has_cont)
self.new_functions.extend((func_true, func_false))
actual_call = self.make_dispatcher(node.test,
func_true, func_false, imported_ids)
# variable modified within the static_if
expected_return = [ast.Name(ii, ast.Store(), None, None)
for ii in assigned_ids]
self.update = True
# name for various variables resulting from the static_if
n = len(self.new_functions)
status_n = "$status{}".format(n)
return_n = "$return{}".format(n)
cont_n = "$cont{}".format(n)
if has_return:
cfg = self.cfgs[-1]
always_return = all(isinstance(x, (ast.Return, ast.Yield))
for x in cfg[node])
always_return &= true_has_return and false_has_return
fast_return = [ast.Name(status_n, ast.Store(), None, None),
ast.Name(return_n, ast.Store(), None, None),
ast.Name(cont_n, ast.Store(), None, None)]
if always_return:
return [ast.Assign([ast.Tuple(fast_return, ast.Store())],
actual_call, None),
ast.Return(ast.Name(return_n, ast.Load(), None, None))]
else:
cont_ass = self.make_control_flow_handlers(cont_n, status_n,
expected_return,
has_cont, has_break)
cmpr = ast.Compare(ast.Name(status_n, ast.Load(), None, None),
[ast.Eq()], [ast.Constant(EARLY_RET, None)])
return [ast.Assign([ast.Tuple(fast_return, ast.Store())],
actual_call, None),
ast.If(cmpr,
[ast.Return(ast.Name(return_n, ast.Load(),
None, None))],
cont_ass)]
elif has_break or has_cont:
cont_ass = self.make_control_flow_handlers(cont_n, status_n,
expected_return,
has_cont, has_break)
fast_return = [ast.Name(status_n, ast.Store(), None, None),
ast.Name(cont_n, ast.Store(), None, None)]
return [ast.Assign([ast.Tuple(fast_return, ast.Store())],
actual_call, None)] + cont_ass
elif expected_return:
return ast.Assign([ast.Tuple(expected_return, ast.Store())],
actual_call, None)
else:
return ast.Expr(actual_call)
def make_fake(stmts):
return ast.If(ast.Num(0), stmts, [])
def run(self, node):
if isinstance(node, list): # so that this pass can be called on list
self.is_list = True
node = ast.If(ast.Constant(1, None), node, [])
return super(ImportedIds, self).run(node)
