From bb5646439af0951bad425c1a11f6ac79e6d4a80b Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 4 Sep 2026 18:06:46 +0200 Subject: [PATCH] fix[next]: collect closure vars from nested code objects `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. Nested code objects are scanned the same way `inspect.getclosurevars` scans the outer one -- `LOAD_GLOBAL` only, so attribute names in `co_names` are not mistaken for references to the enclosing namespace -- and resolved against `__globals__` then the builtins namespace, merged at the same precedence so nonlocals still win. The comprehension target is not collected, being a local of the nested code object. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- src/gt4py/next/ffront/source_utils.py | 49 +++++++++- .../feature_tests/ffront_tests/test_tuples.py | 59 ++++++++++- .../test_tuple_comprehension_closure_vars.py | 98 +++++++++++++++++++ 3 files changed, 204 insertions(+), 2 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..10e40b19fd 100644 --- a/src/gt4py/next/ffront/source_utils.py +++ b/src/gt4py/next/ffront/source_utils.py @@ -8,11 +8,14 @@ from __future__ import annotations +import builtins as _builtins +import dis import functools import inspect import pathlib import symtable import textwrap +import types from collections.abc import Callable, Iterator from dataclasses import dataclass from typing import Any, cast @@ -21,11 +24,55 @@ MISSING_FILENAME = "" +def _iter_nested_code_objects(code: types.CodeType) -> Iterator[types.CodeType]: + """Yield every code object nested in `code`, recursively, excluding `code` itself.""" + for const in code.co_consts: + if isinstance(const, types.CodeType): + yield const + yield from _iter_nested_code_objects(const) + + +def _global_vars_from_nested_code(function: Callable) -> tuple[dict[str, Any], dict[str, Any]]: + """ + Resolve the global and builtin names referenced by `function`'s nested code objects. + + A comprehension or lambda in the body compiles to its own code object, so the names it + references appear in that object's `co_names` rather than the function's, and + `inspect.getclosurevars` (which reads only the function's own) never sees them. Free + variables need no such treatment: closing over one forces a cell, whose name is recorded + on the enclosing code object as `co_freevars`. + """ + global_ns = function.__globals__ + builtin_ns = global_ns.get("__builtins__", _builtins.__dict__) + if inspect.ismodule(builtin_ns): + builtin_ns = builtin_ns.__dict__ + + global_vars: dict[str, Any] = {} + builtin_vars: dict[str, Any] = {} + for code in _iter_nested_code_objects(function.__code__): + # Only `LOAD_GLOBAL` names, mirroring `inspect.getclosurevars`: `co_names` also + # holds attribute names, which are not references to the enclosing namespace. + for name in ( + instruction.argval + for instruction in dis.get_instructions(code) + if instruction.opname == "LOAD_GLOBAL" + ): + if name in global_ns: + global_vars[name] = global_ns[name] + elif name in builtin_ns: + builtin_vars[name] = builtin_ns[name] + + return global_vars, builtin_vars + + def get_closure_vars_from_function(function: Callable) -> dict[str, Any]: (nonlocals, globals, builtins, _unbound) = inspect.getclosurevars(function) # noqa: A001 [builtin-variable-shadowing] + nested_globals, nested_builtins = _global_vars_from_nested_code(function) # nonlocals override globals, sorted for deterministic results - return dict(sorted({**builtins, **globals, **nonlocals}.items())) + return dict( + sorted({**nested_builtins, **builtins, **nested_globals, **globals, **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..d707733bbe --- /dev/null +++ b/tests/next_tests/unit_tests/ffront_tests/test_tuple_comprehension_closure_vars.py @@ -0,0 +1,98 @@ +# 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` walks nested code objects for this reason. + +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)