From 43746f13ab7891adb3145538a1959b758a505759 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 6 Aug 2026 17:16:26 +0200 Subject: [PATCH] proto[next]: ambient values, reimplemented from specification Declared as annotations in a container and bound to plain ContextVars at program-execution time, so a value is reached by bare name inside an operator and never appears in a signature. A declaration becomes a synthesised program parameter when the program is defined, so it travels the ordinary path from there; the two forms differ only in whether that parameter is listed as static. Reimplemented against tmp/ambient_spec.md rather than ported, so it carries none of the residue of the designs it went through. --- src/gt4py/next/__init__.py | 6 + src/gt4py/next/ambient.py | 320 ++++++++++++++++++ src/gt4py/next/ffront/decorator.py | 65 +++- src/gt4py/next/ffront/fbuiltins.py | 6 + src/gt4py/next/ffront/foast_to_gtir.py | 9 +- src/gt4py/next/ffront/foast_to_past.py | 17 +- src/gt4py/next/ffront/func_to_past.py | 14 +- src/gt4py/next/ffront/past_to_itir.py | 19 +- .../next/type_system/type_specifications.py | 17 + .../next/type_system/type_translation.py | 5 +- .../ffront_tests/test_ambient_values.py | 171 ++++++++++ tests/next_tests/unit_tests/test_ambient.py | 158 +++++++++ 12 files changed, 790 insertions(+), 17 deletions(-) create mode 100644 src/gt4py/next/ambient.py create mode 100644 tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py 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 806265a1d1..0e65a7409a 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 Container, Extern, Static, bind from .common import ( CartesianConnectivity, Connectivity, @@ -108,6 +109,11 @@ "iterator", "program_processors", "typing", + # from ambient + "Container", + "Extern", + "Static", + "bind", # from _core.definitions "CUPY_DEVICE_TYPE", "Device", diff --git a/src/gt4py/next/ambient.py b/src/gt4py/next/ambient.py new file mode 100644 index 0000000000..82978dd236 --- /dev/null +++ b/src/gt4py/next/ambient.py @@ -0,0 +1,320 @@ +# 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: values declared once and read from any operator without appearing in a signature. + +A declaration is an annotation in a `Container` subclass:: + + class Grid(gtx.Container): + dx: gtx.Static[float] + nu: gtx.Extern[float] + + + grid = Grid() + +Class access (`Grid.dx`) yields the `contextvars.ContextVar` a value is bound to, instance +access (`grid.dx`) its current value. Values are bound for the dynamic extent of a program +call, either with the `bind` context manager or with the `bind=` argument of a call. +""" + +from __future__ import annotations + +import contextlib +import contextvars +import dataclasses +import hashlib +import typing +import weakref +from typing import Annotated, Any, ClassVar, Iterator, Mapping, TypeAlias + +from gt4py.next.ffront import ( + fbuiltins, + field_operator_ast as foast, + stages as ffront_stages, + type_specifications as ts_ffront, +) +from gt4py.next.type_system import type_specifications as ts, type_translation + + +class _StaticMarker: ... + + +class _ExternMarker: ... + + +type Static[T] = Annotated[T, _StaticMarker] +type Extern[T] = Annotated[T, _ExternMarker] + + +def parameter_name(container_id: str, attr: str) -> str: + """ + Name of the program parameter synthesised for a declaration. + + The name reaches the build-cache key, therefore it is derived from the qualified name of + the declaration instead of anything that changes between interpreter runs. + + Examples: + >>> parameter_name("some.module.Grid", "dx") + 'Grid_dx_1235ba' + """ + digest = hashlib.sha256(f"{container_id}.{attr}".encode()).hexdigest()[:6] + return f"{container_id.rsplit('.', 1)[-1]}_{attr}_{digest}" + + +@dataclasses.dataclass(frozen=True) +class Declaration: + container_id: str + attr: str + type_: ts.TypeSpec + static: bool + var: contextvars.ContextVar[Any] + + @property + def qualified_name(self) -> str: + return f"{self.container_id}.{self.attr}" + + @property + def param_name(self) -> str: + return parameter_name(self.container_id, self.attr) + + def value(self) -> Any: + try: + return self.var.get() + except LookupError: + raise ValueError( + f"Ambient declaration '{self.qualified_name}' is not bound. Bind it with " + f"'gtx.bind(...)' or with the 'bind=' argument of the call." + ) from None + + +def _declaration_type(hint: Any) -> tuple[ts.TypeSpec, bool] | None: + """Deduce type and staticness of a declaration, or `None` if `hint` is not a declaration.""" + origin = typing.get_origin(hint) + if origin is not Static and origin is not Extern: + return None + (value_hint,) = typing.get_args(hint) + return type_translation.from_type_hint(value_hint), origin is Static + + +#: Containers by qualified name, to reject ambiguous declarations at class definition. +_CONTAINERS: weakref.WeakValueDictionary[str, type] = weakref.WeakValueDictionary() + + +class _ContainerMeta(type): + __declarations__: dict[str, Declaration] + + def __getattr__(cls, attr: str) -> contextvars.ContextVar[Any]: + if attr.startswith("__"): + raise AttributeError(attr) + try: + return cls.__declarations__[attr].var + except KeyError: + raise AttributeError(attr) from None + + +class Container(metaclass=_ContainerMeta): + """ + Base class of ambient value declarations. + + An instance without values reads the currently bound values, an instance with values binds + them (`bind(Grid(dx=0.5))`). Values given at construction live in the instance dictionary, + so they never reach `__getattr__` and the two uses do not interfere. + """ + + __declarations__: ClassVar[dict[str, Declaration]] = {} + __container_id__: ClassVar[str] = "" + + def __init_subclass__(cls, /, *, name: str | None = None, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + + if derived_from := [ + base + for base in cls.__mro__[1:] + if isinstance(base, _ContainerMeta) and base is not Container + ]: + raise TypeError( + f"Container '{cls.__name__}' must not derive from container " + f"'{derived_from[0].__name__}', use composition instead." + ) + + container_id = f"{cls.__module__}.{name or cls.__qualname__}" + if container_id in _CONTAINERS: + raise TypeError( + f"Container '{container_id}' is already declared. Give one of them an explicit " + f"name, e.g. 'class {cls.__name__}(gtx.Container, name=\"...\")'." + ) + _CONTAINERS[container_id] = cls + + declarations = {} + for attr, hint in typing.get_type_hints(cls, include_extras=True).items(): + if typing.get_origin(hint) is ClassVar: + continue + if (declaration_type := _declaration_type(hint)) is None: + raise TypeError( + f"Invalid declaration '{cls.__name__}.{attr}', " + f"expected 'Static[...]' or 'Extern[...]'." + ) + type_, static = declaration_type + declarations[attr] = Declaration( + container_id=container_id, + attr=attr, + type_=type_, + static=static, + var=contextvars.ContextVar(f"{container_id}.{attr}"), + ) + cls.__declarations__ = declarations + cls.__container_id__ = container_id + + def __init__(self, **values: Any) -> None: + for attr in values: + if attr not in self.__declarations__: + raise TypeError(f"'{type(self).__name__}' has no declaration '{attr}'.") + self.__dict__.update(values) + + def __getattr__(self, attr: str) -> Any: + try: + declaration = self.__declarations__[attr] + except KeyError: + raise AttributeError(attr) from None + return declaration.value() + + def __gt_type__(self) -> ts.NamespaceType: + return ts.NamespaceType( + qualified_name=self.__container_id__, + element_types=tuple( + (attr, declaration.type_) for attr, declaration in self.__declarations__.items() + ), + ) + + +Binding: TypeAlias = Container | Mapping[Any, Any] + + +def _context_var(key: Any) -> contextvars.ContextVar[Any]: + if isinstance(key, contextvars.ContextVar): + return key + if isinstance(key, fbuiltins.FieldOffset): + return key.ambient_var + raise TypeError(f"Cannot bind '{key}', expected a container declaration or a 'FieldOffset'.") + + +def _pairs(binding: Binding) -> list[tuple[contextvars.ContextVar[Any], Any]]: + if isinstance(binding, Container): + declarations = type(binding).__declarations__ + return [(declarations[attr].var, value) for attr, value in vars(binding).items()] + if isinstance(binding, Mapping): + return [(_context_var(key), value) for key, value in binding.items()] + raise TypeError(f"Cannot bind '{binding}', expected a container or a mapping.") + + +@contextlib.contextmanager +def bind(binding: Binding) -> Iterator[None]: + """ + Bind ambient values for the duration of the context. + + Args: + binding: A container carrying values, e.g. `Grid(dx=0.5)`, or a mapping from + declarations or offsets to values, e.g. `{Grid.dx: 0.5, V2E: connectivity}`. + + Examples: + >>> class Grid(Container): + ... dx: Static[float] + >>> with bind(Grid(dx=0.5)): + ... Grid().dx + 0.5 + """ + tokens = [var.set(value) for var, value in _pairs(binding)] + try: + yield + finally: + for token in reversed(tokens): + token.var.reset(token) + + +def _operator_stage(value: Any) -> ffront_stages.FOASTOperatorDef | None: + """Get the FOAST stage of `value` if it is a field operator, in any of its wrappings.""" + stage = getattr(value, "foast_stage", None) + if stage is None and (definition := getattr(value, "definition", None)) is not None: + stage = getattr(definition, "data", None) + return stage if isinstance(stage, ffront_stages.FOASTOperatorDef) else None + + +def declarations(closure_vars: Mapping[str, Any]) -> dict[str, Declaration]: + """Declarations read by the operators reachable from `closure_vars`, by parameter name.""" + result: dict[str, Declaration] = {} + for value in closure_vars.values(): + if (stage := _operator_stage(value)) is not None: + result |= operator_declarations(stage) + return dict(sorted(result.items())) + + +def operator_declarations(stage: ffront_stages.FOASTOperatorDef) -> dict[str, Declaration]: + """Declarations read by an operator itself and by the operators it calls, by parameter name.""" + result: dict[str, Declaration] = {} + # Attributes are resolved against this operator's own closure variables: the same name may + # refer to different containers in different operators. + for node in stage.foast_node.pre_walk_values().if_isinstance(foast.Attribute): + if isinstance(node.value, foast.Name) and isinstance( + container := stage.closure_vars.get(str(node.value.id)), Container + ): + if (declaration := type(container).__declarations__.get(node.attr)) is not None: + result[declaration.param_name] = declaration + return dict(sorted((result | declarations(stage.closure_vars)).items())) + + +def values(declarations: Mapping[str, Declaration]) -> dict[str, Any]: + """Currently bound value of each declaration, by parameter name.""" + return {name: declaration.value() for name, declaration in declarations.items()} + + +def with_parameters( + program_type: ts_ffront.ProgramType, declarations: Mapping[str, Declaration] +) -> ts_ffront.ProgramType: + """Add the synthesised parameters of `declarations` to a program type.""" + definition = program_type.definition + if not ( + new_params := { + name: declaration.type_ + for name, declaration in declarations.items() + if name not in definition.pos_or_kw_args + } + ): + return program_type + return ts_ffront.ProgramType( + definition=ts.FunctionType( + pos_only_args=definition.pos_only_args, + pos_or_kw_args={**definition.pos_or_kw_args, **new_params}, + kw_only_args=definition.kw_only_args, + returns=definition.returns, + ) + ) + + +def _offsets(closure_vars: Mapping[str, Any]) -> dict[str, fbuiltins.FieldOffset]: + result: dict[str, fbuiltins.FieldOffset] = {} + for value in closure_vars.values(): + if isinstance(value, fbuiltins.FieldOffset): + result[str(value.value)] = value + elif (stage := _operator_stage(value)) is not None: + result |= _offsets(stage.closure_vars) + return result + + +def offset_provider(closure_vars: Mapping[str, Any]) -> dict[str, Any]: + """ + Offset provider assembled from the bound offsets referenced by `closure_vars`. + + Only the offsets reachable from `closure_vars` are considered, so that an unrelated bound + offset does not leak into the offset provider (and hence into the compiled program key). + """ + return { + name: value + for name, offset in sorted(_offsets(closure_vars).items()) + if (value := offset.ambient_var.get(None)) is not None + } diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 8e4c2e3ea0..a9c07a90e9 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, @@ -105,6 +106,12 @@ class _CompilableGTEntryPointMixin(Generic[ffront_stages.DSLDefinitionT]): @abc.abstractmethod def __gt_type__(self) -> ts.CallableType: ... + @property + @abc.abstractmethod + def _ambient_declarations(self) -> dict[str, ambient.Declaration]: + """Ambient declarations read by this program-like, by synthesised parameter name.""" + ... + def with_backend(self, backend: next_backend.Backend | None) -> Self: return dataclasses.replace(self, backend=backend) @@ -137,12 +144,18 @@ def _make_compiled_programs_pool( program_type = ffront_type_info.type_in_program_context(self.__gt_type__()) assert isinstance(program_type, ts_ffront.ProgramType) + ambient_declarations: dict[str, ambient.Declaration] = self._ambient_declarations + program_type = ambient.with_parameters(program_type, ambient_declarations) # The argument descriptor mapping built here must be kept in sync with the descriptors # created in the explicitly-triggered-compilation code path # `CompiledProgramsPool.compile()`. argument_descriptor_mapping: dict[type[arguments.ArgStaticDescriptor], Sequence[str]] = {} + static_params = ( + *static_params, + *(name for name, decl in ambient_declarations.items() if decl.static), + ) if static_params: argument_descriptor_mapping[arguments.StaticArg] = static_params @@ -307,6 +320,10 @@ 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) + @functools.cached_property + def _ambient_declarations(self) -> dict[str, ambient.Declaration]: + return ambient.declarations(self.past_stage.closure_vars) + @functools.cached_property def gtir(self) -> itir.Program: no_args_past = toolchain.ConcreteArtifact( @@ -380,10 +397,17 @@ def __call__( *args: Any, offset_provider: common.OffsetProvider | None = None, enable_jit: bool | None = None, + bind: ambient.Binding | None = None, **kwargs: Any, ) -> None: + if bind is not None: + with ambient.bind(bind): + self(*args, offset_provider=offset_provider, enable_jit=enable_jit, **kwargs) + return + if offset_provider is None: - offset_provider = {} + offset_provider = ambient.offset_provider(self.past_stage.closure_vars) + ambient_kwargs = ambient.values(self._ambient_declarations) enable_jit = self.compilation_options.enable_jit if enable_jit is None else enable_jit with program_call_context( @@ -398,12 +422,19 @@ def __call__( past_process_args._validate_args( self.past_stage.past_node, arg_types=[type_translation.from_value(arg) for arg in args], - kwarg_types={k: type_translation.from_value(v) for k, v in kwargs.items()}, + kwarg_types={ + k: type_translation.from_value(v) + for k, v in {**kwargs, **ambient_kwargs}.items() + }, ) if self.backend is not None: self._compiled_programs( - *args, **kwargs, offset_provider=offset_provider, enable_jit=enable_jit + *args, + **kwargs, + **ambient_kwargs, + offset_provider=offset_provider, + enable_jit=enable_jit, ) else: # Embedded execution. @@ -654,10 +685,29 @@ def __gt_gtir__(self) -> itir.FunctionDefinition: 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: + @functools.cached_property + def _ambient_declarations(self) -> dict[str, ambient.Declaration]: + return ambient.operator_declarations(self.foast_stage) + + def _offset_provider(self, offset_provider: common.OffsetProvider | None) -> dict[str, Any]: + if offset_provider is None: + return ambient.offset_provider(self.foast_stage.closure_vars) + return {**offset_provider} + + def __call__( + self, + *args: Any, + enable_jit: bool | None = None, + bind: ambient.Binding | None = None, + **kwargs: Any, + ) -> Any: + if bind is not None: + with ambient.bind(bind): + return self(*args, enable_jit=enable_jit, **kwargs) + 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 = self._offset_provider(kwargs.pop("offset_provider", None)) if "out" not in kwargs: raise errors.MissingArgumentError(None, "out", True) out = kwargs.pop("out") @@ -670,6 +720,7 @@ def __call__(self, *args: Any, enable_jit: bool | None = None, **kwargs: Any) -> return self._compiled_programs( *args, **kwargs, + **ambient.values(self._ambient_declarations), out=out, offset_provider=offset_provider, enable_jit=self.compilation_options.enable_jit @@ -679,7 +730,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"] = self._offset_provider( + kwargs.pop("offset_provider", None) + ) 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..2cf4786a91 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[common.Connectivity]: + """Context variable the connectivity of this offset is bound to, see `gt4py.next.ambient`.""" + 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..266d50c304 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, @@ -236,10 +236,15 @@ 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: if isinstance(node.type, ts.DimensionType): return itir.AxisLiteral(value=node.type.dim.value, kind=node.type.dim.kind) + if isinstance(namespace_type := node.value.type, ts.NamespaceType): + # An ambient value: a free symbol resolving against the parameter the enclosing + # program synthesised for this declaration. + return im.ref(ambient.parameter_name(namespace_type.qualified_name, node.attr)) + if isinstance(named_tup_type := node.value.type, ts.NamedCollectionType): ind = named_tup_type.keys.index(node.attr) return im.tuple_get(ind, self.visit(node.value, **kwargs)) diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index 9a560b7ff8..f176e26d36 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -9,6 +9,7 @@ import dataclasses from typing import Any, Optional +from gt4py.next import ambient from gt4py.next.ffront import ( dialect_ast_enums, foast_to_gtir, @@ -113,10 +114,15 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: *partial_program_type.definition.kw_only_args.keys(), ] assert isinstance(type_, ts.CallableType) - assert arg_types[-1] == type_info.return_type( - type_, with_args=list(arg_types), with_kwargs=kwarg_types - ) assert args_names[-1] == "out" + # Like an explicitly written program, the generated one carries the parameters + # synthesised for the ambient values the operator reads. + ambient_declarations = ambient.operator_declarations(inp.data) + operator_arg_types = arg_types[: len(arg_types) - len(ambient_declarations)] + assert operator_arg_types[-1] == type_info.return_type( + type_, with_args=list(operator_arg_types), with_kwargs=kwarg_types + ) + args_names += [*ambient_declarations] params_decl: list[past.Symbol] = [ past.DataSymbol( @@ -131,7 +137,10 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: strict=True, ) ] - params_ref = [past.Name(id=pdecl.id, location=loc) for pdecl in params_decl[:-1]] + params_ref = [ + past.Name(id=pdecl.id, location=loc) + for pdecl in params_decl[: len(operator_arg_types) - 1] + ] out_ref = past.Name(id="out", location=loc) if inp.data.foast_node.id in inp.data.closure_vars: diff --git a/src/gt4py/next/ffront/func_to_past.py b/src/gt4py/next/ffront/func_to_past.py index 292a56767b..77dc30fcaa 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, @@ -112,6 +112,14 @@ def _postprocess_dialect_ast( def visit_FunctionDef(self, node: ast.FunctionDef) -> past.Program: self._check_not_a_reserved_name(node.name, self.get_location(node)) + loc = self.get_location(node) + params: list[past.DataSymbol] = [ + *self.visit(node.args), + *( + past.DataSymbol(id=name, type=declaration.type_, location=loc) + for name, declaration in ambient.declarations(self.closure_vars).items() + ), + ] closure_symbols: list[past.Symbol] = [ past.Symbol( id=name, @@ -125,10 +133,10 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> past.Program: return past.Program( id=node.name, type=ts.DeferredType(constraint=ts_ffront.ProgramType), - params=self.visit(node.args), + params=params, body=[self.visit(node) for node in node.body], closure_vars=closure_symbols, - location=self.get_location(node), + location=loc, ) def visit_arguments(self, node: ast.arguments) -> list[past.DataSymbol]: diff --git a/src/gt4py/next/ffront/past_to_itir.py b/src/gt4py/next/ffront/past_to_itir.py index 3febb910ef..077ce504c1 100644 --- a/src/gt4py/next/ffront/past_to_itir.py +++ b/src/gt4py/next/ffront/past_to_itir.py @@ -112,9 +112,26 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> definitions.CompilableProgramDe if not any(el is None for el in utils.flatten_nested_tuple(descr)) # type: ignore[arg-type] } body = remap_symbols.RemapSymbolRefs().visit(itir_program.body, symbol_map=static_args) + # Ambient values appear as free symbols inside the function definitions, resolving + # against the program parameters, so the substitution has to reach them there. + function_definitions = [ + itir.FunctionDefinition( + id=fun.id, + params=fun.params, + expr=remap_symbols.RemapSymbolRefs().visit( + fun.expr, + symbol_map={ + name: value + for name, value in static_args.items() + if name not in {str(param.id) for param in fun.params} + }, + ), + ) + for fun in itir_program.function_definitions + ] itir_program = itir.Program( id=itir_program.id, - function_definitions=itir_program.function_definitions, + function_definitions=function_definitions, params=itir_program.params, declarations=itir_program.declarations, body=body, diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 59ac40f0f3..01060f281b 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -79,6 +79,23 @@ def __str__(self) -> str: return f"Offset[{self.source}, {self.target}]" +class NamespaceType(TypeSpec): + """Type of an object whose attributes are named typed values, e.g. an ambient container.""" + + #: '__module__.__qualname__' of the object the namespace stands for. + qualified_name: str + element_types: tuple[tuple[str, TypeSpec], ...] + + def __getattr__(self, name: str) -> TypeSpec: + for key, type_ in object.__getattribute__(self, "element_types"): + if key == name: + return type_ + raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'") + + def __str__(self) -> str: + return f"Namespace[{self.qualified_name}]" + + class ScalarKind(eve_types.IntEnum): BOOL = 1 INT8 = 2 diff --git a/src/gt4py/next/type_system/type_translation.py b/src/gt4py/next/type_system/type_translation.py index aa003003f5..0dce9fe699 100644 --- a/src/gt4py/next/type_system/type_translation.py +++ b/src/gt4py/next/type_system/type_translation.py @@ -345,7 +345,10 @@ 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)): + if isinstance( + symbol_type, + (ts.DataType, ts.CallableType, ts.OffsetType, ts.DimensionType, ts.NamespaceType), + ): 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 new file mode 100644 index 0000000000..4fc66a4eae --- /dev/null +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py @@ -0,0 +1,171 @@ +# 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 numpy as np +import pytest + +from gt4py import next as gtx +from gt4py.next import common +from gt4py.next.ffront.fbuiltins import neighbor_sum + +from next_tests.integration_tests import cases +from next_tests.integration_tests.cases import ( + V2E, + V2EDim, + cartesian_case, + mesh_descriptor, + unstructured_case, +) +from next_tests.integration_tests.cases_utils import exec_alloc_descriptor + + +class Grid(gtx.Container): + dx: gtx.Static[float] + nu: gtx.Extern[float] + + +grid = Grid() + + +@gtx.field_operator +def scale_by_dx(a: cases.IFloatField) -> cases.IFloatField: + return grid.dx * a + + +@gtx.program +def scale_by_dx_program(a: cases.IFloatField, out: cases.IFloatField): + scale_by_dx(a, out=out) + + +@gtx.field_operator +def scale_by_nu(a: cases.IFloatField) -> cases.IFloatField: + return grid.nu * a + + +@gtx.program +def scale_by_nu_program(a: cases.IFloatField, out: cases.IFloatField): + scale_by_nu(a, out=out) + + +@gtx.field_operator +def scale_by_dx_twice(a: cases.IFloatField) -> cases.IFloatField: + return scale_by_dx(a) + scale_by_dx(a) + + +@gtx.program +def nested_program(a: cases.IFloatField, out: cases.IFloatField): + scale_by_dx_twice(a, out=out) + + +@gtx.field_operator +def sum_neighbors(a: cases.EField) -> cases.VField: + return neighbor_sum(a(V2E), axis=V2EDim) + + +@gtx.program +def sum_neighbors_program(a: cases.EField, out: cases.VField): + sum_neighbors(a, out=out) + + +def _inout(case, program): + return ( + cases.allocate(case, program, "a")(), + cases.allocate(case, program, "out").zeros()(), + ) + + +@pytest.mark.parametrize("spacing", [0.5, 2.0]) +def test_value_bound_at_call(cartesian_case, spacing): + a, out = _inout(cartesian_case, scale_by_dx_program) + + scale_by_dx_program.with_backend(cartesian_case.backend)(a, out, bind=Grid(dx=spacing)) + + np.testing.assert_allclose(out.asnumpy(), spacing * a.asnumpy()) + + +def test_value_bound_in_region(cartesian_case): + a, out = _inout(cartesian_case, scale_by_dx_program) + + with gtx.bind(Grid(dx=0.5)): + cases.verify(cartesian_case, scale_by_dx_program, a, out=out, ref=0.5 * a.asnumpy()) + + +def test_extern_bound_at_call(cartesian_case): + a, out = _inout(cartesian_case, scale_by_nu_program) + + scale_by_nu_program.with_backend(cartesian_case.backend)(a, out, bind=Grid(nu=1e-3)) + + np.testing.assert_allclose(out.asnumpy(), 1e-3 * a.asnumpy()) + + +def test_value_reaches_nested_operator(cartesian_case): + a, out = _inout(cartesian_case, nested_program) + + nested_program.with_backend(cartesian_case.backend)(a, out, bind=Grid(dx=0.5)) + + np.testing.assert_allclose(out.asnumpy(), a.asnumpy()) + + +def test_distinct_values_do_not_share_a_compiled_program(cartesian_case): + a, out = _inout(cartesian_case, scale_by_dx_program) + testee = scale_by_dx_program.with_backend(cartesian_case.backend) + + for spacing in (0.5, 2.0): + testee(a, out, bind=Grid(dx=spacing)) + np.testing.assert_allclose(out.asnumpy(), spacing * a.asnumpy()) + + +@pytest.mark.parametrize( + "program, variants", + [(scale_by_dx_program, 2), (scale_by_nu_program, 1)], + ids=["static", "extern"], +) +def test_static_specialises_and_extern_does_not(cartesian_case, program, variants): + if cartesian_case.backend is None: + pytest.skip("Embedded execution does not compile programs.") + a, out = _inout(cartesian_case, program) + testee = program.with_backend(cartesian_case.backend) + + for value in (0.5, 2.0): + # Both declarations are bound although each program reads only one: a static + # declaration a program does not read must not specialise it either. + testee(a, out, bind=Grid(dx=value, nu=value)) + + assert len(testee._compiled_programs.compiled_programs) == variants + + +def test_bind_does_not_leak_past_the_call(cartesian_case): + a, out = _inout(cartesian_case, scale_by_dx_program) + + scale_by_dx_program.with_backend(cartesian_case.backend)(a, out, bind=Grid(dx=0.5)) + + with pytest.raises(ValueError, match="not bound"): + grid.dx + + +@pytest.mark.uses_unstructured_shift +@pytest.mark.parametrize( + "testee", [sum_neighbors, sum_neighbors_program], ids=["field-operator", "program"] +) +def test_offset_spellings_agree(unstructured_case, testee): + connectivity = unstructured_case.offset_provider[V2E.value] + a = cases.allocate(unstructured_case, sum_neighbors, "a")() + outs = [ + cases.allocate(unstructured_case, sum_neighbors, cases.RETURN).zeros()() for _ in range(3) + ] + testee = testee.with_backend(unstructured_case.backend) + + testee(a, out=outs[0], offset_provider=unstructured_case.offset_provider) + with gtx.bind({V2E: connectivity}): + testee(a, out=outs[1]) + testee(a, out=outs[2], bind={V2E: connectivity}) + + v2e_table = connectivity.asnumpy() + ref = np.sum(a.asnumpy()[v2e_table], axis=1, where=v2e_table != common._DEFAULT_SKIP_VALUE) + for out in outs: + np.testing.assert_allclose(out.asnumpy(), ref) 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..80b8132c12 --- /dev/null +++ b/tests/next_tests/unit_tests/test_ambient.py @@ -0,0 +1,158 @@ +# 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 contextvars +import re +import sys +import types + +import pytest + +from gt4py import next as gtx +from gt4py.next import ambient +from gt4py.next.type_system import type_specifications as ts, type_translation + + +IDim = gtx.Dimension("IDim") +KDim = gtx.Dimension("KDim", kind=gtx.DimensionKind.VERTICAL) +Koff = gtx.FieldOffset("Koff", source=KDim, target=(KDim,)) +Ioff = gtx.FieldOffset("Ioff", source=IDim, target=(IDim,)) + +KField = gtx.Field[gtx.Dims[KDim], gtx.float64] + + +class Grid(gtx.Container): + dx: gtx.Static[float] + nu: gtx.Extern[float] + + +grid = Grid() + + +def container_in_module(module_name: str) -> type: + """Declare a container in a module of its own, as two unrelated user modules would.""" + module = types.ModuleType(module_name) + module.gtx = gtx + sys.modules[module_name] = module + + class Grid(gtx.Container): + __module__ = module_name + + dx: gtx.Static[float] + + return Grid + + +def test_class_access_is_the_context_variable(): + assert isinstance(Grid.dx, contextvars.ContextVar) + assert Grid.dx is not Grid.nu + + +def test_instance_access_is_the_value(): + with gtx.bind({Grid.dx: 0.5}): + assert grid.dx == 0.5 + assert 1.0 / grid.dx == 2.0 # a plain scalar, no arithmetic protocol involved + + +def test_container_binds_all_its_values(): + with gtx.bind(Grid(dx=0.5, nu=1e-3)): + assert (grid.dx, grid.nu) == (0.5, 1e-3) + + +def test_partial_container_binds_what_it_carries(): + with gtx.bind(Grid(dx=0.5)): + assert grid.dx == 0.5 + with pytest.raises(ValueError, match="Grid.nu' is not bound"): + grid.nu + + +def test_undeclared_keyword_is_rejected(): + with pytest.raises(TypeError, match="no declaration 'dz'"): + Grid(dz=0.5) + + +def test_bindings_nest_and_unwind(): + with gtx.bind({Grid.dx: 0.5}): + with gtx.bind({Grid.dx: 0.25}): + assert grid.dx == 0.25 + assert grid.dx == 0.5 + + +def test_bindings_are_context_local(): + def bound(): + with gtx.bind({Grid.dx: 0.5}): + return grid.dx + + assert contextvars.copy_context().run(bound) == 0.5 + with pytest.raises(ValueError, match="Grid.dx' is not bound"): + grid.dx + + +def test_unbound_declaration_names_the_declaration(): + with pytest.raises(ValueError, match=re.escape(f"'{__name__}.Grid.dx'")): + grid.dx + + +def test_same_name_in_different_modules_gives_distinct_parameters(): + first = container_in_module("next_tests_ambient_module_a") + second = container_in_module("next_tests_ambient_module_b") + + assert first.__declarations__["dx"].param_name != second.__declarations__["dx"].param_name + + +def test_parameter_name_depends_only_on_the_qualified_name(): + assert Grid.__declarations__["dx"].param_name == ambient.parameter_name( + f"{__name__}.Grid", "dx" + ) + + +def test_ambiguous_container_is_rejected(): + def declare(**kwargs): + class Ambiguous(gtx.Container, **kwargs): + dx: gtx.Static[float] + + return Ambiguous + + first = declare() + with pytest.raises(TypeError, match="already declared"): + declare() + assert declare(name="explicit").__container_id__ != first.__container_id__ + + +def test_subclassing_is_rejected(): + with pytest.raises(TypeError, match="must not derive from container"): + + class Derived(Grid): + dz: gtx.Static[float] + + +def test_declaration_without_type_is_rejected(): + with pytest.raises(TypeError, match="Invalid declaration"): + + class Untagged(gtx.Container): + dx: float + + +def test_container_types_itself_without_binding(): + type_ = type_translation.from_value(grid) + + assert isinstance(type_, ts.NamespaceType) + assert type_.dx == ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + + +def test_offset_provider_is_scoped_to_the_referenced_offsets(): + @gtx.field_operator + def shift_k(f: KField) -> KField: + return f(Koff[1]) + + @gtx.program + def testee(f: KField, out: KField) -> None: + shift_k(f, out=out) + + with gtx.bind({Koff: "k-connectivity", Ioff: "unrelated"}): + assert ambient.offset_provider(testee.past_stage.closure_vars) == {"Koff": "k-connectivity"}