fix[next]: collect closure vars via the compiler's scope analysis - #2866
Merged
Conversation
`get_closure_vars_from_function` used `inspect.getclosurevars`, which disassembles the enclosing function's own code object and collects the names loaded by its `LOAD_GLOBAL` instructions. A generator expression, lambda or nested function compiles to a separate code object, so a global referenced only inside one is recorded there and nowhere else, and was never collected. Free variables are unaffected, because closing over one forces a cell whose name is recorded on the enclosing code object as `co_freevars`. Take the global names from the compiler's own scope analysis instead, via the `symtable` module this file already imports. Every scope the compiler creates is a child `symtable.Function` whose `get_globals()` is exactly the set of names that compile to `LOAD_GLOBAL` there, so locals shadowing a global and comprehension targets are excluded by construction. The source is analyzed under `from __future__ import annotations` so annotations contribute no names on any supported version. Free variables keep coming from `inspect.getclosurevars`. No construct the frontend accepts today creates a nested scope, so this changes nothing for existing programs; it closes the mechanism ahead of tuple comprehensions (GridTools#2833), which reach it. Same source change as GridTools#2864, rebased onto main; the tests that need tuple comprehensions stay with GridTools#2833. Claude-Session: https://claude.ai/code/session_01Sr9xdMdZgG4wjLLAKbweYb
This was referenced Sep 9, 2026
havogt
marked this pull request as ready for review
September 9, 2026 08:51
Contributor
Author
|
fyi, here is a memoization addition, but not sure if it's worth taking it havogt#79 |
egparedes
approved these changes
Sep 9, 2026
egparedes
left a comment
egparedes
left a comment
Contributor
There was a problem hiding this comment.
Looks good to me. I would merge here first the memoization optimization and then merge this.
`get_closure_vars_from_function` re-read and re-analyzed the source on every call, and the toolchain calls it many times per definition: every in-memory stage cache fingerprints its input with `semantic_fingerprinter`, which deconstructs each DSL definition function it meets, including those reached through closure variables. For a chain of 16 operators that is 186 collections at decoration and 65 more on the first call. The global names are a property of the code object, so cache them per code object in a `WeakKeyDictionary`; values are still looked up on every call, so a rebound global is seen. Free variables are read from the closure cells directly instead of through `inspect.getclosurevars`, which disassembled the function only to reach them. Per call this is 3 us instead of 284 us for a 5-line function and 94 us instead of 12.6 ms for a 200-line one, against 92 us and 6.6 ms for the `getclosurevars` scan on main.
havogt
added a commit
to havogt/gt4py
that referenced
this pull request
Sep 9, 2026
…sions GridTools#2866 fixed closure variable collection across nested scopes and landed the collection-level tests in 'test_source_utils.py'. These are the parser- and execution-level counterparts, which exercise the same fix through 'FieldOperatorParser' and through a compiled program. 'test_free_variables_still_resolve_through_the_parser' is the parser-level counterpart of the identically-shaped collection-level test in 'test_source_utils.py'; the name differs so the tree holds no two same-named tests. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh
2 tasks
havogt
added a commit
that referenced
this pull request
Sep 9, 2026
`get_closure_vars_from_function` calls `inspect.getclosurevars` only for its `nonlocals`, but `getclosurevars` disassembles the whole function to also find the global names, which since #2866 come from the cached symbol-table analysis anyway. Read the free variables from `co_freevars` and the closure cells directly, in `_free_variables_from_closure_cells`.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
get_closure_vars_from_functionusedinspect.getclosurevars, which reads only the enclosing function's own code object. A generator expression, lambda or nested function compiles to its own code object, so a global referenced only inside one was never collected. Free variables are unaffected, since closing over one creates a cell on the enclosing code object.Take the global names from the compiler's scope analysis instead: every scope is a child
symtable.Functionwhoseget_globals()is exactly the set of names that compile toLOAD_GLOBALthere, so locals shadowing a global and comprehension targets are excluded by construction. The source is analyzed underfrom __future__ import annotationsso annotations contribute no names on any supported Python version.The name analysis is cached per code object, because the toolchain collects the closure variables of a function many times over (every stage fingerprint does). Values are still looked up on every call.
No construct the frontend accepts today creates a nested scope, so existing programs are unaffected. Tuple comprehensions (#2833) reach it; the tests that need them stay with that PR.
Requirements
https://claude.ai/code/session_01Sr9xdMdZgG4wjLLAKbweYb