提交 1d94ed68 authored 作者: ricardoV94's avatar ricardoV94 提交者: Ricardo Vieira

Less myopic check for protected inputs

上级 d9e8728a
...@@ -244,17 +244,38 @@ def inplace_candidates(fgraph, inputs, protected_inputs=None): ...@@ -244,17 +244,38 @@ def inplace_candidates(fgraph, inputs, protected_inputs=None):
protected_inputs.update(fgraph.outputs) protected_inputs.update(fgraph.outputs)
has_destroyers = fgraph.has_destroyers has_destroyers = fgraph.has_destroyers
view_i = fgraph.destroy_handler.view_i
candidate_roots = {}
candidate_inputs = []
for inp in inputs:
if isinstance(inp, Constant):
# Can't inplace on constants.
continue
return [ # Find the root of the view chain, and while traversing check if it passes on any protected inputs.
inp view_of_protected = False
# Remove duplicates, while preserving order by using dict.fromkeys root = inp
for inp in dict.fromkeys(inputs) try:
if ( while True:
not isinstance(inp, Constant) if root in protected_inputs:
and inp not in protected_inputs view_of_protected = True
and not has_destroyers([inp]) root = view_i[root]
) except KeyError:
] pass
if root in candidate_roots:
# Another input views on the same root, we can't destroy either
if (invalid_candidate := candidate_roots[root]) is not None:
# Invalidate the previous candidate
candidate_inputs.remove(invalid_candidate)
candidate_roots[root] = None
elif not view_of_protected and not has_destroyers([inp]):
candidate_inputs.append(inp)
candidate_roots[root] = inp
else:
candidate_roots[root] = None
return candidate_inputs
class DestroyHandler(Bookkeeper): class DestroyHandler(Bookkeeper):
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论