Skip to content

Commit

Permalink
don't use Loop._invariants in Loop._node
Browse files Browse the repository at this point in the history
The `Loop._invariants` are quite expensive to build. This patch aims to reduce
the usage of `Loop._invariants` by reimplementing `Loop._node` without
`Loop._invariants`.
  • Loading branch information
joostvanzwieten committed May 6, 2024
1 parent 5cf7ea4 commit 28a3453
Showing 1 changed file with 10 additions and 2 deletions.
12 changes: 10 additions & 2 deletions nutils/evaluable.py
Original file line number Diff line number Diff line change
Expand Up @@ -4521,8 +4521,16 @@ def evalf_loop_body_withtimes(self, times, output, body_arg):
def _node(self, cache, subgraph, times, unique_loop_ids):
if (cached := cache.get(self)) is not None:
return cached
for arg in itertools.chain(self._invariants, (self.init_arg,)):
arg._node(cache, subgraph, times, unique_loop_ids)

# Populate the `cache` with objects that do not depend on `self.index`.
stack = [self.init_arg, self.body_arg]
while stack:
func = stack.pop()
if self.index in func.arguments:
stack.extend(func._Evaluable__args)
else:
func._node(cache, subgraph, times, unique_loop_ids)

if unique_loop_ids:
loopcache = cache
loopgraph = cache.setdefault(('subgraph', self.loop_id), Subgraph('Loop', subgraph))
Expand Down

0 comments on commit 28a3453

Please sign in to comment.