Skip to content

compiler: Lazy IET visitors + Search #2621

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jul 7, 2025
Merged

Conversation

enwask
Copy link
Contributor

@enwask enwask commented May 30, 2025

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.

Copy link
Contributor

@mloubout mloubout left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some minor comments.

I think it needs a rebase on main too

@FabioLuporini
Copy link
Contributor

@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
Copy link
Contributor

also note @enwask that CI is currently failing

@enwask
Copy link
Contributor Author

enwask commented May 31, 2025

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
Copy link
Contributor Author

enwask commented Jun 1, 2025

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.

Copy link

codecov bot commented Jun 1, 2025

Codecov Report

Attention: Patch coverage is 92.85714% with 9 lines in your changes missing coverage. Please review.

Project coverage is 87.60%. Comparing base (77ffae7) to head (c67bb40).
Report is 6 commits behind head on main.

Files with missing lines Patch % Lines
devito/symbolics/search.py 82.50% 5 Missing and 2 partials ⚠️
devito/ir/iet/visitors.py 97.67% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2621      +/-   ##
==========================================
- Coverage   92.00%   87.60%   -4.40%     
==========================================
  Files         245      245              
  Lines       48813    48790      -23     
  Branches     4310     4305       -5     
==========================================
- Hits        44908    42744    -2164     
- Misses       3211     5313    +2102     
- Partials      694      733      +39     
Flag Coverage Δ
pytest-gpu-aomp-amdgpuX ?
pytest-gpu-nvc-nvidiaX ?

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.

@enwask enwask force-pushed the lazy-visitors branch 2 times, most recently from 4f92b56 to 60aaf00 Compare June 5, 2025 21:58
@ggorman ggorman requested a review from Copilot June 5, 2025 23:15
Copilot

This comment was marked as outdated.

@enwask enwask requested a review from mloubout June 6, 2025 06:58
@enwask enwask changed the title compiler: Lazy IET visitors compiler: Lazy IET visitors + Search Jun 9, 2025
@enwask
Copy link
Contributor Author

enwask commented Jun 9, 2025

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

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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lol

@enwask enwask force-pushed the lazy-visitors branch 3 times, most recently from 26fcb48 to 528d8f5 Compare June 19, 2025 10:29
@enwask enwask requested a review from Copilot June 19, 2025 10:40
Copy link

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

@enwask enwask requested review from EdCaunt and mloubout June 20, 2025 12:59
@enwask enwask requested a review from FabioLuporini June 30, 2025 10:33
@enwask enwask force-pushed the lazy-visitors branch 3 times, most recently from b582035 to d09d33a Compare July 7, 2025 11:15
@mloubout mloubout merged commit ac5097e into devitocodes:main Jul 7, 2025
31 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants