compiler: Lazy IET visitors + Search #2621

enwask · 2025-05-30T10:37:47Z

This PR implements lazy visitors in the IET layer which operate via generators instead of flattening lists of children's results at every node. Replaces FindSymbols, FindNodes, FindWithin and FindApplications with such versions, as well as introducing type hints for the affected visitors.

From my testing this results in a speedup in IET lowering of up to 50%.

Also reimplements Search similarly, benchmarks forthcoming.

devito/ir/iet/visitors.py

mloubout

Some minor comments.

I think it needs a rebase on main too

FabioLuporini · 2025-05-31T14:28:03Z

@enwask can you give us one or more command lines to reproduce such a compilation speed improvement? perhaps something using the available examples/seismic/... ? or a script of your own? thanks!

FabioLuporini · 2025-05-31T14:28:26Z

also note @enwask that CI is currently failing

enwask · 2025-05-31T14:58:43Z

Yeah will cook up a demo. Did run the test suite yesterday but some of my typing changes for the review must have broken things, will take a look

enwask · 2025-06-01T17:16:39Z

Profiled with this script on dewback through numactl, getting a less substantial but measurable 5% improvement overall: you can view the cProfile output from both the upstream head and my fork with lazy visitors.

I also separately ran some of the examples through cProfile; this gist has my test script as well as the profiling output across three examples (namely acoustic, elastic and TTI), where I'm seeing a similar 5% improvement in operator construction. Also in Python 3.10 on dewback.

codecov · 2025-06-01T21:18:43Z

Codecov Report

Attention: Patch coverage is 94.07407% with 8 lines in your changes missing coverage. Please review.

Project coverage is 87.60%. Comparing base (bca0b10) to head (4f5a6ec).

Files with missing lines	Patch %	Lines
devito/symbolics/search.py	82.50%	5 Missing and 2 partials ⚠️
devito/ir/iet/visitors.py	98.94%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2621      +/-   ##
==========================================
- Coverage   91.30%   87.60%   -3.71%     
==========================================
  Files         245      245              
  Lines       48752    48739      -13     
  Branches     4298     4299       +1     
==========================================
- Hits        44515    42698    -1817     
- Misses       3529     5311    +1782     
- Partials      708      730      +22

Flag	Coverage Δ
pytest-gpu-aomp-amdgpuX	`?`

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

devito/symbolics/search.py

enwask · 2025-06-09T15:43:52Z

Noticed that the Search rewrite is actually a regression; trying to work out why...

Found out why! Turns out bfs_first_hit is neither a BFS nor does it return only the first hit

EdCaunt · 2025-06-12T09:12:26Z

devito/symbolics/search.py

-            found.update(searcher.bfs_first_hit(e))
-        else:
-            raise ValueError("Unknown visit type `%s`" % visit)
+    # Search doesn't actually use a BFS (rather, a preorder DFS), but the terminology


commit 2079dda Author: enwask <enwask@ucf.edu> Date: Fri Jun 6 10:34:58 2025 +0200 Remove FindApplications.visit_Node override for coverage commit 60aaf00 Author: enwask <enwask@ucf.edu> Date: Thu Jun 5 00:17:10 2025 +0200 compiler: Lazy IET visitors commit 167d26d Author: enwask <enwask@ucf.edu> Date: Mon Jun 2 12:09:46 2025 +0200 Fix FindSymbols commit 871e42a Author: enwask <enwask@ucf.edu> Date: Sun Jun 1 23:59:20 2025 +0200 Remove TypeAlias for Python 3.9 commit 38230cf Author: enwask <enwask@ucf.edu> Date: Sat May 31 19:22:04 2025 +0100 Fix all the tests failing commit f7bdfbd Author: enwask <enwask@ucf.edu> Date: Fri May 30 16:27:12 2025 +0100 Remove redundant key=id, rename type parameters commit 2b13898 Author: enwask <enwask@ucf.edu> Date: Fri May 30 15:56:10 2025 +0100 Move typing import, add RulesDict type aliases commit 82cc017 Author: enwask <enwask@ucf.edu> Date: Fri May 30 12:52:22 2025 +0100 Move visit_tuple up + remove redundant set casts commit 4d1b64b Author: enwask <enwask@ucf.edu> Date: Fri May 30 11:38:02 2025 +0100 Formatting fixes commit 9181005 Author: enwask <enwask@ucf.edu> Date: Thu May 29 11:25:35 2025 +0100 Lazy FindApplications commit 508d41c Author: enwask <enwask@ucf.edu> Date: Thu May 29 11:23:53 2025 +0100 LazyVisitor generic typing commit 82a11c9 Author: enwask <enwask@ucf.edu> Date: Wed May 28 18:30:40 2025 +0100 Cleanup + faster FindSymbols commit 3d9fe74 Author: enwask <enwask@ucf.edu> Date: Wed May 28 17:25:51 2025 +0100 FindSymbols tweaks commit e56e768 Author: enwask <enwask@ucf.edu> Date: Wed May 28 12:03:16 2025 +0100 Lazy FindNodes commit de94ac3 Author: enwask <enwask@ucf.edu> Date: Tue May 27 18:57:18 2025 +0100 Lazy FindSymbols commit 921125b Author: enwask <enwask@ucf.edu> Date: Tue May 27 18:55:49 2025 +0100 Add LazyVisitor

Copilot

Pull Request Overview

This PR replaces eager tree traversal in the IET layer with generator-based (lazy) visitors to boost performance, and refactors the Search utility similarly while adding type hints. It also updates Python version support across project configs and CI workflows.

Bump supported Python versions to >=3.10,<3.13 in project metadata and CI
Refactor devito/symbolics/search.py to generator-based Search, adding type hints
Introduce LazyVisitor in devito/ir/iet/visitors.py and replace Find* visitors accordingly

Reviewed Changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
pyproject.toml	Updated Python version requirement
environment-dev.yml	Aligned dev environment Python constraint
devito/symbolics/search.py	Refactored `Search` to use generators; added type annotations
devito/ir/iet/visitors.py	Introduced `LazyVisitor` and rewrote `FindSymbols`, `FindNodes`, `FindWithin`, `FindApplications`
benchmarks/user/advisor/README.md	Updated required Python version in documentation
.github/workflows/tutorials.yml	Updated CI matrix for Python 3.10 and added image version check
.github/workflows/pytest-core-nompi.yml	Adjusted CI matrix to drop Python 3.9 and align versions
.github/workflows/pytest-core-mpi.yml	Adjusted MPI CI matrix for Python 3.10
.github/workflows/flake8.yml	Updated linter job to use Python 3.10

Comments suppressed due to low confidence (3)

devito/symbolics/search.py:90

[nitpick] The return annotation uses List (the custom class) which may be confused with the built-in list or typing.List. Consider annotating as list[Expression] or renaming the custom class to avoid ambiguity.

def search(exprs: Expression | Iterable[Expression],

devito/symbolics/search.py:116

as_tuple is used here but not imported, causing a NameError. Please add the appropriate import (e.g., from devito.symbolics.search import as_tuple).

    exprs = filter(lambda e: isinstance(e, sympy.Basic), as_tuple(exprs))

devito/ir/iet/visitors.py:1144

Accessing ret[-1] without checking if ret is non-empty can raise an IndexError. Add a guard like if ret and ret[-1] is self.STOP: before this line.

        if ret[-1] is self.STOP:

.github/workflows/pytest-core-nompi.yml

devito/symbolics/search.py

devito/ir/iet/visitors.py

devito/symbolics/search.py

FabioLuporini · 2025-06-24T07:56:42Z

devito/ir/iet/visitors.py

    """
    Like FindNodes, but given an additional parameter `within=(start, stop)`,
    it starts collecting matching nodes only after `start` is found, and stops
    collecting matching nodes after `stop` is found.
    """

-    def __init__(self, match, start, stop=None):
+    # Sentinel values to signal the start/end of a matching window
+    SET_FLAG = object()


I have several concerns about the increased complexity in FindWithin

FabioLuporini · 2025-06-24T08:02:22Z

devito/ir/iet/visitors.py

+            for i in self._visit(el, flag=flag):
+                # New flag state is yielded at the end of child results
+                if i is self.SET_FLAG:
+                    flag = True


just flag = self.SET_FLAG, and same below, so that you can eventually just return flag

FabioLuporini · 2025-06-24T08:03:32Z

devito/ir/iet/visitors.py

+                    flag = True
+                    continue
+                if i is self.UNSET_FLAG:
+                    if flag:


why is this different from what happens in visit_tuple?

EdCaunt reviewed May 30, 2025

View reviewed changes

mloubout added the compiler label May 30, 2025

EdCaunt reviewed May 30, 2025

View reviewed changes

devito/ir/iet/visitors.py Outdated Show resolved Hide resolved

mloubout reviewed May 30, 2025

View reviewed changes

devito/ir/iet/visitors.py Show resolved Hide resolved

mloubout reviewed May 30, 2025

View reviewed changes

devito/ir/iet/visitors.py Outdated Show resolved Hide resolved

mloubout reviewed May 30, 2025

View reviewed changes

devito/ir/iet/visitors.py Outdated Show resolved Hide resolved

mloubout reviewed May 30, 2025

View reviewed changes

enwask force-pushed the lazy-visitors branch from 25b668d to f7bdfbd Compare May 30, 2025 15:30

enwask force-pushed the lazy-visitors branch 2 times, most recently from 4f92b56 to 60aaf00 Compare June 5, 2025 21:58

ggorman requested a review from Copilot June 5, 2025 23:15

This comment was marked as outdated.

Sign in to view

enwask requested a review from mloubout June 6, 2025 06:58

enwask force-pushed the lazy-visitors branch from 2079dda to 8e3752e Compare June 6, 2025 08:35

enwask changed the title ~~compiler: Lazy IET visitors~~ compiler: Lazy IET visitors + Search Jun 9, 2025

mloubout reviewed Jun 9, 2025

View reviewed changes

devito/symbolics/search.py Outdated Show resolved Hide resolved

enwask force-pushed the lazy-visitors branch from fbf0d32 to 8c02095 Compare June 9, 2025 16:46

EdCaunt reviewed Jun 12, 2025

View reviewed changes

EdCaunt mentioned this pull request Jun 12, 2025

compiler: Revamp MPI hoisting and merging #2629

Merged

enwask force-pushed the lazy-visitors branch 2 times, most recently from 8426db2 to 26fcb48 Compare June 19, 2025 10:28

enwask added 3 commits June 19, 2025 11:29

deps: Drop Python 3.9

09de795

ci: Fix flake8 workflow configuration

bb123c4

ci: Add Docker Python version checks

2886742

enwask added 13 commits June 19, 2025 11:29

deps: Add Python version upper bound

7314651

ci: Oops

a6b6124

temp: Do CI

9f4b8c5

Add version check logging

9bef19c

ci: Not anymore

252eaa0

ci: Oops

14ff62a

Lazy search

130efc1

Better lazy search

a3b4048

Fix search regression + typing

8e1bc1b

Slightly faster search

d69462b

misc: Cleanup + flake8 fixes

1cdf3e8

misc: LazyVisitor tweaks + lazy FindWithin

528d8f5

enwask force-pushed the lazy-visitors branch from 26fcb48 to 528d8f5 Compare June 19, 2025 10:29

enwask requested a review from Copilot June 19, 2025 10:40

Copilot AI reviewed Jun 19, 2025

View reviewed changes

misc: Fix FindWithin

48373f5

EdCaunt reviewed Jun 19, 2025

View reviewed changes

.github/workflows/pytest-core-nompi.yml Outdated Show resolved Hide resolved

devito/symbolics/search.py Outdated Show resolved Hide resolved

enwask added 3 commits June 19, 2025 16:13

misc: Fix FindWithin (again)

f0bc509

misc: flake8 fix + match/case for search

d5c4aab

Merge branch 'main' into lazy-visitors

4f5a6ec

enwask requested review from EdCaunt and mloubout June 20, 2025 12:59

EdCaunt reviewed Jun 23, 2025

View reviewed changes

devito/ir/iet/visitors.py Show resolved Hide resolved

devito/symbolics/search.py Show resolved Hide resolved

EdCaunt approved these changes Jun 23, 2025

View reviewed changes

FabioLuporini reviewed Jun 24, 2025

View reviewed changes

compiler: Lazy IET visitors + Search #2621

Are you sure you want to change the base?

compiler: Lazy IET visitors + Search #2621

Uh oh!

Conversation

enwask commented May 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mloubout left a comment

Choose a reason for hiding this comment

Uh oh!

FabioLuporini commented May 31, 2025

Uh oh!

FabioLuporini commented May 31, 2025

Uh oh!

enwask commented May 31, 2025

Uh oh!

enwask commented Jun 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Jun 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

enwask commented Jun 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

EdCaunt Jun 12, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

FabioLuporini Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

FabioLuporini Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

FabioLuporini Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

enwask commented May 30, 2025 •

edited

Loading

enwask commented Jun 1, 2025 •

edited

Loading

codecov bot commented Jun 1, 2025 •

edited

Loading

enwask commented Jun 9, 2025 •

edited

Loading

FabioLuporini Jun 24, 2025 •

edited

Loading