diff --git a/loopy/tools.py b/loopy/tools.py index 8e8453602..2a119eae2 100644 --- a/loopy/tools.py +++ b/loopy/tools.py @@ -1004,7 +1004,7 @@ def from_root(root: T): return Tree(pmap({root: frozenset()}), pmap({root: None})) - @property + @cached_property def root(self) -> T: guess = set(self._child_to_parent).pop() while self.parent(guess) is not None: @@ -1012,6 +1012,7 @@ def root(self) -> T: return guess + @memoize_method def ancestors(self, node: T) -> "FrozenSet[T]": """ Returns a :class:`frozenset` of nodes that are ancestors of *node*. @@ -1039,6 +1040,7 @@ def children(self, node: T) -> "FrozenSet[T]": return self._parent_to_children[node] + @memoize_method def depth(self, node: T) -> int: if not self.is_a_node(node): raise ValueError(f"'{node}' not in tree.") diff --git a/loopy/transform/loop_fusion.py b/loopy/transform/loop_fusion.py index 16df080f7..584fc3f50 100644 --- a/loopy/transform/loop_fusion.py +++ b/loopy/transform/loop_fusion.py @@ -275,18 +275,17 @@ def _get_ldg_nodes_from_loopy_insn(kernel, insn, candidates, non_candidates, for candidate in ((insn.within_inames | insn.reduction_inames()) & candidates)] - elif {loop_nest - for loop_nest in non_candidates - if (loop_nest & insn.within_inames)}: + else: non_candidate, = {loop_nest for loop_nest in non_candidates if (loop_nest & insn.within_inames)} - return [NonCandidateLoop(non_candidate)] - else: - assert ((insn.within_inames & just_outer_loop_nest) - or (insn.within_inames == just_outer_loop_nest)) - return [OuterLoopNestStatement(insn.id)] + if non_candidate: + return [NonCandidateLoop(non_candidate)] + else: + assert ((insn.within_inames & just_outer_loop_nest) + or (insn.within_inames == just_outer_loop_nest)) + return [OuterLoopNestStatement(insn.id)] @memoize_on_first_arg