提交 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
# Find the root of the view chain, and while traversing check if it passes on any protected inputs.
view_of_protected = False
root = inp
try:
while True:
if root in protected_inputs:
view_of_protected = True
root = view_i[root]
except KeyError:
pass
return [ if root in candidate_roots:
inp # Another input views on the same root, we can't destroy either
# Remove duplicates, while preserving order by using dict.fromkeys if (invalid_candidate := candidate_roots[root]) is not None:
for inp in dict.fromkeys(inputs) # Invalidate the previous candidate
if ( candidate_inputs.remove(invalid_candidate)
not isinstance(inp, Constant) candidate_roots[root] = None
and inp not in protected_inputs elif not view_of_protected and not has_destroyers([inp]):
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 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论