diff --git a/src/gt4py/next/__init__.py b/src/gt4py/next/__init__.py index 806265a1d1..0f729fdbcd 100644 --- a/src/gt4py/next/__init__.py +++ b/src/gt4py/next/__init__.py @@ -20,7 +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 . import ambient, common, ffront, iterator, program_processors, typing +from .ambient import Container, Extern, Static, bind, bindings, freeze from .common import ( CartesianConnectivity, Connectivity, @@ -136,6 +137,14 @@ "full", "as_field", "as_connectivity", + # from ambient + "ambient", + "bind", + "bindings", + "Container", + "Extern", + "Static", + "freeze", # 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..ae9f76306c --- /dev/null +++ b/src/gt4py/next/ambient.py @@ -0,0 +1,392 @@ +# 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, 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. + +A declaration is an *annotation* in a container, and the thing it binds to is a +plain `contextvars.ContextVar`:: + + class Grid(Container): + dx: Static[float] + nu: Extern[float] + + + grid = Grid() + + + @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 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 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 + 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 hashlib +import typing +import weakref +from collections.abc import Generator, Mapping +from typing import Annotated, Any, ClassVar + +import numpy as np + +from gt4py.eve import utils as eve_utils +from gt4py.next import common +from gt4py.next.ffront import fbuiltins + + +_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.""" + + +class _ExternMarker: + """Annotation marker: pass the value as a runtime argument.""" + + +#: `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] + + +@dataclasses.dataclass(frozen=True) +class Declaration: + """What a container annotation declares: a name, a type, a kind, a variable.""" + + #: 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 + + 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: + value = self.var.get(_UNSET) + if value is _UNSET: + raise ValueError( + f"Ambient value '{self.qualname}' is not bound." + " Pass 'bind={: }' at the call, or use 'gtx.bind'." + ) + return value + + +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 + + +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) + + +class Container(metaclass=_ContainerMeta): + """ + Base for a container of ambient declarations. + + Declarations are annotations, so they create no class attribute and instance + access reaches `__getattr__` -- which is where the `ContextVar` is read. + """ + + _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, *, 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"{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] + cls._declarations[attr] = Declaration( + name=f"{cls.__name__}_{attr}_{digest}", + qualname=qualname, + 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 __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] + except KeyError: + raise AttributeError(name) from None + return declaration.value + + 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 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: + """ + 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 + if not hasattr(elem, "ndarray"): + return elem + if readonly: + # 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 + + +def as_bindings(spec: Any) -> dict[Any, Any]: + """ + 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) + 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 +def bindings(mapping: Mapping[Any, Any]) -> Generator[None, None, None]: + """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) + var = variable_for(declaration) + tokens.append((var, var.set(value))) + try: + yield + finally: + for var, token in reversed(tokens): + var.reset(token) + + +@contextlib.contextmanager +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 + + +def offset_provider_for(closure_vars: Mapping[str, Any]) -> common.OffsetProvider: + """ + Assemble the offset provider from the offsets *this* program references. + + 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(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 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( + closure_vars: Mapping[str, Any], +) -> 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 type(value)._declarations.items() + } + + +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/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/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 8e4c2e3ea0..4582f0dcaa 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 @@ -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, @@ -125,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: @@ -307,6 +319,19 @@ 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, ...]: + # 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.referenced_declarations(self._all_closure_vars).items() + if decl.static and name in synthesised + ) + ) + @functools.cached_property def gtir(self) -> itir.Program: no_args_past = toolchain.ConcreteArtifact( @@ -376,15 +401,37 @@ def with_bound_args(self, **kwargs: Any) -> ProgramWithBoundArgs: ) def __call__( + self, + *args: Any, + offset_provider: common.OffsetProvider | 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(ambient.as_bindings(bind) if bind else {}): + 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: - if offset_provider is None: - 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 + # 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.referenced_declarations(self._all_closure_vars).items() + if name in synthesised + } + kwargs = {**kwargs, **ambient_args} with program_call_context( program=self, @@ -414,9 +461,12 @@ def __call__( 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: @@ -432,11 +482,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, self._all_closure_vars) type_ = self.past_stage.past_node.type assert isinstance(type_, ts_ffront.ProgramType) new_type = ts_ffront.ProgramType( @@ -485,7 +534,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( @@ -654,10 +703,21 @@ 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 {}): + 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), self._all_closure_vars) + } if "out" not in kwargs: raise errors.MissingArgumentError(None, "out", True) out = kwargs.pop("out") @@ -679,7 +739,9 @@ 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), 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/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 292a56767b..96a40e4ff2 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.referenced_declarations( + 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/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_binding.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py new file mode 100644 index 0000000000..1b5e63410e --- /dev/null +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_binding.py @@ -0,0 +1,90 @@ +# 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 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)() + + # 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) + 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(V2E, bound[V2E]): + callee(*args, **kwargs) + else: + callee(*args, **kwargs, bind=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)() + bound = {V2E: unstructured_case.offset_provider["V2E"]} + + sum_edges_program.with_backend(unstructured_case.backend)(inp, out, bind=bound) + + 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 new file mode 100644 index 0000000000..e7020cc1e6 --- /dev/null +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -0,0 +1,231 @@ +# 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 contextvars + +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] + + +class Grid(gtx.Container): + """Declarations live in a container; `grid.dx` reads the bound value.""" + + dx: gtx.Static[float] + dx_extern: gtx.Extern[float] + + +grid = Grid() + + +@gtx.field_operator +def delta_x(f: IJFloatField) -> IJFloatField: + """Forward difference in x.""" + return (1.0 / grid.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.field_operator +def delta_x_extern(f: IJFloatField) -> IJFloatField: + """Same, but supplied as a runtime argument instead of folded in.""" + return (1.0 / grid.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) + return data, out + + +def _binding(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): + 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=_binding(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=_binding(0.5)) + np.testing.assert_allclose(out.asnumpy(), _reference(data, 0.5)) + + prog(data, out, bind=_binding(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=_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(_binding(0.5)): + run_delta_x.with_backend(cartesian_case.backend)(data, out) + 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=_binding(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=_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 + + +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 Grid._declarations["dx"].__gt_type__() == expected + 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_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_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 + 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) + assert Grid.dx.get(None) is None + + +def test_unbound_declaration_reports_itself(): + with pytest.raises(ValueError, match="not bound"): + grid.dx + + +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.bind(Grid.dx, 0.5): + assert grid.dx == 0.5 + 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 new file mode 100644 index 0000000000..74e435ccd5 --- /dev/null +++ b/tests/next_tests/unit_tests/test_ambient.py @@ -0,0 +1,166 @@ +# 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, neighbor_sum + + +Vertex = common.Dimension("Vertex") +Edge = common.Dimension("Edge") +V2EDim = common.Dimension("V2E", kind=common.DimensionKind.LOCAL) + + +V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim)) + + +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_for({"V2E": V2E}) == {} + + +def test_bound_offset_declaration_provides_its_connectivity(): + 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_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_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_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_conn([[0, 1], [1, 2]]), make_conn([[0, 1], [1, 2]]) + unfrozen = ( + 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) + gtx.freeze(m2) + assert common.hash_offset_provider_items_by_id( + {"V2E": m1} + ) == common.hash_offset_provider_items_by_id({"V2E": m2}) + + +def test_differing_frozen_connectivities_are_keyed_apart(): + 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_conn([[0, 1], [1, 2]]), readonly=True) + with pytest.raises(ValueError): + conn.ndarray[0, 0] = 7 + + +# --- embedded end-to-end (backend-free) -------------------------------------- + + +@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_conn([[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) + if mechanism == "offset_provider": + callee(*args, **kwargs, offset_provider={"V2E": m}) + elif mechanism == "context_manager": + with gtx.bind(V2E, m): + callee(*args, **kwargs) + else: + callee(*args, **kwargs, bind={V2E: m}) + + np.testing.assert_array_equal(out.asnumpy(), expected) + + +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={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 + + class Scoped(gtx.Container): + dx: gtx.Static[float] + + scoped = Scoped() + seen = {} + + def bind_and_record(): + with gtx.bind(Scoped.dx, 0.25): + seen["inner"] = scoped.dx + + with gtx.bind(Scoped.dx, 0.5): + contextvars.copy_context().run(bind_and_record) + seen["outer"] = scoped.dx + + assert seen == {"inner": 0.25, "outer": 0.5}