From fa6f9bed811944224123e7c406fafb691ec821ec Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 4 Sep 2026 18:10:02 +0200 Subject: [PATCH] fix[next]: collect closure vars via the compiler's scope analysis `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 compiles to a separate code object, so a name referenced only inside a tuple comprehension body is recorded there and nowhere else, and was never collected. Type deduction then failed with `UndefinedSymbolError` on a name that is visibly imported: @gtx.field_operator def testee(tracers: tuple[EField, ...]) -> tuple[CField, ...]: return tuple(neighbor_sum(t(C2E), axis=C2EDim) for t in tracers) # UndefinedSymbolError: Undeclared symbol 'neighbor_sum' This affected every module-level name: builtins, `FieldOffset`s, `Dimension`s and module-level field operators, so `tuple(_inner(t, ...) for t in tracers)` did not work either. Free variables were unaffected, because closing over one forces a cell whose name is recorded on the enclosing code object as `co_freevars`. Globals need no cell, so nothing links them back. That is also why the existing comprehension tests pass: they all define their helpers inside the test function. 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, the generator expression's included, 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`. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- src/gt4py/next/ffront/source_utils.py | 42 +++++++- .../feature_tests/ffront_tests/test_tuples.py | 59 ++++++++++- .../test_tuple_comprehension_closure_vars.py | 99 +++++++++++++++++++ 3 files changed, 197 insertions(+), 3 deletions(-) create mode 100644 tests/next_tests/unit_tests/ffront_tests/test_tuple_comprehension_closure_vars.py diff --git a/src/gt4py/next/ffront/source_utils.py b/src/gt4py/next/ffront/source_utils.py index cc02c1d89b..44d59dff28 100644 --- a/src/gt4py/next/ffront/source_utils.py +++ b/src/gt4py/next/ffront/source_utils.py @@ -8,6 +8,7 @@ from __future__ import annotations +import builtins import functools import inspect import pathlib @@ -21,11 +22,48 @@ MISSING_FILENAME = "" +def _global_names_from_source(source: str) -> set[str]: + """ + Names referenced as globals in the function defined by `source`, including its nested scopes. + + The compiler's own scope analysis decides what is a global, so locals shadowing a + global name and comprehension targets are never reported. + """ + + def walk(table: symtable.SymbolTable) -> Iterator[str]: + if isinstance(table, symtable.Function): + yield from table.get_globals() + for child in table.get_children(): + yield from walk(child) + + # Analyzed as if under PEP 563 so that annotations contribute no names: Python evaluates + # parameter and return annotations in the enclosing scope and never evaluates local + # variable annotations, but only Python 3.14's symtable stops reporting the latter. + source = "from __future__ import annotations\n" + source + return set(walk(symtable.symtable(source, MISSING_FILENAME, "exec"))) + + def get_closure_vars_from_function(function: Callable) -> dict[str, Any]: - (nonlocals, globals, builtins, _unbound) = inspect.getclosurevars(function) # noqa: A001 [builtin-variable-shadowing] + # `inspect.getclosurevars` only sees the names of the function's own code object, which + # misses names referenced only inside a nested scope such as a generator expression. + # Free variables are unaffected (they are cells of the function itself), so only the + # global names are taken from the source instead. + nonlocals = inspect.getclosurevars(function).nonlocals + global_ns = function.__globals__ + builtin_ns = global_ns.get("__builtins__", builtins.__dict__) + if inspect.ismodule(builtin_ns): + builtin_ns = builtin_ns.__dict__ + + source = make_source_definition_from_function(function).source + closure_vars: dict[str, Any] = {} + for name in _global_names_from_source(source): + if name in global_ns: + closure_vars[name] = global_ns[name] + elif name in builtin_ns: + closure_vars[name] = builtin_ns[name] # nonlocals override globals, sorted for deterministic results - return dict(sorted({**builtins, **globals, **nonlocals}.items())) + return dict(sorted({**closure_vars, **nonlocals}.items())) def make_source_definition_from_function(func: Callable) -> SourceDefinition: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 39c53e0073..61896c3ac4 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -10,7 +10,7 @@ import pytest import gt4py.next as gtx -from gt4py.next import broadcast, errors, float64, int32, neighbor_sum, utils as gt_utils +from gt4py.next import broadcast, common, errors, float64, int32, neighbor_sum, utils as gt_utils from next_tests.integration_tests import cases from next_tests.integration_tests.cases import ( @@ -199,6 +199,63 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie ) +# Defined at module level, unlike `inner` in 'test_tuple_comprehension_other_fo' above. +# That difference alone decides whether the comprehension below can see it. +@gtx.field_operator +def _module_level_inner(tracer: cases.IField, factor: int32) -> cases.IField: + return tracer * factor + + +@pytest.mark.uses_tuple_args +def test_tuple_comprehension_module_level_fo(cartesian_case): + # Identical to 'test_tuple_comprehension_other_fo', except that the called operator is + # a module-level global rather than a local of the test function. + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(_module_level_inner(tracer, factor) for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(el * f for el in t), + ) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_unstructured_shift +def test_tuple_comprehension_module_level_builtin(unstructured_case): + # Same cause, but for an imported builtin -- the shape every real stencil has. + @gtx.field_operator + def testee(tracers: tuple[cases.EField, ...]) -> tuple[cases.VField, ...]: + return tuple(neighbor_sum(tracer(V2E), axis=V2EDim) for tracer in tracers) + + v2e = unstructured_case.offset_provider["V2E"].asnumpy() + valid = v2e != common._DEFAULT_SKIP_VALUE + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda t: tuple(np.sum(el[v2e], axis=1, where=valid) for el in t), + ) + + +@pytest.mark.uses_tuple_args +def test_tuple_comprehension_module_level_fo_used_outside_too(cartesian_case): + # The workaround, and why this is easy to miss: the comprehension body is unchanged; + # only the unrelated statement above it makes '_module_level_inner' collectable. + @gtx.field_operator + def testee( + tracers: tuple[cases.IField, ...], a: cases.IField, factor: int32 + ) -> tuple[cases.IField, ...]: + unrelated = _module_level_inner(a, factor) + return tuple(_module_level_inner(tracer, factor) for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, a, f: tuple(el * f for el in t), + ) + + @pytest.mark.uses_tuple_args def test_nested_tuple_comprehension(cartesian_case): @gtx.field_operator diff --git a/tests/next_tests/unit_tests/ffront_tests/test_tuple_comprehension_closure_vars.py b/tests/next_tests/unit_tests/ffront_tests/test_tuple_comprehension_closure_vars.py new file mode 100644 index 0000000000..90d827c171 --- /dev/null +++ b/tests/next_tests/unit_tests/ffront_tests/test_tuple_comprehension_closure_vars.py @@ -0,0 +1,99 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +"""Closure variable resolution inside tuple comprehension bodies. + +A generator expression compiles to its own code object, so names referenced only +inside it appear in that object's `co_names` and never in the enclosing function's. +Collecting closure variables from the enclosing function alone therefore misses +them; `get_closure_vars_from_function` takes the global names from the compiler's +symbol table of the source instead, which covers nested scopes. + +Free variables need no such treatment: the enclosing code object carries a cell +for them, which is why comprehensions over locally defined helpers always worked. +""" + +import pytest + +import gt4py.next as gtx +from gt4py.next import Dims, Dimension, float64, neighbor_sum +from gt4py.next.ffront.func_to_foast import FieldOperatorParser +from gt4py.next.ffront.source_utils import get_closure_vars_from_function + + +Cell = Dimension("Cell") +Edge = Dimension("Edge") +C2EDim = Dimension("C2E", kind=gtx.DimensionKind.LOCAL) +C2E = gtx.FieldOffset("C2E", source=Edge, target=(Cell, C2EDim)) + +CField = gtx.Field[Dims[Cell], float64] +EField = gtx.Field[Dims[Edge], float64] + + +@gtx.field_operator +def scale(f: CField, factor: float64) -> CField: + return f * factor + + +def _builtin_and_offset(tracers: tuple[EField, ...]) -> tuple[CField, ...]: + return tuple(neighbor_sum(t(C2E), axis=C2EDim) for t in tracers) + + +def _module_level_operator(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]: + return tuple(scale(t, factor) for t in tracers) + + +@pytest.mark.parametrize( + "func", [_builtin_and_offset, _module_level_operator], ids=["builtin_and_offset", "operator"] +) +def test_module_level_name_used_only_in_comprehension(func): + """Module-level names are reachable from inside a comprehension body.""" + parsed = FieldOperatorParser.apply_to_function(func) + assert parsed.type is not None + + +def test_names_are_collected_from_the_nested_code_object(): + """The names live in the generator expression's code object, not the function's.""" + assert "neighbor_sum" not in _builtin_and_offset.__code__.co_names + nested = [c for c in _builtin_and_offset.__code__.co_consts if hasattr(c, "co_names")] + assert any("neighbor_sum" in c.co_names for c in nested) + + collected = get_closure_vars_from_function(_builtin_and_offset) + assert {"neighbor_sum", "C2E", "C2EDim"} <= set(collected) + + +def test_comprehension_target_is_not_collected(): + """The loop target is a local of the nested code object, not a global reference.""" + assert "t" not in get_closure_vars_from_function(_module_level_operator) + + +def test_free_variables_still_resolve(): + """The path that already worked, kept as a guard.""" + + @gtx.field_operator + def local_scale(f: CField, factor: float64) -> CField: + return f * factor + + def uses_freevar(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]: + return tuple(local_scale(t, factor) for t in tracers) + + parsed = FieldOperatorParser.apply_to_function(uses_freevar) + assert parsed.type is not None + + +def test_local_name_shadowing_a_global_is_not_collected_as_global(): + """A comprehension referencing an enclosing local must bind the local, not the global.""" + + def shadows(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]: + scale = local_helper # noqa: F841 shadows the module-level 'scale' + return tuple(scale(t, factor) for t in tracers) + + def local_helper(t, factor): + return t + + assert "scale" not in get_closure_vars_from_function(shadows)