forked from M-Labs/artiq
transform: add intermediate assignment removal
This commit is contained in:
parent
f012151506
commit
70cc0d1766
|
@ -2,6 +2,7 @@ import os
|
|||
|
||||
from artiq.transforms.inline import inline
|
||||
from artiq.transforms.lower_units import lower_units
|
||||
from artiq.transforms.remove_inter_assigns import remove_inter_assigns
|
||||
from artiq.transforms.fold_constants import fold_constants
|
||||
from artiq.transforms.unroll_loops import unroll_loops
|
||||
from artiq.transforms.interleave import interleave
|
||||
|
@ -51,6 +52,9 @@ class Core:
|
|||
lower_units(func_def, rpc_map)
|
||||
_debug_unparse("lower_units", func_def)
|
||||
|
||||
remove_inter_assigns(func_def)
|
||||
_debug_unparse("remove_inter_assigns_1", func_def)
|
||||
|
||||
fold_constants(func_def)
|
||||
_debug_unparse("fold_constants_1", func_def)
|
||||
|
||||
|
@ -65,6 +69,9 @@ class Core:
|
|||
self.runtime_env.ref_period)
|
||||
_debug_unparse("lower_time", func_def)
|
||||
|
||||
remove_inter_assigns(func_def)
|
||||
_debug_unparse("remove_inter_assigns_2", func_def)
|
||||
|
||||
fold_constants(func_def)
|
||||
_debug_unparse("fold_constants_2", func_def)
|
||||
|
||||
|
|
|
@ -0,0 +1,145 @@
|
|||
import ast
|
||||
from copy import copy, deepcopy
|
||||
|
||||
|
||||
_replaceable_funcs = {
|
||||
"bool", "int", "float", "round",
|
||||
"int64", "round64", "Fraction",
|
||||
}
|
||||
|
||||
|
||||
def _is_replaceable(value):
|
||||
if isinstance(value, (ast.NameConstant, ast.Num, ast.Str)):
|
||||
return True
|
||||
elif isinstance(value, ast.BinOp):
|
||||
return _is_replaceable(value.left) and _is_replaceable(value.right)
|
||||
elif isinstance(value, ast.BoolOp):
|
||||
return all(_is_replaceable(v) for v in value.values)
|
||||
elif isinstance(value, ast.Call) and isinstance(value.func, ast.Name):
|
||||
return value.func.id in _replaceable_funcs
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
class _TargetLister(ast.NodeVisitor):
|
||||
def __init__(self):
|
||||
self.targets = set()
|
||||
|
||||
def visit_Name(self, node):
|
||||
if isinstance(node.ctx, ast.Store):
|
||||
self.targets.add(node.id)
|
||||
|
||||
|
||||
class _InterAssignRemover(ast.NodeTransformer):
|
||||
def __init__(self):
|
||||
self.replacements = dict()
|
||||
self.modified_names = set()
|
||||
|
||||
def visit_Name(self, node):
|
||||
if isinstance(node.ctx, ast.Load):
|
||||
try:
|
||||
return self.replacements[node.id]
|
||||
except KeyError:
|
||||
return node
|
||||
else:
|
||||
self.modified_names.add(node.id)
|
||||
return node
|
||||
|
||||
def visit_Assign(self, node):
|
||||
self.generic_visit(node)
|
||||
if _is_replaceable(node.value):
|
||||
for target in node.targets:
|
||||
if isinstance(target, ast.Name):
|
||||
self.replacements[target.id] = node.value
|
||||
else:
|
||||
for target in node.targets:
|
||||
try:
|
||||
del self.replacements[target.id]
|
||||
except KeyError:
|
||||
pass
|
||||
return node
|
||||
|
||||
def visit_AugAssign(self, node):
|
||||
left = deepcopy(node.target)
|
||||
left.ctx = ast.Load()
|
||||
newnode = ast.copy_location(
|
||||
ast.Assign(
|
||||
targets=[node.target],
|
||||
value=ast.BinOp(left=left, op=node.op, right=node.value)
|
||||
),
|
||||
node
|
||||
)
|
||||
return self.visit_Assign(newnode)
|
||||
|
||||
def modified_names_push(self):
|
||||
prev_modified_names = self.modified_names
|
||||
self.modified_names = set()
|
||||
return prev_modified_names
|
||||
|
||||
def modified_names_pop(self, prev_modified_names):
|
||||
for name in self.modified_names:
|
||||
try:
|
||||
del self.replacements[name]
|
||||
except KeyError:
|
||||
pass
|
||||
self.modified_names |= prev_modified_names
|
||||
|
||||
def visit_Try(self, node):
|
||||
prev_modified_names = self.modified_names_push()
|
||||
self.generic_visit(node)
|
||||
self.modified_names_pop(prev_modified_names)
|
||||
return node
|
||||
|
||||
def visit_If(self, node):
|
||||
node.test = self.visit(node.test)
|
||||
|
||||
prev_modified_names = self.modified_names_push()
|
||||
|
||||
prev_replacements = self.replacements
|
||||
self.replacements = copy(prev_replacements)
|
||||
node.body = [self.visit(n) for n in node.body]
|
||||
self.replacements = copy(prev_replacements)
|
||||
node.orelse = [self.visit(n) for n in node.orelse]
|
||||
self.replacements = prev_replacements
|
||||
|
||||
self.modified_names_pop(prev_modified_names)
|
||||
|
||||
return node
|
||||
|
||||
def visit_loop(self, node):
|
||||
prev_modified_names = self.modified_names_push()
|
||||
prev_replacements = self.replacements
|
||||
|
||||
self.replacements = copy(prev_replacements)
|
||||
tl = _TargetLister()
|
||||
for n in node.body:
|
||||
tl.visit(n)
|
||||
for name in tl.targets:
|
||||
try:
|
||||
del self.replacements[name]
|
||||
except KeyError:
|
||||
pass
|
||||
node.body = [self.visit(n) for n in node.body]
|
||||
|
||||
self.replacements = copy(prev_replacements)
|
||||
node.orelse = [self.visit(n) for n in node.orelse]
|
||||
|
||||
self.replacements = prev_replacements
|
||||
self.modified_names_pop(prev_modified_names)
|
||||
|
||||
def visit_For(self, node):
|
||||
prev_modified_names = self.modified_names_push()
|
||||
node.target = self.visit(node.target)
|
||||
self.modified_names_pop(prev_modified_names)
|
||||
node.iter = self.visit(node.iter)
|
||||
self.visit_loop(node)
|
||||
return node
|
||||
|
||||
def visit_While(self, node):
|
||||
self.visit_loop(node)
|
||||
node.test = self.visit(node.test)
|
||||
return node
|
||||
|
||||
|
||||
def remove_inter_assigns(func_def):
|
||||
_InterAssignRemover().visit(func_def)
|
Loading…
Reference in New Issue