提交 73f7b1ce authored 作者: Frederic's avatar Frederic

Make the merge op in canonicalize eq a local opt to now make all merge work there

上级 599995d0
...@@ -1773,6 +1773,8 @@ class NavigatorOptimizer(Optimizer): ...@@ -1773,6 +1773,8 @@ class NavigatorOptimizer(Optimizer):
raise raise
if replacements is False or replacements is None: if replacements is False or replacements is None:
return False return False
if replacements is True:
return True
old_vars = node.outputs old_vars = node.outputs
if isinstance(replacements, dict): if isinstance(replacements, dict):
old_vars = list(replacements.keys()) old_vars = list(replacements.keys())
......
...@@ -504,9 +504,22 @@ def register_specialize_device(lopt, *tags, **kwargs): ...@@ -504,9 +504,22 @@ def register_specialize_device(lopt, *tags, **kwargs):
return lopt return lopt
# Register merge_optimizer as a global opt during canonicalize # Register in the canonizer Equilibrium as a local opt the merge opt.
compile.optdb['canonicalize'].register('canon_merge', merge_optimizer, # Without this, as the equilibrium have ignore_newtrees=False, we
'fast_run', final_opt=True) # won't merge all nodes if it is set as a global optimizer with
# final_opt=True.
#
# This work due to those properties:
# 1) the EQ will execute first the optimizer that trac all nodes.
# 2) after an local optimization being applied, if the
# current node is still in the graph, it will continue to the next
# local optimizer. So this won't trigger more iteration.
@register_canonicalize('fast_compile', 'merge')
@gof.local_optimizer(None)
def local_merge_optimizer(node):
if node.fgraph.merge_feature.scheduled:
ret = merge_optimizer(node.fgraph)
return ret[5] > 0
##################### #####################
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论