From a178c91457183b6ec4e9aec05414f741312a99d7 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 18:12:18 +0200 Subject: [PATCH 01/14] proto[next]: ambient binding of the offset provider Prototype for the 'ambient values' idea in the knowledge base: a program with no 'offset_provider=' takes its connectivities from whatever is bound to an ambient namespace at call time. Binding half only; ambient *fields* referenced by name inside an operator are not implemented. --- src/gt4py/next/__init__.py | 4 ++ src/gt4py/next/ambient.py | 96 ++++++++++++++++++++++++++++++ src/gt4py/next/ffront/decorator.py | 3 +- 3 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 src/gt4py/next/ambient.py diff --git a/src/gt4py/next/__init__.py b/src/gt4py/next/__init__.py index 806265a1d1..4bf10036e7 100644 --- a/src/gt4py/next/__init__.py +++ b/src/gt4py/next/__init__.py @@ -21,6 +21,7 @@ # ruff: noqa: F401 from .._core.definitions import CUPY_DEVICE_TYPE, Device, DeviceType, is_scalar_type from . import common, ffront, iterator, program_processors, typing +from .ambient import Namespace, bind from .common import ( CartesianConnectivity, Connectivity, @@ -136,6 +137,9 @@ "full", "as_field", "as_connectivity", + # from ambient + "Namespace", + "bind", # from ffront "FieldOffset", "field_operator", diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py new file mode 100644 index 0000000000..7af06ce082 --- /dev/null +++ b/src/gt4py/next/ambient.py @@ -0,0 +1,96 @@ +# 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 + +""" +Prototype: ambient values, bound at program-execution time. + +An ambient value is declared once and reached from any program without +appearing in a signature. Binding happens at execution time (JIT time for +compiled backends), so the same programs can run against a second mesh in one +process. + +This module implements the *binding* half only: the connectivities a program +needs are taken from the ambient context when the caller passes no +``offset_provider``. Referring to ambient *fields* by name inside an operator +(``mesh.edge_length``) is not implemented here. +""" + +from __future__ import annotations + +import contextlib +import contextvars +from collections.abc import Generator, Mapping +from typing import Any + +from gt4py.next import common + + +_bindings: contextvars.ContextVar[Mapping[Namespace, Any]] = contextvars.ContextVar( + "_ambient_bindings" +) + + +class Namespace: + """ + A named collection of ambient values, resolved by attribute access. + + The declaration is the namespace object itself; ``mesh.e2v`` is a reference + that only becomes a value once something is bound to ``mesh``. Attribute + names are not declared up front in this prototype. + """ + + def __init__(self, name: str) -> None: + self._name = name + + def __repr__(self) -> str: + return f"Namespace('{self._name}')" + + @property + def bound(self) -> Any: + """The object currently bound to this namespace.""" + binding = _bindings.get({}).get(self, None) + if binding is None: + raise ValueError( + f"Nothing is bound to ambient namespace '{self._name}'." + " Use 'gtx.bind(, )' around the call." + ) + return binding + + +@contextlib.contextmanager +def bind(namespace: Namespace, value: Any) -> Generator[None, None, None]: + """Bind `value` to `namespace` for the duration of the context.""" + token = _bindings.set({**_bindings.get({}), namespace: value}) + try: + yield + finally: + _bindings.reset(token) + + +def offset_provider() -> common.OffsetProvider: + """ + Collect the offset provider from all bound namespaces. + + Every `common.Connectivity` reachable as an attribute of a bound object + contributes under its attribute name. Names must not collide across + namespaces — an ambiguous offset would silently pick one mesh's table. + """ + collected: dict[str, Any] = {} + for namespace, value in _bindings.get({}).items(): + for key in dir(value): + if key.startswith("_"): + continue + elem = getattr(value, key) + if isinstance(elem, (common.Connectivity, common.Dimension)): + if key in collected and collected[key] is not elem: + raise ValueError( + f"Ambient offset '{key}' is provided by more than one namespace;" + f" '{namespace}' conflicts with an earlier binding." + ) + collected[key] = elem + return collected diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 8e4c2e3ea0..cc1392fd58 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -27,6 +27,7 @@ from gt4py.eve import extended_typing as xtyping from gt4py.eve.extended_typing import Self, Unpack, override from gt4py.next import ( + ambient, backend as next_backend, common, custom_layout_allocators as next_allocators, @@ -383,7 +384,7 @@ def __call__( **kwargs: Any, ) -> None: if offset_provider is None: - offset_provider = {} + offset_provider = ambient.offset_provider() enable_jit = self.compilation_options.enable_jit if enable_jit is None else enable_jit with program_call_context( From 9529d9f56e142abca193cce85a89c4d1454d387c Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 18:23:28 +0200 Subject: [PATCH 02/14] proto[next]: content-hash frozen offset provider elements A frozen element carries a content hash computed once at freeze time, and 'hash_offset_provider_items_by_id' prefers it over 'id(v)'. Two semantically equal meshes then share a compiled program instead of triggering a recompile. Unfrozen elements keep the previous id-based behaviour. --- src/gt4py/next/__init__.py | 6 +- src/gt4py/next/ambient.py | 59 +++++++++--- src/gt4py/next/common.py | 32 +++++-- tests/next_tests/unit_tests/test_ambient.py | 101 ++++++++++++++++++++ 4 files changed, 177 insertions(+), 21 deletions(-) create mode 100644 tests/next_tests/unit_tests/test_ambient.py diff --git a/src/gt4py/next/__init__.py b/src/gt4py/next/__init__.py index 4bf10036e7..d3dbde3ddc 100644 --- a/src/gt4py/next/__init__.py +++ b/src/gt4py/next/__init__.py @@ -20,8 +20,8 @@ # ruff: noqa: F401 from .._core.definitions import CUPY_DEVICE_TYPE, Device, DeviceType, is_scalar_type -from . import common, ffront, iterator, program_processors, typing -from .ambient import Namespace, bind +from . import ambient, common, ffront, iterator, program_processors, typing +from .ambient import Namespace, bind, freeze from .common import ( CartesianConnectivity, Connectivity, @@ -138,8 +138,10 @@ "as_field", "as_connectivity", # from ambient + "ambient", "Namespace", "bind", + "freeze", # from ffront "FieldOffset", "field_operator", diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index 7af06ce082..a49640943c 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -27,6 +27,9 @@ from collections.abc import Generator, Mapping from typing import Any +import numpy as np + +from gt4py.eve import utils as eve_utils from gt4py.next import common @@ -62,9 +65,37 @@ def bound(self) -> Any: return binding +def freeze(elem: Any, *, readonly: bool = False) -> Any: + """ + Give an offset provider element a content hash, so it identifies by value. + + An ambient value is static for the jitted programs that see it, so it may + identify itself by *content* rather than by `id`. The hash is computed once, + here — the O(size) cost is paid at freeze time, never per call. + + `readonly` additionally marks the buffer immutable, which is what makes the + cached hash trustworthy. It is **off by default because it breaks the gtfn + bindings**: they are generated with mutable `ndarray` parameters and reject + a non-writeable array outright. Until that is fixed, the hash is only as + stable as the caller's discipline. + """ + if common.frozen_content_hash(elem) is not None: + return elem + buffer = getattr(elem, "ndarray", None) + if buffer is None: + return elem + if readonly: + buffer.flags.writeable = False + digest = int(eve_utils.content_hash(np.asarray(buffer)), 16) + object.__setattr__(elem, common.FROZEN_HASH_ATTR, digest) + return elem + + @contextlib.contextmanager def bind(namespace: Namespace, value: Any) -> Generator[None, None, None]: """Bind `value` to `namespace` for the duration of the context.""" + for elem in offset_provider_of(value).values(): + freeze(elem) token = _bindings.set({**_bindings.get({}), namespace: value}) try: yield @@ -72,6 +103,16 @@ def bind(namespace: Namespace, value: Any) -> Generator[None, None, None]: _bindings.reset(token) +def offset_provider_of(value: Any) -> dict[str, Any]: + """The connectivities and dimensions reachable as public attributes of `value`.""" + return { + key: elem + for key in dir(value) + if not key.startswith("_") + and isinstance(elem := getattr(value, key), (common.Connectivity, common.Dimension)) + } + + def offset_provider() -> common.OffsetProvider: """ Collect the offset provider from all bound namespaces. @@ -82,15 +123,11 @@ def offset_provider() -> common.OffsetProvider: """ collected: dict[str, Any] = {} for namespace, value in _bindings.get({}).items(): - for key in dir(value): - if key.startswith("_"): - continue - elem = getattr(value, key) - if isinstance(elem, (common.Connectivity, common.Dimension)): - if key in collected and collected[key] is not elem: - raise ValueError( - f"Ambient offset '{key}' is provided by more than one namespace;" - f" '{namespace}' conflicts with an earlier binding." - ) - collected[key] = elem + for key, elem in offset_provider_of(value).items(): + if key in collected and collected[key] is not elem: + raise ValueError( + f"Ambient offset '{key}' is provided by more than one namespace;" + f" '{namespace}' conflicts with an earlier binding." + ) + collected[key] = elem return collected diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index 65b564c14f..4a3ef29f8e 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -1222,17 +1222,33 @@ def has_offset(offset_provider: OffsetProvider | OffsetProviderType, offset_tag: return True +#: Attribute carrying the content hash of a frozen (immutable) offset provider element. +FROZEN_HASH_ATTR: Final[str] = "__gt_frozen_hash__" + + +def frozen_content_hash(elem: OffsetProviderElem) -> int | None: + """Content hash of a frozen offset provider element, or `None` if it is not frozen.""" + return getattr(elem, FROZEN_HASH_ATTR, None) + + def hash_offset_provider_items_by_id(offset_provider: OffsetProvider) -> int: """ - Compute hash of an offset provider on the tuples of key and value id. - - This function is unsafe since it uses the `id` of the values in the - offset provider, which could generate different hashes for two - offset providers that are semantically equal. It additionally relies - on the ordering of the items in the mapping, which could also lead to - different hashes for semantically equal offset providers. + Compute hash of an offset provider on the tuples of key and value. + + A *frozen* element (immutable buffer, content hash computed once at freeze + time) contributes that content hash, so two semantically equal offset + providers share a compiled program. Any other element falls back to its + `id`, which is unsafe: it could generate different hashes for two offset + providers that are semantically equal. This additionally relies on the + ordering of the items in the mapping, which could also lead to different + hashes for semantically equal offset providers. """ - return hash(tuple((k, id(v)) for k, v in offset_provider.items())) + return hash( + tuple( + (k, h if (h := frozen_content_hash(v)) is not None else id(v)) + for k, v in offset_provider.items() + ) + ) DomainDimT = TypeVar("DomainDimT", bound="Dimension") diff --git a/tests/next_tests/unit_tests/test_ambient.py b/tests/next_tests/unit_tests/test_ambient.py new file mode 100644 index 0000000000..ec6584f1a2 --- /dev/null +++ b/tests/next_tests/unit_tests/test_ambient.py @@ -0,0 +1,101 @@ +# 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 + +import dataclasses + +import numpy as np +import pytest + +import gt4py.next as gtx +from gt4py.next import common + + +Vertex = common.Dimension("Vertex") +Edge = common.Dimension("Edge") +V2EDim = common.Dimension("V2E", kind=common.DimensionKind.LOCAL) + + +@dataclasses.dataclass(frozen=True) +class Mesh: + V2E: common.Connectivity + + +def make_mesh(table) -> Mesh: + return Mesh( + V2E=gtx.as_connectivity( + domain={Vertex: 2, V2EDim: 2}, + codomain=Edge, + data=np.asarray(table, dtype=np.int32), + skip_value=None, + ) + ) + + +def test_nothing_bound_yields_empty_offset_provider(): + assert gtx.ambient.offset_provider() == {} + + +def test_bound_namespace_provides_its_connectivities(): + mesh = gtx.Namespace("mesh") + m = make_mesh([[0, 1], [1, 2]]) + with gtx.bind(mesh, m): + assert gtx.ambient.offset_provider() == {"V2E": m.V2E} + assert gtx.ambient.offset_provider() == {} + + +def test_bindings_nest_and_unwind(): + mesh = gtx.Namespace("mesh") + m1, m2 = make_mesh([[0, 1], [1, 2]]), make_mesh([[1, 2], [0, 1]]) + with gtx.bind(mesh, m1): + with gtx.bind(mesh, m2): + assert gtx.ambient.offset_provider()["V2E"] is m2.V2E + assert gtx.ambient.offset_provider()["V2E"] is m1.V2E + + +def test_colliding_offset_names_are_rejected(): + a, b = gtx.Namespace("a"), gtx.Namespace("b") + with gtx.bind(a, make_mesh([[0, 1], [1, 2]])): + with gtx.bind(b, make_mesh([[1, 2], [0, 1]])): + with pytest.raises(ValueError, match="provided by more than one namespace"): + gtx.ambient.offset_provider() + + +def test_freeze_gives_a_content_hash(): + conn = make_mesh([[0, 1], [1, 2]]).V2E + assert common.frozen_content_hash(conn) is None + gtx.freeze(conn) + assert common.frozen_content_hash(conn) is not None + + +def test_equal_frozen_connectivities_share_a_cache_key(): + m1, m2 = make_mesh([[0, 1], [1, 2]]), make_mesh([[0, 1], [1, 2]]) + unfrozen = ( + common.hash_offset_provider_items_by_id({"V2E": m1.V2E}), + common.hash_offset_provider_items_by_id({"V2E": m2.V2E}), + ) + assert unfrozen[0] != unfrozen[1], "distinct objects are keyed apart before freezing" + + gtx.freeze(m1.V2E) + gtx.freeze(m2.V2E) + assert common.hash_offset_provider_items_by_id( + {"V2E": m1.V2E} + ) == common.hash_offset_provider_items_by_id({"V2E": m2.V2E}) + + +def test_differing_frozen_connectivities_are_keyed_apart(): + m1 = gtx.freeze(make_mesh([[0, 1], [1, 2]]).V2E) + m2 = gtx.freeze(make_mesh([[1, 2], [0, 1]]).V2E) + assert common.hash_offset_provider_items_by_id( + {"V2E": m1} + ) != common.hash_offset_provider_items_by_id({"V2E": m2}) + + +def test_readonly_freeze_marks_the_buffer_immutable(): + conn = gtx.freeze(make_mesh([[0, 1], [1, 2]]).V2E, readonly=True) + with pytest.raises(ValueError): + conn.ndarray[0, 0] = 7 From 2473af83ec0c9dc87b3d305059cbc78c04b75389 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 19:49:48 +0200 Subject: [PATCH 03/14] proto[next]: call-time 'bind=' and ambient resolution for field operators Adds a second spelling next to the 'gtx.bind' context manager: a 'bind={ns: value}' kwarg where 'offset_provider=' goes today, scoped to that one call. Both program and field operator entry points now resolve the offset provider from the ambient context when the caller passes none, so direct field-operator calls work too. '__call__' becomes a thin wrapper around '_invoke' so the existing bodies are untouched; 'ProgramWithBoundArgs' overrides '_invoke' accordingly. --- src/gt4py/next/ambient.py | 26 ++++++++-- src/gt4py/next/ffront/decorator.py | 33 ++++++++---- tests/next_tests/unit_tests/test_ambient.py | 57 ++++++++++++++++++++- 3 files changed, 101 insertions(+), 15 deletions(-) diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index a49640943c..91d34e41c4 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -92,17 +92,33 @@ def freeze(elem: Any, *, readonly: bool = False) -> Any: @contextlib.contextmanager -def bind(namespace: Namespace, value: Any) -> Generator[None, None, None]: - """Bind `value` to `namespace` for the duration of the context.""" - for elem in offset_provider_of(value).values(): - freeze(elem) - token = _bindings.set({**_bindings.get({}), namespace: value}) +def bindings(mapping: Mapping[Namespace, Any]) -> Generator[None, None, None]: + """Bind several namespaces at once, for the duration of the context.""" + if not mapping: + yield + return + for value in mapping.values(): + for elem in offset_provider_of(value).values(): + freeze(elem) + token = _bindings.set({**_bindings.get({}), **mapping}) try: yield finally: _bindings.reset(token) +@contextlib.contextmanager +def bind(namespace: Namespace, value: Any) -> Generator[None, None, None]: + """Bind `value` to `namespace` for the duration of the context.""" + with bindings({namespace: value}): + yield + + +def resolve(explicit: common.OffsetProvider | None) -> common.OffsetProvider: + """The caller's offset provider if given, otherwise the ambient one.""" + return offset_provider() if explicit is None else explicit + + def offset_provider_of(value: Any) -> dict[str, Any]: """The connectivities and dimensions reachable as public attributes of `value`.""" return { diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index cc1392fd58..6800b6fc08 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -19,7 +19,7 @@ import types import typing import warnings -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Any, Generic, Optional, Sequence, TypeAlias from gt4py import eve @@ -380,11 +380,22 @@ def __call__( self, *args: Any, offset_provider: common.OffsetProvider | None = None, + bind: Mapping[ambient.Namespace, Any] | None = None, enable_jit: bool | None = None, **kwargs: Any, ) -> None: - if offset_provider is None: - offset_provider = ambient.offset_provider() + """Call the program; `bind` scopes ambient bindings to this call.""" + with ambient.bindings(bind or {}): + self._invoke(*args, offset_provider=offset_provider, enable_jit=enable_jit, **kwargs) + + def _invoke( + self, + *args: Any, + offset_provider: common.OffsetProvider | None = None, + enable_jit: bool | None = None, + **kwargs: Any, + ) -> None: + offset_provider = ambient.resolve(offset_provider) enable_jit = self.compilation_options.enable_jit if enable_jit is None else enable_jit with program_call_context( @@ -433,11 +444,10 @@ class ProgramWithBoundArgs(Program): bound_args: dict[str, float | int | bool] = dataclasses.field(default_factory=dict) @override - def __call__( + def _invoke( self, *args: Any, offset_provider: common.OffsetProvider | None = None, **kwargs: Any ) -> None: - if offset_provider is None: - offset_provider = {} + offset_provider = ambient.resolve(offset_provider) type_ = self.past_stage.past_node.type assert isinstance(type_, ts_ffront.ProgramType) new_type = ts_ffront.ProgramType( @@ -486,7 +496,7 @@ def __call__( else: full_kwargs[str(param.id)] = self.bound_args[param.id] - return super().__call__(*tuple(full_args), offset_provider=offset_provider, **full_kwargs) + return super()._invoke(*tuple(full_args), offset_provider=offset_provider, **full_kwargs) @override def compile( @@ -656,9 +666,14 @@ def __gt_closure_vars__(self) -> dict[str, Any]: return self.foast_stage.closure_vars def __call__(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> Any: + """Call the field operator; `bind` scopes ambient bindings to this call.""" + with ambient.bindings(kwargs.pop("bind", None) or {}): + return self._invoke(*args, enable_jit=enable_jit, **kwargs) + + def _invoke(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> Any: if not next_embedded.context.within_valid_context() and self.backend is not None: # non embedded execution - offset_provider = {**kwargs.pop("offset_provider", {})} + offset_provider = {**ambient.resolve(kwargs.pop("offset_provider", None))} if "out" not in kwargs: raise errors.MissingArgumentError(None, "out", True) out = kwargs.pop("out") @@ -680,7 +695,7 @@ def __call__(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> else: if not next_embedded.context.within_valid_context(): # field_operator as program - kwargs["offset_provider"] = {**kwargs.pop("offset_provider", {})} + kwargs["offset_provider"] = {**ambient.resolve(kwargs.pop("offset_provider", None))} attributes = ( self.definition_stage.attributes if self.definition_stage diff --git a/tests/next_tests/unit_tests/test_ambient.py b/tests/next_tests/unit_tests/test_ambient.py index ec6584f1a2..d1d1c8b06d 100644 --- a/tests/next_tests/unit_tests/test_ambient.py +++ b/tests/next_tests/unit_tests/test_ambient.py @@ -12,7 +12,7 @@ import pytest import gt4py.next as gtx -from gt4py.next import common +from gt4py.next import common, neighbor_sum Vertex = common.Dimension("Vertex") @@ -99,3 +99,58 @@ def test_readonly_freeze_marks_the_buffer_immutable(): conn = gtx.freeze(make_mesh([[0, 1], [1, 2]]).V2E, readonly=True) with pytest.raises(ValueError): conn.ndarray[0, 0] = 7 + + +# --- embedded end-to-end (backend-free) -------------------------------------- + +V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim)) + + +@gtx.field_operator +def sum_edges(a: gtx.Field[gtx.Dims[Edge], gtx.int32]) -> gtx.Field[gtx.Dims[Vertex], gtx.int32]: + return neighbor_sum(a(V2E), axis=V2EDim) + + +@gtx.program +def run( + a: gtx.Field[gtx.Dims[Edge], gtx.int32], out: gtx.Field[gtx.Dims[Vertex], gtx.int32] +) -> None: + sum_edges(a, out=out) + + +@pytest.fixture +def inputs(): + return ( + gtx.as_field([Edge], np.arange(3, dtype=np.int32)), + make_mesh([[0, 1], [1, 2]]), + np.asarray([1, 3], dtype=np.int32), + ) + + +@pytest.mark.parametrize("entry_point", ["program", "field_operator"]) +@pytest.mark.parametrize("mechanism", ["offset_provider", "context_manager", "bind_kwarg"]) +def test_embedded_execution_via_every_mechanism(inputs, entry_point, mechanism): + a, m, expected = inputs + out = gtx.zeros(gtx.domain({Vertex: 2}), dtype=np.int32) + callee = run if entry_point == "program" else sum_edges + kwargs = {"out": out} if entry_point == "field_operator" else {} + args = (a,) if entry_point == "field_operator" else (a, out) + mesh = gtx.Namespace("mesh") + + if mechanism == "offset_provider": + callee(*args, **kwargs, offset_provider={"V2E": m.V2E}) + elif mechanism == "context_manager": + with gtx.bind(mesh, m): + callee(*args, **kwargs) + else: + callee(*args, **kwargs, bind={mesh: m}) + + np.testing.assert_array_equal(out.asnumpy(), expected) + + +def test_bind_kwarg_is_scoped_to_the_call(inputs): + a, m, _ = inputs + mesh = gtx.Namespace("mesh") + out = gtx.zeros(gtx.domain({Vertex: 2}), dtype=np.int32) + run(a, out, bind={mesh: m}) + assert gtx.ambient.offset_provider() == {} From df750afb77e79386302b78acafc3e6ffb8383d12 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 20:36:19 +0200 Subject: [PATCH 04/14] proto[next]: fix ambient offset provider ordering and GPU freeze Two defects the backend matrix exposed, neither visible in a single-connectivity demo or in embedded-only tests: - 'offset_provider_of' built the mapping from 'dir()' (alphabetical) rather than the bound object's own '__dict__' (insertion order). gt4py's offset provider is order-sensitive, so a multi-connectivity mesh got wrong results on gtfn and a segfault elsewhere. - 'freeze' hashed via 'np.asarray', which refuses device arrays; every GPU backend failed. Uses 'asnumpy()' now. Adds a backend-matrix test covering {program, field_operator} x {offset_provider=, context manager, bind=}. --- src/gt4py/next/ambient.py | 52 ++++++++-- .../ffront_tests/test_ambient_binding.py | 94 +++++++++++++++++++ 2 files changed, 139 insertions(+), 7 deletions(-) create mode 100644 tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index 91d34e41c4..c73a156f7c 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -24,6 +24,7 @@ import contextlib import contextvars +import dataclasses from collections.abc import Generator, Mapping from typing import Any @@ -53,6 +54,13 @@ def __init__(self, name: str) -> None: def __repr__(self) -> str: return f"Namespace('{self._name}')" + def __getattr__(self, name: str) -> AmbientRef: + # only reached when normal lookup fails; dunder/private probes (pickle, + # copy, the DSL frontend) must not be turned into references + if name.startswith("_"): + raise AttributeError(name) + return AmbientRef(self, name) + @property def bound(self) -> Any: """The object currently bound to this namespace.""" @@ -65,6 +73,26 @@ def bound(self) -> Any: return binding +@dataclasses.dataclass(frozen=True) +class AmbientRef: + """ + A deferred reference to `namespace.name`, resolved when something is bound. + + Attribute access on a `Namespace` yields one of these instead of a value, + because at the point an operator is *defined* nothing is bound yet. + """ + + namespace: Namespace + name: str + + def __repr__(self) -> str: + return f"{self.namespace._name}.{self.name}" + + @property + def value(self) -> Any: + return getattr(self.namespace.bound, self.name) + + def freeze(elem: Any, *, readonly: bool = False) -> Any: """ Give an offset provider element a content hash, so it identifies by value. @@ -81,13 +109,15 @@ def freeze(elem: Any, *, readonly: bool = False) -> Any: """ if common.frozen_content_hash(elem) is not None: return elem - buffer = getattr(elem, "ndarray", None) - if buffer is None: + if not hasattr(elem, "ndarray"): return elem if readonly: - buffer.flags.writeable = False - digest = int(eve_utils.content_hash(np.asarray(buffer)), 16) - object.__setattr__(elem, common.FROZEN_HASH_ATTR, digest) + # numpy-only; device buffers have no writeable flag + elem.ndarray.flags.writeable = False + # asnumpy, not np.asarray: the latter refuses device arrays outright, so + # hashing through it fails on every GPU backend. + host = elem.asnumpy() if hasattr(elem, "asnumpy") else np.asarray(elem.ndarray) + object.__setattr__(elem, common.FROZEN_HASH_ATTR, int(eve_utils.content_hash(host), 16)) return elem @@ -120,10 +150,18 @@ def resolve(explicit: common.OffsetProvider | None) -> common.OffsetProvider: def offset_provider_of(value: Any) -> dict[str, Any]: - """The connectivities and dimensions reachable as public attributes of `value`.""" + """ + The connectivities and dimensions reachable as public attributes of `value`. + + Attribute *order* is preserved from the bound object's own `__dict__`, not + taken from `dir()`: `dir()` sorts alphabetically, and gt4py's offset + provider is order-sensitive (see `hash_offset_provider_items_by_id`), so + reordering silently hands a compiled program the wrong tables. + """ + names = vars(value).keys() if hasattr(value, "__dict__") else dir(value) return { key: elem - for key in dir(value) + for key in names if not key.startswith("_") and isinstance(elem := getattr(value, key), (common.Connectivity, common.Dimension)) } diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py new file mode 100644 index 0000000000..ab74e0d58e --- /dev/null +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py @@ -0,0 +1,94 @@ +# 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 + +"""Ambient binding of the offset provider, across the backend matrix.""" + +import types + +import numpy as np +import pytest + +import gt4py.next as gtx +from gt4py.next import common, neighbor_sum + +from next_tests.integration_tests import cases +from next_tests.integration_tests.cases import ( + V2E, + Edge, + V2EDim, + Vertex, + unstructured_case, +) +from next_tests.integration_tests.cases_utils import ( + exec_alloc_descriptor, + mesh_descriptor, +) + + +@gtx.field_operator +def sum_edges(edge_f: cases.EField) -> cases.VField: + return neighbor_sum(edge_f(V2E), axis=V2EDim) + + +@gtx.program +def sum_edges_program(edge_f: cases.EField, out: cases.VField) -> None: + sum_edges(edge_f, out=out) + + +def _reference(case, inp): + v2e_table = case.offset_provider["V2E"].asnumpy() + return np.sum( + inp.asnumpy()[v2e_table], + axis=1, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ) + + +@pytest.mark.uses_unstructured_shift +@pytest.mark.parametrize("entry_point", ["program", "field_operator"]) +@pytest.mark.parametrize("mechanism", ["offset_provider", "context_manager", "bind_kwarg"]) +def test_ambient_binding_matches_explicit_offset_provider( + unstructured_case, entry_point, mechanism +): + """Every spelling produces the same result on every backend in the matrix.""" + inp = cases.allocate(unstructured_case, sum_edges, "edge_f")() + out = cases.allocate(unstructured_case, sum_edges, cases.RETURN)() + + mesh = gtx.Namespace("mesh") + # the fixture hands out a plain mapping; an ambient namespace binds an + # *object* whose public attributes are the connectivities + bound = types.SimpleNamespace(**unstructured_case.offset_provider) + + if entry_point == "program": + callee = sum_edges_program.with_backend(unstructured_case.backend) + args, kwargs = (inp, out), {} + else: + callee = sum_edges.with_backend(unstructured_case.backend) + args, kwargs = (inp,), {"out": out} + + if mechanism == "offset_provider": + callee(*args, **kwargs, offset_provider=unstructured_case.offset_provider) + elif mechanism == "context_manager": + with gtx.bind(mesh, bound): + callee(*args, **kwargs) + else: + callee(*args, **kwargs, bind={mesh: bound}) + + np.testing.assert_allclose(out.asnumpy(), _reference(unstructured_case, inp)) + + +@pytest.mark.uses_unstructured_shift +def test_bind_kwarg_does_not_leak_past_the_call(unstructured_case): + inp = cases.allocate(unstructured_case, sum_edges, "edge_f")() + out = cases.allocate(unstructured_case, sum_edges, cases.RETURN)() + mesh = gtx.Namespace("mesh") + bound = types.SimpleNamespace(**unstructured_case.offset_provider) + + sum_edges_program.with_backend(unstructured_case.backend)(inp, out, bind={mesh: bound}) + + assert gtx.ambient.offset_provider() == {} From 5b39d47ad08b03efabf1ef8ae2a7bcefbf48ed6f Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 21:10:45 +0200 Subject: [PATCH 05/14] proto[next]: Extern[T] / Static[T] ambient value declarations A declaration carries the type, so an operator referring to an ambient value types at decoration without the value being bound: 'from_value' dispatches on '__gt_type__', so no type-system change is needed. Execution is not wired up yet. The reference survives FOAST but dies at 'itir.Program' construction ('Symbols {SymbolRef(dx)} not found'), because eve validates symbol refs before any remap can run. Both forms therefore need the reference to become a real parameter first. --- src/gt4py/next/__init__.py | 4 ++- src/gt4py/next/ambient.py | 65 +++++++++++++++++++++++++++++++++++--- 2 files changed, 64 insertions(+), 5 deletions(-) diff --git a/src/gt4py/next/__init__.py b/src/gt4py/next/__init__.py index d3dbde3ddc..cfcebc515d 100644 --- a/src/gt4py/next/__init__.py +++ b/src/gt4py/next/__init__.py @@ -21,7 +21,7 @@ # ruff: noqa: F401 from .._core.definitions import CUPY_DEVICE_TYPE, Device, DeviceType, is_scalar_type from . import ambient, common, ffront, iterator, program_processors, typing -from .ambient import Namespace, bind, freeze +from .ambient import Extern, Namespace, Static, bind, freeze from .common import ( CartesianConnectivity, Connectivity, @@ -141,6 +141,8 @@ "ambient", "Namespace", "bind", + "Extern", + "Static", "freeze", # from ffront "FieldOffset", diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index c73a156f7c..c89062edc6 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -34,9 +34,10 @@ from gt4py.next import common -_bindings: contextvars.ContextVar[Mapping[Namespace, Any]] = contextvars.ContextVar( - "_ambient_bindings" -) +_UNBOUND: Any = object() + +#: keyed by declaration object: a `Namespace` or an `AmbientValue` +_bindings: contextvars.ContextVar[Mapping[Any, Any]] = contextvars.ContextVar("_ambient_bindings") class Namespace: @@ -93,6 +94,62 @@ def value(self) -> Any: return getattr(self.namespace.bound, self.name) +class AmbientValue: + """ + A value declared here and bound later: `dx = Extern[float]`. + + The declaration carries the *type*, which is all the frontend needs when the + operator is defined; the *value* arrives at bind time. Use the declaration + object itself as the binding key: `bind={dx: 0.5}`. + + `Extern[T]` is supplied to the compiled program as a runtime argument. + `Static[T]` is folded into it as a literal, so each distinct value gets its + own compiled variant. + """ + + def __init__(self, type_hint: Any, *, static: bool, name: str = "?") -> None: + self.type_hint = type_hint + self.static = static + self.name = name + + def __repr__(self) -> str: + return f"{'Static' if self.static else 'Extern'}[{getattr(self.type_hint, '__name__', self.type_hint)}]" + + def __gt_type__(self) -> Any: + from gt4py.next.type_system import type_translation + + return type_translation.from_type_hint(self.type_hint) + + @property + def value(self) -> Any: + binding = _bindings.get({}).get(self, _UNBOUND) + if binding is _UNBOUND: + raise ValueError( + f"Ambient value '{self!r}' is not bound." + f" Pass 'bind={{: }}' at the call, or use 'gtx.bind'." + ) + return binding + + +class _Declarator: + def __init__(self, static: bool) -> None: + self._static = static + + def __getitem__(self, type_hint: Any) -> AmbientValue: + return AmbientValue(type_hint, static=self._static) + + +#: `dx = Extern[float]` — supplied as a runtime argument. +Extern = _Declarator(static=False) +#: `dx = Static[float]` — folded in as a literal; one compiled variant per value. +Static = _Declarator(static=True) + + +def ambient_values_in(closure_vars: Mapping[str, Any]) -> dict[str, AmbientValue]: + """The ambient declarations referenced by a set of closure variables.""" + return {k: v for k, v in closure_vars.items() if isinstance(v, AmbientValue)} + + def freeze(elem: Any, *, readonly: bool = False) -> Any: """ Give an offset provider element a content hash, so it identifies by value. @@ -122,7 +179,7 @@ def freeze(elem: Any, *, readonly: bool = False) -> Any: @contextlib.contextmanager -def bindings(mapping: Mapping[Namespace, Any]) -> Generator[None, None, None]: +def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: """Bind several namespaces at once, for the duration of the context.""" if not mapping: yield From fbf5638c5f37edcf7ecfb413a4f1827a47542f01 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 21:52:26 +0200 Subject: [PATCH 06/14] proto[next]: ambient Static[T] values reachable by bare name An operator can refer to 'dx = Static[float]' without it being a parameter, so it need not be threaded through nested operators. Three pieces: - embedded: a bound declaration behaves as the scalar it stands for. - compiled: '_SubstituteAmbientValues' replaces the reference by its value before lowering, since 'ClosureVarFolding' runs at decoration when nothing is bound yet. - caching: bound values enter the compiled-program key, and are mirrored onto the declaration so the lowering cache (which fingerprints closure variables) cannot serve one value's code for another. 'Extern[T]' currently folds like 'Static[T]'; making it a runtime argument needs a synthesised program parameter, as does the ambient-field case. --- src/gt4py/next/ambient.py | 111 +++++++++++++++++- src/gt4py/next/ffront/foast_to_gtir.py | 27 ++++- src/gt4py/next/otf/compiled_program.py | 11 +- .../ffront_tests/test_ambient_values.py | 108 +++++++++++++++++ 4 files changed, 248 insertions(+), 9 deletions(-) create mode 100644 tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index c89062edc6..554b444b2b 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -14,10 +14,21 @@ compiled backends), so the same programs can run against a second mesh in one process. -This module implements the *binding* half only: the connectivities a program -needs are taken from the ambient context when the caller passes no -``offset_provider``. Referring to ambient *fields* by name inside an operator -(``mesh.edge_length``) is not implemented here. +Two things are ambient: + +- **Connectivities**, via a `Namespace`: a program called without an + ``offset_provider`` takes one from whatever is bound. +- **Values**, via a declaration ``dx = Static[float]`` referenced by bare name + inside an operator. It never appears in a signature, so it does not have to be + threaded through nested operators. + +``Static[T]`` is folded into the generated code, so each distinct value gets its +own compiled variant. ``Extern[T]`` is meant to be supplied as a runtime +argument instead — **not yet implemented**: it currently behaves like +``Static[T]``, because making it a runtime argument requires synthesising a +program parameter (the reference cannot stay a free symbol: eve validates symbol +refs when ``itir.Program`` is constructed). Ambient *fields* (``mesh.edge_length``) +need that same parameter machinery. """ from __future__ import annotations @@ -130,6 +141,53 @@ def value(self) -> Any: ) return binding + # Embedded execution runs the operator body as plain Python, so a bound + # declaration has to behave like the scalar it stands for. + def _v(self) -> Any: + return self.value + + def __float__(self) -> float: + return float(self.value) + + def __int__(self) -> int: + return int(self.value) + + def __bool__(self) -> bool: + return bool(self.value) + + def __neg__(self) -> Any: + return -self.value + + def __add__(self, other: Any) -> Any: + return self.value + other + + def __radd__(self, other: Any) -> Any: + return other + self.value + + def __sub__(self, other: Any) -> Any: + return self.value - other + + def __rsub__(self, other: Any) -> Any: + return other - self.value + + def __mul__(self, other: Any) -> Any: + return self.value * other + + def __rmul__(self, other: Any) -> Any: + return other * self.value + + def __truediv__(self, other: Any) -> Any: + return self.value / other + + def __rtruediv__(self, other: Any) -> Any: + return other / self.value + + def __pow__(self, other: Any) -> Any: + return self.value**other + + def __rpow__(self, other: Any) -> Any: + return other**self.value + class _Declarator: def __init__(self, static: bool) -> None: @@ -188,9 +246,27 @@ def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: for elem in offset_provider_of(value).values(): freeze(elem) token = _bindings.set({**_bindings.get({}), **mapping}) + # The lowered code bakes a `Static[T]` value in, and the lowering cache + # fingerprints the stage's closure variables — i.e. the declaration object. + # Mirroring the binding into the instance makes that fingerprint vary with + # the value, so two values cannot share a cache entry. (Prototype + # limitation: this mirror is process-wide, unlike the ContextVar itself.) + previous = [ + (decl, decl.__dict__.get("_bound", _UNBOUND)) + for decl in mapping + if isinstance(decl, AmbientValue) + ] + for decl, value in mapping.items(): + if isinstance(decl, AmbientValue): + decl.__dict__["_bound"] = value try: yield finally: + for decl, old_value in previous: + if old_value is _UNBOUND: + decl.__dict__.pop("_bound", None) + else: + decl.__dict__["_bound"] = old_value _bindings.reset(token) @@ -242,3 +318,30 @@ def offset_provider() -> common.OffsetProvider: ) collected[key] = elem return collected + + +def bound_values_in(closure_vars: Mapping[str, Any]) -> dict[str, Any]: + """ + Resolved values of the ambient declarations referenced by `closure_vars`. + + Unbound declarations are skipped rather than raising: a program may close + over declarations it does not use on this path, and the ones it does use + surface later as a missing symbol. + """ + current = _bindings.get({}) + return { + name: current[decl] + for name, decl in ambient_values_in(closure_vars).items() + if decl in current + } + + +def current_static_key() -> tuple[Any, ...]: + """A hashable summary of the bound `Static[T]` values, for compiled-program keys.""" + return tuple( + sorted( + (id(decl), eve_utils.content_hash(value)) + for decl, value in _bindings.get({}).items() + if isinstance(decl, AmbientValue) and decl.static + ) + ) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 10bc754526..6b47e305f3 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -12,7 +12,7 @@ from gt4py import eve from gt4py.eve.extended_typing import Never, cast -from gt4py.next import common, utils +from gt4py.next import ambient, common, utils from gt4py.next.ffront import ( dialect_ast_enums, experimental as experimental_builtins, @@ -38,7 +38,30 @@ def foast_to_gtir(inp: ffront_stages.FOASTOperatorDef) -> itir.FunctionDefinitio See the docstring of `FieldOperatorLowering` for details. """ - return FieldOperatorLowering.apply(inp.foast_node) + node = inp.foast_node + bound = ambient.bound_values_in(inp.closure_vars) + if bound: + node = _SubstituteAmbientValues(bound).visit(node) + return FieldOperatorLowering.apply(node) + + +class _SubstituteAmbientValues(eve.PreserveLocationVisitor, eve.NodeTranslator): + """ + Replace references to bound ambient declarations by their value. + + Runs here rather than in `ClosureVarFolding` because nothing is bound when + the operator is defined; the value only exists at call time. + """ + + def __init__(self, values: dict[str, Any]) -> None: + self.values = values + + def visit_Name(self, node: foast.Name, **kwargs: Any) -> foast.Name | foast.Constant: + if node.id in self.values: + return foast.Constant( + value=self.values[node.id], type=node.type, location=node.location + ) + return node def foast_to_gtir_factory( diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 0784ca0f73..36592fda98 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -24,7 +24,7 @@ from gt4py._core import definitions as core_defs from gt4py.eve import extended_typing as xtyping, utils as eve_utils -from gt4py.next import backend as gtx_backend, common, errors, utils as gtx_utils +from gt4py.next import ambient, backend as gtx_backend, common, errors, utils as gtx_utils from gt4py.next.ffront import ( stages as ffront_stages, type_info as ffront_type_info, @@ -41,8 +41,9 @@ ScalarOrTupleOfScalars: TypeAlias = xtyping.MaybeNestedInTuple[core_defs.Scalar] -#: Content of the key: (*hashable_arg_descriptors, id(offset_provider), concrete_instantation_if_generic) -CompiledProgramsKey: TypeAlias = tuple[tuple[Hashable, ...], int, None | str] +#: Content of the key: (*hashable_arg_descriptors, id(offset_provider), +#: concrete_instantation_if_generic, bound_ambient_static_values) +CompiledProgramsKey: TypeAlias = tuple[tuple[Hashable, ...], int, None | str, tuple[Any, ...]] ArgStaticDescriptorsByType: TypeAlias = dict[ type[arguments.ArgStaticDescriptor], dict[str, arguments.ArgStaticDescriptor] @@ -425,6 +426,9 @@ def __call__( static_args_values, common.hash_offset_provider_items_by_id(offset_provider), arg_specialization_key, + # ambient `Static[T]` values are folded into the generated code, so + # each distinct value needs its own compiled variant + ambient.current_static_key(), ) try: @@ -615,6 +619,7 @@ def _compile_variant( self._argument_descriptor_cache_key_from_descriptors(argument_descriptor_contexts), common.hash_offset_provider_items_by_id(offset_provider), eve_utils.content_hash(arg_specialization_info) if self._is_generic else None, + ambient.current_static_key(), ) assert call_key is None or call_key == key diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py new file mode 100644 index 0000000000..0339f0908c --- /dev/null +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -0,0 +1,108 @@ +# 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 + +"""Ambient values referenced by name inside an operator, bound at call time.""" + +import numpy as np +import pytest + +import gt4py.next as gtx +from gt4py.next.type_system import type_specifications as ts + +from next_tests.integration_tests import cases +from next_tests.integration_tests.cases import IDim, JDim, cartesian_case +from next_tests.integration_tests.cases_utils import exec_alloc_descriptor + + +IJFloatField = gtx.Field[gtx.Dims[IDim, JDim], gtx.float64] + +#: declared here, bound at the call — never a parameter, so it does not have to +#: be threaded through nested operators +dx = gtx.Static[gtx.float64] + + +@gtx.field_operator +def delta_x(f: IJFloatField) -> IJFloatField: + """Forward difference in x.""" + return (1.0 / dx) * (f(IDim + 1) - f) + + +@gtx.field_operator +def delta_x_twice(f: IJFloatField) -> IJFloatField: + """`dx` is used one level down, and still never appears in a signature.""" + return delta_x(f) + delta_x(f) + + +@gtx.program +def run_delta_x(f: IJFloatField, out: IJFloatField) -> None: + delta_x(f, out=out) + + +@gtx.program +def run_delta_x_twice(f: IJFloatField, out: IJFloatField) -> None: + delta_x_twice(f, out=out) + + +def _inputs(case): + data = gtx.as_field([IDim, JDim], np.arange(20.0).reshape(5, 4), allocator=case.allocator) + out = gtx.zeros(gtx.domain({IDim: 4, JDim: 4}), dtype=np.float64, allocator=case.allocator) + return data, out + + +def _reference(data, spacing, factor=1): + a = data.asnumpy() + return factor * (1.0 / spacing) * (a[1:5, :] - a[0:4, :]) + + +@pytest.mark.uses_cartesian_shift +@pytest.mark.parametrize("spacing", [0.5, 0.25]) +def test_static_value_is_bound_at_call(cartesian_case, spacing): + data, out = _inputs(cartesian_case) + run_delta_x.with_backend(cartesian_case.backend)(data, out, bind={dx: spacing}) + np.testing.assert_allclose(out.asnumpy(), _reference(data, spacing)) + + +@pytest.mark.uses_cartesian_shift +def test_distinct_values_do_not_share_a_compiled_program(cartesian_case): + """The second binding must not reuse the first one's folded literal.""" + data, out = _inputs(cartesian_case) + prog = run_delta_x.with_backend(cartesian_case.backend) + + prog(data, out, bind={dx: 0.5}) + np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5)) + + prog(data, out, bind={dx: 0.25}) + np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.25)) + + +@pytest.mark.uses_cartesian_shift +def test_value_reaches_a_nested_operator(cartesian_case): + """`dx` is referenced two levels down without being passed as an argument.""" + data, out = _inputs(cartesian_case) + run_delta_x_twice.with_backend(cartesian_case.backend)(data, out, bind={dx: 0.5}) + np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5, factor=2)) + + +@pytest.mark.uses_cartesian_shift +def test_context_manager_binding(cartesian_case): + data, out = _inputs(cartesian_case) + with gtx.bind(dx, 0.5): + run_delta_x.with_backend(cartesian_case.backend)(data, out) + np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5)) + + +def test_declaration_types_itself_without_a_binding(): + """The frontend only needs the type at decoration; the value comes later.""" + expected = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + assert gtx.Static[gtx.float64].__gt_type__() == expected + assert gtx.Extern[gtx.float64].__gt_type__() == expected + + +def test_unbound_declaration_reports_itself(): + with pytest.raises(ValueError, match="not bound"): + gtx.Static[gtx.float64].value From c89780a73cc2893e46851a27fc14373579aacea2 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 21:57:55 +0200 Subject: [PATCH 07/14] proto[next]: fingerprint ambient bindings instead of mirroring them The lowering cache did not distinguish two bindings of the same declaration because its fingerprinter walks closure variables and sees only the declaration object. That was worked around by mirroring the bound value onto the declaration's '__dict__', which made the binding process-wide. Register a deconstructor for 'AmbientValue' with the frontend fingerprinter instead: it reads the current binding from the ContextVar, so the cache distinguishes values and bindings stay context-local. --- src/gt4py/next/ambient.py | 36 ++++++++++----------- src/gt4py/next/ffront/stages.py | 15 +++++++-- tests/next_tests/unit_tests/test_ambient.py | 18 +++++++++++ 3 files changed, 49 insertions(+), 20 deletions(-) diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index 554b444b2b..a07e784d1d 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -246,27 +246,9 @@ def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: for elem in offset_provider_of(value).values(): freeze(elem) token = _bindings.set({**_bindings.get({}), **mapping}) - # The lowered code bakes a `Static[T]` value in, and the lowering cache - # fingerprints the stage's closure variables — i.e. the declaration object. - # Mirroring the binding into the instance makes that fingerprint vary with - # the value, so two values cannot share a cache entry. (Prototype - # limitation: this mirror is process-wide, unlike the ContextVar itself.) - previous = [ - (decl, decl.__dict__.get("_bound", _UNBOUND)) - for decl in mapping - if isinstance(decl, AmbientValue) - ] - for decl, value in mapping.items(): - if isinstance(decl, AmbientValue): - decl.__dict__["_bound"] = value try: yield finally: - for decl, old_value in previous: - if old_value is _UNBOUND: - decl.__dict__.pop("_bound", None) - else: - decl.__dict__["_bound"] = old_value _bindings.reset(token) @@ -345,3 +327,21 @@ def current_static_key() -> tuple[Any, ...]: if isinstance(decl, AmbientValue) and decl.static ) ) + + +def fingerprint_declaration(decl: AmbientValue) -> Any: + """ + What an ambient declaration contributes to a stage fingerprint. + + A `Static[T]` value is baked into the lowered code, so the *current binding* + has to be part of the fingerprint — otherwise the lowering cache serves one + value's code for another. Reads the binding rather than mirroring it onto + the declaration, so it stays context-local. + """ + current = _bindings.get({}) + bound = current.get(decl, _UNBOUND) + return ( + decl.static, + decl.type_hint, + None if bound is _UNBOUND else eve_utils.content_hash(bound), + ) diff --git a/src/gt4py/next/ffront/stages.py b/src/gt4py/next/ffront/stages.py index 0651a69739..bbf4134ee4 100644 --- a/src/gt4py/next/ffront/stages.py +++ b/src/gt4py/next/ffront/stages.py @@ -27,7 +27,7 @@ import typing from typing import Any, Optional, TypeVar -from gt4py.next import common, fingerprinting +from gt4py.next import ambient, common, fingerprinting from gt4py.next.ffront import field_operator_ast as foast, program_ast as past, source_utils from gt4py.next.otf import arguments, toolchain @@ -59,9 +59,20 @@ def _deconstruct_definition_function(func: types.FunctionType) -> fingerprinting #: identical operators at different locations must not share an entry) and #: fingerprints DSL definition functions by their source code and closure #: variables (instead of by qualified name). +#: An ambient declaration contributes its *current binding*: a `Static[T]` value +#: is baked into the lowered code, so two bindings must not share a cache entry. +def _deconstruct_ambient_declaration(decl: ambient.AmbientValue) -> fingerprinting.Deconstruction: + return fingerprinting.Deconstruction.from_pieces( + ambient.fingerprint_declaration(decl), state=b"ambient_declaration" + ) + + semantic_fingerprinter: fingerprinting.Fingerprinter = fingerprinting.make_fingerprinter( deconstructor=fingerprinting.make_lenient_data_deconstructor( - {types.FunctionType: _deconstruct_definition_function} + { + types.FunctionType: _deconstruct_definition_function, + ambient.AmbientValue: _deconstruct_ambient_declaration, + } ), ) diff --git a/tests/next_tests/unit_tests/test_ambient.py b/tests/next_tests/unit_tests/test_ambient.py index d1d1c8b06d..6a09bb2b9a 100644 --- a/tests/next_tests/unit_tests/test_ambient.py +++ b/tests/next_tests/unit_tests/test_ambient.py @@ -154,3 +154,21 @@ def test_bind_kwarg_is_scoped_to_the_call(inputs): out = gtx.zeros(gtx.domain({Vertex: 2}), dtype=np.int32) run(a, out, bind={mesh: m}) assert gtx.ambient.offset_provider() == {} + + +def test_bindings_are_context_local(): + """The binding must not leak between contexts (the fingerprint reads it, not a mirror).""" + import contextvars + + decl = gtx.Static[gtx.float64] + seen = {} + + def bind_and_record(): + with gtx.bind(decl, 0.25): + seen["inner"] = decl.value + + with gtx.bind(decl, 0.5): + contextvars.copy_context().run(bind_and_record) + seen["outer"] = decl.value + + assert seen == {"inner": 0.25, "outer": 0.5} From 1a8052a5c0b9ef341f30b1769ab30d217b46a0dc Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 22:13:01 +0200 Subject: [PATCH 08/14] proto[next]: synthesise a program parameter for ambient values MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the bespoke machinery with a parameter, added once when the program is defined ('func_to_past'). From there ambient values travel the ordinary path: type checking, lowering, 'static_params' and the compiled-program key all treat them as arguments, and only the value is supplied per call. The two forms now differ in one place only — whether the synthesised parameter is listed as static: - 'Extern[T]': ordinary runtime argument, one compiled program for all values. - 'Static[T]': static argument, so the existing fold bakes it in and the existing key specialises on it — one program per distinct value. Measured on gtfn and dace, 2 values: Static 2 variants, Extern 1, both correct. Deletes '_SubstituteAmbientValues', 'current_static_key' and the ambient fingerprint deconstructor: the operator's IR no longer holds the value, so the lowering cache needs no help. A free symbol in a lowered operator resolves against the program's parameters, so no threading into operator signatures is needed. --- src/gt4py/next/ambient.py | 66 ++++--------------- src/gt4py/next/ffront/decorator.py | 39 ++++++++++- src/gt4py/next/ffront/foast_to_gtir.py | 27 +------- src/gt4py/next/ffront/func_to_past.py | 54 ++++++++++++++- src/gt4py/next/ffront/stages.py | 15 +---- src/gt4py/next/otf/compiled_program.py | 11 +--- .../ffront_tests/test_ambient_values.py | 41 +++++++++++- 7 files changed, 148 insertions(+), 105 deletions(-) diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index a07e784d1d..8aa643f57c 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -22,13 +22,20 @@ inside an operator. It never appears in a signature, so it does not have to be threaded through nested operators. -``Static[T]`` is folded into the generated code, so each distinct value gets its -own compiled variant. ``Extern[T]`` is meant to be supplied as a runtime -argument instead — **not yet implemented**: it currently behaves like -``Static[T]``, because making it a runtime argument requires synthesising a -program parameter (the reference cannot stay a free symbol: eve validates symbol -refs when ``itir.Program`` is constructed). Ambient *fields* (``mesh.edge_length``) -need that same parameter machinery. +A declaration becomes a **synthesised program parameter** when the program is +defined (`func_to_past`), so from there on it travels the ordinary path: type +checking, lowering, `static_params` and the compiled-program key all treat it +like any other argument, and only the *value* is supplied per call. The caller +never names it. + +The two forms differ in one place only — whether that parameter is listed as a +static one: + +- ``Extern[T]`` is an ordinary runtime argument: one compiled program serves + every value. +- ``Static[T]`` is a static argument, so the existing static-argument machinery + folds it into the generated code and keys the compiled variant on it: one + compiled program per distinct value. """ from __future__ import annotations @@ -300,48 +307,3 @@ def offset_provider() -> common.OffsetProvider: ) collected[key] = elem return collected - - -def bound_values_in(closure_vars: Mapping[str, Any]) -> dict[str, Any]: - """ - Resolved values of the ambient declarations referenced by `closure_vars`. - - Unbound declarations are skipped rather than raising: a program may close - over declarations it does not use on this path, and the ones it does use - surface later as a missing symbol. - """ - current = _bindings.get({}) - return { - name: current[decl] - for name, decl in ambient_values_in(closure_vars).items() - if decl in current - } - - -def current_static_key() -> tuple[Any, ...]: - """A hashable summary of the bound `Static[T]` values, for compiled-program keys.""" - return tuple( - sorted( - (id(decl), eve_utils.content_hash(value)) - for decl, value in _bindings.get({}).items() - if isinstance(decl, AmbientValue) and decl.static - ) - ) - - -def fingerprint_declaration(decl: AmbientValue) -> Any: - """ - What an ambient declaration contributes to a stage fingerprint. - - A `Static[T]` value is baked into the lowered code, so the *current binding* - has to be part of the fingerprint — otherwise the lowering cache serves one - value's code for another. Reads the binding rather than mirroring it onto - the declaration, so it stays context-local. - """ - current = _bindings.get({}) - bound = current.get(decl, _UNBOUND) - return ( - decl.static, - decl.type_hint, - None if bound is _UNBOUND else eve_utils.content_hash(bound), - ) diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 6800b6fc08..3c10c0232a 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -126,10 +126,21 @@ def _compiled_programs(self) -> compiled_program.CompiledProgramsPool: # calling `compile()`, the pool is initialized with the options passed # to `compile()` instead of re-using the existing compilations options. return self._make_compiled_programs_pool( - static_params=self.compilation_options.static_params or (), + static_params=(*(self.compilation_options.static_params or ()), *self._ambient_statics), static_domains=self.compilation_options.static_domains, ) + @property + def _ambient_statics(self) -> tuple[str, ...]: + """Synthesised parameters for ambient `Static[T]` declarations. + + Listing them as static parameters is the whole difference between + `Static[T]` and `Extern[T]`: the existing static-argument machinery then + folds the value into the generated code and keys the compiled variant on + it, while an `Extern[T]` stays an ordinary runtime argument. + """ + return () + def _make_compiled_programs_pool( self, static_params: Sequence[str], static_domains: bool ) -> compiled_program.CompiledProgramsPool: @@ -308,6 +319,16 @@ def _frontend_transforms(self) -> next_backend.Transforms: def _all_closure_vars(self) -> dict[str, Any]: return transform_utils._get_closure_vars_recursively(self.past_stage.closure_vars) + @property + def _ambient_statics(self) -> tuple[str, ...]: + return tuple( + sorted( + name + for name, decl in ambient.ambient_values_in(self._all_closure_vars).items() + if decl.static + ) + ) + @functools.cached_property def gtir(self) -> itir.Program: no_args_past = toolchain.ConcreteArtifact( @@ -397,6 +418,15 @@ def _invoke( ) -> None: offset_provider = ambient.resolve(offset_provider) enable_jit = self.compilation_options.enable_jit if enable_jit is None else enable_jit + # Ambient declarations became synthesised parameters when the program was + # defined, so from here on they are ordinary arguments and the caller + # never names them. Embedded execution is the exception: it runs the + # Python function, which reads them from its own closure. + ambient_args = { + name: decl.value + for name, decl in ambient.ambient_values_in(self._all_closure_vars).items() + } + kwargs = {**kwargs, **ambient_args} with program_call_context( program=self, @@ -426,9 +456,12 @@ def _invoke( stacklevel=2, ) + embedded_kwargs = {k: v for k, v in kwargs.items() if k not in ambient_args} with next_embedded.context.update(offset_provider=offset_provider): - with embedded_program_call_context(self, args, offset_provider, kwargs): - self.definition_stage.definition(*args, **kwargs) + with embedded_program_call_context( + self, args, offset_provider, embedded_kwargs + ): + self.definition_stage.definition(*args, **embedded_kwargs) try: diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 6b47e305f3..10bc754526 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -12,7 +12,7 @@ from gt4py import eve from gt4py.eve.extended_typing import Never, cast -from gt4py.next import ambient, common, utils +from gt4py.next import common, utils from gt4py.next.ffront import ( dialect_ast_enums, experimental as experimental_builtins, @@ -38,30 +38,7 @@ def foast_to_gtir(inp: ffront_stages.FOASTOperatorDef) -> itir.FunctionDefinitio See the docstring of `FieldOperatorLowering` for details. """ - node = inp.foast_node - bound = ambient.bound_values_in(inp.closure_vars) - if bound: - node = _SubstituteAmbientValues(bound).visit(node) - return FieldOperatorLowering.apply(node) - - -class _SubstituteAmbientValues(eve.PreserveLocationVisitor, eve.NodeTranslator): - """ - Replace references to bound ambient declarations by their value. - - Runs here rather than in `ClosureVarFolding` because nothing is bound when - the operator is defined; the value only exists at call time. - """ - - def __init__(self, values: dict[str, Any]) -> None: - self.values = values - - def visit_Name(self, node: foast.Name, **kwargs: Any) -> foast.Name | foast.Constant: - if node.id in self.values: - return foast.Constant( - value=self.values[node.id], type=node.type, location=node.location - ) - return node + return FieldOperatorLowering.apply(inp.foast_node) def foast_to_gtir_factory( diff --git a/src/gt4py/next/ffront/func_to_past.py b/src/gt4py/next/ffront/func_to_past.py index 292a56767b..a2a99c4627 100644 --- a/src/gt4py/next/ffront/func_to_past.py +++ b/src/gt4py/next/ffront/func_to_past.py @@ -14,7 +14,7 @@ from typing import Any, cast from gt4py._core import definitions as core_defs -from gt4py.next import errors +from gt4py.next import ambient, errors from gt4py.next.ffront import ( dialect_ast_enums, experimental, @@ -22,6 +22,7 @@ program_ast as past, source_utils, stages as ffront_stages, + transform_utils, type_specifications as ts_ffront, ) from gt4py.next.ffront.dialect_parser import DialectParser @@ -65,13 +66,62 @@ def func_to_past(inp: DSLProgramDef) -> PASTProgramDef: closure_vars = source_utils.get_closure_vars_from_function(inp.definition) annotations = typing.get_type_hints(inp.definition) return ffront_stages.PASTProgramDef( - past_node=ProgramParser.apply(source_def, closure_vars, annotations), + past_node=_with_ambient_params( + ProgramParser.apply(source_def, closure_vars, annotations), closure_vars + ), closure_vars=closure_vars, grid_type=inp.grid_type, debug=inp.debug, ) +def _with_ambient_params(node: past.Program, closure_vars: dict[str, typing.Any]) -> past.Program: + """ + Give every ambient declaration the program references a synthesised parameter. + + Ambient values are referenced by bare name inside operators, so at the IR + level they are free symbols. Turning them into parameters here — once, when + the program is defined — puts them on the ordinary path: type checking, + lowering, `static_params` and the compiled-program key all treat them like + any other argument, and only the *value* has to be supplied per call. + """ + declarations = ambient.ambient_values_in( + transform_utils._get_closure_vars_recursively(closure_vars) + ) + if not declarations: + return node + assert isinstance(node.type, ts_ffront.ProgramType) + extra = [ + past.DataSymbol( + id=name, + type=decl.__gt_type__(), + namespace=dialect_ast_enums.Namespace.LOCAL, + location=node.location, + ) + # deterministic order: the parameter list is part of the compiled signature + for name, decl in sorted(declarations.items()) + ] + definition = node.type.definition + return past.Program( + id=node.id, + type=ts_ffront.ProgramType( + definition=ts.FunctionType( + pos_only_args=definition.pos_only_args, + pos_or_kw_args={ + **definition.pos_or_kw_args, + **{s.id: s.type for s in extra}, + }, + kw_only_args=definition.kw_only_args, + returns=definition.returns, + ) + ), + params=[*node.params, *extra], + body=node.body, + closure_vars=node.closure_vars, + location=node.location, + ) + + def func_to_past_factory(cached: bool = True) -> workflow.Workflow[DSLProgramDef, PASTProgramDef]: """ Wrap `func_to_past` in a chainable and optionally cached workflow step. diff --git a/src/gt4py/next/ffront/stages.py b/src/gt4py/next/ffront/stages.py index bbf4134ee4..0651a69739 100644 --- a/src/gt4py/next/ffront/stages.py +++ b/src/gt4py/next/ffront/stages.py @@ -27,7 +27,7 @@ import typing from typing import Any, Optional, TypeVar -from gt4py.next import ambient, common, fingerprinting +from gt4py.next import common, fingerprinting from gt4py.next.ffront import field_operator_ast as foast, program_ast as past, source_utils from gt4py.next.otf import arguments, toolchain @@ -59,20 +59,9 @@ def _deconstruct_definition_function(func: types.FunctionType) -> fingerprinting #: identical operators at different locations must not share an entry) and #: fingerprints DSL definition functions by their source code and closure #: variables (instead of by qualified name). -#: An ambient declaration contributes its *current binding*: a `Static[T]` value -#: is baked into the lowered code, so two bindings must not share a cache entry. -def _deconstruct_ambient_declaration(decl: ambient.AmbientValue) -> fingerprinting.Deconstruction: - return fingerprinting.Deconstruction.from_pieces( - ambient.fingerprint_declaration(decl), state=b"ambient_declaration" - ) - - semantic_fingerprinter: fingerprinting.Fingerprinter = fingerprinting.make_fingerprinter( deconstructor=fingerprinting.make_lenient_data_deconstructor( - { - types.FunctionType: _deconstruct_definition_function, - ambient.AmbientValue: _deconstruct_ambient_declaration, - } + {types.FunctionType: _deconstruct_definition_function} ), ) diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 36592fda98..0784ca0f73 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -24,7 +24,7 @@ from gt4py._core import definitions as core_defs from gt4py.eve import extended_typing as xtyping, utils as eve_utils -from gt4py.next import ambient, backend as gtx_backend, common, errors, utils as gtx_utils +from gt4py.next import backend as gtx_backend, common, errors, utils as gtx_utils from gt4py.next.ffront import ( stages as ffront_stages, type_info as ffront_type_info, @@ -41,9 +41,8 @@ ScalarOrTupleOfScalars: TypeAlias = xtyping.MaybeNestedInTuple[core_defs.Scalar] -#: Content of the key: (*hashable_arg_descriptors, id(offset_provider), -#: concrete_instantation_if_generic, bound_ambient_static_values) -CompiledProgramsKey: TypeAlias = tuple[tuple[Hashable, ...], int, None | str, tuple[Any, ...]] +#: Content of the key: (*hashable_arg_descriptors, id(offset_provider), concrete_instantation_if_generic) +CompiledProgramsKey: TypeAlias = tuple[tuple[Hashable, ...], int, None | str] ArgStaticDescriptorsByType: TypeAlias = dict[ type[arguments.ArgStaticDescriptor], dict[str, arguments.ArgStaticDescriptor] @@ -426,9 +425,6 @@ def __call__( static_args_values, common.hash_offset_provider_items_by_id(offset_provider), arg_specialization_key, - # ambient `Static[T]` values are folded into the generated code, so - # each distinct value needs its own compiled variant - ambient.current_static_key(), ) try: @@ -619,7 +615,6 @@ def _compile_variant( self._argument_descriptor_cache_key_from_descriptors(argument_descriptor_contexts), common.hash_offset_provider_items_by_id(offset_provider), eve_utils.content_hash(arg_specialization_info) if self._is_generic else None, - ambient.current_static_key(), ) assert call_key is None or call_key == key diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py index 0339f0908c..e354425bad 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -21,9 +21,10 @@ IJFloatField = gtx.Field[gtx.Dims[IDim, JDim], gtx.float64] -#: declared here, bound at the call — never a parameter, so it does not have to -#: be threaded through nested operators +#: declared here, bound at the call — never a parameter in the user's source, so +#: it does not have to be threaded through nested operators dx = gtx.Static[gtx.float64] +dx_extern = gtx.Extern[gtx.float64] @gtx.field_operator @@ -43,11 +44,22 @@ def run_delta_x(f: IJFloatField, out: IJFloatField) -> None: delta_x(f, out=out) +@gtx.field_operator +def delta_x_extern(f: IJFloatField) -> IJFloatField: + """Same, but supplied as a runtime argument instead of folded in.""" + return (1.0 / dx_extern) * (f(IDim + 1) - f) + + @gtx.program def run_delta_x_twice(f: IJFloatField, out: IJFloatField) -> None: delta_x_twice(f, out=out) +@gtx.program +def run_delta_x_extern(f: IJFloatField, out: IJFloatField) -> None: + delta_x_extern(f, out=out) + + def _inputs(case): data = gtx.as_field([IDim, JDim], np.arange(20.0).reshape(5, 4), allocator=case.allocator) out = gtx.zeros(gtx.domain({IDim: 4, JDim: 4}), dtype=np.float64, allocator=case.allocator) @@ -96,6 +108,31 @@ def test_context_manager_binding(cartesian_case): np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5)) +@pytest.mark.uses_cartesian_shift +@pytest.mark.parametrize("spacing", [0.5, 0.25]) +def test_extern_value_is_bound_at_call(cartesian_case, spacing): + data, out = _inputs(cartesian_case) + run_delta_x_extern.with_backend(cartesian_case.backend)(data, out, bind={dx_extern: spacing}) + np.testing.assert_allclose(out.asnumpy(), _reference(data, spacing)) + + +@pytest.mark.uses_cartesian_shift +def test_static_specializes_but_extern_does_not(cartesian_case): + """The one difference between the two forms: how many programs get compiled.""" + if cartesian_case.backend is None: + pytest.skip("compiled-program pool only exists for compiled backends") + data, out = _inputs(cartesian_case) + + static_prog = run_delta_x.with_backend(cartesian_case.backend) + extern_prog = run_delta_x_extern.with_backend(cartesian_case.backend) + for spacing in (0.5, 0.25): + static_prog(data, out, bind={dx: spacing}) + extern_prog(data, out, bind={dx_extern: spacing}) + + assert len(static_prog._compiled_programs.compiled_programs) == 2 + assert len(extern_prog._compiled_programs.compiled_programs) == 1 + + def test_declaration_types_itself_without_a_binding(): """The frontend only needs the type at decoration; the value comes later.""" expected = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) From e10a3d9b93e104f6e7e545eba7619a24e71b1844 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 22:45:55 +0200 Subject: [PATCH 09/14] proto[next]: bind offsets by declaration, retiring Namespace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Connectivities and values were bound by two different rules: a 'Namespace' harvested connectivities from a bound object by attribute *name*, while values were keyed by declaration identity. Now there is one rule — the declaration is the key: prog(a, out, bind={V2E: connectivity, dx: 0.5}) A 'FieldOffset' already names the offset and fixes its source and target, so it is the declaration; the offset provider is assembled from the bound ones. This retires 'Namespace', 'offset_provider_of' and the name-collision check it needed, and lets a container supply the very offset an operator refers to rather than one that merely shares its name. The offset provider itself is unchanged: it is still assembled and passed as today, so only the binding surface moves. --- src/gt4py/next/__init__.py | 4 +- src/gt4py/next/ambient.py | 161 +++++++----------- src/gt4py/next/ffront/decorator.py | 6 +- .../ffront_tests/test_ambient_binding.py | 18 +- tests/next_tests/unit_tests/test_ambient.py | 49 +++--- 5 files changed, 100 insertions(+), 138 deletions(-) diff --git a/src/gt4py/next/__init__.py b/src/gt4py/next/__init__.py index cfcebc515d..011d5d8f8e 100644 --- a/src/gt4py/next/__init__.py +++ b/src/gt4py/next/__init__.py @@ -21,7 +21,7 @@ # ruff: noqa: F401 from .._core.definitions import CUPY_DEVICE_TYPE, Device, DeviceType, is_scalar_type from . import ambient, common, ffront, iterator, program_processors, typing -from .ambient import Extern, Namespace, Static, bind, freeze +from .ambient import Extern, Static, bind, bindings, freeze from .common import ( CartesianConnectivity, Connectivity, @@ -139,8 +139,8 @@ "as_connectivity", # from ambient "ambient", - "Namespace", "bind", + "bindings", "Extern", "Static", "freeze", diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index 8aa643f57c..dbd855d345 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -14,13 +14,27 @@ compiled backends), so the same programs can run against a second mesh in one process. -Two things are ambient: +Everything ambient is bound the same way: **the declaration is the key**. -- **Connectivities**, via a `Namespace`: a program called without an - ``offset_provider`` takes one from whatever is bound. -- **Values**, via a declaration ``dx = Static[float]`` referenced by bare name - inside an operator. It never appears in a signature, so it does not have to be - threaded through nested operators. + prog(a, out, bind={V2E: connectivity, dx: 0.5}) + +A `FieldOffset` is already a declaration — it names the offset and fixes its +source and target — so it binds exactly like a `Static[T]` / `Extern[T]` value, +and a program called without an ``offset_provider`` assembles one from the bound +offsets. A container may declare what it supplies, in which case the class +attribute is the declaration and the instance attribute the value:: + + class Mesh: + V2E = V2E # this mesh supplies the V2E connectivity + dx = physics.dx + +Binding by declaration rather than by attribute *name* is what lets a container +supply the very offset an operator refers to, instead of something that merely +shares its name. + +A value declared this way is referenced by bare name inside an operator and +never appears in a signature, so it does not have to be threaded through nested +operators. A declaration becomes a **synthesised program parameter** when the program is defined (`func_to_past`), so from there on it travels the ordinary path: type @@ -42,7 +56,6 @@ import contextlib import contextvars -import dataclasses from collections.abc import Generator, Mapping from typing import Any @@ -50,6 +63,7 @@ from gt4py.eve import utils as eve_utils from gt4py.next import common +from gt4py.next.ffront import fbuiltins _UNBOUND: Any = object() @@ -58,60 +72,6 @@ _bindings: contextvars.ContextVar[Mapping[Any, Any]] = contextvars.ContextVar("_ambient_bindings") -class Namespace: - """ - A named collection of ambient values, resolved by attribute access. - - The declaration is the namespace object itself; ``mesh.e2v`` is a reference - that only becomes a value once something is bound to ``mesh``. Attribute - names are not declared up front in this prototype. - """ - - def __init__(self, name: str) -> None: - self._name = name - - def __repr__(self) -> str: - return f"Namespace('{self._name}')" - - def __getattr__(self, name: str) -> AmbientRef: - # only reached when normal lookup fails; dunder/private probes (pickle, - # copy, the DSL frontend) must not be turned into references - if name.startswith("_"): - raise AttributeError(name) - return AmbientRef(self, name) - - @property - def bound(self) -> Any: - """The object currently bound to this namespace.""" - binding = _bindings.get({}).get(self, None) - if binding is None: - raise ValueError( - f"Nothing is bound to ambient namespace '{self._name}'." - " Use 'gtx.bind(, )' around the call." - ) - return binding - - -@dataclasses.dataclass(frozen=True) -class AmbientRef: - """ - A deferred reference to `namespace.name`, resolved when something is bound. - - Attribute access on a `Namespace` yields one of these instead of a value, - because at the point an operator is *defined* nothing is bound yet. - """ - - namespace: Namespace - name: str - - def __repr__(self) -> str: - return f"{self.namespace._name}.{self.name}" - - @property - def value(self) -> Any: - return getattr(self.namespace.bound, self.name) - - class AmbientValue: """ A value declared here and bound later: `dx = Extern[float]`. @@ -243,6 +203,36 @@ def freeze(elem: Any, *, readonly: bool = False) -> Any: return elem +def as_bindings(spec: Any) -> dict[Any, Any]: + """ + Normalise what `bind=` accepts into a declaration -> value mapping. + + A mapping is taken as-is. Anything else is treated as a *container*: its + class attributes name the declarations it supplies, and the instance + attribute of the same name carries the value. The declaration object itself + is the key, so a container binds the very `Extern` an operator refers to + rather than something that merely shares its name. + + class Grid: + dx = physics.dx # the declaration this grid supplies + + grid = Grid(); grid.dx = 0.5 + prog(f, out, bind=grid) + """ + if isinstance(spec, Mapping): + return dict(spec) + resolved: dict[Any, Any] = {} + for name, decl in vars(type(spec)).items(): + if not isinstance(decl, (AmbientValue, fbuiltins.FieldOffset)): + continue + if name not in vars(spec): + raise ValueError( + f"'{type(spec).__name__}' declares '{name}' but the instance does not set it." + ) + resolved[decl] = vars(spec)[name] + return resolved + + @contextlib.contextmanager def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: """Bind several namespaces at once, for the duration of the context.""" @@ -250,8 +240,8 @@ def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: yield return for value in mapping.values(): - for elem in offset_provider_of(value).values(): - freeze(elem) + if isinstance(value, common.Connectivity): + freeze(value) token = _bindings.set({**_bindings.get({}), **mapping}) try: yield @@ -260,9 +250,9 @@ def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: @contextlib.contextmanager -def bind(namespace: Namespace, value: Any) -> Generator[None, None, None]: - """Bind `value` to `namespace` for the duration of the context.""" - with bindings({namespace: value}): +def bind(declaration: Any, value: Any) -> Generator[None, None, None]: + """Bind `value` to `declaration` for the duration of the context.""" + with bindings({declaration: value}): yield @@ -271,39 +261,16 @@ def resolve(explicit: common.OffsetProvider | None) -> common.OffsetProvider: return offset_provider() if explicit is None else explicit -def offset_provider_of(value: Any) -> dict[str, Any]: +def offset_provider() -> common.OffsetProvider: """ - The connectivities and dimensions reachable as public attributes of `value`. + Assemble the offset provider from the bound offset declarations. - Attribute *order* is preserved from the bound object's own `__dict__`, not - taken from `dir()`: `dir()` sorts alphabetically, and gt4py's offset - provider is order-sensitive (see `hash_offset_provider_items_by_id`), so - reordering silently hands a compiled program the wrong tables. + A `FieldOffset` is itself the declaration — it already names the offset and + fixes its source and target — so binding one is the same act as binding an + ambient value, and no attribute of the bound object has to be inspected. """ - names = vars(value).keys() if hasattr(value, "__dict__") else dir(value) return { - key: elem - for key in names - if not key.startswith("_") - and isinstance(elem := getattr(value, key), (common.Connectivity, common.Dimension)) + str(decl.value): value + for decl, value in _bindings.get({}).items() + if isinstance(decl, fbuiltins.FieldOffset) } - - -def offset_provider() -> common.OffsetProvider: - """ - Collect the offset provider from all bound namespaces. - - Every `common.Connectivity` reachable as an attribute of a bound object - contributes under its attribute name. Names must not collide across - namespaces — an ambiguous offset would silently pick one mesh's table. - """ - collected: dict[str, Any] = {} - for namespace, value in _bindings.get({}).items(): - for key, elem in offset_provider_of(value).items(): - if key in collected and collected[key] is not elem: - raise ValueError( - f"Ambient offset '{key}' is provided by more than one namespace;" - f" '{namespace}' conflicts with an earlier binding." - ) - collected[key] = elem - return collected diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 3c10c0232a..77e0a978ca 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -401,12 +401,12 @@ def __call__( self, *args: Any, offset_provider: common.OffsetProvider | None = None, - bind: Mapping[ambient.Namespace, Any] | None = None, + bind: Mapping[Any, Any] | Any | None = None, enable_jit: bool | None = None, **kwargs: Any, ) -> None: """Call the program; `bind` scopes ambient bindings to this call.""" - with ambient.bindings(bind or {}): + with ambient.bindings(ambient.as_bindings(bind) if bind else {}): self._invoke(*args, offset_provider=offset_provider, enable_jit=enable_jit, **kwargs) def _invoke( @@ -700,7 +700,7 @@ def __gt_closure_vars__(self) -> dict[str, Any]: def __call__(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> Any: """Call the field operator; `bind` scopes ambient bindings to this call.""" - with ambient.bindings(kwargs.pop("bind", None) or {}): + with ambient.bindings(ambient.as_bindings(b) if (b := kwargs.pop("bind", None)) else {}): return self._invoke(*args, enable_jit=enable_jit, **kwargs) def _invoke(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> Any: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py index ab74e0d58e..27a0c154cb 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py @@ -8,8 +8,6 @@ """Ambient binding of the offset provider, across the backend matrix.""" -import types - import numpy as np import pytest @@ -59,10 +57,9 @@ def test_ambient_binding_matches_explicit_offset_provider( inp = cases.allocate(unstructured_case, sum_edges, "edge_f")() out = cases.allocate(unstructured_case, sum_edges, cases.RETURN)() - mesh = gtx.Namespace("mesh") - # the fixture hands out a plain mapping; an ambient namespace binds an - # *object* whose public attributes are the connectivities - bound = types.SimpleNamespace(**unstructured_case.offset_provider) + # one rule for everything ambient: the declaration is the key. The fixture + # hands out a name-keyed mapping, so re-key it on the offset declarations. + bound = {V2E: unstructured_case.offset_provider["V2E"]} if entry_point == "program": callee = sum_edges_program.with_backend(unstructured_case.backend) @@ -74,10 +71,10 @@ def test_ambient_binding_matches_explicit_offset_provider( if mechanism == "offset_provider": callee(*args, **kwargs, offset_provider=unstructured_case.offset_provider) elif mechanism == "context_manager": - with gtx.bind(mesh, bound): + with gtx.bind(V2E, bound[V2E]): callee(*args, **kwargs) else: - callee(*args, **kwargs, bind={mesh: bound}) + callee(*args, **kwargs, bind=bound) np.testing.assert_allclose(out.asnumpy(), _reference(unstructured_case, inp)) @@ -86,9 +83,8 @@ def test_ambient_binding_matches_explicit_offset_provider( def test_bind_kwarg_does_not_leak_past_the_call(unstructured_case): inp = cases.allocate(unstructured_case, sum_edges, "edge_f")() out = cases.allocate(unstructured_case, sum_edges, cases.RETURN)() - mesh = gtx.Namespace("mesh") - bound = types.SimpleNamespace(**unstructured_case.offset_provider) + bound = {V2E: unstructured_case.offset_provider["V2E"]} - sum_edges_program.with_backend(unstructured_case.backend)(inp, out, bind={mesh: bound}) + sum_edges_program.with_backend(unstructured_case.backend)(inp, out, bind=bound) assert gtx.ambient.offset_provider() == {} diff --git a/tests/next_tests/unit_tests/test_ambient.py b/tests/next_tests/unit_tests/test_ambient.py index 6a09bb2b9a..bb23267cae 100644 --- a/tests/next_tests/unit_tests/test_ambient.py +++ b/tests/next_tests/unit_tests/test_ambient.py @@ -20,14 +20,21 @@ V2EDim = common.Dimension("V2E", kind=common.DimensionKind.LOCAL) -@dataclasses.dataclass(frozen=True) +V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim)) + + class Mesh: - V2E: common.Connectivity + """A container declaring what it supplies; the class attribute is the declaration.""" + + V2E = V2E + + def __init__(self, connectivity): + self.V2E = connectivity def make_mesh(table) -> Mesh: return Mesh( - V2E=gtx.as_connectivity( + gtx.as_connectivity( domain={Vertex: 2, V2EDim: 2}, codomain=Edge, data=np.asarray(table, dtype=np.int32), @@ -40,31 +47,28 @@ def test_nothing_bound_yields_empty_offset_provider(): assert gtx.ambient.offset_provider() == {} -def test_bound_namespace_provides_its_connectivities(): - mesh = gtx.Namespace("mesh") +def test_bound_offset_declaration_provides_its_connectivity(): m = make_mesh([[0, 1], [1, 2]]) - with gtx.bind(mesh, m): + with gtx.bind(V2E, m.V2E): assert gtx.ambient.offset_provider() == {"V2E": m.V2E} assert gtx.ambient.offset_provider() == {} +def test_container_binds_by_declaration_not_by_name(): + """The container's class attribute names the declaration it supplies.""" + m = make_mesh([[0, 1], [1, 2]]) + with gtx.bindings(gtx.ambient.as_bindings(m)): + assert gtx.ambient.offset_provider() == {"V2E": m.V2E} + + def test_bindings_nest_and_unwind(): - mesh = gtx.Namespace("mesh") m1, m2 = make_mesh([[0, 1], [1, 2]]), make_mesh([[1, 2], [0, 1]]) - with gtx.bind(mesh, m1): - with gtx.bind(mesh, m2): + with gtx.bind(V2E, m1.V2E): + with gtx.bind(V2E, m2.V2E): assert gtx.ambient.offset_provider()["V2E"] is m2.V2E assert gtx.ambient.offset_provider()["V2E"] is m1.V2E -def test_colliding_offset_names_are_rejected(): - a, b = gtx.Namespace("a"), gtx.Namespace("b") - with gtx.bind(a, make_mesh([[0, 1], [1, 2]])): - with gtx.bind(b, make_mesh([[1, 2], [0, 1]])): - with pytest.raises(ValueError, match="provided by more than one namespace"): - gtx.ambient.offset_provider() - - def test_freeze_gives_a_content_hash(): conn = make_mesh([[0, 1], [1, 2]]).V2E assert common.frozen_content_hash(conn) is None @@ -103,8 +107,6 @@ def test_readonly_freeze_marks_the_buffer_immutable(): # --- embedded end-to-end (backend-free) -------------------------------------- -V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim)) - @gtx.field_operator def sum_edges(a: gtx.Field[gtx.Dims[Edge], gtx.int32]) -> gtx.Field[gtx.Dims[Vertex], gtx.int32]: @@ -135,24 +137,21 @@ def test_embedded_execution_via_every_mechanism(inputs, entry_point, mechanism): callee = run if entry_point == "program" else sum_edges kwargs = {"out": out} if entry_point == "field_operator" else {} args = (a,) if entry_point == "field_operator" else (a, out) - mesh = gtx.Namespace("mesh") - if mechanism == "offset_provider": callee(*args, **kwargs, offset_provider={"V2E": m.V2E}) elif mechanism == "context_manager": - with gtx.bind(mesh, m): + with gtx.bind(V2E, m.V2E): callee(*args, **kwargs) else: - callee(*args, **kwargs, bind={mesh: m}) + callee(*args, **kwargs, bind=m) np.testing.assert_array_equal(out.asnumpy(), expected) def test_bind_kwarg_is_scoped_to_the_call(inputs): a, m, _ = inputs - mesh = gtx.Namespace("mesh") out = gtx.zeros(gtx.domain({Vertex: 2}), dtype=np.int32) - run(a, out, bind={mesh: m}) + run(a, out, bind=m) assert gtx.ambient.offset_provider() == {} From e2e3871f8cfd2543a20562a32bf55b396599f14b Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 6 Aug 2026 12:36:35 +0200 Subject: [PATCH 10/14] proto[next]: read ambient values through a container, dropping the arithmetic protocol MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Declarations now live in a container and are read as 'grid.dx'. The declaration is a descriptor: class access ('Grid.dx') yields the declaration, which is what 'bind=' takes as a key; instance access yields the bound value. Embedded execution therefore sees a plain scalar, and the 12 arithmetic dunders that made 'AmbientValue' impersonate one are gone — they were also incomplete (no comparisons, no '%', no numpy interop). A container-qualified name ('Grid_dx') gives the synthesised parameter an identifier that cannot collide across modules, which was the known gap. Only declarations an operator actually *reads* become parameters: a container is a place to declare things, and an operator that never reads 'grid.dx' must not acquire it — for a 'Static[T]' that would specialise the compiled program on an unused value. The frontend needed two hooks: a container types itself as a namespace over its own class (so 'grid.dx' resolves to the declared type with nothing bound), and the lowering rewrites the attribute to a reference to the synthesised parameter. --- src/gt4py/next/__init__.py | 3 +- src/gt4py/next/ambient.py | 107 +++++++++++------- src/gt4py/next/ffront/decorator.py | 11 +- src/gt4py/next/ffront/foast_to_gtir.py | 25 +++- src/gt4py/next/ffront/func_to_past.py | 25 +++- .../next/type_system/type_translation.py | 10 +- .../ffront_tests/test_ambient_values.py | 48 +++++--- 7 files changed, 158 insertions(+), 71 deletions(-) diff --git a/src/gt4py/next/__init__.py b/src/gt4py/next/__init__.py index 011d5d8f8e..0f729fdbcd 100644 --- a/src/gt4py/next/__init__.py +++ b/src/gt4py/next/__init__.py @@ -21,7 +21,7 @@ # ruff: noqa: F401 from .._core.definitions import CUPY_DEVICE_TYPE, Device, DeviceType, is_scalar_type from . import ambient, common, ffront, iterator, program_processors, typing -from .ambient import Extern, Static, bind, bindings, freeze +from .ambient import Container, Extern, Static, bind, bindings, freeze from .common import ( CartesianConnectivity, Connectivity, @@ -141,6 +141,7 @@ "ambient", "bind", "bindings", + "Container", "Extern", "Static", "freeze", diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index dbd855d345..09bb688290 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -98,6 +98,20 @@ def __gt_type__(self) -> Any: return type_translation.from_type_hint(self.type_hint) + def __set_name__(self, owner: type, name: str) -> None: + # a declaration inside a container knows where it lives, which gives the + # synthesised parameter a name that cannot collide across modules + self.name = f"{owner.__name__}_{name}" + + def __get__(self, obj: Any, objtype: type | None = None) -> Any: + """Class access yields the declaration, instance access the bound value. + + That is what lets a container be both the thing you *declare* with + (`Grid.dx` as a `bind=` key) and the thing you *read* inside an operator + (`grid.dx`), without the declaration having to impersonate a scalar. + """ + return self if obj is None else self.value + @property def value(self) -> Any: binding = _bindings.get({}).get(self, _UNBOUND) @@ -108,52 +122,24 @@ def value(self) -> Any: ) return binding - # Embedded execution runs the operator body as plain Python, so a bound - # declaration has to behave like the scalar it stands for. - def _v(self) -> Any: - return self.value - - def __float__(self) -> float: - return float(self.value) - - def __int__(self) -> int: - return int(self.value) - - def __bool__(self) -> bool: - return bool(self.value) - - def __neg__(self) -> Any: - return -self.value - def __add__(self, other: Any) -> Any: - return self.value + other - - def __radd__(self, other: Any) -> Any: - return other + self.value - - def __sub__(self, other: Any) -> Any: - return self.value - other - - def __rsub__(self, other: Any) -> Any: - return other - self.value - - def __mul__(self, other: Any) -> Any: - return self.value * other - - def __rmul__(self, other: Any) -> Any: - return other * self.value +class Container: + """ + Base for a container of ambient declarations: `class Grid(Container): dx = Static[float]`. - def __truediv__(self, other: Any) -> Any: - return self.value / other + Reading `grid.dx` inside an operator goes through the declaration's + descriptor, so in embedded execution it is simply the bound value — the + declaration never has to impersonate a scalar. `Grid.dx` (class access) + stays the declaration, which is what `bind=` takes as its key. - def __rtruediv__(self, other: Any) -> Any: - return other / self.value + The container types itself as a *namespace* over its own class, so the + frontend resolves `grid.dx` to the declared type without needing a value. + """ - def __pow__(self, other: Any) -> Any: - return self.value**other + def __gt_type__(self) -> Any: + from gt4py.next.type_system import type_translation - def __rpow__(self, other: Any) -> Any: - return other**self.value + return type_translation.NamespaceProxy(type(self)) class _Declarator: @@ -170,9 +156,22 @@ def __getitem__(self, type_hint: Any) -> AmbientValue: Static = _Declarator(static=True) -def ambient_values_in(closure_vars: Mapping[str, Any]) -> dict[str, AmbientValue]: - """The ambient declarations referenced by a set of closure variables.""" - return {k: v for k, v in closure_vars.items() if isinstance(v, AmbientValue)} +def declarations_in(closure_vars: Mapping[str, Any]) -> dict[str, AmbientValue]: + """ + Ambient declarations reachable from a set of closure variables, by parameter name. + + Declarations live in containers, so the key is the declaration's own + container-qualified name (`Grid_dx`) rather than whatever the referring + scope happens to call the container. That is what keeps two modules from + colliding when both declare a `dx`. + """ + found: dict[str, AmbientValue] = {} + for value in closure_vars.values(): + if isinstance(value, Container): + for decl in vars(type(value)).values(): + if isinstance(decl, AmbientValue): + found[decl.name] = decl + return found def freeze(elem: Any, *, readonly: bool = False) -> Any: @@ -274,3 +273,23 @@ def offset_provider() -> common.OffsetProvider: for decl, value in _bindings.get({}).items() if isinstance(decl, fbuiltins.FieldOffset) } + + +def attribute_declarations( + closure_vars: Mapping[str, Any], +) -> dict[tuple[str | None, str], AmbientValue]: + """Map `(container closure-var name, attribute)` to the declaration it reads.""" + return { + (var_name, attr): decl + for var_name, value in closure_vars.items() + if isinstance(value, Container) + for attr, decl in vars(type(value)).items() + if isinstance(decl, AmbientValue) + } + + +def attribute_parameter_names( + closure_vars: Mapping[str, Any], +) -> dict[tuple[str | None, str], str]: + """Map `(container closure-var name, attribute)` to the synthesised parameter name.""" + return {key: decl.name for key, decl in attribute_declarations(closure_vars).items()} diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 77e0a978ca..f60c51865d 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -321,11 +321,14 @@ def _all_closure_vars(self) -> dict[str, Any]: @property def _ambient_statics(self) -> tuple[str, ...]: + # only declarations that actually became parameters of *this* program; + # a container may hold others that it never reads + synthesised = {p.id for p in self.past_stage.past_node.params} return tuple( sorted( name - for name, decl in ambient.ambient_values_in(self._all_closure_vars).items() - if decl.static + for name, decl in ambient.declarations_in(self._all_closure_vars).items() + if decl.static and name in synthesised ) ) @@ -422,9 +425,11 @@ def _invoke( # defined, so from here on they are ordinary arguments and the caller # never names them. Embedded execution is the exception: it runs the # Python function, which reads them from its own closure. + synthesised = {p.id for p in self.past_stage.past_node.params} ambient_args = { name: decl.value - for name, decl in ambient.ambient_values_in(self._all_closure_vars).items() + for name, decl in ambient.declarations_in(self._all_closure_vars).items() + if name in synthesised } kwargs = {**kwargs, **ambient_args} diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 10bc754526..4e76816e9e 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -12,7 +12,7 @@ from gt4py import eve from gt4py.eve.extended_typing import Never, cast -from gt4py.next import common, utils +from gt4py.next import ambient, common, utils from gt4py.next.ffront import ( dialect_ast_enums, experimental as experimental_builtins, @@ -38,7 +38,9 @@ def foast_to_gtir(inp: ffront_stages.FOASTOperatorDef) -> itir.FunctionDefinitio See the docstring of `FieldOperatorLowering` for details. """ - return FieldOperatorLowering.apply(inp.foast_node) + return FieldOperatorLowering.apply( + inp.foast_node, ambient_names=ambient.attribute_parameter_names(inp.closure_vars) + ) def foast_to_gtir_factory( @@ -68,6 +70,9 @@ def promote_to_list(node_type: ts.TypeSpec) -> Callable[[itir.Expr], itir.Expr]: @dataclasses.dataclass class FieldOperatorLowering(eve.PreserveLocationVisitor, eve.NodeTranslator): + #: (container closure-var name, attribute) -> synthesised parameter name + ambient_names: dict[tuple[str | None, str], str] = dataclasses.field(default_factory=dict) + """ Lower FieldOperator AST (FOAST) to GTIR. @@ -98,8 +103,12 @@ class FieldOperatorLowering(eve.PreserveLocationVisitor, eve.NodeTranslator): uid_generator: utils.IDGeneratorPool = dataclasses.field(default_factory=utils.IDGeneratorPool) @classmethod - def apply(cls, node: foast.LocatedNode) -> itir.FunctionDefinition: - result = cls().visit(node) + def apply( + cls, + node: foast.LocatedNode, + ambient_names: dict[tuple[str | None, str], str] | None = None, + ) -> itir.FunctionDefinition: + result = cls(ambient_names=ambient_names or {}).visit(node) assert isinstance(result, itir.FunctionDefinition) return result @@ -236,7 +245,13 @@ def visit_Name(self, node: foast.Name, **kwargs: Any) -> itir.SymRef | itir.Axis return itir.AxisLiteral(value=node.type.dim.value, kind=node.type.dim.kind) return im.ref(node.id) - def visit_Attribute(self, node: foast.Attribute, **kwargs: Any) -> itir.AxisLiteral: + def visit_Attribute(self, node: foast.Attribute, **kwargs: Any) -> itir.Expr: + # `container.declaration` becomes a reference to the parameter + # synthesised for that declaration; it stays free here and is bound by + # the enclosing program's parameter list. + key = (getattr(node.value, "id", None), node.attr) + if key in self.ambient_names: + return im.ref(self.ambient_names[key]) if isinstance(node.type, ts.DimensionType): return itir.AxisLiteral(value=node.type.dim.value, kind=node.type.dim.kind) diff --git a/src/gt4py/next/ffront/func_to_past.py b/src/gt4py/next/ffront/func_to_past.py index a2a99c4627..6884f645d4 100644 --- a/src/gt4py/next/ffront/func_to_past.py +++ b/src/gt4py/next/ffront/func_to_past.py @@ -19,6 +19,7 @@ dialect_ast_enums, experimental, fbuiltins, + field_operator_ast as foast, program_ast as past, source_utils, stages as ffront_stages, @@ -75,6 +76,28 @@ def func_to_past(inp: DSLProgramDef) -> PASTProgramDef: ) +def _referenced_declarations(closure_vars: dict[str, typing.Any]) -> dict[str, typing.Any]: + """ + Declarations the program's operators actually read, by parameter name. + + Only these get a synthesised parameter: a container is a convenient place to + declare things, but an operator that never reads `grid.dx` must not acquire + it as an argument — for a `Static[T]` that would specialise the compiled + program on a value it does not use. + """ + referenced: dict[str, typing.Any] = {} + for value in closure_vars.values(): + foast_stage = getattr(value, "foast_stage", None) + if foast_stage is None: + continue + by_attribute = ambient.attribute_declarations(foast_stage.closure_vars) + for node in foast_stage.foast_node.walk_values().if_isinstance(foast.Attribute): + decl = by_attribute.get((getattr(node.value, "id", None), node.attr), None) + if decl is not None: + referenced[decl.name] = decl + return referenced + + def _with_ambient_params(node: past.Program, closure_vars: dict[str, typing.Any]) -> past.Program: """ Give every ambient declaration the program references a synthesised parameter. @@ -85,7 +108,7 @@ def _with_ambient_params(node: past.Program, closure_vars: dict[str, typing.Any] lowering, `static_params` and the compiled-program key all treat them like any other argument, and only the *value* has to be supplied per call. """ - declarations = ambient.ambient_values_in( + declarations = _referenced_declarations( transform_utils._get_closure_vars_recursively(closure_vars) ) if not declarations: diff --git a/src/gt4py/next/type_system/type_translation.py b/src/gt4py/next/type_system/type_translation.py index aa003003f5..3530421025 100644 --- a/src/gt4py/next/type_system/type_translation.py +++ b/src/gt4py/next/type_system/type_translation.py @@ -260,7 +260,7 @@ def from_type_hint( ConstantPythonNamespaceObject: TypeAlias = eve_utils.FrozenNamespace | enum.EnumMeta -PythonNamespaceObject: TypeAlias = ConstantPythonNamespaceObject | types.ModuleType +PythonNamespaceObject: TypeAlias = ConstantPythonNamespaceObject | types.ModuleType | type class NamespaceProxy(ts.TypeSpec): @@ -345,7 +345,13 @@ def from_value(value: Any) -> ts.TypeSpec: type_ = xtyping.infer_type(value, annotate_callable_kwargs=True) symbol_type = from_type_hint(type_) - if isinstance(symbol_type, (ts.DataType, ts.CallableType, ts.OffsetType, ts.DimensionType)): + # `NamespaceProxy` is accepted here so an object can declare itself a + # namespace through `__gt_type__`, not only by being one of the built-in + # namespace kinds above. + if isinstance( + symbol_type, + (ts.DataType, ts.CallableType, ts.OffsetType, ts.DimensionType, NamespaceProxy), + ): return symbol_type else: raise ValueError(f"Impossible to map '{value}' value to a 'Symbol'.") diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py index e354425bad..0fdcb1c2fa 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -21,16 +21,21 @@ IJFloatField = gtx.Field[gtx.Dims[IDim, JDim], gtx.float64] -#: declared here, bound at the call — never a parameter in the user's source, so -#: it does not have to be threaded through nested operators -dx = gtx.Static[gtx.float64] -dx_extern = gtx.Extern[gtx.float64] + +class Grid(gtx.Container): + """Declarations live in a container; `grid.dx` reads the bound value.""" + + dx = gtx.Static[gtx.float64] + dx_extern = gtx.Extern[gtx.float64] + + +grid = Grid() @gtx.field_operator def delta_x(f: IJFloatField) -> IJFloatField: """Forward difference in x.""" - return (1.0 / dx) * (f(IDim + 1) - f) + return (1.0 / grid.dx) * (f(IDim + 1) - f) @gtx.field_operator @@ -47,7 +52,7 @@ def run_delta_x(f: IJFloatField, out: IJFloatField) -> None: @gtx.field_operator def delta_x_extern(f: IJFloatField) -> IJFloatField: """Same, but supplied as a runtime argument instead of folded in.""" - return (1.0 / dx_extern) * (f(IDim + 1) - f) + return (1.0 / grid.dx_extern) * (f(IDim + 1) - f) @gtx.program @@ -66,6 +71,11 @@ def _inputs(case): return data, out +def _binding(spacing): + """A container is bound as a unit, so every declaration it holds gets a value.""" + return {Grid.dx: spacing, Grid.dx_extern: spacing} + + def _reference(data, spacing, factor=1): a = data.asnumpy() return factor * (1.0 / spacing) * (a[1:5, :] - a[0:4, :]) @@ -75,7 +85,7 @@ def _reference(data, spacing, factor=1): @pytest.mark.parametrize("spacing", [0.5, 0.25]) def test_static_value_is_bound_at_call(cartesian_case, spacing): data, out = _inputs(cartesian_case) - run_delta_x.with_backend(cartesian_case.backend)(data, out, bind={dx: spacing}) + run_delta_x.with_backend(cartesian_case.backend)(data, out, bind=_binding(spacing)) np.testing.assert_allclose(out.asnumpy(), _reference(data, spacing)) @@ -85,10 +95,10 @@ def test_distinct_values_do_not_share_a_compiled_program(cartesian_case): data, out = _inputs(cartesian_case) prog = run_delta_x.with_backend(cartesian_case.backend) - prog(data, out, bind={dx: 0.5}) + prog(data, out, bind=_binding(0.5)) np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5)) - prog(data, out, bind={dx: 0.25}) + prog(data, out, bind=_binding(0.25)) np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.25)) @@ -96,14 +106,14 @@ def test_distinct_values_do_not_share_a_compiled_program(cartesian_case): def test_value_reaches_a_nested_operator(cartesian_case): """`dx` is referenced two levels down without being passed as an argument.""" data, out = _inputs(cartesian_case) - run_delta_x_twice.with_backend(cartesian_case.backend)(data, out, bind={dx: 0.5}) + run_delta_x_twice.with_backend(cartesian_case.backend)(data, out, bind=_binding(0.5)) np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5, factor=2)) @pytest.mark.uses_cartesian_shift def test_context_manager_binding(cartesian_case): data, out = _inputs(cartesian_case) - with gtx.bind(dx, 0.5): + with gtx.bindings(_binding(0.5)): run_delta_x.with_backend(cartesian_case.backend)(data, out) np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5)) @@ -112,7 +122,7 @@ def test_context_manager_binding(cartesian_case): @pytest.mark.parametrize("spacing", [0.5, 0.25]) def test_extern_value_is_bound_at_call(cartesian_case, spacing): data, out = _inputs(cartesian_case) - run_delta_x_extern.with_backend(cartesian_case.backend)(data, out, bind={dx_extern: spacing}) + run_delta_x_extern.with_backend(cartesian_case.backend)(data, out, bind=_binding(spacing)) np.testing.assert_allclose(out.asnumpy(), _reference(data, spacing)) @@ -126,8 +136,8 @@ def test_static_specializes_but_extern_does_not(cartesian_case): static_prog = run_delta_x.with_backend(cartesian_case.backend) extern_prog = run_delta_x_extern.with_backend(cartesian_case.backend) for spacing in (0.5, 0.25): - static_prog(data, out, bind={dx: spacing}) - extern_prog(data, out, bind={dx_extern: spacing}) + static_prog(data, out, bind=_binding(spacing)) + extern_prog(data, out, bind=_binding(spacing)) assert len(static_prog._compiled_programs.compiled_programs) == 2 assert len(extern_prog._compiled_programs.compiled_programs) == 1 @@ -142,4 +152,12 @@ def test_declaration_types_itself_without_a_binding(): def test_unbound_declaration_reports_itself(): with pytest.raises(ValueError, match="not bound"): - gtx.Static[gtx.float64].value + Grid.dx.value + + +def test_class_access_is_the_declaration_instance_access_the_value(): + """The descriptor is what lets embedded execution see a plain scalar.""" + assert Grid.dx is not None and not isinstance(Grid.dx, float) + with gtx.bindings({Grid.dx: 0.5}): + assert grid.dx == 0.5 + assert 1.0 / grid.dx == 2.0 # no arithmetic protocol on the declaration From 0224db1126c80ac85728f73f2e4458a6d96c8a94 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 6 Aug 2026 13:54:32 +0200 Subject: [PATCH 11/14] proto[next]: replace AmbientValue with annotations and plain ContextVars MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Declarations are annotations in a container ('dx: Static[float]') and bind to a 'contextvars.ContextVar'. 'Static[T]'/'Extern[T]' are PEP 695 aliases over 'Annotated', so a type checker sees plain 'T' while the binding machinery reads the marker — the bespoke 'AmbientValue' object is gone, and with it the descriptor it needed. 'Grid.dx' (class access) is the variable, which 'bind=' takes as a key; 'grid.dx' (instance access) is its value. Binding is stdlib set/reset, so the dict-in-a-ContextVar store disappears; a 'FieldOffset' carries its own variable instead of a central registry. The offset provider is now assembled from the offsets *this* program references rather than from everything currently bound. That removes global state and fixes a real defect: an unrelated bound mesh leaked into every program's offset provider, where it also perturbed the compiled-program key and forced spurious recompiles. --- src/gt4py/next/ambient.py | 320 +++++++++--------- src/gt4py/next/ffront/decorator.py | 16 +- src/gt4py/next/ffront/fbuiltins.py | 6 + .../ffront_tests/test_ambient_binding.py | 2 +- .../ffront_tests/test_ambient_values.py | 25 +- tests/next_tests/unit_tests/test_ambient.py | 101 +++--- 6 files changed, 244 insertions(+), 226 deletions(-) diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index 09bb688290..f7d166df00 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -9,55 +9,61 @@ """ Prototype: ambient values, bound at program-execution time. -An ambient value is declared once and reached from any program without -appearing in a signature. Binding happens at execution time (JIT time for -compiled backends), so the same programs can run against a second mesh in one -process. +An ambient value is declared once and reached from any program without appearing +in a signature, so it does not have to be threaded through nested operators. +Binding happens at execution time (JIT time for compiled backends), which is what +lets the same programs run against a second mesh in one process. -Everything ambient is bound the same way: **the declaration is the key**. +A declaration is an *annotation* in a container, and the thing it binds to is a +plain `contextvars.ContextVar`:: - prog(a, out, bind={V2E: connectivity, dx: 0.5}) + class Grid(Container): + dx: Static[float] + nu: Extern[float] -A `FieldOffset` is already a declaration — it names the offset and fixes its -source and target — so it binds exactly like a `Static[T]` / `Extern[T]` value, -and a program called without an ``offset_provider`` assembles one from the bound -offsets. A container may declare what it supplies, in which case the class -attribute is the declaration and the instance attribute the value:: - class Mesh: - V2E = V2E # this mesh supplies the V2E connectivity - dx = physics.dx + grid = Grid() -Binding by declaration rather than by attribute *name* is what lets a container -supply the very offset an operator refers to, instead of something that merely -shares its name. -A value declared this way is referenced by bare name inside an operator and -never appears in a signature, so it does not have to be threaded through nested -operators. + @gtx.field_operator + def delta_x(f: IJField) -> IJField: + return (1.0 / grid.dx) * (f(I + 1) - f) + + + prog(f, out, bind={Grid.dx: 0.5, Grid.nu: 1e-3}) + +`Grid.dx` (class access) *is* the `ContextVar`, which is what `bind=` takes as a +key; `grid.dx` (instance access) is its current value, so embedded execution -- +which runs the operator body as plain Python -- sees an ordinary float. A +`FieldOffset` binds the same way, carrying its own `ContextVar`. + +`Static[T]` and `Extern[T]` are `Annotated` aliases, so a type checker sees plain +`T` and only the binding machinery reads the marker. A declaration becomes a **synthesised program parameter** when the program is -defined (`func_to_past`), so from there on it travels the ordinary path: type +defined (`func_to_past`), so from there it travels the ordinary path: type checking, lowering, `static_params` and the compiled-program key all treat it -like any other argument, and only the *value* is supplied per call. The caller -never names it. - -The two forms differ in one place only — whether that parameter is listed as a -static one: +like any other argument, and only the *value* is supplied per call. The two forms +differ in one place only -- whether that parameter is listed as static: -- ``Extern[T]`` is an ordinary runtime argument: one compiled program serves - every value. -- ``Static[T]`` is a static argument, so the existing static-argument machinery +- `Extern[T]` is an ordinary runtime argument: one compiled program serves every + value. +- `Static[T]` is a static argument, so the existing static-argument machinery folds it into the generated code and keys the compiled variant on it: one compiled program per distinct value. + +Nothing here is global: what a program can see is decided by *its own* closure +variables, never by a registry of everything that happens to be bound. """ from __future__ import annotations import contextlib import contextvars +import dataclasses +import typing from collections.abc import Generator, Mapping -from typing import Any +from typing import Annotated, Any, ClassVar import numpy as np @@ -66,112 +72,121 @@ class Mesh: from gt4py.next.ffront import fbuiltins -_UNBOUND: Any = object() +_UNSET: Any = object() -#: keyed by declaration object: a `Namespace` or an `AmbientValue` -_bindings: contextvars.ContextVar[Mapping[Any, Any]] = contextvars.ContextVar("_ambient_bindings") +class _StaticMarker: + """Annotation marker: fold the value into the generated code.""" -class AmbientValue: - """ - A value declared here and bound later: `dx = Extern[float]`. - The declaration carries the *type*, which is all the frontend needs when the - operator is defined; the *value* arrives at bind time. Use the declaration - object itself as the binding key: `bind={dx: 0.5}`. +class _ExternMarker: + """Annotation marker: pass the value as a runtime argument.""" - `Extern[T]` is supplied to the compiled program as a runtime argument. - `Static[T]` is folded into it as a literal, so each distinct value gets its - own compiled variant. - """ - def __init__(self, type_hint: Any, *, static: bool, name: str = "?") -> None: - self.type_hint = type_hint - self.static = static - self.name = name +#: `dx: Static[float]` -- folded in; one compiled variant per distinct value. +type Static[T] = Annotated[T, _StaticMarker] +#: `nu: Extern[float]` -- a runtime argument; one compiled program for all values. +type Extern[T] = Annotated[T, _ExternMarker] + - def __repr__(self) -> str: - return f"{'Static' if self.static else 'Extern'}[{getattr(self.type_hint, '__name__', self.type_hint)}]" +@dataclasses.dataclass(frozen=True) +class Declaration: + """What a container annotation declares: a name, a type, a kind, a variable.""" + + #: container-qualified, so two modules declaring a `dx` cannot collide + name: str + type_hint: Any + static: bool + var: contextvars.ContextVar def __gt_type__(self) -> Any: from gt4py.next.type_system import type_translation return type_translation.from_type_hint(self.type_hint) - def __set_name__(self, owner: type, name: str) -> None: - # a declaration inside a container knows where it lives, which gives the - # synthesised parameter a name that cannot collide across modules - self.name = f"{owner.__name__}_{name}" - - def __get__(self, obj: Any, objtype: type | None = None) -> Any: - """Class access yields the declaration, instance access the bound value. - - That is what lets a container be both the thing you *declare* with - (`Grid.dx` as a `bind=` key) and the thing you *read* inside an operator - (`grid.dx`), without the declaration having to impersonate a scalar. - """ - return self if obj is None else self.value - @property def value(self) -> Any: - binding = _bindings.get({}).get(self, _UNBOUND) - if binding is _UNBOUND: + value = self.var.get(_UNSET) + if value is _UNSET: raise ValueError( - f"Ambient value '{self!r}' is not bound." - f" Pass 'bind={{: }}' at the call, or use 'gtx.bind'." + f"Ambient value '{self.name}' is not bound." + " Pass 'bind={: }' at the call, or use 'gtx.bind'." ) - return binding + return value -class Container: - """ - Base for a container of ambient declarations: `class Grid(Container): dx = Static[float]`. +def _declared(hint: Any) -> tuple[Any, bool] | None: + """Split a `Static[T]` / `Extern[T]` annotation into `(T, is_static)`.""" + alias = getattr(hint, "__origin__", None) + markers = getattr(getattr(alias, "__value__", None), "__metadata__", ()) + if _StaticMarker in markers: + static = True + elif _ExternMarker in markers: + static = False + else: + return None + (base,) = typing.get_args(hint) + return base, static - Reading `grid.dx` inside an operator goes through the declaration's - descriptor, so in embedded execution it is simply the bound value — the - declaration never has to impersonate a scalar. `Grid.dx` (class access) - stays the declaration, which is what `bind=` takes as its key. - The container types itself as a *namespace* over its own class, so the - frontend resolves `grid.dx` to the declared type without needing a value. - """ +class _ContainerMeta(type): + def __getattr__(cls, name: str) -> contextvars.ContextVar: + # class access yields the variable, which is the `bind=` key + declarations: dict[str, Declaration] = cls.__dict__.get("_declarations", {}) + if name in declarations: + return declarations[name].var + raise AttributeError(name) - def __gt_type__(self) -> Any: - from gt4py.next.type_system import type_translation - - return type_translation.NamespaceProxy(type(self)) +class Container(metaclass=_ContainerMeta): + """ + Base for a container of ambient declarations. -class _Declarator: - def __init__(self, static: bool) -> None: - self._static = static + Declarations are annotations, so they create no class attribute and instance + access reaches `__getattr__` -- which is where the `ContextVar` is read. + """ - def __getitem__(self, type_hint: Any) -> AmbientValue: - return AmbientValue(type_hint, static=self._static) + _declarations: ClassVar[dict[str, Declaration]] = {} + #: class whose attributes carry the declared *types*, for the frontend + _type_view: ClassVar[type] + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + cls._declarations = {} + for attr, hint in typing.get_type_hints(cls, include_extras=True).items(): + if (declared := _declared(hint)) is None: + continue + type_hint, static = declared + cls._declarations[attr] = Declaration( + name=f"{cls.__name__}_{attr}", + type_hint=type_hint, + static=static, + var=contextvars.ContextVar(f"{cls.__name__}.{attr}"), + ) + cls._type_view = type(f"{cls.__name__}_types", (), dict(cls._declarations)) + def __getattr__(self, name: str) -> Any: + try: + declaration = type(self)._declarations[name] + except KeyError: + raise AttributeError(name) from None + return declaration.value -#: `dx = Extern[float]` — supplied as a runtime argument. -Extern = _Declarator(static=False) -#: `dx = Static[float]` — folded in as a literal; one compiled variant per value. -Static = _Declarator(static=True) + def __gt_type__(self) -> Any: + from gt4py.next.type_system import type_translation + # a container types itself as a namespace over its declared types, so + # `grid.dx` resolves at definition time with nothing bound + return type_translation.NamespaceProxy(type(self)._type_view) -def declarations_in(closure_vars: Mapping[str, Any]) -> dict[str, AmbientValue]: - """ - Ambient declarations reachable from a set of closure variables, by parameter name. - Declarations live in containers, so the key is the declaration's own - container-qualified name (`Grid_dx`) rather than whatever the referring - scope happens to call the container. That is what keeps two modules from - colliding when both declare a `dx`. - """ - found: dict[str, AmbientValue] = {} - for value in closure_vars.values(): - if isinstance(value, Container): - for decl in vars(type(value)).values(): - if isinstance(decl, AmbientValue): - found[decl.name] = decl - return found +def variable_for(declaration: Any) -> contextvars.ContextVar: + """The `ContextVar` a declaration binds to.""" + if isinstance(declaration, contextvars.ContextVar): + return declaration + if isinstance(declaration, fbuiltins.FieldOffset): + return declaration._ambient_var + raise TypeError(f"'{declaration!r}' is not an ambient declaration.") def freeze(elem: Any, *, readonly: bool = False) -> Any: @@ -180,13 +195,13 @@ def freeze(elem: Any, *, readonly: bool = False) -> Any: An ambient value is static for the jitted programs that see it, so it may identify itself by *content* rather than by `id`. The hash is computed once, - here — the O(size) cost is paid at freeze time, never per call. + here -- the O(size) cost is paid at freeze time, never per call. `readonly` additionally marks the buffer immutable, which is what makes the cached hash trustworthy. It is **off by default because it breaks the gtfn - bindings**: they are generated with mutable `ndarray` parameters and reject - a non-writeable array outright. Until that is fixed, the hash is only as - stable as the caller's discipline. + bindings**: they are generated with mutable `ndarray` parameters and reject a + non-writeable array outright. Until that is fixed, the hash is only as stable + as the caller's discipline. """ if common.frozen_content_hash(elem) is not None: return elem @@ -203,49 +218,26 @@ def freeze(elem: Any, *, readonly: bool = False) -> Any: def as_bindings(spec: Any) -> dict[Any, Any]: - """ - Normalise what `bind=` accepts into a declaration -> value mapping. - - A mapping is taken as-is. Anything else is treated as a *container*: its - class attributes name the declarations it supplies, and the instance - attribute of the same name carries the value. The declaration object itself - is the key, so a container binds the very `Extern` an operator refers to - rather than something that merely shares its name. - - class Grid: - dx = physics.dx # the declaration this grid supplies - - grid = Grid(); grid.dx = 0.5 - prog(f, out, bind=grid) - """ + """Normalise what `bind=` accepts into a declaration -> value mapping.""" if isinstance(spec, Mapping): return dict(spec) - resolved: dict[Any, Any] = {} - for name, decl in vars(type(spec)).items(): - if not isinstance(decl, (AmbientValue, fbuiltins.FieldOffset)): - continue - if name not in vars(spec): - raise ValueError( - f"'{type(spec).__name__}' declares '{name}' but the instance does not set it." - ) - resolved[decl] = vars(spec)[name] - return resolved + raise TypeError(f"'{spec!r}' is not a mapping of declarations to values.") @contextlib.contextmanager def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: - """Bind several namespaces at once, for the duration of the context.""" - if not mapping: - yield - return - for value in mapping.values(): + """Bind declarations to values for the duration of the context.""" + tokens: list[tuple[contextvars.ContextVar, contextvars.Token]] = [] + for declaration, value in mapping.items(): if isinstance(value, common.Connectivity): freeze(value) - token = _bindings.set({**_bindings.get({}), **mapping}) + var = variable_for(declaration) + tokens.append((var, var.set(value))) try: yield finally: - _bindings.reset(token) + for var, token in reversed(tokens): + var.reset(token) @contextlib.contextmanager @@ -255,36 +247,48 @@ def bind(declaration: Any, value: Any) -> Generator[None, None, None]: yield -def resolve(explicit: common.OffsetProvider | None) -> common.OffsetProvider: - """The caller's offset provider if given, otherwise the ambient one.""" - return offset_provider() if explicit is None else explicit - - -def offset_provider() -> common.OffsetProvider: +def offset_provider_for(closure_vars: Mapping[str, Any]) -> common.OffsetProvider: """ - Assemble the offset provider from the bound offset declarations. + Assemble the offset provider from the offsets *this* program references. - A `FieldOffset` is itself the declaration — it already names the offset and - fixes its source and target — so binding one is the same act as binding an - ambient value, and no attribute of the bound object has to be inspected. + Scoped to the given closure variables rather than to everything currently + bound: an unrelated mesh must not leak into a program's offset provider, + where it would also perturb the compiled-program key. """ return { - str(decl.value): value - for decl, value in _bindings.get({}).items() - if isinstance(decl, fbuiltins.FieldOffset) + str(offset.value): value + for offset in closure_vars.values() + if isinstance(offset, fbuiltins.FieldOffset) + if (value := offset._ambient_var.get(_UNSET)) is not _UNSET + } + + +def resolve( + explicit: common.OffsetProvider | None, closure_vars: Mapping[str, Any] +) -> common.OffsetProvider: + """The caller's offset provider if given, otherwise the ambient one.""" + return offset_provider_for(closure_vars) if explicit is None else explicit + + +def declarations_in(closure_vars: Mapping[str, Any]) -> dict[str, Declaration]: + """Ambient declarations reachable from a set of closure variables, by parameter name.""" + return { + decl.name: decl + for value in closure_vars.values() + if isinstance(value, Container) + for decl in type(value)._declarations.values() } def attribute_declarations( closure_vars: Mapping[str, Any], -) -> dict[tuple[str | None, str], AmbientValue]: +) -> dict[tuple[str | None, str], Declaration]: """Map `(container closure-var name, attribute)` to the declaration it reads.""" return { (var_name, attr): decl for var_name, value in closure_vars.items() if isinstance(value, Container) - for attr, decl in vars(type(value)).items() - if isinstance(decl, AmbientValue) + for attr, decl in type(value)._declarations.items() } diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index f60c51865d..1500516052 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -419,7 +419,7 @@ def _invoke( enable_jit: bool | None = None, **kwargs: Any, ) -> None: - offset_provider = ambient.resolve(offset_provider) + offset_provider = ambient.resolve(offset_provider, self._all_closure_vars) enable_jit = self.compilation_options.enable_jit if enable_jit is None else enable_jit # Ambient declarations became synthesised parameters when the program was # defined, so from here on they are ordinary arguments and the caller @@ -485,7 +485,7 @@ class ProgramWithBoundArgs(Program): def _invoke( self, *args: Any, offset_provider: common.OffsetProvider | None = None, **kwargs: Any ) -> None: - offset_provider = ambient.resolve(offset_provider) + offset_provider = ambient.resolve(offset_provider, self._all_closure_vars) type_ = self.past_stage.past_node.type assert isinstance(type_, ts_ffront.ProgramType) new_type = ts_ffront.ProgramType( @@ -703,6 +703,10 @@ def __gt_gtir__(self) -> itir.FunctionDefinition: def __gt_closure_vars__(self) -> dict[str, Any]: return self.foast_stage.closure_vars + @functools.cached_property + def _all_closure_vars(self) -> dict[str, Any]: + return transform_utils._get_closure_vars_recursively(self.foast_stage.closure_vars) + def __call__(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> Any: """Call the field operator; `bind` scopes ambient bindings to this call.""" with ambient.bindings(ambient.as_bindings(b) if (b := kwargs.pop("bind", None)) else {}): @@ -711,7 +715,9 @@ def __call__(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> def _invoke(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> Any: if not next_embedded.context.within_valid_context() and self.backend is not None: # non embedded execution - offset_provider = {**ambient.resolve(kwargs.pop("offset_provider", None))} + offset_provider = { + **ambient.resolve(kwargs.pop("offset_provider", None), self._all_closure_vars) + } if "out" not in kwargs: raise errors.MissingArgumentError(None, "out", True) out = kwargs.pop("out") @@ -733,7 +739,9 @@ def _invoke(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> else: if not next_embedded.context.within_valid_context(): # field_operator as program - kwargs["offset_provider"] = {**ambient.resolve(kwargs.pop("offset_provider", None))} + kwargs["offset_provider"] = { + **ambient.resolve(kwargs.pop("offset_provider", None), self._all_closure_vars) + } attributes = ( self.definition_stage.attributes if self.definition_stage diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 37ddf9183a..8d791ac729 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -6,6 +6,7 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import contextvars import dataclasses import functools import inspect @@ -471,6 +472,11 @@ class FieldOffset(runtime.Offset): def _cache(self) -> dict: return {} + @functools.cached_property + def _ambient_var(self) -> contextvars.ContextVar: + """The variable this offset binds to, so it needs no central registry.""" + return contextvars.ContextVar(f"offset.{self.value}") + def __post_init__(self) -> None: if len(self.target) == 2 and self.target[1].kind != common.DimensionKind.LOCAL: raise ValueError("Second dimension in offset must be a local dimension.") diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py index 27a0c154cb..1b5e63410e 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py @@ -87,4 +87,4 @@ def test_bind_kwarg_does_not_leak_past_the_call(unstructured_case): sum_edges_program.with_backend(unstructured_case.backend)(inp, out, bind=bound) - assert gtx.ambient.offset_provider() == {} + assert gtx.ambient.offset_provider_for({"V2E": V2E}) == {} diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py index 0fdcb1c2fa..6fc55b3b86 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -8,6 +8,8 @@ """Ambient values referenced by name inside an operator, bound at call time.""" +import contextvars + import numpy as np import pytest @@ -25,8 +27,8 @@ class Grid(gtx.Container): """Declarations live in a container; `grid.dx` reads the bound value.""" - dx = gtx.Static[gtx.float64] - dx_extern = gtx.Extern[gtx.float64] + dx: gtx.Static[float] + dx_extern: gtx.Extern[float] grid = Grid() @@ -146,18 +148,23 @@ def test_static_specializes_but_extern_does_not(cartesian_case): def test_declaration_types_itself_without_a_binding(): """The frontend only needs the type at decoration; the value comes later.""" expected = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - assert gtx.Static[gtx.float64].__gt_type__() == expected - assert gtx.Extern[gtx.float64].__gt_type__() == expected + assert Grid._declarations["dx"].__gt_type__() == expected + assert Grid._declarations["dx_extern"].__gt_type__() == expected + + +def test_bind_key_is_a_plain_contextvar(): + """No bespoke declaration object: binding is stdlib set/reset.""" + assert isinstance(Grid.dx, contextvars.ContextVar) + assert Grid.dx.get(None) is None def test_unbound_declaration_reports_itself(): with pytest.raises(ValueError, match="not bound"): - Grid.dx.value + grid.dx -def test_class_access_is_the_declaration_instance_access_the_value(): - """The descriptor is what lets embedded execution see a plain scalar.""" - assert Grid.dx is not None and not isinstance(Grid.dx, float) +def test_class_access_is_the_variable_instance_access_the_value(): + """What lets embedded execution see a plain scalar, with no arithmetic protocol.""" with gtx.bindings({Grid.dx: 0.5}): assert grid.dx == 0.5 - assert 1.0 / grid.dx == 2.0 # no arithmetic protocol on the declaration + assert 1.0 / grid.dx == 2.0 diff --git a/tests/next_tests/unit_tests/test_ambient.py b/tests/next_tests/unit_tests/test_ambient.py index bb23267cae..74e435ccd5 100644 --- a/tests/next_tests/unit_tests/test_ambient.py +++ b/tests/next_tests/unit_tests/test_ambient.py @@ -23,84 +23,74 @@ V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim)) -class Mesh: - """A container declaring what it supplies; the class attribute is the declaration.""" - - V2E = V2E - - def __init__(self, connectivity): - self.V2E = connectivity - - -def make_mesh(table) -> Mesh: - return Mesh( - gtx.as_connectivity( - domain={Vertex: 2, V2EDim: 2}, - codomain=Edge, - data=np.asarray(table, dtype=np.int32), - skip_value=None, - ) +def make_conn(table): + return gtx.as_connectivity( + domain={Vertex: 2, V2EDim: 2}, + codomain=Edge, + data=np.asarray(table, dtype=np.int32), + skip_value=None, ) def test_nothing_bound_yields_empty_offset_provider(): - assert gtx.ambient.offset_provider() == {} + assert gtx.ambient.offset_provider_for({"V2E": V2E}) == {} def test_bound_offset_declaration_provides_its_connectivity(): - m = make_mesh([[0, 1], [1, 2]]) - with gtx.bind(V2E, m.V2E): - assert gtx.ambient.offset_provider() == {"V2E": m.V2E} - assert gtx.ambient.offset_provider() == {} + m = make_conn([[0, 1], [1, 2]]) + with gtx.bind(V2E, m): + assert gtx.ambient.offset_provider_for({"V2E": V2E}) == {"V2E": m} + assert gtx.ambient.offset_provider_for({"V2E": V2E}) == {} -def test_container_binds_by_declaration_not_by_name(): - """The container's class attribute names the declaration it supplies.""" - m = make_mesh([[0, 1], [1, 2]]) - with gtx.bindings(gtx.ambient.as_bindings(m)): - assert gtx.ambient.offset_provider() == {"V2E": m.V2E} +def test_offset_provider_is_scoped_to_the_referenced_offsets(): + """An unrelated bound offset must not leak into a program's offset provider.""" + other = gtx.FieldOffset("OTHER", source=Edge, target=(Vertex, V2EDim)) + m1, m2 = make_conn([[0, 1], [1, 2]]), make_conn([[1, 2], [0, 1]]) + with gtx.bindings({V2E: m1, other: m2}): + assert gtx.ambient.offset_provider_for({"V2E": V2E}) == {"V2E": m1} def test_bindings_nest_and_unwind(): - m1, m2 = make_mesh([[0, 1], [1, 2]]), make_mesh([[1, 2], [0, 1]]) - with gtx.bind(V2E, m1.V2E): - with gtx.bind(V2E, m2.V2E): - assert gtx.ambient.offset_provider()["V2E"] is m2.V2E - assert gtx.ambient.offset_provider()["V2E"] is m1.V2E + m1, m2 = make_conn([[0, 1], [1, 2]]), make_conn([[1, 2], [0, 1]]) + with gtx.bind(V2E, m1): + with gtx.bind(V2E, m2): + assert gtx.ambient.offset_provider_for({"V2E": V2E})["V2E"] is m2 + assert gtx.ambient.offset_provider_for({"V2E": V2E})["V2E"] is m1 def test_freeze_gives_a_content_hash(): - conn = make_mesh([[0, 1], [1, 2]]).V2E + conn = make_conn([[0, 1], [1, 2]]) assert common.frozen_content_hash(conn) is None gtx.freeze(conn) assert common.frozen_content_hash(conn) is not None def test_equal_frozen_connectivities_share_a_cache_key(): - m1, m2 = make_mesh([[0, 1], [1, 2]]), make_mesh([[0, 1], [1, 2]]) + m1, m2 = make_conn([[0, 1], [1, 2]]), make_conn([[0, 1], [1, 2]]) unfrozen = ( - common.hash_offset_provider_items_by_id({"V2E": m1.V2E}), - common.hash_offset_provider_items_by_id({"V2E": m2.V2E}), + common.hash_offset_provider_items_by_id({"V2E": m1}), + common.hash_offset_provider_items_by_id({"V2E": m2}), ) assert unfrozen[0] != unfrozen[1], "distinct objects are keyed apart before freezing" - gtx.freeze(m1.V2E) - gtx.freeze(m2.V2E) + gtx.freeze(m1) + gtx.freeze(m2) assert common.hash_offset_provider_items_by_id( - {"V2E": m1.V2E} - ) == common.hash_offset_provider_items_by_id({"V2E": m2.V2E}) + {"V2E": m1} + ) == common.hash_offset_provider_items_by_id({"V2E": m2}) def test_differing_frozen_connectivities_are_keyed_apart(): - m1 = gtx.freeze(make_mesh([[0, 1], [1, 2]]).V2E) - m2 = gtx.freeze(make_mesh([[1, 2], [0, 1]]).V2E) + m1 = gtx.freeze(make_conn([[0, 1], [1, 2]])) + m2 = gtx.freeze(make_conn([[1, 2], [0, 1]])) assert common.hash_offset_provider_items_by_id( {"V2E": m1} ) != common.hash_offset_provider_items_by_id({"V2E": m2}) def test_readonly_freeze_marks_the_buffer_immutable(): - conn = gtx.freeze(make_mesh([[0, 1], [1, 2]]).V2E, readonly=True) + conn = gtx.freeze(make_conn([[0, 1], [1, 2]]), readonly=True) with pytest.raises(ValueError): conn.ndarray[0, 0] = 7 @@ -124,7 +114,7 @@ def run( def inputs(): return ( gtx.as_field([Edge], np.arange(3, dtype=np.int32)), - make_mesh([[0, 1], [1, 2]]), + make_conn([[0, 1], [1, 2]]), np.asarray([1, 3], dtype=np.int32), ) @@ -138,12 +128,12 @@ def test_embedded_execution_via_every_mechanism(inputs, entry_point, mechanism): kwargs = {"out": out} if entry_point == "field_operator" else {} args = (a,) if entry_point == "field_operator" else (a, out) if mechanism == "offset_provider": - callee(*args, **kwargs, offset_provider={"V2E": m.V2E}) + callee(*args, **kwargs, offset_provider={"V2E": m}) elif mechanism == "context_manager": - with gtx.bind(V2E, m.V2E): + with gtx.bind(V2E, m): callee(*args, **kwargs) else: - callee(*args, **kwargs, bind=m) + callee(*args, **kwargs, bind={V2E: m}) np.testing.assert_array_equal(out.asnumpy(), expected) @@ -151,23 +141,26 @@ def test_embedded_execution_via_every_mechanism(inputs, entry_point, mechanism): def test_bind_kwarg_is_scoped_to_the_call(inputs): a, m, _ = inputs out = gtx.zeros(gtx.domain({Vertex: 2}), dtype=np.int32) - run(a, out, bind=m) - assert gtx.ambient.offset_provider() == {} + run(a, out, bind={V2E: m}) + assert gtx.ambient.offset_provider_for({"V2E": V2E}) == {} def test_bindings_are_context_local(): """The binding must not leak between contexts (the fingerprint reads it, not a mirror).""" import contextvars - decl = gtx.Static[gtx.float64] + class Scoped(gtx.Container): + dx: gtx.Static[float] + + scoped = Scoped() seen = {} def bind_and_record(): - with gtx.bind(decl, 0.25): - seen["inner"] = decl.value + with gtx.bind(Scoped.dx, 0.25): + seen["inner"] = scoped.dx - with gtx.bind(decl, 0.5): + with gtx.bind(Scoped.dx, 0.5): contextvars.copy_context().run(bind_and_record) - seen["outer"] = decl.value + seen["outer"] = scoped.dx assert seen == {"inner": 0.25, "outer": 0.5} From fea8af5539250941be710eaf1ba4b132e1f2b1fe Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 6 Aug 2026 14:12:39 +0200 Subject: [PATCH 12/14] proto[next]: bind a filled container as a whole MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 'Grid(dx=0.5, nu=1e-3)' carries values and binds every declaration it holds, so the caller provides the grid (or the mesh) as one thing and each program picks the parts it needs — rather than the caller tracking which program reads what. The two uses of a container instance do not collide: one constructed with values keeps them in its instance dict, so attribute access finds them directly; one constructed empty has nothing there, so access falls through to '__getattr__' and reads the bound variable. That is the instance an operator reads through. 'gtx.bind' now also takes containers, and an undeclared keyword is rejected at construction rather than silently ignored. --- src/gt4py/next/ambient.py | 54 +++++++++++++++++-- .../ffront_tests/test_ambient_values.py | 19 +++++-- 2 files changed, 64 insertions(+), 9 deletions(-) diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index f7d166df00..9af9f60939 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -165,6 +165,24 @@ def __init_subclass__(cls, **kwargs: Any) -> None: ) cls._type_view = type(f"{cls.__name__}_types", (), dict(cls._declarations)) + def __init__(self, **values: Any) -> None: + """ + Carry values for this container's declarations: `Grid(dx=0.5, nu=1e-3)`. + + A container instance is used two ways, and the two do not collide: one + constructed *with* values carries them in its instance dict and is what + `bind=` takes; one constructed empty carries nothing, so attribute access + falls through to `__getattr__` and reads the bound value. That is the + instance an operator reads through. + """ + unknown = set(values) - set(type(self)._declarations) + if unknown: + raise TypeError( + f"'{type(self).__name__}' does not declare {sorted(unknown)};" + f" it declares {sorted(type(self)._declarations)}." + ) + self.__dict__.update(values) + def __getattr__(self, name: str) -> Any: try: declaration = type(self)._declarations[name] @@ -218,10 +236,27 @@ def freeze(elem: Any, *, readonly: bool = False) -> Any: def as_bindings(spec: Any) -> dict[Any, Any]: - """Normalise what `bind=` accepts into a declaration -> value mapping.""" + """ + Normalise what `bind=` accepts into a declaration -> value mapping. + + A filled container binds everything it carries, which is usually what you + want: a grid or a mesh is one thing semantically, and each program picks the + parts it needs rather than the caller tracking which those are. + """ + if isinstance(spec, Container): + declarations = type(spec)._declarations + return {declarations[attr].var: value for attr, value in vars(spec).items()} if isinstance(spec, Mapping): return dict(spec) - raise TypeError(f"'{spec!r}' is not a mapping of declarations to values.") + if isinstance(spec, (list, tuple)): + merged: dict[Any, Any] = {} + for element in spec: + merged.update(as_bindings(element)) + return merged + raise TypeError( + f"'{spec!r}' is not a container, a mapping of declarations to values," + " or a sequence of those." + ) @contextlib.contextmanager @@ -241,9 +276,18 @@ def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: @contextlib.contextmanager -def bind(declaration: Any, value: Any) -> Generator[None, None, None]: - """Bind `value` to `declaration` for the duration of the context.""" - with bindings({declaration: value}): +def bind(*specs: Any, **values: Any) -> Generator[None, None, None]: + """ + Bind for the duration of the context. + + with gtx.bind(Grid(dx=0.5, nu=1e-3)): # a filled container + with gtx.bind(Grid.dx, 0.5): # one declaration + """ + if len(specs) == 2 and not isinstance(specs[0], (Container, Mapping, list, tuple)): + mapping = {specs[0]: specs[1]} + else: + mapping = as_bindings(list(specs) + list(values.values())) + with bindings(mapping): yield diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py index 6fc55b3b86..3708d9d1cb 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -74,8 +74,8 @@ def _inputs(case): def _binding(spacing): - """A container is bound as a unit, so every declaration it holds gets a value.""" - return {Grid.dx: spacing, Grid.dx_extern: spacing} + """The whole grid is provided; each program picks the parts it needs.""" + return Grid(dx=spacing, dx_extern=spacing) def _reference(data, spacing, factor=1): @@ -115,7 +115,7 @@ def test_value_reaches_a_nested_operator(cartesian_case): @pytest.mark.uses_cartesian_shift def test_context_manager_binding(cartesian_case): data, out = _inputs(cartesian_case) - with gtx.bindings(_binding(0.5)): + with gtx.bind(_binding(0.5)): run_delta_x.with_backend(cartesian_case.backend)(data, out) np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5)) @@ -152,6 +152,17 @@ def test_declaration_types_itself_without_a_binding(): assert Grid._declarations["dx_extern"].__gt_type__() == expected +def test_whole_container_binds_every_value_it_carries(): + """A grid is one thing semantically; the caller need not track who uses what.""" + with gtx.bind(Grid(dx=0.5, dx_extern=1.5)): + assert (grid.dx, grid.dx_extern) == (0.5, 1.5) + + +def test_container_rejects_an_undeclared_value(): + with pytest.raises(TypeError, match="does not declare"): + Grid(dz=1.0) + + def test_bind_key_is_a_plain_contextvar(): """No bespoke declaration object: binding is stdlib set/reset.""" assert isinstance(Grid.dx, contextvars.ContextVar) @@ -165,6 +176,6 @@ def test_unbound_declaration_reports_itself(): def test_class_access_is_the_variable_instance_access_the_value(): """What lets embedded execution see a plain scalar, with no arithmetic protocol.""" - with gtx.bindings({Grid.dx: 0.5}): + with gtx.bind(Grid.dx, 0.5): assert grid.dx == 0.5 assert 1.0 / grid.dx == 2.0 From 9b91f43fc678005c9a99877e78872beb36224829 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 6 Aug 2026 15:01:08 +0200 Subject: [PATCH 13/14] proto[next]: disambiguate declarations across modules Two containers with the same class name in different modules produced the same synthesised parameter, and the result was silently wrong rather than an error: both declarations collapsed onto one parameter and one binding won. A container class name is not unique, so the parameter now carries a stable digest of the fully qualified name, and the declaration keeps the qualified name for diagnostics. A second collision sat behind it: '_get_closure_vars_recursively' merges by name, so two modules that both call their container 'grid' shadowed one another. Every caller now uses one per-operator walk that looks at each operator's own closure variables instead of a merged mapping. --- src/gt4py/next/ambient.py | 48 ++++++++++++++----- src/gt4py/next/ffront/decorator.py | 4 +- src/gt4py/next/ffront/func_to_past.py | 25 +--------- .../ffront_tests/test_ambient_values.py | 21 ++++++++ 4 files changed, 61 insertions(+), 37 deletions(-) diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index 9af9f60939..4260d8f7d7 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -61,6 +61,7 @@ def delta_x(f: IJField) -> IJField: import contextlib import contextvars import dataclasses +import hashlib import typing from collections.abc import Generator, Mapping from typing import Annotated, Any, ClassVar @@ -93,8 +94,13 @@ class _ExternMarker: class Declaration: """What a container annotation declares: a name, a type, a kind, a variable.""" - #: container-qualified, so two modules declaring a `dx` cannot collide + #: the synthesised parameter name. Readable, but disambiguated by a digest of + #: the fully qualified name: a container class name is not unique, and two + #: modules each declaring a `Grid.dx` would otherwise share one parameter and + #: silently take one another's value. name: str + #: fully qualified, for diagnostics + qualname: str type_hint: Any static: bool var: contextvars.ContextVar @@ -109,7 +115,7 @@ def value(self) -> Any: value = self.var.get(_UNSET) if value is _UNSET: raise ValueError( - f"Ambient value '{self.name}' is not bound." + f"Ambient value '{self.qualname}' is not bound." " Pass 'bind={: }' at the call, or use 'gtx.bind'." ) return value @@ -157,8 +163,13 @@ def __init_subclass__(cls, **kwargs: Any) -> None: if (declared := _declared(hint)) is None: continue type_hint, static = declared + qualname = f"{cls.__module__}.{cls.__qualname__}.{attr}" + # a stable digest, not a counter: the name lands in the compiled + # signature, so it must not shift with import order + digest = hashlib.sha1(qualname.encode()).hexdigest()[:6] cls._declarations[attr] = Declaration( - name=f"{cls.__name__}_{attr}", + name=f"{cls.__name__}_{attr}_{digest}", + qualname=qualname, type_hint=type_hint, static=static, var=contextvars.ContextVar(f"{cls.__name__}.{attr}"), @@ -314,14 +325,29 @@ def resolve( return offset_provider_for(closure_vars) if explicit is None else explicit -def declarations_in(closure_vars: Mapping[str, Any]) -> dict[str, Declaration]: - """Ambient declarations reachable from a set of closure variables, by parameter name.""" - return { - decl.name: decl - for value in closure_vars.values() - if isinstance(value, Container) - for decl in type(value)._declarations.values() - } +def referenced_declarations(closure_vars: Mapping[str, Any]) -> dict[str, Declaration]: + """ + Declarations the operators reachable from `closure_vars` actually read. + + Walks each operator's *own* closure variables rather than a merged mapping: + merging is keyed by name, so two modules that both call their container + `grid` would shadow one another. Only read declarations are returned — an + operator that never reads `grid.dx` must not acquire it as a parameter, or a + `Static[T]` would specialise the compiled program on a value it does not use. + """ + from gt4py.next.ffront import field_operator_ast as foast + + referenced: dict[str, Declaration] = {} + for value in closure_vars.values(): + foast_stage = getattr(value, "foast_stage", None) + if foast_stage is None: + continue + by_attribute = attribute_declarations(foast_stage.closure_vars) + for node in foast_stage.foast_node.walk_values().if_isinstance(foast.Attribute): + decl = by_attribute.get((getattr(node.value, "id", None), node.attr), None) + if decl is not None: + referenced[decl.name] = decl + return referenced def attribute_declarations( diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 1500516052..4582f0dcaa 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -327,7 +327,7 @@ def _ambient_statics(self) -> tuple[str, ...]: return tuple( sorted( name - for name, decl in ambient.declarations_in(self._all_closure_vars).items() + for name, decl in ambient.referenced_declarations(self._all_closure_vars).items() if decl.static and name in synthesised ) ) @@ -428,7 +428,7 @@ def _invoke( synthesised = {p.id for p in self.past_stage.past_node.params} ambient_args = { name: decl.value - for name, decl in ambient.declarations_in(self._all_closure_vars).items() + for name, decl in ambient.referenced_declarations(self._all_closure_vars).items() if name in synthesised } kwargs = {**kwargs, **ambient_args} diff --git a/src/gt4py/next/ffront/func_to_past.py b/src/gt4py/next/ffront/func_to_past.py index 6884f645d4..96a40e4ff2 100644 --- a/src/gt4py/next/ffront/func_to_past.py +++ b/src/gt4py/next/ffront/func_to_past.py @@ -19,7 +19,6 @@ dialect_ast_enums, experimental, fbuiltins, - field_operator_ast as foast, program_ast as past, source_utils, stages as ffront_stages, @@ -76,28 +75,6 @@ def func_to_past(inp: DSLProgramDef) -> PASTProgramDef: ) -def _referenced_declarations(closure_vars: dict[str, typing.Any]) -> dict[str, typing.Any]: - """ - Declarations the program's operators actually read, by parameter name. - - Only these get a synthesised parameter: a container is a convenient place to - declare things, but an operator that never reads `grid.dx` must not acquire - it as an argument — for a `Static[T]` that would specialise the compiled - program on a value it does not use. - """ - referenced: dict[str, typing.Any] = {} - for value in closure_vars.values(): - foast_stage = getattr(value, "foast_stage", None) - if foast_stage is None: - continue - by_attribute = ambient.attribute_declarations(foast_stage.closure_vars) - for node in foast_stage.foast_node.walk_values().if_isinstance(foast.Attribute): - decl = by_attribute.get((getattr(node.value, "id", None), node.attr), None) - if decl is not None: - referenced[decl.name] = decl - return referenced - - def _with_ambient_params(node: past.Program, closure_vars: dict[str, typing.Any]) -> past.Program: """ Give every ambient declaration the program references a synthesised parameter. @@ -108,7 +85,7 @@ def _with_ambient_params(node: past.Program, closure_vars: dict[str, typing.Any] lowering, `static_params` and the compiled-program key all treat them like any other argument, and only the *value* has to be supplied per call. """ - declarations = _referenced_declarations( + declarations = ambient.referenced_declarations( transform_utils._get_closure_vars_recursively(closure_vars) ) if not declarations: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py index 3708d9d1cb..47c2fdb02c 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -163,6 +163,27 @@ def test_container_rejects_an_undeclared_value(): Grid(dz=1.0) +def test_same_named_containers_in_different_modules_do_not_collide(): + """A class name is not unique; two modules may both declare a `Grid.dx`. + + Sharing a synthesised parameter is silently wrong rather than an error -- + one binding simply wins -- so the name is disambiguated by module. + """ + elsewhere = type( + "Grid", + (gtx.Container,), + {"__annotations__": {"dx": gtx.Static[float]}, "__module__": "some.other.module"}, + ) + assert Grid._declarations["dx"].name != elsewhere._declarations["dx"].name + assert Grid._declarations["dx"].qualname != elsewhere._declarations["dx"].qualname + + +def test_declaration_names_are_stable_across_runs(): + """The name lands in the compiled signature, so it must not shift.""" + assert Grid._declarations["dx"].name == Grid._declarations["dx"].name + assert Grid._declarations["dx"].name.startswith("Grid_dx_") + + def test_bind_key_is_a_plain_contextvar(): """No bespoke declaration object: binding is stdlib set/reset.""" assert isinstance(Grid.dx, contextvars.ContextVar) From ce280d7ca0ea5ce01a43ff8f9931466e4ed34900 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 6 Aug 2026 15:15:27 +0200 Subject: [PATCH 14/14] proto[next]: reject containers that cannot be told apart across runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Container identity ends up in the synthesised parameter name, which changes the stage fingerprint and therefore the build-cache key. 'id()' would be unique but would differ on every interpreter restart, so the cache would never hit; module and qualified name are stable but not always unique. Two containers built by the same factory share both and nothing stable distinguishes them, so they are now rejected at definition rather than silently sharing a parameter. 'class Grid(Container, name=...)' separates them when that is intended. The registry backing the check is weak, write-once at class definition, and never consulted on the execution path — unlike the offset registry removed earlier, which sat on the lookup path and leaked. --- src/gt4py/next/ambient.py | 27 +++++++++++++++-- .../ffront_tests/test_ambient_values.py | 29 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py index 4260d8f7d7..ae9f76306c 100644 --- a/src/gt4py/next/ambient.py +++ b/src/gt4py/next/ambient.py @@ -63,6 +63,7 @@ def delta_x(f: IJField) -> IJField: import dataclasses import hashlib import typing +import weakref from collections.abc import Generator, Mapping from typing import Annotated, Any, ClassVar @@ -75,6 +76,10 @@ def delta_x(f: IJField) -> IJField: _UNSET: Any = object() +#: live ambient containers by stable key, so indistinguishable ones are rejected +#: at definition. Weak, write-once, and never consulted on the execution path. +_containers: weakref.WeakValueDictionary[str, type] = weakref.WeakValueDictionary() + class _StaticMarker: """Annotation marker: fold the value into the generated code.""" @@ -153,17 +158,35 @@ class Container(metaclass=_ContainerMeta): """ _declarations: ClassVar[dict[str, Declaration]] = {} + #: stable identity of this container, unique among live containers + _key: ClassVar[str] = "" #: class whose attributes carry the declared *types*, for the frontend _type_view: ClassVar[type] - def __init_subclass__(cls, **kwargs: Any) -> None: + def __init_subclass__(cls, *, name: str | None = None, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) + # Identity has to be *stable across interpreter restarts*: it ends up in + # the synthesised parameter name, which changes the stage fingerprint and + # therefore the build-cache key. `id()` would be unique but would miss the + # cache on every run. Module and qualified name are stable, but not always + # unique -- two containers built by the same factory are indistinguishable + # -- so that case is rejected rather than silently sharing a parameter. + key = name or f"{cls.__module__}.{cls.__qualname__}" + if (previous := _containers.get(key)) is not None and previous is not cls: + raise TypeError( + f"Ambient container '{key}' is already defined." + " Two containers that share a module and qualified name cannot be" + " told apart across runs; pass a distinct" + " 'class MyGrid(Container, name=...)' to separate them." + ) + _containers[key] = cls + cls._key = key cls._declarations = {} for attr, hint in typing.get_type_hints(cls, include_extras=True).items(): if (declared := _declared(hint)) is None: continue type_hint, static = declared - qualname = f"{cls.__module__}.{cls.__qualname__}.{attr}" + qualname = f"{key}.{attr}" # a stable digest, not a counter: the name lands in the compiled # signature, so it must not shift with import order digest = hashlib.sha1(qualname.encode()).hexdigest()[:6] diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py index 47c2fdb02c..e7020cc1e6 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -178,6 +178,35 @@ def test_same_named_containers_in_different_modules_do_not_collide(): assert Grid._declarations["dx"].qualname != elsewhere._declarations["dx"].qualname +def test_indistinguishable_containers_are_rejected(): + """Two containers built by one factory share a module and qualified name. + + Nothing stable tells them apart, and `id()` cannot be used -- it would change + the parameter name every run and miss the build cache -- so they are rejected + rather than silently sharing a parameter. + """ + + def make(): + class Twin(gtx.Container): + dx: gtx.Static[float] + + return Twin + + make() + with pytest.raises(TypeError, match="already defined"): + make() + + +def test_explicit_name_separates_them(): + class Source(gtx.Container, name="ambient-test-source"): + dx: gtx.Static[float] + + class Target(gtx.Container, name="ambient-test-target"): + dx: gtx.Static[float] + + assert Source._declarations["dx"].name != Target._declarations["dx"].name + + def test_declaration_names_are_stable_across_runs(): """The name lands in the compiled signature, so it must not shift.""" assert Grid._declarations["dx"].name == Grid._declarations["dx"].name