diff --git a/docs/development/ADRs/next/0028-Homogeneous_Tuple_Comprehensions.md b/docs/development/ADRs/next/0028-Homogeneous_Tuple_Comprehensions.md new file mode 100644 index 0000000000..9b3b34e4f3 --- /dev/null +++ b/docs/development/ADRs/next/0028-Homogeneous_Tuple_Comprehensions.md @@ -0,0 +1,41 @@ +--- +tags: [] +--- + +# Homogeneous Tuple Comprehensions + +- **Status**: valid +- **Authors**: Till Ehrengruber (@tehrengruber), Sara Faghih-Naini (@SF-N) +- **Created**: 2026-08-27 +- **Updated**: 2026-08-27 + +In the context of tuple comprehensions in the field-view frontend, facing the constraint that every FOAST node carries exactly one type, we decided to support only homogeneous iterables — all elements of the same type — and reject heterogeneous ones in the type deduction. + +## Context + +Tuple comprehensions, e.g. `tuple(2.0 * el for el in (a, b))`, are typed and lowered with a single mapper: one target (`el`) and one element expression (`2.0 * el`), shared by all elements of the iterable. In FOAST every node has a single `type` attribute, so the target symbol — and consequently every node in the element expression — can only be typed once. If the iterable's elements had different types, the mapper would need a different type per element, i.e. per-element re-typing (monomorphization) of the element expression, which the FOAST type system does not support. + +The same constraint exists at the GTIR level: the ITIR type inference also stores a single type per node (and asserts on conflicting re-assignment), so the single `map_tuple` lambda used to lower comprehensions over variable-length tuples can only have one function type. For variable-length iterables heterogeneity cannot occur in the first place, since `VarArgType` describes all elements with a single element type. Rejecting heterogeneous iterables in the FOAST type deduction just surfaces the error earliest, with a source location. + +## Decision + +Only homogeneous iterables are supported, both fixed-length and variable-length. Heterogeneous ones are rejected in the type deduction. E.g., with `a`, `b`, `c`, `d` of equal type: + +```python +tuple(2.0 * el for el in (a, b)) # supported +tuple(2.0 * el for el in (a(V2E), b(V2E))) # supported +tuple(local_el + el for local_el, el in ((a(V2E), b), (c(V2E), d))) # supported +tuple(2.0 * el for el in (a(V2E), b)) # rejected: local vs. non-local element +``` + +Note that homogeneity applies to the iterable's elements as a whole: in the third example each element is a pair of a local and a non-local field, but all elements share that same tuple type, so each target symbol still has a single consistent type. + +## Consequences + +- Typing and lowering stay simple: the element expression is visited once, with one type per node. +- Computations over differently-typed elements cannot be written as a comprehension; they must be spelled out per element. +- The restriction could be lifted for fixed-length iterables by typing the mapper generically: the target symbol gets a type variable bounded by the valid element types, so a single type per node still suffices during type deduction. After `map_tuple` expansion the mapper is instantiated once per element, and each instance can then be specialized to its concrete element type. This is possible as a follow-up without breaking existing code, since it only widens the set of accepted programs. For variable-length iterables there is nothing to lift: heterogeneity cannot occur, as `VarArgType` has a single element type by construction. + +## References + +- PR [#2833](https://github.com/GridTools/gt4py/pull/2833) (tuple comprehension support) diff --git a/docs/development/ADRs/next/README.md b/docs/development/ADRs/next/README.md index f19167ef9b..e2f26b4854 100644 --- a/docs/development/ADRs/next/README.md +++ b/docs/development/ADRs/next/README.md @@ -29,6 +29,7 @@ Writing a new ADR is simple: - [0002 - Field View Lowering](0002-Field_View_Lowering.md) - [0010 - Domain in Field View](0010-Domain_in_Field_View.md) - [0013 - Scalar vs 0d-Fields](0013-Scalar_vs_0d_Fields.md) +- [0028 - Homogeneous Tuple Comprehensions](0028-Homogeneous_Tuple_Comprehensions.md) ### Iterator IR #iterator diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index bf63aa692a..e13a85b9b2 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -214,6 +214,12 @@ def _field_domain_descriptor_mapping_from_func_type(func_type: ts.FunctionType) for name, type_ in param_types.items(): for el_type_, path in type_info.primitive_constituents(type_, with_path_arg=True): if isinstance(el_type_, ts.FieldType): + if any(isinstance(idx, type_info.VarLenTupleElementKey) for idx in path): + raise NotImplementedError( + f"Static domains are not supported for parameter '{name}': the " + "elements of a variable-length tuple can not be addressed before " + "the concrete length is known at the call." + ) path_as_expr = "".join(f"[{idx}]" for idx in path) static_domain_args.append(f"{name}{path_as_expr}") return static_domain_args @@ -306,15 +312,20 @@ def _all_closure_vars(self) -> dict[str, Any]: @functools.cached_property def gtir(self) -> itir.Program: - no_args_past = workflow.ConcreteArtifact( + # The program is lowered without a call here, so use the declared parameter + # types as compile-time arguments. + declared_args_past = workflow.ConcreteArtifact( data=ffront_stages.PASTProgramDef( past_node=self.past_stage.past_node, closure_vars=self.past_stage.closure_vars, grid_type=self.definition_stage.grid_type, ), - args=arguments.CompileTimeArgs.empty(), + args=dataclasses.replace( + arguments.CompileTimeArgs.empty(), + args=tuple(param.type for param in self.past_stage.past_node.params), + ), ) - return self._frontend_transforms.past_to_itir(no_args_past).data + return self._frontend_transforms.past_to_itir(declared_args_past).data def with_grid_type(self, grid_type: common.GridType) -> Program: return dataclasses.replace( diff --git a/src/gt4py/next/ffront/dialect_parser.py b/src/gt4py/next/ffront/dialect_parser.py index 14155b6f31..7e2417783c 100644 --- a/src/gt4py/next/ffront/dialect_parser.py +++ b/src/gt4py/next/ffront/dialect_parser.py @@ -54,7 +54,10 @@ ), ast.SetComp: ("set comprehension", ()), ast.DictComp: ("dictionary comprehension", ()), - ast.GeneratorExp: ("generator expression", ()), + ast.GeneratorExp: ( + "generator expression", + ("Generator expressions are only supported as the argument of 'tuple(...)'.",), + ), ast.Lambda: ( "'lambda' expression", ("Define a separate function decorated with '@field_operator' instead.",), diff --git a/src/gt4py/next/ffront/field_operator_ast.py b/src/gt4py/next/ffront/field_operator_ast.py index f84dd53d62..ba99476d23 100644 --- a/src/gt4py/next/ffront/field_operator_ast.py +++ b/src/gt4py/next/ffront/field_operator_ast.py @@ -14,6 +14,7 @@ from gt4py.eve import Coerced, Node, SourceLocation, SymbolName, SymbolRef, datamodels from gt4py.eve.traits import SymbolTableTrait from gt4py.eve.type_definitions import StrEnum +from gt4py.eve.xtyping import MaybeNestedInTuple from gt4py.next import utils from gt4py.next.ffront import dialect_ast_enums, type_specifications as ts_ffront from gt4py.next.type_system import type_specifications as ts @@ -98,6 +99,35 @@ class TupleExpr(Expr): elts: list[Expr] +# TODO(tehrengruber): extend this to supported nested tuple comprehension. +# e.g. `tuple(element_expr for child in nested_tuple for grand_child in child)` +# would be represented by: +# ``` +# class TupleComprehension(Expr): # ruff: noqa: ERA001 +# inner: TupleComprehensionMapper | NestedTupleCompr # ruff: noqa: ERA001 +# class NestedTupleCompr(Expr, SymbolTableTrait): # ruff: noqa: ERA001 +# params: tuple[DataSymbol] # ruff: noqa: ERA001 +# body: TupleComprehension # ruff: noqa: ERA001 +# ``` +class TupleComprehension(Expr): + """ + tuple(element_expr for target in iterable) + Note: The structure here differs from the one in the Python AST. Here we group target and + element expression in order to cleanly nest by the symbols being introduced, whereas in + the Python AST target and iterable are grouped into generator nodes. + """ + + inner: TupleComprehensionMapper + iterable: Expr + + +# This is essentially a lambda. The difference is that for a lambda we might not know the type of +# the args; therefore this is named differently at the moment. +class TupleComprehensionMapper(LocatedNode, SymbolTableTrait): + target: MaybeNestedInTuple[DataSymbol] + element_expr: Expr + + class UnaryOp(Expr): op: dialect_ast_enums.UnaryOperator operand: Expr diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 738b105a61..f21c85f2c4 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -12,6 +12,7 @@ import gt4py.next.ffront.field_operator_ast as foast from gt4py import eve from gt4py.eve import NodeTranslator, NodeVisitor, traits +from gt4py.eve.xtyping import MaybeNestedInTuple from gt4py.next import common, errors from gt4py.next.common import Dimension, DimensionKind, promote_dims from gt4py.next.ffront import ( @@ -24,6 +25,7 @@ from gt4py.next.ffront.foast_passes import utils as foast_utils from gt4py.next.iterator import builtins from gt4py.next.type_system import type_info, type_specifications as ts, type_translation +from gt4py.next.utils import tree_map OperatorNodeT = TypeVar("OperatorNodeT", bound=foast.LocatedNode) @@ -756,6 +758,132 @@ def visit_TupleExpr(self, node: foast.TupleExpr, **kwargs: Any) -> foast.TupleEx new_type = ts.TupleType(types=[element.type for element in new_elts]) return foast.TupleExpr(elts=new_elts, type=new_type, location=node.location) + def _deduce_tuple_comprehension_target_type( + self, + target: MaybeNestedInTuple[foast.Symbol], + element_type: ts.TypeSpec, + **kwargs: Any, + ) -> MaybeNestedInTuple[foast.Symbol]: + """ + Deduce the types of the loop target, e.g. `(a, b)` in `tuple(a + b for (a, b) in it)`. + + Each target symbol starts out with a deferred type; revisiting it with `refine_type` + replaces that by the constituent of `element_type` at the symbol's unpacking position. + """ + + @tree_map(with_path_arg=True) + def process_target(target_el: foast.Symbol, path: tuple[int, ...]) -> foast.Symbol: + # Walk the unpacking pattern and the element type side-by-side so that the + # number of names and values can be compared at every nesting level. + type_ = element_type + target_tree: MaybeNestedInTuple[foast.Symbol] = target + for i in path: + assert isinstance(target_tree, tuple) + if not isinstance(type_, ts.TupleType): + raise errors.DSLError( + target_el.location, f"Cannot unpack non-iterable '{type_}' object." + ) + if len(type_.types) > len(target_tree): + raise errors.DSLError( + target_el.location, + f"Too many values to unpack (expected {len(target_tree)}).", + ) + if len(type_.types) < len(target_tree): + raise errors.DSLError( + target_el.location, + f"Not enough values to unpack (expected {len(target_tree)}, " + f"got {len(type_.types)}).", + ) + type_ = type_.types[i] + target_tree = target_tree[i] + return self.visit(target_el, refine_type=type_, **kwargs) + + return process_target(target) + + def _deduce_tuple_comprehension_mapper( + self, + node: foast.TupleComprehension, + target: MaybeNestedInTuple[foast.Symbol], + element_type: ts.DataType, + **kwargs: Any, + ) -> foast.TupleComprehensionMapper: + """ + Deduce the per-element part, e.g. `a + b for (a, b)` in `tuple(a + b for (a, b) in it)`. + + Refines the target symbols against `element_type` and then types the element + expression with those symbols in scope (a fresh child symbol table). + """ + inner_kwargs = {**kwargs, "symtable": kwargs["symtable"].new_child()} + new_target = self._deduce_tuple_comprehension_target_type( + target, element_type, **inner_kwargs + ) + return foast.TupleComprehensionMapper( + target=new_target, + element_expr=self.visit(node.inner.element_expr, **inner_kwargs), + location=node.location, + ) + + def visit_TupleComprehension( + self, node: foast.TupleComprehension, **kwargs: Any + ) -> foast.TupleComprehension: + # The target symbols are visited in `_deduce_tuple_comprehension_target_type`, + # where their types are refined against the iterable's element type. + target = node.inner.target + iterable = self.visit(node.iterable, **kwargs) + + if isinstance(iterable.type, ts.TupleType): + if len(iterable.type.types) == 0: + raise errors.DSLError( + iterable.location, + "Cannot iterate over an empty tuple in a tuple comprehension.", + ) + if not all( + isinstance(element_type, ts.DataType) for element_type in iterable.type.types + ): + raise errors.DSLError( + iterable.location, + "Tuple comprehension iterable elements must be data types.", + ) + + element_types = cast(list[ts.DataType], iterable.type.types) + # Only homogeneous iterables are supported, see ADR 0028. + if not all(element_type == element_types[0] for element_type in element_types): + raise errors.DSLError( + iterable.location, + "Tuple comprehensions over fixed-length tuples with differently typed " + "iterable elements are not implemented (see ADR 0028).", + ) + new_mapper = self._deduce_tuple_comprehension_mapper( + node, target, element_types[0], **kwargs + ) + result = foast.TupleComprehension( + inner=new_mapper, + iterable=iterable, + location=node.location, + type=ts.TupleType(types=[new_mapper.element_expr.type for _ in element_types]), + ) + return result + elif isinstance(iterable.type, ts.VarLenTupleType): + element_type = iterable.type.element_type + assert isinstance(element_type, ts.DataType) # 'DeferredType' can not occur here + new_mapper = self._deduce_tuple_comprehension_mapper( + node, target, element_type, **kwargs + ) + element_expr = new_mapper.element_expr + return_type = ts.VarLenTupleType(element_type=element_expr.type) + + return foast.TupleComprehension( + inner=new_mapper, + iterable=iterable, + location=node.location, + type=return_type, + ) + else: + raise errors.DSLError( + iterable.location, + f"Iterable in generator expression must be a tuple, got '{iterable.type}'.", + ) + def visit_Call(self, node: foast.Call, **kwargs: Any) -> foast.Call: new_func = self.visit(node.func, **kwargs) new_args = self.visit(node.args, **kwargs) diff --git a/src/gt4py/next/ffront/foast_pretty_printer.py b/src/gt4py/next/ffront/foast_pretty_printer.py index 8b2e369501..8c4203be6d 100644 --- a/src/gt4py/next/ffront/foast_pretty_printer.py +++ b/src/gt4py/next/ffront/foast_pretty_printer.py @@ -120,6 +120,22 @@ def apply(cls, node: foast.LocatedNode, **kwargs: Any) -> str: # type: ignore[o UnaryOp = as_fmt("{op}{operand}") + def visit_TupleComprehensionMapper( + self, node: foast.TupleComprehensionMapper, **kwargs: Any + ) -> str: + def format_target(target: Any) -> str: + if isinstance(target, tuple): + return f"({', '.join(format_target(el) for el in target)})" + return self.visit(target, **kwargs) + + element_expr = self.visit(node.element_expr, **kwargs) + return f"{element_expr} for {format_target(node.target)}" + + def visit_TupleComprehension(self, node: foast.TupleComprehension, **kwargs: Any) -> str: + mapper = self.visit(node.inner, **kwargs) + iterable = self.visit(node.iterable, **kwargs) + return f"tuple(({mapper} in {iterable}))" + def visit_UnaryOp(self, node: foast.UnaryOp, **kwargs: Any) -> str: if node.op is dialect_ast_enums.UnaryOperator.NOT: op = "not " diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 6913b5d6be..08ba691f7e 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -8,6 +8,7 @@ import dataclasses +import functools from typing import Any, Callable, Never, Optional from gt4py import eve @@ -258,6 +259,75 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> itir.Expr: def visit_TupleExpr(self, node: foast.TupleExpr, **kwargs: Any) -> itir.Expr: return im.make_tuple(*[self.visit(el, **kwargs) for el in node.elts]) + def _bind_tuple_comprehension_target( + self, + comprehension_target: itir.Sym | tuple, + element_expr: itir.Expr, + iterable_element: itir.Expr | str, + ) -> itir.Expr: + """ + Wrap `element_expr` in a `let` binding the `comprehension_target` to `iterable_element`. + + For `2.0 * a + b for (a, b) in iterable`: + - `comprehension_target`: `(a, b)` + - `element_expr`: `2.0 * a + b` + - `iterable_element`: the current element of `iterable` + + returns + `let a = iterable_element[0], b = iterable_element[1] in element_expr`. + """ + if isinstance(comprehension_target, itir.Sym): + return im.let(comprehension_target, iterable_element)(element_expr) + + flat_targets = utils.flatten_nested_tuple(comprehension_target) + nested_target_values = utils.tree_map( + lambda _, path: functools.reduce( + lambda element, index: im.tuple_get(index, element), path, iterable_element + ), + with_path_arg=True, + )(comprehension_target) + + flat_target_values = utils.flatten_nested_tuple(nested_target_values) # type: ignore[arg-type] + + target_bindings = tuple(zip(flat_targets, flat_target_values, strict=True)) + return im.let(*target_bindings)(element_expr) # type: ignore[arg-type] + + def visit_TupleComprehension(self, node: foast.TupleComprehension, **kwargs: Any) -> itir.Expr: + # Only homogeneous iterables — all elements of the same type — are supported; + # heterogeneous ones are rejected in the type deduction (see ADR 0028). + comprehension_target = self.visit(node.inner.target, **kwargs) + element_expr = self.visit(node.inner.element_expr, **kwargs) + iterable_expr = self.visit(node.iterable, **kwargs) + iterable_type = node.iterable.type + + lower_body_for_iterable_element = functools.partial( + self._bind_tuple_comprehension_target, comprehension_target, element_expr + ) + + if isinstance(iterable_type, ts.TupleType): + assert isinstance(node.type, ts.TupleType) + iterable_value_name = next(self.uid_generator["__tuple_comprh"]) + + fixed_tuple_elements = [ + lower_body_for_iterable_element(im.tuple_get(element_index, iterable_value_name)) + for element_index in range(len(iterable_type.types)) + ] + + result_tuple = im.make_tuple(*fixed_tuple_elements) + return im.let(iterable_value_name, iterable_expr)(result_tuple) + + assert isinstance(iterable_type, ts.VarLenTupleType) + assert isinstance(node.type, ts.VarLenTupleType) + if isinstance(comprehension_target, itir.Sym): + map_tuple_lambda = im.lambda_(comprehension_target)(element_expr) + else: + iterable_element_param = next(self.uid_generator["__tuple_comprh"]) + map_tuple_lambda = im.lambda_(iterable_element_param)( + lower_body_for_iterable_element(iterable_element_param) + ) + + return im.call(im.call("map_tuple")(map_tuple_lambda))(iterable_expr) + def visit_UnaryOp(self, node: foast.UnaryOp, **kwargs: Any) -> itir.Expr: # TODO(tehrengruber): extend iterator ir to support unary operators dtype = type_info.extract_dtype(node.type) diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index beabe53c71..fa44a99bbd 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -12,6 +12,7 @@ from gt4py.next.ffront import ( dialect_ast_enums, foast_to_gtir, + past_to_itir, program_ast as past, stages as ffront_stages, type_info as ffront_type_info, @@ -21,7 +22,7 @@ from gt4py.next.ffront.stages import ConcreteFOASTOperatorDef, ConcretePASTProgramDef from gt4py.next.iterator import ir as itir from gt4py.next.otf import workflow -from gt4py.next.type_system import type_info, type_specifications as ts +from gt4py.next.type_system import type_specifications as ts @dataclasses.dataclass(frozen=True) @@ -99,37 +100,28 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: # of arg and kwarg types # TODO(tehrengruber): check foast operator has no out argument that clashes # with the out argument of the program we generate here. - arg_types, kwarg_types = inp.args.args, inp.args.kwargs - assert not kwarg_types - - type_ = inp.data.foast_node.type loc = inp.data.foast_node.location assert isinstance(inp.data.foast_node.type, ts.CallableType) - partial_program_type = ffront_type_info.type_in_program_context(inp.data.foast_node.type) - assert isinstance(partial_program_type, ts_ffront.ProgramType) - args_names = [ - *partial_program_type.definition.pos_only_args, - *partial_program_type.definition.pos_or_kw_args.keys(), - *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 + program_type = ffront_type_info.type_in_program_context(inp.data.foast_node.type) + assert isinstance(program_type, ts_ffront.ProgramType) + param_types: tuple[tuple[str, ts.TypeSpec], ...] = ( + *( + (f"__pos_arg{i}", type_) + for i, type_ in enumerate(program_type.definition.pos_only_args) + ), + *program_type.definition.pos_or_kw_args.items(), + *program_type.definition.kw_only_args.items(), ) - assert args_names[-1] == "out" + assert param_types[-1][0] == "out" params_decl: list[past.Symbol] = [ past.DataSymbol( id=name, - type=arg_type, + type=type_, # note: this might be a generic type like VarLenTupleType namespace=dialect_ast_enums.Namespace.LOCAL, location=loc, ) - for name, arg_type in zip( - args_names, - arg_types, - strict=True, - ) + for name, type_ in param_types ] params_ref = [past.Name(id=pdecl.id, location=loc) for pdecl in params_decl[:-1]] out_ref = past.Name(id="out", location=loc) @@ -163,6 +155,25 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: closure_vars=closure_symbols, location=loc, ) + # Specialize the parameter types to the concrete argument types of the call before + # the type deduction: the declared types may be generic — e.g. `DeferredType` for + # scan operator parameters, which are only promoted to fields at the call — and the + # body call can not be type checked against them. The specialization happens again + # later, in `past_to_gtir`, which also covers hand-written programs that are not + # specialized here. With support for generic types in the program type deduction + # this step would become obsolete, as the generic variant could be type checked + # directly. + # TODO(tehrengruber): Remove again. This specialization is only needed as the `program_type` + # above contains `DeferredType` for direct scan operator calls due to the lack of + # generics. + untyped_past_node = past_to_itir._specialize_program_params( + untyped_past_node, + inp.args.args, + inp.args.kwargs, + # argument validation has not run yet on this path; incompatible arguments are + # rejected right below with the proper error message + check_compatibility=False, + ) untyped_past_node = closure_var_type_deduction.ClosureVarTypeDeduction.apply( untyped_past_node, fieldop_itir_closure_vars, # type: ignore[arg-type] diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py index 14dceb25d1..1820c3203b 100644 --- a/src/gt4py/next/ffront/func_to_foast.py +++ b/src/gt4py/next/ffront/func_to_foast.py @@ -11,9 +11,10 @@ import ast import textwrap import typing -from typing import Any, Type +from typing import Any import gt4py.eve as eve +from gt4py.eve.xtyping import MaybeNestedInTuple from gt4py.next import errors from gt4py.next.ffront import ( dialect_ast_enums, @@ -324,7 +325,7 @@ def visit_Assign( if not isinstance(target, ast.Name): raise errors.DSLError(self.get_location(node), "Can only assign to names.") new_value = self.visit(node.value) - constraint_type: Type[ts.DataType] = ts.DataType + constraint_type: type[ts.DataType] = ts.DataType if isinstance(new_value, foast.TupleExpr): constraint_type = ts.TupleType elif ( @@ -542,24 +543,94 @@ def visit_NotEq(self, node: ast.NotEq, **kwargs: Any) -> foast.CompareOperator: return foast.CompareOperator.NOTEQ def _verify_builtin_type_constructor(self, node: ast.Call) -> None: - if len(node.args) > 0: - arg = node.args[0] - if not ( - isinstance(arg, ast.Constant) - or (isinstance(arg, ast.UnaryOp) and isinstance(arg.operand, ast.Constant)) - ): - raise errors.DSLError( - self.get_location(node), - f"'{self._func_name(node)}()' only takes literal arguments.", - ) + assert isinstance(node.func, ast.Name) + if len(node.args) != 1: + raise errors.DSLError( + self.get_location(node), + f"'{self._func_name(node)}()' takes exactly one argument, got {len(node.args)}.", + ) + (arg,) = node.args + if not ( + isinstance(arg, ast.Constant) + or (isinstance(arg, ast.UnaryOp) and isinstance(arg.operand, ast.Constant)) + or (node.func.id == "tuple" and isinstance(arg, ast.GeneratorExp)) + ): + allowed = ( + "literal arguments or a generator expression" + if node.func.id == "tuple" + else "literal arguments" + ) + raise errors.DSLError( + self.get_location(node), + f"'{self._func_name(node)}()' only takes {allowed}.", + ) def _func_name(self, node: ast.Call) -> str: return node.func.id # type: ignore[attr-defined] # We want this to fail if the attribute does not exist unexpectedly. - def visit_Call(self, node: ast.Call, **kwargs: Any) -> foast.Call: - # TODO(tehrengruber): is this still needed or redundant with the checks in type deduction? + def visit_Call(self, node: ast.Call, **kwargs: Any) -> foast.Call | foast.TupleComprehension: if isinstance(node.func, ast.Name): func_name = self._func_name(node) + + if ( + func_name == "tuple" + and len(node.args) == 1 + and isinstance(gen_expr := node.args[0], ast.GeneratorExp) + ): + if len(gen_expr.generators) != 1: + raise errors.DSLError( + self.get_location(gen_expr), + "Nested generator expressions are not supported.", + hints=["Use a single 'for' clause iterating over one tuple."], + ) + if gen_expr.generators[0].ifs != []: + raise errors.DSLError( + self.get_location(gen_expr.generators[0].ifs[0]), + "Conditionals are not supported in generator expressions.", + notes=[ + ( + "The length of the resulting tuple must be known at compile time, " + "but an 'if' filter makes it depend on runtime values." + ) + ], + ) + + seen_target_names: set[str] = set() + + def parse_target(target: ast.expr) -> MaybeNestedInTuple[foast.DataSymbol]: + if isinstance(target, ast.Tuple): + return tuple(parse_target(el) for el in target.elts) + if not (isinstance(target, ast.Name) and isinstance(target.ctx, ast.Store)): + raise errors.UnsupportedPythonFeatureError( + self.get_location(target), + "generator expression targets other than names " + "and (nested) tuples of names", + ) + if target.id in seen_target_names: + raise errors.DSLError( + self.get_location(target), + f"Duplicate name '{target.id}' in generator expression target.", + ) + seen_target_names.add(target.id) + return foast.DataSymbol( + id=target.id, + location=self.get_location(target), + type=ts.DeferredType(constraint=None), + ) + + target = parse_target(gen_expr.generators[0].target) + + return foast.TupleComprehension( + inner=foast.TupleComprehensionMapper( + target=target, + element_expr=self.visit(gen_expr.elt, **kwargs), + location=self.get_location(node), + ), + iterable=self.visit(gen_expr.generators[0].iter, **kwargs), + location=self.get_location(node), + ) + + # TODO(tehrengruber): is this still needed or redundant with the checks in type deduction? if func_name in fbuiltins.TYPE_BUILTIN_NAMES: self._verify_builtin_type_constructor(node) diff --git a/src/gt4py/next/ffront/lowering_utils.py b/src/gt4py/next/ffront/lowering_utils.py index c4e35c9e18..f7a5d26991 100644 --- a/src/gt4py/next/ffront/lowering_utils.py +++ b/src/gt4py/next/ffront/lowering_utils.py @@ -20,6 +20,7 @@ def process_elements( objs: itir.Expr | Iterable[itir.Expr], current_el_type: ts.TypeSpec, arg_types: Optional[Iterable[ts.TypeSpec]] = None, + with_path_arg: bool = False, ) -> itir.FunCall: """ Recursively applies a processing function to all primitive constituents of a tuple or @@ -34,6 +35,7 @@ def process_elements( arg_types: If provided, a tuple of the type of each argument is passed to `process_func` as last argument. Note, that `arg_types` might coincide with `(current_el_type,)*len(objs)`, but not necessarily, in case of implicit broadcasts. + with_path_arg: If true, the index path to the current leaf is passed as last argument. """ if isinstance(objs, itir.Expr): objs = (objs,) @@ -47,6 +49,8 @@ def process_elements( tuple(im.ref(var_name) for var_name in var_names), current_el_type, arg_types=arg_types, + path=(), + with_path_arg=with_path_arg, ) return im.let(*bound_vars.items())(body) @@ -60,6 +64,8 @@ def _process_elements_impl( _current_el_exprs: Iterable[T], current_el_type: ts.TypeSpec, arg_types: Optional[Iterable[ts.TypeSpec]], + path: tuple[int, ...], + with_path_arg: bool, ) -> itir.Expr: if isinstance(current_el_type, (ts.TupleType, ts.NamedCollectionType)): result = im.make_tuple( @@ -70,17 +76,23 @@ def _process_elements_impl( im.tuple_get(i, current_el_expr) for current_el_expr in _current_el_exprs ), current_el_type.types[i], - arg_types=tuple(arg_t.types[i] for arg_t in arg_types) # type: ignore[attr-defined] # guaranteed by the requirement that `current_el_type` and each element of `arg_types` have the same tuple structure - if arg_types is not None - else None, + arg_types=( + tuple(arg_t.types[i] for arg_t in arg_types) # type: ignore[attr-defined] # guaranteed by the requirement that `current_el_type` and each element of `arg_types` have the same tuple structure + if arg_types is not None + else None + ), + path=(*path, i), + with_path_arg=with_path_arg, ) for i in range(len(current_el_type.types)) ) ) else: if arg_types is not None: - result = process_func(*_current_el_exprs, arg_types) + result = process_func( + *_current_el_exprs, arg_types, *((path,) if with_path_arg else ()) + ) else: - result = process_func(*_current_el_exprs) + result = process_func(*_current_el_exprs, *((path,) if with_path_arg else ())) return result diff --git a/src/gt4py/next/ffront/past_passes/type_deduction.py b/src/gt4py/next/ffront/past_passes/type_deduction.py index 9d021ceb51..a735833084 100644 --- a/src/gt4py/next/ffront/past_passes/type_deduction.py +++ b/src/gt4py/next/ffront/past_passes/type_deduction.py @@ -57,6 +57,21 @@ def _validate_domain_out( f"Only 'Dimension' allowed in domain dictionary keys, got '{dim}' which is of type '{dim.type}'." ) + dict_dims = [cast(ts.DimensionType, dim.type).dim for dim in dom.keys_] + + # the result is irrelevant, the mapped function only validates + @type_info.tree_map_type(result_collection_constructor=lambda _, elts: None) + def _validate_field_dims(type_: ts.TypeSpec) -> None: + # non-field leaves are invalid `out` types, but only rejected later, by the + # check of the `out` type against the operator return type + if isinstance(type_, ts.FieldType) and list(type_.dims) != dict_dims: + raise ValueError( + f"Dimensions in out field and field domain are not equivalent: " + f"expected '{type_.dims}', got '{dict_dims}'." + ) + + _validate_field_dims(out) + for domain_values in dom.values_: if len(domain_values.elts) != 2: raise ValueError( @@ -68,11 +83,18 @@ def _validate_domain_out( ) elif isinstance(dom, past.TupleExpr): + if isinstance(out, ts.VarLenTupleType): + # The fixed-length domain tuple would pin the length of the variable-length + # output, so any call with a different length would break the program. + raise ValueError( + f"Per-element domain tuples are not supported for the variable-length tuple " + f"output '{out}'; use a single domain that applies to every element." + ) if isinstance(out, ts.CollectionTypeSpec): # TODO(havogt): for NamedCollectionType, we assume the domain tuple is in element order out_elts = out.types else: - raise ValueError(f"Tuple domain requires tuple output, got {type(out)}.") + raise ValueError(f"Tuple domain requires tuple output, got '{out}'.") if len(dom.elts) != len(out_elts): raise ValueError("Mismatched tuple lengths between domain and output.") @@ -142,6 +164,14 @@ def visit_Slice(self, node: past.Slice, **kwargs: Any) -> past.Slice: def visit_Subscript(self, node: past.Subscript, **kwargs: Any) -> past.Subscript: value = self.visit(node.value, **kwargs) + if isinstance(value.type, (*ts.COLLECTION_TYPE_SPECS, ts.VarLenTupleType)): + # A subscript on a program argument is a domain slice; applying one to a + # (possibly variable-length) tuple is not supported, and element selection is + # not supported for any tuple argument. + raise errors.DSLError( + node.location, + f"Subscripting a tuple is not supported, got '{value.type}'.", + ) return past.Subscript( value=value, slice_=self.visit(node.slice_, **kwargs), @@ -248,7 +278,9 @@ def visit_Call(self, node: past.Call, **kwargs: Any) -> past.Call: operator_return_type = type_info.return_type( new_func.type, with_args=arg_types, with_kwargs=kwarg_types ) - if operator_return_type != new_kwargs["out"].type: + # Not equality: variadic comprehensions give a variable-length tuple return type, + # while 'out' is a concrete tuple. + if not type_info.is_compatible_type(operator_return_type, new_kwargs["out"].type): raise ValueError( "Expected keyword argument 'out' to be of " f"type '{operator_return_type}', got " diff --git a/src/gt4py/next/ffront/past_to_itir.py b/src/gt4py/next/ffront/past_to_itir.py index 05af5fa937..911af82a43 100644 --- a/src/gt4py/next/ffront/past_to_itir.py +++ b/src/gt4py/next/ffront/past_to_itir.py @@ -9,7 +9,6 @@ from __future__ import annotations import dataclasses -import functools from typing import Any, Optional, Sequence, cast import devtools @@ -27,7 +26,7 @@ ) from gt4py.next.ffront.stages import ConcretePASTProgramDef from gt4py.next.iterator import ir as itir -from gt4py.next.iterator.ir_utils import ir_makers as im +from gt4py.next.iterator.ir_utils import ir_makers as im, misc as ir_misc from gt4py.next.iterator.transforms import remap_symbols, replace_get_domain_range_with_constants from gt4py.next.otf import arguments, stages, workflow from gt4py.next.type_system import type_info, type_specifications as ts @@ -93,8 +92,10 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> stages.CompilableProgramDef: for gt_callable in gt_callables: lowered_funcs.append(gt_callable.__gt_gtir__()) + past_node = _specialize_program_params(inp.data.past_node, inp.args.args, inp.args.kwargs) + itir_program = ProgramLowering.apply( - inp.data.past_node, function_definitions=lowered_funcs, grid_type=grid_type + past_node, function_definitions=lowered_funcs, grid_type=grid_type ) # TODO(tehrengruber): Put this in a dedicated transformation step. @@ -209,29 +210,55 @@ def _compute_field_slice(node: past.Subscript) -> list[past.Slice]: return out_field_slice_ -def _get_element_from_tuple_expr(node: past.Expr, path: tuple[int, ...]) -> past.Expr: - """Get element from a (nested) TupleExpr by following the given path. - - Pre-condition: `node` is a `past.TupleExpr` (if `path ! = ()`) - and `path` is a valid path through the nested tuple structure. - """ - return functools.reduce(lambda e, i: e.elts[i], path, node) # type: ignore[attr-defined] # see pre-condition - - -def _unwrap_tuple_expr(expr: past.Expr, path: tuple[int, ...]) -> tuple[past.Expr, Sequence[int]]: +def _unwrap_tuple_expr( + expr: past.Expr, path: tuple[int | type_info.VarLenTupleElementKey, ...] +) -> tuple[past.Expr, Sequence[int | type_info.VarLenTupleElementKey]]: """Unwrap (nested) TupleExpr by following the given path as long as possible. If a non-tuple expression is encountered, the current expression and the remaining path are returned. """ - path_remainder: Sequence[int] = path + path_remainder: Sequence[int | type_info.VarLenTupleElementKey] = path while isinstance(expr, past.TupleExpr): idx, *path_remainder = path_remainder + # a `TupleExpr` always has a fixed-length tuple type, so the path element is an index + assert isinstance(idx, int) expr = expr.elts[idx] return expr, path_remainder +def _specialize_program_params( + past_node: past.Program, + arg_types: Sequence[ts.TypeSpec], + kwarg_types: dict[str, ts.TypeSpec], + check_compatibility: bool = True, +) -> past.Program: + """ + Replace the program parameter types by the concrete argument types of the call. + + In particular, parameters declared with a generic type, e.g. a variable-length tuple, + only have a concrete type at the call. The arguments are expected to be canonicalized + against the signature already (see `transform_program_args`). + """ + param_types: dict[str, ts.TypeSpec] = { + **{str(param.id): arg_type for param, arg_type in zip(past_node.params, arg_types)}, + **kwarg_types, + } + new_params = [] + for param in past_node.params: + param_type = param_types[str(param.id)] + # incompatible argument types have already been rejected in `transform_program_args` + assert not check_compatibility or type_info.is_concretizable( + param.type, to_type=param_type + ), ( + f"Argument of type '{param_type}' is not compatible with parameter " + f"'{param.id}' of type '{param.type}'." + ) + new_params.append(param.copy(update={"type": param_type})) + return past_node.copy(update={"params": new_params}) + + @dataclasses.dataclass class ProgramLowering( traits.PreserveLocationVisitor, traits.VisitorWithSymbolTableTrait, NodeTranslator @@ -352,17 +379,18 @@ def _visit_slice_bound( lowered_bound.location = slice_bound.location return lowered_bound - def _construct_itir_out_arg(self, node: past.Expr) -> itir.Expr: + def _construct_itir_out_arg(self, node: past.Expr, allow_slicing: bool) -> itir.Expr: if isinstance(node, past.Name): return itir.SymRef(id=node.id, location=node.location) elif isinstance(node, past.Subscript): - itir_node = self._construct_itir_out_arg(node.value) + assert allow_slicing + itir_node = self._construct_itir_out_arg(node.value, allow_slicing) itir_node.location = node.location return itir_node elif isinstance(node, past.TupleExpr): return itir.FunCall( fun=itir.SymRef(id="make_tuple"), - args=[self._construct_itir_out_arg(el) for el in node.elts], + args=[self._construct_itir_out_arg(el, allow_slicing) for el in node.elts], location=node.location, ) else: @@ -375,71 +403,48 @@ def _construct_itir_domain_arg( self, out_expr: itir.Expr, out_type: ts.FieldType, - node_domain: Optional[past.Expr], slices: Optional[list[past.Slice]] = None, ) -> itir.FunCall: - domain_args = [] + ranges: dict[common.Dimension, tuple[itir.Expr, itir.Expr]] = {} for dim_i, dim in enumerate(out_type.dims): + if dim.kind == common.DimensionKind.LOCAL: + raise ValueError(f"common.Dimension '{dim.value}' must not be local.") # an expression for the range of a dimension - dim_range = im.call("get_domain_range")( - out_expr, itir.AxisLiteral(value=dim.value, kind=dim.kind) - ) + dim_range = im.call("get_domain_range")(out_expr, im.axis_literal(dim)) dim_start, dim_stop = im.tuple_get(0, dim_range), im.tuple_get(1, dim_range) - # bounds - lower: itir.Expr - upper: itir.Expr - if node_domain is not None: - assert isinstance(node_domain, past.Dict) - lower, upper = self._construct_itir_initialized_domain_arg(dim_i, dim, node_domain) - else: - lower = self._visit_slice_bound( - slices[dim_i].lower if slices else None, - dim_start, - dim_start, - dim_stop, - ) - upper = self._visit_slice_bound( - slices[dim_i].upper if slices else None, - dim_stop, - dim_start, - dim_stop, - ) - - if dim.kind == common.DimensionKind.LOCAL: - raise ValueError(f"common.Dimension '{dim.value}' must not be local.") - domain_args.append( - itir.FunCall( - fun=itir.SymRef(id="named_range"), - args=[itir.AxisLiteral(value=dim.value, kind=dim.kind), lower, upper], - ) + lower = self._visit_slice_bound( + slices[dim_i].lower if slices else None, + dim_start, + dim_start, + dim_stop, ) - - if self.grid_type == common.GridType.CARTESIAN: - domain_builtin = "cartesian_domain" - elif self.grid_type == common.GridType.UNSTRUCTURED: - domain_builtin = "unstructured_domain" - else: - raise AssertionError() - - return itir.FunCall( - fun=itir.SymRef(id=domain_builtin), - args=domain_args, - location=(node_domain or out_expr).location, - ) - - def _construct_itir_initialized_domain_arg( - self, dim_i: int, dim: common.Dimension, node_domain: past.Dict - ) -> list[itir.FunCall]: - assert len(node_domain.values_[dim_i].elts) == 2 - keys_dims_types = cast(ts.DimensionType, node_domain.keys_[dim_i].type).dim - if keys_dims_types != dim: - raise ValueError( - "common.Dimensions in out field and field domain are not equivalent:" - f"expected '{dim}', got '{keys_dims_types}'." + upper = self._visit_slice_bound( + slices[dim_i].upper if slices else None, + dim_stop, + dim_start, + dim_stop, ) - - return [self.visit(bound) for bound in node_domain.values_[dim_i].elts] + ranges[dim] = (lower, upper) + + domain_expr = im.domain(self.grid_type, ranges) + domain_expr.location = out_expr.location + return domain_expr + + def _lower_domain_arg(self, domain_arg: past.Expr) -> itir.Expr: + if isinstance(domain_arg, past.TupleExpr): + return im.make_tuple(*(self._lower_domain_arg(el) for el in domain_arg.elts)) + assert isinstance(domain_arg, past.Dict) + ranges = { + cast(ts.DimensionType, dim_name.type).dim: ( + self.visit(bounds.elts[0]), + self.visit(bounds.elts[1]), + ) + for dim_name, bounds in zip(domain_arg.keys_, domain_arg.values_, strict=True) + } + domain_expr = im.domain(self.grid_type, ranges) + domain_expr.location = domain_arg.location + return domain_expr def _split_field_and_slice( self, field: past.Name | past.Subscript @@ -452,40 +457,43 @@ def _split_field_and_slice( def _visit_stencil_call_out_arg( self, out_arg: past.Expr, domain_arg: Optional[past.Expr], **kwargs: Any - ) -> tuple[itir.Expr, itir.FunCall]: + ) -> tuple[itir.Expr, itir.Expr]: assert isinstance(out_arg, (past.Subscript, past.Name, past.TupleExpr)), ( "Unexpected 'out' argument. Must be a 'past.Subscript', 'past.Name' or 'past.TupleExpr' node." ) + assert out_arg.type is not None + + if domain_arg is not None: + # Mutual exclusivity of a `domain` argument and slices is ensured in the type + # deduction, so the target is just the plain lowered `out` argument. Note that + # a single domain may apply to an entire tuple subtree of the target (see + # `itir.SetAt`). + return self._construct_itir_out_arg( + out_arg, allow_slicing=False + ), self._lower_domain_arg(domain_arg) + else: - @type_info.tree_map_type( - with_path_arg=True, - unpack=True, - result_collection_constructor=lambda _, elts: im.make_tuple(*elts), - ) - def impl(out_type: ts.FieldType, path: tuple[int, ...]) -> tuple[itir.Expr, itir.Expr]: - out_field, path_remainder = _unwrap_tuple_expr(out_arg, path) - - assert isinstance(out_field, (past.Name, past.Subscript)) - out_field, slice_info = self._split_field_and_slice(out_field) + @ir_misc.tree_map_type_into_expr(with_path_arg=True) + def impl( + out_type: ts.FieldType, + current: itir.Expr, + path: tuple[int | type_info.VarLenTupleElementKey, ...], + ) -> itir.Expr: + out_field, path_remainder = _unwrap_tuple_expr(out_arg, path) + + # A subscript (domain slice) can only appear on a single field, i.e. on an + # expression that does not decompose into further tuple elements. + assert isinstance(out_field, past.Name) or ( + isinstance(out_field, past.Subscript) and len(path_remainder) == 0 + ) + _, slice_info = self._split_field_and_slice(out_field) - domain_element = ( - _get_element_from_tuple_expr(domain_arg, path) - if isinstance(domain_arg, past.TupleExpr) - else domain_arg - ) + return self._construct_itir_domain_arg(current, out_type, slice_info) - lowered_out_field = functools.reduce( - lambda expr, i: im.tuple_get(i, expr), path_remainder, self.visit(out_field) - ) - lowered_domain = self._construct_itir_domain_arg( - lowered_out_field, - out_type, - domain_element, - slice_info, - ) - return lowered_out_field, lowered_domain + lowered_out = self._construct_itir_out_arg(out_arg, allow_slicing=True) + lowered_domain = impl(out_arg.type, lowered_out) - return impl(out_arg.type) + return lowered_out, lowered_domain def visit_Constant(self, node: past.Constant, **kwargs: Any) -> itir.Literal: if isinstance(node.type, ts.ScalarType) and node.type.shape is None: diff --git a/src/gt4py/next/ffront/type_info.py b/src/gt4py/next/ffront/type_info.py index e7280d51b2..b26cea2cc7 100644 --- a/src/gt4py/next/ffront/type_info.py +++ b/src/gt4py/next/ffront/type_info.py @@ -153,9 +153,9 @@ def function_signature_incompatibilities_fieldop( def _tree_map_type_constructor_drop_python_type( - value: ts.CollectionTypeSpecT, + value: ts.CollectionTypeSpecT | ts.VarLenTupleType, elems: Iterable[ts.DataType | ts.DimensionType | ts.DeferredType], -) -> ts.CollectionTypeSpecT: +) -> ts.CollectionTypeSpecT | ts.VarLenTupleType: result = type_info.tree_map_type_constructor(value, elems) if isinstance(value, ts.NamedCollectionType): result = datamodels.evolve(result, original_python_type=ts.ANY_PYTHON_TYPE_NAME) @@ -164,7 +164,7 @@ def _tree_map_type_constructor_drop_python_type( def _scan_param_promotion( param: ts.TypeSpec, arg: ts.TypeSpec -) -> ts.FieldType | ts.TupleType | ts.NamedCollectionType: +) -> ts.FieldType | ts.TupleType | ts.NamedCollectionType | ts.VarLenTupleType: """ Promote parameter of a scan pass to match dimensions of respective scan operator argument. @@ -183,14 +183,24 @@ def _scan_param_promotion( FieldType(dims=[Dimension(value='I', kind=)], dtype=ScalarType(kind=, shape=None)) """ - def _as_field(dtype: ts.TypeSpec, path: tuple[int, ...]) -> ts.FieldType: + def _index_arg(type_: ts.TypeSpec, idx: int | type_info.VarLenTupleElementKey) -> ts.TypeSpec: + if isinstance(idx, type_info.VarLenTupleElementKey): + if isinstance(type_, ts.TupleType): + # case where `param` is variable length, but `arg` is actual tuple + if not all(el == type_.types[0] for el in type_.types): + raise IndexError() + return type_.types[0] + elif isinstance(type_, ts.VarLenTupleType): + return type_.element_type + raise AttributeError + return type_.types[idx] # type: ignore[attr-defined] + + def _as_field( + dtype: ts.TypeSpec, path: tuple[int | type_info.VarLenTupleElementKey, ...] + ) -> ts.FieldType: assert isinstance(dtype, ts.ScalarType) try: - el_type = functools.reduce( - lambda type_, idx: type_.types[idx], # type: ignore[attr-defined] - path, - arg, - ) + el_type = functools.reduce(_index_arg, path, arg) return ts.FieldType(dims=type_info.extract_dims(el_type), dtype=dtype) except (IndexError, AttributeError): # The structure of the scan passes argument and the requested @@ -208,7 +218,7 @@ def _as_field(dtype: ts.TypeSpec, path: tuple[int, ...]) -> ts.FieldType: result_collection_constructor=_tree_map_type_constructor_drop_python_type, with_path_arg=True, )(param) - assert isinstance(res, (ts.FieldType, ts.TupleType, ts.NamedCollectionType)) + assert isinstance(res, (ts.FieldType, ts.TupleType, ts.NamedCollectionType, ts.VarLenTupleType)) return res diff --git a/src/gt4py/next/iterator/ir.py b/src/gt4py/next/iterator/ir.py index f024ef6168..6a237ffd39 100644 --- a/src/gt4py/next/iterator/ir.py +++ b/src/gt4py/next/iterator/ir.py @@ -125,6 +125,16 @@ class Stmt(Node): ... class SetAt(Stmt): # from JAX array.at[...].set() + """ + Write `expr`, computed on `domain`, into `target`. + + For a tuple-valued `target` the `domain` is a matching tuple of domains (built with + `make_tuple`), except that a single domain may also apply to an entire tuple subtree + of the target, i.e. to each of its elements. This shallower structure is needed for + variable-length tuple outputs, whose concrete length is unknown when the domain + expression is constructed. + """ + expr: Expr # only `as_fieldop(stencil)(inp0, ...)` in first refactoring domain: Expr target: Expr # `make_tuple` or SymRef @@ -153,6 +163,13 @@ class Program(Node, ValidatedSymbolTableTrait): Sym(id=name) for name in sorted(BUILTINS) ] # sorted for serialization stability + def __hash__(self) -> int: + # Only here do types contribute to the hash as they are the starting point for + # type inference, whereas in all other places the types can be recovered. Note + # that equality still ignores types, so equal programs may have different hashes; + # the users only consume the hash value itself, never using programs as dict keys. + return hash((Node.__hash__(self), tuple(str(param.type) for param in self.params))) + # TODO(fthaler): just use hashable types in nodes (tuples instead of lists) Sym.__hash__ = Node.__hash__ # type: ignore[method-assign] @@ -165,7 +182,6 @@ class Program(Node, ValidatedSymbolTableTrait): Lambda.__hash__ = Node.__hash__ # type: ignore[method-assign] FunCall.__hash__ = Node.__hash__ # type: ignore[method-assign] FunctionDefinition.__hash__ = Node.__hash__ # type: ignore[method-assign] -Program.__hash__ = Node.__hash__ # type: ignore[method-assign] SetAt.__hash__ = Node.__hash__ # type: ignore[method-assign] IfStmt.__hash__ = Node.__hash__ # type: ignore[method-assign] InfinityLiteral.__hash__ = Node.__hash__ # type: ignore[method-assign] diff --git a/src/gt4py/next/iterator/ir_utils/misc.py b/src/gt4py/next/iterator/ir_utils/misc.py index d9a6e85f57..7b8a9d75fd 100644 --- a/src/gt4py/next/iterator/ir_utils/misc.py +++ b/src/gt4py/next/iterator/ir_utils/misc.py @@ -7,8 +7,9 @@ # SPDX-License-Identifier: BSD-3-Clause import dataclasses +import functools from collections import ChainMap -from typing import Callable, Iterable, TypeVar, cast +from typing import Any, Callable, Iterable, Optional, TypeVar, cast from gt4py import eve from gt4py._core import definitions as core_defs @@ -17,7 +18,7 @@ from gt4py.next.iterator import embedded, ir as itir from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, ir_makers as im from gt4py.next.iterator.transforms import inline_lambdas -from gt4py.next.type_system import type_specifications as ts +from gt4py.next.type_system import type_info, type_specifications as ts @dataclasses.dataclass(frozen=True) @@ -279,3 +280,86 @@ def value_from_literal(literal: itir.Literal) -> core_defs.Scalar: assert literal.value in ["True", "False"] return literal.value == "True" return getattr(embedded, str(literal.type))(literal.value) + + +#: Name bound to the value of the current element in `tree_map_type_into_expr`; only the +#: (shadowed) binders introduced by `map_tuple` lambdas remain visible in the result. +_CURRENT_ELEMENT = "__tmtie_el" + + +def tree_map_type_into_expr( + fun: Optional[Callable[..., itir.Expr]] = None, + *, + with_path_arg: bool = False, + unpack: bool = False, +) -> Callable[..., Any]: + """ + Map `fun` over the leaves of a type tree, building an ITIR expression over a + corresponding value expression. + + Like `type_info.tree_map_type`, but `fun` additionally receives a `current` argument — + a reference to the value of the current element — and the mapped callable is called + with the value expression alongside the type. Fixed-length collection nodes recombine + the element results with `make_tuple`, rebinding `current` per element; `VarLenTupleType` + nodes map their elements with the `map_tuple` builtin, so the result is also valid for + types whose concrete length is only known later. The bindings introduced along the way + are inlined unconditionally before returning — the value expression should hence be + cheap (e.g. a symbol reference or tuple accesses thereon), as it may be duplicated. + With `with_path_arg`, paths contain `type_info.VAR_LENGTH_TUPLE_ELEMENT` below variable-length + tuples. + + Can be used directly or as a (possibly parametrized) decorator. + + >>> field = ts.FieldType(dims=[], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)) + >>> type_ = ts.TupleType(types=[field, ts.VarLenTupleType(element_type=field)]) + >>> expr = tree_map_type_into_expr(lambda t, current: im.call("get_dom")(current))( + ... type_, im.ref("out") + ... ) + >>> str(expr) + '{get_dom(out[0]), map_tuple(λ(__tmtie_el_) → get_dom(__tmtie_el_))(out[1])}' + """ + if fun is None: + return functools.partial( + tree_map_type_into_expr, with_path_arg=with_path_arg, unpack=unpack + ) + + @functools.wraps(fun) + def impl(type_: ts.TypeSpec, expr: itir.Expr) -> itir.Expr | tuple[itir.Expr, ...]: + current = im.ref(_CURRENT_ELEMENT) + + def collection_constructor(value: ts.TypeSpec, elts: Iterable[itir.Expr]) -> itir.Expr: + if isinstance(value, ts.VarLenTupleType): + (element_expr,) = elts + return im.call(im.call("map_tuple")(im.lambda_(_CURRENT_ELEMENT)(element_expr)))( + current + ) + return im.make_tuple( + *( + im.let(_CURRENT_ELEMENT, im.tuple_get(i, current))(el) + for i, el in enumerate(elts) + ) + ) + + leaf_fun = ( + (lambda type_, path: fun(type_, current, path)) + if with_path_arg + else (lambda type_: fun(type_, current)) + ) + mapped = type_info.tree_map_type( + leaf_fun, + with_path_arg=with_path_arg, + unpack=unpack, + result_collection_constructor=collection_constructor, + ) + + def bind(tree: itir.Expr) -> itir.Expr: + # inline unconditionally so consumers see plain tuple accesses instead of + # `let`-bound references (see docstring) + return inline_lambdas.InlineLambdas.apply(im.let(_CURRENT_ELEMENT, expr)(tree)) + + result = mapped(type_) + if unpack: + return tuple(bind(tree) for tree in cast(tuple[itir.Expr, ...], result)) + return bind(result) + + return impl diff --git a/src/gt4py/next/named_collections.py b/src/gt4py/next/named_collections.py index e2c7c09a52..31dbab9b83 100644 --- a/src/gt4py/next/named_collections.py +++ b/src/gt4py/next/named_collections.py @@ -165,6 +165,16 @@ def make_extractor_expr(named_collection_type_hint: xtyping.TypeAnnotation, valu elif issubclass(actual_type, tuple) and ( tuple_arg_hints := typing.get_args(named_collection_type_hint) ): + if len(tuple_arg_hints) == 2 and tuple_arg_hints[1] is Ellipsis: + # Variable-length tuple: the length is only known at run time, extract + # with a comprehension, or pass the value through if the elements do not + # need extraction. Note that reusing the variable name in nested + # comprehensions is fine: a comprehension's iterable is evaluated in the + # enclosing scope and the element expression only uses its own variable. + element_expr = make_extractor_expr(tuple_arg_hints[0], "__el") + if element_expr == "__el": + return value_expr + return f"tuple({element_expr} for __el in {value_expr})" # This is a `tuple` with type arguments expr_parts = [ make_extractor_expr(value, f"{value_expr}[{i}]") @@ -202,6 +212,11 @@ def make_extractor_expr_from_type_spec(type_: ts.TypeSpec, value_expr: str) -> s for i, t in enumerate(type_.types) ] return f"({', '.join(args)})" + case ts.VarLenTupleType() if type_info.needs_value_extraction(type_): + # the length is only known at run time, extract with a comprehension (see + # the note on nested comprehensions in `make_extractor_expr`) + element_expr = make_extractor_expr_from_type_spec(type_.element_type, "__el") + return f"tuple({element_expr} for __el in {value_expr})" return value_expr diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index b71fdd9df9..c9700ce8b1 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -507,13 +507,22 @@ def _is_generic(self) -> bool: Is the operator or program generic in the sense that it can be called for different argument types. - Right now this is only the case for scan operators. + Right now this is the case for scan operators and for parameters containing a + variable-length tuple, whose concrete length is only known at the call. """ + # TODO(tehrengruber): This concept does not exist elsewhere and is not properly reflected # in the type system. For now we just use `DeferredType` to communicate between # here and `type_info.type_in_program_context`. + def is_generic_type(t: ts.TypeSpec) -> bool: + if isinstance(t, (ts.DeferredType, ts.VarLenTupleType)): + return True + if isinstance(t, ts.TupleType): + return any(is_generic_type(el) for el in t.types) + return False + return any( - isinstance(t, ts.DeferredType) + is_generic_type(t) for t in itertools.chain( self.program_type.definition.pos_only_args, self.program_type.definition.pos_or_kw_args.values(), @@ -724,6 +733,12 @@ def compile( pool.compile(static_arg0=[0], static_arg1=[2]).compile(static_arg=[1], static_arg1=[3]) will compile for (0,2), (1,3) """ + if self._is_generic: + raise ValueError( + f"Cannot compile '{self.definition_stage.definition.__name__}' ahead of time: " + "generic programs (e.g. with variable-length tuple parameters or direct scan " + "operator calls) require the concrete argument types of a call." + ) for offset_provider in offset_providers: # not included in product for better type checking for static_values in itertools.product(*static_args.values()): # The argument descriptors built here must be kept in sync with the diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py index e645068f64..2442939bbf 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py +++ b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py @@ -1151,6 +1151,10 @@ def _visit_target( # Visit the field operator expression. source_tree = self._visit_expression(stmt.expr, ctx) + + # A single domain may apply to an entire tuple subtree of the target (see + # `gtir.SetAt`); promote it so all trees have the same structure. + _, domain = gtx_utils.equalize_tuple_structure(source_tree, domain, bidirectional=False) ctx_input_data = ctx.get_input_data() # In order to support inout argument, write the result in separate next state diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index 414d665012..2d616f6c5c 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import functools +import itertools import types from collections.abc import Callable, Collection, Iterable, Iterator from typing import Any, Final, Literal, Sequence, Type, TypeGuard, TypeVar, overload @@ -83,6 +84,19 @@ def type_class(symbol_type: ts.TypeSpec) -> Type[ts.TypeSpec]: ) +class VarLenTupleElementKey: + """Path key of the representative element of a `VarLenTupleType` in type traversals.""" + + def __repr__(self) -> str: + return "VAR_LENGTH_TUPLE_ELEMENT" + + +#: Type traversals (`tree_map_type`, `primitive_constituents`) visit a `VarLenTupleType` as +#: a single representative element (its `element_type`), addressed by this sentinel +#: instead of a position. +VAR_LENGTH_TUPLE_ELEMENT: Final = VarLenTupleElementKey() + + @overload def primitive_constituents( symbol_type: ts.TypeSpec, with_path_arg: Literal[False] = False @@ -92,15 +106,21 @@ def primitive_constituents( @overload def primitive_constituents( symbol_type: ts.TypeSpec, with_path_arg: Literal[True] -) -> utils.XIterable[tuple[ts.TypeSpec, tuple[int, ...]]]: ... +) -> utils.XIterable[tuple[ts.TypeSpec, tuple[int | VarLenTupleElementKey, ...]]]: ... def primitive_constituents( symbol_type: ts.TypeSpec, with_path_arg: bool = False -) -> utils.XIterable[ts.TypeSpec] | utils.XIterable[tuple[ts.TypeSpec, tuple[int, ...]]]: +) -> ( + utils.XIterable[ts.TypeSpec] + | utils.XIterable[tuple[ts.TypeSpec, tuple[int | VarLenTupleElementKey, ...]]] +): """ Return the primitive types contained in a composite type. + A `VarLenTupleType` is traversed as its single representative element, addressed by + `VAR_LENGTH_TUPLE_ELEMENT` in paths. + >>> from gt4py.next import common >>> I = common.Dimension(value="I") >>> int_type = ts.ScalarType(kind=ts.ScalarKind.INT64) @@ -116,9 +136,16 @@ def primitive_constituents( """ def constituents_yielder( - symbol_type: ts.TypeSpec, path: tuple[int, ...] - ) -> Iterator[ts.TypeSpec] | Iterator[tuple[ts.TypeSpec, tuple[int, ...]]]: - if isinstance(symbol_type, ts.COLLECTION_TYPE_SPECS): + symbol_type: ts.TypeSpec, path: tuple[int | VarLenTupleElementKey, ...] + ) -> ( + Iterator[ts.TypeSpec] + | Iterator[tuple[ts.TypeSpec, tuple[int | VarLenTupleElementKey, ...]]] + ): + if isinstance(symbol_type, ts.VarLenTupleType): + yield from constituents_yielder( + symbol_type.element_type, (*path, VAR_LENGTH_TUPLE_ELEMENT) + ) + elif isinstance(symbol_type, ts.COLLECTION_TYPE_SPECS): for i, el_type in enumerate(symbol_type.types): yield from constituents_yielder(el_type, (*path, i)) else: @@ -134,10 +161,27 @@ def constituents_yielder( _C = TypeVar("_C") +def _tree_map_type_collection_elements(value: ts.TypeSpec) -> Iterable[ts.TypeSpec]: + if isinstance(value, ts.VarLenTupleType): + return (value.element_type,) + assert isinstance(value, ts.COLLECTION_TYPE_SPECS) + return value.types + + +def _tree_map_type_collection_keys(value: ts.TypeSpec) -> Iterable[Any]: + if isinstance(value, ts.VarLenTupleType): + return (VAR_LENGTH_TUPLE_ELEMENT,) + return itertools.count() + + def tree_map_type_constructor( - value: ts.CollectionTypeSpecT, + value: ts.CollectionTypeSpecT | ts.VarLenTupleType, elems: Iterable[ts.DataType | ts.DimensionType | ts.DeferredType], -) -> ts.CollectionTypeSpecT: +) -> ts.CollectionTypeSpecT | ts.VarLenTupleType: + if isinstance(value, ts.VarLenTupleType): + (element_type,) = elems + assert isinstance(element_type, (ts.DataType, ts.DeferredType)) + return ts.VarLenTupleType(element_type=element_type) return ( ts.NamedCollectionType( keys=value.keys, original_python_type=value.original_python_type, types=list(elems) @@ -194,8 +238,9 @@ def tree_map_type( ) return next_utils.tree_map( fun, - collection_type=ts.COLLECTION_TYPE_SPECS, - collection_elements=lambda t: t.types, + collection_type=(*ts.COLLECTION_TYPE_SPECS, ts.VarLenTupleType), + collection_elements=_tree_map_type_collection_elements, + collection_keys=_tree_map_type_collection_keys, result_collection_constructor=result_collection_constructor, with_path_arg=with_path_arg, unpack=unpack, @@ -574,7 +619,22 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool: or issubclass(type_class(to_type), symbol_type.constraint) ): return True - elif is_concrete(symbol_type): + if isinstance(symbol_type, ts.VarLenTupleType) and isinstance(to_type, ts.VarLenTupleType): + return is_concretizable(symbol_type.element_type, to_type.element_type) + if isinstance(symbol_type, ts.VarLenTupleType) and isinstance(to_type, ts.TupleType): + # The empty tuple is a valid concretization of any variable-length tuple: there + # are no elements whose type could conflict with the element type. + return len(to_type.types) == 0 or ( + all(type_ == to_type.types[0] for type_ in to_type.types) + and is_concretizable(symbol_type.element_type, to_type.types[0]) + ) + if isinstance(symbol_type, ts.TupleType) and isinstance(to_type, ts.TupleType): + # Recurse elementwise so that e.g. a variable-length tuple nested inside a + # fixed-length tuple can be concretized. + return len(symbol_type.types) == len(to_type.types) and all( + is_concretizable(el, to_el) for el, to_el in zip(symbol_type.types, to_type.types) + ) + if is_concrete(symbol_type): return symbol_type == to_type return False @@ -990,8 +1050,15 @@ def accepts_args( def needs_value_extraction( type_spec: ts.TypeSpec, -) -> typing_extensions.TypeIs[ts.NamedCollectionType | ts.TupleType]: - return isinstance(type_spec, ts.NamedCollectionType) or ( - isinstance(type_spec, ts.TupleType) - and any(needs_value_extraction(t) for t in type_spec.types) +) -> typing_extensions.TypeIs[ts.NamedCollectionType | ts.TupleType | ts.VarLenTupleType]: + return ( + isinstance(type_spec, ts.NamedCollectionType) + or ( + isinstance(type_spec, ts.TupleType) + and any(needs_value_extraction(t) for t in type_spec.types) + ) + or ( + isinstance(type_spec, ts.VarLenTupleType) + and needs_value_extraction(type_spec.element_type) + ) ) diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 077a008045..3261ab0e36 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -142,6 +142,21 @@ def __str__(self) -> str: return f"tuple[{', '.join(map(str, self.types))}]" +class VarLenTupleType(DataType): + """ + A homogeneous tuple whose length is not part of the type. + + The length is only known at the call, where the type is specialized to a concrete + `TupleType` (or stays variable-length in length-agnostic code paths). + """ + + # TODO(tehrengruber): `DeferredType` should be removed again, see `TupleType.types`. + element_type: DataType | DeferredType + + def __str__(self) -> str: + return f"tuple[{self.element_type}, ...]" + + class AnyPythonType: """Marker type representing any Python type which cannot be used for instantiation. @@ -185,6 +200,8 @@ def __str__(self) -> str: return f"NamedTuple{{{', '.join(f'{k}: {v}' for k, v in zip(self.keys, self.types))}}}" +# TODO(tehrengruber): Add VarLenTupleType to this list, which requires checking all users of it +# that potentially depend on the length being known. CollectionTypeSpecT = TypeVar("CollectionTypeSpecT", TupleType, NamedCollectionType) CollectionTypeSpec = TupleType | NamedCollectionType COLLECTION_TYPE_SPECS: Final[tuple[type[CollectionTypeSpec], ...]] = typing.get_args( diff --git a/src/gt4py/next/type_system/type_translation.py b/src/gt4py/next/type_system/type_translation.py index df721da0e7..afeaf254e2 100644 --- a/src/gt4py/next/type_system/type_translation.py +++ b/src/gt4py/next/type_system/type_translation.py @@ -19,7 +19,7 @@ import sys import types import typing -from typing import Any, ForwardRef, Optional, TypeAlias +from typing import Any, ForwardRef, Optional, TypeAlias, cast import numpy as np import numpy.typing as npt @@ -75,6 +75,16 @@ def make_constructor_type(type_spec: ts.TypeSpec) -> ts.ConstructorType: ) ) + case ts.DeferredType(constraint=ts.TupleType): + return ts.ConstructorType( + definition=ts.FunctionType( + pos_only_args=[ts.DeferredType(constraint=None)], + pos_or_kw_args={}, + kw_only_args={}, + returns=ts.DeferredType(constraint=ts.VarLenTupleType), + ) + ) + case ts.NamedCollectionType() as named_tuple_type: type_ = pkgutil.resolve_name(named_tuple_type.original_python_type) pos_or_kw_args = {k: t for k, t in zip(type_spec.keys, type_spec.types)} @@ -147,7 +157,7 @@ def canonicalize_type_hint( *, globalns: Optional[dict[str, Any]] = None, localns: Optional[dict[str, Any]] = None, -) -> tuple[Any, tuple[Any, ...], tuple[Any, ...]]: +) -> tuple[Any, tuple[Any, ...] | None, tuple[Any, ...]]: """ Canonicalize python type annotations as a tuple of (canonical_type, type_args, annotated_extra_args). """ @@ -178,7 +188,14 @@ def canonicalize_type_hint( type_hint = _resolve_type_alias(type_hint) canonical_type = typing.get_origin(type_hint) or type_hint - args = typing.get_args(type_hint) + # `None` (rather than `()`) marks an unparametrized annotation, distinguishing bare + # `tuple` / `typing.Tuple` from the empty-tuple annotation `tuple[()]`. Unsubscripted + # `typing` aliases have a truthy origin but no `__args__` attribute. + args = ( + typing.get_args(type_hint) + if typing.get_origin(type_hint) and hasattr(type_hint, "__args__") + else None + ) return canonical_type, args, tuple(extra_args) @@ -199,17 +216,46 @@ def from_type_hint( match canonical_type: case builtins.tuple: - if not args: - raise ValueError(f"Tuple annotation '{type_hint}' requires at least one argument.") - if Ellipsis in args: - raise ValueError(f"Unbound tuples '{type_hint}' are not allowed.") - tuple_types = [from_type_hint_same_ns(arg) for arg in args] - assert all(isinstance(elem, ts.DataType) for elem in tuple_types) - return ts.TupleType(types=tuple_types) + # Fixed-length tuple, e.g. `tuple[int32, float64]` + if ( + isinstance(args, tuple) + and len(args) > 0 + and not any(arg is Ellipsis for arg in args) + ): + tuple_types = [from_type_hint_same_ns(arg) for arg in args] + if not all(isinstance(elem, ts.DataType) for elem in tuple_types): + raise ValueError( + f"Tuple annotation '{type_hint}' contains invalid element types." + ) + return ts.TupleType(types=tuple_types) + # Variable-length tuple, e.g. `tuple[int32, ...]` + elif isinstance(args, tuple) and len(args) == 2 and args[1] is Ellipsis: + element_type = from_type_hint_same_ns(args[0]) + if not isinstance(element_type, ts.DataType): + raise ValueError( + f"Tuple annotation '{type_hint}' contains an invalid element type." + ) + return ts.VarLenTupleType(element_type=element_type) + # Unparametrized annotation, i.e. bare `tuple` / `typing.Tuple` (the + # empty-tuple annotation `tuple[()]` has `args == ()` and is rejected below) + elif args is None: + # TODO(tehrengruber): We use `DeferredType` until we have an actual representation + # for a generic type. + return ts.DeferredType(constraint=ts.TupleType) + else: + raise ValueError( + f"Tuple annotation '{type_hint}' must either " + f"be a list of concrete arguments (e.g. 'tuple[int]'), " + f"be a variadic tuple (e.g. 'tuple[int, ...]'), " + f"or have no arguments (e.g. 'tuple')." + ) case common.Field: - if (n_args := len(args)) != 2: - raise ValueError(f"Field type requires two arguments, got {n_args}: '{type_hint}'.") + # `args` is `None` for an unparametrized `Field` annotation + if args is None or len(args) != 2: + raise ValueError( + f"Field type requires two arguments, got {len(args or ())}: '{type_hint}'." + ) dims: list[common.Dimension] = [] dim_arg, dtype_arg = args dim_arg = ( @@ -359,6 +405,30 @@ def from_value(value: Any) -> ts.TypeSpec: elems = [from_value(el) for el in value] assert all(isinstance(elem, ts.DataType) for elem in elems) return ts.TupleType(types=elems) + elif isinstance(value, named_collections.CUSTOM_NAMED_COLLECTION_TYPES): + # Take the member types from the runtime values — the annotated types may be + # generic, e.g. a variable-length tuple, whose concrete type is only known from + # the value. Members whose type is incompatible with the annotation are rejected + # later, by the signature checks against the (annotated) parameter types; in + # debug mode reject them already here, with a more precise error. + annotated_type = from_type_hint(xtyping.infer_type(value)) + assert isinstance(annotated_type, ts.NamedCollectionType) + member_types = [from_value(getattr(value, key)) for key in annotated_type.keys] + if __debug__: + for key, annotated_member_type, member_type in zip( + annotated_type.keys, annotated_type.types, member_types, strict=True + ): + if not type_info.is_concretizable(annotated_member_type, to_type=member_type): + raise ValueError( + f"Member '{key}' of '{type(value).__name__}' has type '{member_type}'," + f" which is incompatible with its annotation '{annotated_member_type}'." + ) + assert all(isinstance(t, (ts.DataType, ts.DimensionType)) for t in member_types) + return ts.NamedCollectionType( + types=cast(list[ts.DataType], member_types), + keys=list(annotated_type.keys), + original_python_type=annotated_type.original_python_type, + ) elif isinstance(value, PythonNamespaceObject): return NamespaceProxy(value) else: diff --git a/src/gt4py/next/utils.py b/src/gt4py/next/utils.py index 54491d3663..7bf6ec5a5f 100644 --- a/src/gt4py/next/utils.py +++ b/src/gt4py/next/utils.py @@ -26,7 +26,7 @@ overload, ) -from gt4py.eve import utils as eve_utils +from gt4py.eve import NOTHING, utils as eve_utils _T = TypeVar("_T") @@ -153,6 +153,9 @@ def tree_map( reconstruct a collection, since `result_collection_constructor` receives the original collection and can recover them from there. result_collection_constructor: Type of the collection to be returned. If `None` the same type as `collection_type` is used. + `fun` is guaranteed to have been applied to every element before the constructor + is called (the mapped elements are materialized eagerly), so a constructor may + discard the results, e.g. for validation-only maps. unpack: Replicate tuple structure returned from `fun` to the mapped result, i.e. return tuple of result collections instead of result collections of tuples. with_path_arg: Pass the path to access the current element to `fun`. @@ -298,13 +301,14 @@ def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]: def equalize_tuple_structure( - d1: _T1, d2: _T2, *, fill_value: Any, bidirectional: bool = True + d1: _T1, d2: _T2, *, fill_value: Any = NOTHING, bidirectional: bool = True ) -> tuple[_T1, _T2]: """ Given two values or tuples thereof equalize their structure. If one of the arguments is a tuple the argument will be promoted to a tuple of the same - structure. + structure. By default only the structure is equalized and differing tuple lengths are + an error; pass `fill_value` to instead pad the shorter tuple with it. >>> equalize_tuple_structure((1,), (2, 3), fill_value=None) == ( ... (1, None), @@ -342,6 +346,8 @@ def equalize_tuple_structure( if d1_is_tuple and d2_is_tuple: if not bidirectional and len(d1) < len(d2): # type: ignore[arg-type] # ensured by condition above raise ValueError(f"Expected `{d2!s}` to be longer than or equal to `{d1!s}`.") + if fill_value is NOTHING and len(d1) != len(d2): # type: ignore[arg-type] # ensured by condition above + raise ValueError(f"Expected `{d1!s}` and `{d2!s}` to have equal lengths.") return tuple( # type: ignore[return-value] # mypy not smart enough zip( *( diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index b5d99e2eb5..a1c39e6f90 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -24,7 +24,7 @@ import inspect import types import typing -from typing import Any, Callable, Literal, Optional, Mapping, Protocol, Self, TypeAlias +from typing import Any, Callable, Final, Literal, Optional, Mapping, Protocol, Self, TypeAlias import numpy as np import pytest @@ -347,6 +347,7 @@ def allocate( strategy: Optional[DataInitializer] = None, dtype: Optional[np.typing.DTypeLike] = None, extend: Optional[dict[gtx.Dimension, tuple[int, int]]] = None, + tuple_length: Optional[int] = None, ) -> FieldViewArg: """ Allocate a parameter or return value from a fieldview code object. @@ -364,6 +365,8 @@ def allocate( extend: Lower and upper size extension per dimension. Useful for shifted fields, which must start off bigger than the output field in the shifted dimension. + tuple_length: Concrete length to allocate for variable-length tuple + parameters (default: ``DEFAULT_TUPLE_LENGTH``). """ if sizes: assert not domain and all(dim in case.default_sizes for dim in sizes) @@ -393,6 +396,7 @@ def allocate( domain=domain, dtype=dtype, strategy=strategy.from_case(case=case, fieldview_prog=fieldview_prog, arg_name=name), + tuple_length=tuple_length, ) @@ -588,6 +592,11 @@ def unstructured_case_3d(unstructured_case): ) +#: Default concrete length used for variable-length tuple parameters, e.g. when +#: allocating arguments for them. +DEFAULT_TUPLE_LENGTH: Final = 3 + + def _allocate_from_type( case: Case, arg_type: ts.TypeSpec, @@ -595,6 +604,7 @@ def _allocate_from_type( strategy: DataInitializer, dtype: Optional[np.typing.DTypeLike] = None, tuple_start: Optional[int] = None, + tuple_length: Optional[int] = None, ) -> FieldViewArg: """Allocate data based on the type or a (nested) tuple thereof.""" match arg_type: @@ -610,11 +620,34 @@ def _allocate_from_type( return tuple( ( _allocate_from_type( - case=case, arg_type=t, domain=domain, dtype=dtype, strategy=strategy + case=case, + arg_type=t, + domain=domain, + dtype=dtype, + strategy=strategy, + tuple_length=tuple_length, ) for t in types ) ) + case ts.VarLenTupleType(element_type=element_type): + length = tuple_length if tuple_length is not None else DEFAULT_TUPLE_LENGTH + # the unique-value offsets are computed from the declared types using the + # default length (see `get_param_size`) + assert length <= DEFAULT_TUPLE_LENGTH + return tuple( + ( + _allocate_from_type( + case=case, + arg_type=t, + domain=domain, + dtype=dtype, + strategy=strategy, + tuple_length=tuple_length, + ) + for t in [element_type] * length + ) + ) case ts.NamedCollectionType(types=types) as named_collection_type_spec: container_constructor = ( named_collections.make_named_collection_constructor_from_type_spec( @@ -625,7 +658,12 @@ def _allocate_from_type( tuple( ( _allocate_from_type( - case=case, arg_type=t, domain=domain, dtype=dtype, strategy=strategy + case=case, + arg_type=t, + domain=domain, + dtype=dtype, + strategy=strategy, + tuple_length=tuple_length, ) for t in types ) @@ -660,6 +698,8 @@ def get_param_size(param_type: ts.TypeSpec, sizes: dict[gtx.Dimension, int]) -> return sum([get_param_size(t, sizes=sizes) for t in types]) case ts.NamedCollectionType(types=types): return sum([get_param_size(t, sizes=sizes) for t in types]) + case ts.VarLenTupleType(element_type=element_type): + return get_param_size(ts.TupleType(types=[element_type] * DEFAULT_TUPLE_LENGTH), sizes) case _: raise TypeError(f"Can not get size for parameter of type '{param_type}'.") diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py index 8ba6f3091a..734d734a91 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py @@ -84,3 +84,15 @@ def scan(inp: int32) -> int32: ).strip() assert pretty_format(scan.foast_stage.foast_node) == expected + + +def test_tuple_comprehension_multi_target(): + I = Dimension("I") + + @field_operator + def foo( + b: tuple[tuple[Field[[I], int64], Field[[I], int64]], ...], + ) -> tuple[Field[[I], int64], ...]: + return tuple(x + y for x, y in b) + + assert "for (x, y) in b" in pretty_format(foo.foast_stage.foast_node) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py index 911ace0804..3ef33405fd 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py @@ -110,6 +110,99 @@ def test_named_tuple_like_constructed_outside(cartesian_case, testee): ) +class VarLenMemberNamedTuple(NamedTuple): + tracers: tuple[gtx.Field[[IDim, JDim], gtx.float32], ...] + rho: gtx.Field[[IDim, JDim], gtx.float32] + + +@dataclasses.dataclass +class VarLenMemberDataclass: + tracers: tuple[gtx.Field[[IDim, JDim], gtx.float32], ...] + rho: gtx.Field[[IDim, JDim], gtx.float32] + + +@gtx.field_operator +def add_rho_named_tuple( + c: VarLenMemberNamedTuple, +) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(tracer + c.rho for tracer in c.tracers) + + +@gtx.field_operator +def add_rho_dataclass( + c: VarLenMemberDataclass, +) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(tracer + c.rho for tracer in c.tracers) + + +@pytest.mark.parametrize("testee", [add_rho_named_tuple, add_rho_dataclass]) +@pytest.mark.uses_tuple_args +def test_named_collection_with_var_len_tuple_member(cartesian_case, testee): + c = cases.allocate(cartesian_case, testee, "c")() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + cases.verify( + cartesian_case, + testee, + c, + out=out, + ref=tuple(t.asnumpy() + c.rho.asnumpy() for t in c.tracers), + ) + + +@gtx.field_operator +def sum_named_tuple_elements( + pairs: tuple[NamedTupleNamedCollection, ...], +) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(p.u + p.v for p in pairs) + + +@gtx.field_operator +def sum_dataclass_elements( + pairs: tuple[DataclassNamedCollection, ...], +) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(p.u + p.v for p in pairs) + + +@pytest.mark.parametrize("testee", [sum_named_tuple_elements, sum_dataclass_elements]) +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_var_len_tuple_of_named_collections(cartesian_case, testee): + pairs = cases.allocate(cartesian_case, testee, "pairs")() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + cases.verify( + cartesian_case, + testee, + pairs, + out=out, + ref=tuple(p.u.asnumpy() + p.v.asnumpy() for p in pairs), + ) + + +@dataclasses.dataclass +class PairsContainer: + pairs: tuple[DataclassNamedCollection, ...] + offset: gtx.Field[[IDim, JDim], gtx.float32] + + +@gtx.field_operator +def sum_nested_pairs(c: PairsContainer) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(p.u + p.v + c.offset for p in c.pairs) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_named_collection_with_var_len_tuple_of_named_collections_member(cartesian_case): + c = cases.allocate(cartesian_case, sum_nested_pairs, "c")() + out = cases.allocate(cartesian_case, sum_nested_pairs, cases.RETURN)() + cases.verify( + cartesian_case, + sum_nested_pairs, + c, + out=out, + ref=tuple(p.u.asnumpy() + p.v.asnumpy() + c.offset.asnumpy() for p in c.pairs), + ) + + @gtx.field_operator def constructed_inside_named_tuple( vel: tuple[gtx.Field[[IDim, JDim], gtx.float32], gtx.Field[[IDim, JDim], gtx.float32]], @@ -307,7 +400,7 @@ class StateNamedTuple(NamedTuple): value: gtx.float32 -@gtx.scan_operator(axis=cases.KDim, forward=True, init=StateDataclass(value=0.0)) +@gtx.scan_operator(axis=cases.KDim, forward=True, init=StateDataclass(value=gtx.float32(0.0))) def scan_dataclass( state: StateDataclass, inp: gtx.float32, @@ -322,7 +415,7 @@ def scan_dataclass_wrapper(inp: gtx.Field[[KDim], gtx.float32]) -> gtx.Field[[KD return scan_dataclass(inp).value -@gtx.scan_operator(axis=cases.KDim, forward=True, init=StateNamedTuple(value=0.0)) +@gtx.scan_operator(axis=cases.KDim, forward=True, init=StateNamedTuple(value=gtx.float32(0.0))) def scan_named_tuple( state: StateNamedTuple, inp: gtx.float32, diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py index d882a88ec4..62d80a8e47 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py @@ -250,17 +250,13 @@ def test_dimensions_domain(cartesian_case): def empty_domain_fieldop(a: cases.IJField): return a - @gtx.program - def empty_domain_program(a: cases.IJField, out_field: cases.IJField): - empty_domain_fieldop(a, out=out_field, domain={JDim: (0, 1), IDim: (0, 1)}) - - a = cases.allocate(cartesian_case, empty_domain_program, "a")() - out_field = cases.allocate(cartesian_case, empty_domain_program, "out_field")() - with pytest.raises( - ValueError, match=(r"Dimensions in out field and field domain are not equivalent") + errors.DSLError, match=(r"Dimensions in out field and field domain are not equivalent") ): - cases.run(cartesian_case, empty_domain_program, a, out_field, offset_provider={}) + + @gtx.program + def empty_domain_program(a: cases.IJField, out_field: cases.IJField): + empty_domain_fieldop(a, out=out_field, domain={JDim: (0, 1), IDim: (0, 1)}) @pytest.mark.uses_origin diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 91025f156f..487558cb91 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -10,14 +10,7 @@ import pytest import gt4py.next as gtx -from gt4py.next import ( - broadcast, - errors, - float64, - int32, - neighbor_sum, - utils as gt_utils, -) +from gt4py.next import broadcast, common, errors, float64, int32, neighbor_sum, utils as gt_utils from next_tests.integration_tests import cases from next_tests.integration_tests.cases import ( @@ -120,6 +113,367 @@ def testee(a: tuple[cases.IField, cases.IJField]) -> cases.IJField: ) +@pytest.mark.uses_tuple_args +def test_fixed_len_tuple_comprehension(cartesian_case): + @gtx.field_operator + def testee( + tracers: tuple[cases.IField, cases.IField], factor: int32 + ) -> tuple[cases.IField, cases.IField]: + return tuple(tracer * factor for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(el * f for el in t), + ) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_comprehension(cartesian_case): + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(el * f for el in t), + ) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_comprehension_empty(cartesian_case): + # The empty tuple is a valid concretization of a variable-length tuple; the program + # then has nothing to write and must simply run without error. + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + if cartesian_case.backend is not None: + pytest.xfail( + "The compiled toolchain cannot execute a program with an empty tuple output: " + "a 'SetAt' for an empty 'out' has no target and no domain to derive." + ) + cases.run(cartesian_case, testee, (), 42, out=()) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_var_len_tuple_comprehension_in_fixed_tuple_return(cartesian_case): + # A comprehension result nested inside a fixed-length return tuple, called directly. + @gtx.field_operator + def testee( + tracers: tuple[cases.IField, ...], rho: cases.IField + ) -> tuple[tuple[cases.IField, ...], cases.IField]: + return tuple(tracer * 2 for tracer in tracers), rho + + tracers = cases.allocate(cartesian_case, testee, "tracers")() + rho = cases.allocate(cartesian_case, testee, "rho")() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + cases.verify( + cartesian_case, + testee, + tracers, + rho, + out=out, + ref=(tuple(t.asnumpy() * 2 for t in tracers), rho.asnumpy()), + ) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_comprehension_different_lengths(cartesian_case): + # An operator with variable-length tuple parameters is generic: every call must be + # specialized (and compiled) for the concrete tuple length of its arguments. + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + for length in (3, 2): + tracers = cases.allocate(cartesian_case, testee, "tracers", tuple_length=length)() + out = cases.allocate(cartesian_case, testee, cases.RETURN, tuple_length=length)() + cases.verify( + cartesian_case, + testee, + tracers, + 42, + out=out, + ref=tuple(t.asnumpy() * 42 for t in tracers), + ) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_program(cartesian_case): + @gtx.field_operator + def scale(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + @gtx.program + def testee(tracers: tuple[cases.IField, ...], factor: int32, out: tuple[cases.IField, ...]): + scale(tracers, factor, out=out) + + for length in (3, 2): # ensure each concrete length gets its own specialization + tracers = cases.allocate(cartesian_case, testee, "tracers", tuple_length=length)() + out = cases.allocate(cartesian_case, testee, "out", tuple_length=length)() + cases.run(cartesian_case, testee, tracers, 42, out) + assert all(np.allclose(o.asnumpy(), t.asnumpy() * 42) for o, t in zip(out, tracers)) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_program_domain(cartesian_case): + # A single `domain` argument applies to (is broadcast over) every element of the + # variable-length tuple output. + @gtx.field_operator + def scale(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + @gtx.program + def testee( + tracers: tuple[cases.IField, ...], + factor: int32, + size: int32, + out: tuple[cases.IField, ...], + ): + scale(tracers, factor, out=out, domain={IDim: (1, size - 1)}) + + size = cartesian_case.default_sizes[IDim] + tracers = cases.allocate(cartesian_case, testee, "tracers")() + out = cases.allocate(cartesian_case, testee, "out")() + cases.run(cartesian_case, testee, tracers, 42, int32(size), out) + for o, t in zip(out, tracers): + assert np.allclose(o.asnumpy()[1:-1], t.asnumpy()[1:-1] * 42) + assert np.all(o.asnumpy()[[0, -1]] == 0) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_program_fixed_in_vararg_out(cartesian_case): + @gtx.field_operator + def scale_pairs( + pairs: tuple[tuple[cases.IField, cases.IField], ...], factor: int32 + ) -> tuple[tuple[cases.IField, cases.IField], ...]: + return tuple((a * factor, b * factor) for a, b in pairs) + + @gtx.program + def testee( + pairs: tuple[tuple[cases.IField, cases.IField], ...], + factor: int32, + out: tuple[tuple[cases.IField, cases.IField], ...], + ): + scale_pairs(pairs, factor, out=out) + + pairs = cases.allocate(cartesian_case, testee, "pairs")() + out = cases.allocate(cartesian_case, testee, "out")() + cases.run(cartesian_case, testee, pairs, 42, out) + assert all( + np.allclose(o.asnumpy(), p.asnumpy() * 42) + for o_pair, p_pair in zip(out, pairs) + for o, p in zip(o_pair, p_pair) + ) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_var_len_tuple_program_vararg_in_fixed_out(cartesian_case): + @gtx.field_operator + def scale_and_copy( + tracers: tuple[cases.IField, ...], rho: cases.IField, factor: int32 + ) -> tuple[tuple[cases.IField, ...], cases.IField]: + return tuple(tracer * factor for tracer in tracers), rho + + @gtx.program + def testee( + tracers: tuple[cases.IField, ...], + rho: cases.IField, + factor: int32, + out: tuple[tuple[cases.IField, ...], cases.IField], + ): + scale_and_copy(tracers, rho, factor, out=out) + + tracers = cases.allocate(cartesian_case, testee, "tracers")() + rho = cases.allocate(cartesian_case, testee, "rho")() + out = cases.allocate(cartesian_case, testee, "out")() + cases.run(cartesian_case, testee, tracers, rho, 42, out) + assert all(np.allclose(o.asnumpy(), t.asnumpy() * 42) for o, t in zip(out[0], tracers)) + assert np.allclose(out[1].asnumpy(), rho.asnumpy()) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_var_len_tuple_program_composed_out(cartesian_case): + # The 'out' argument is composed from several program parameters inside the program. + @gtx.field_operator + def scale_and_copy( + tracers: tuple[cases.IField, ...], rho: cases.IField, factor: int32 + ) -> tuple[tuple[cases.IField, ...], cases.IField]: + return tuple(tracer * factor for tracer in tracers), rho + + @gtx.program + def testee( + tracers: tuple[cases.IField, ...], + rho: cases.IField, + factor: int32, + out_tracers: tuple[cases.IField, ...], + out_rho: cases.IField, + ): + scale_and_copy(tracers, rho, factor, out=(out_tracers, out_rho)) + + tracers = cases.allocate(cartesian_case, testee, "tracers")() + rho = cases.allocate(cartesian_case, testee, "rho")() + out_tracers = cases.allocate( + cartesian_case, testee, "out_tracers", strategy=cases.ZeroInitializer() + )() + out_rho = cases.allocate(cartesian_case, testee, "out_rho", strategy=cases.ZeroInitializer())() + cases.run(cartesian_case, testee, tracers, rho, 42, out_tracers, out_rho) + assert all(np.allclose(o.asnumpy(), t.asnumpy() * 42) for o, t in zip(out_tracers, tracers)) + assert np.allclose(out_rho.asnumpy(), rho.asnumpy()) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_program_compile_rejected(cartesian_case): + @gtx.field_operator + def scale(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + @gtx.program(backend=cartesian_case.backend) + def testee(tracers: tuple[cases.IField, ...], factor: int32, out: tuple[cases.IField, ...]): + scale(tracers, factor, out=out) + + if cartesian_case.backend is None: + pytest.skip("Ahead-of-time compilation requires a backend.") + with pytest.raises(ValueError, match="generic programs"): + testee.compile(offset_provider={}) + + +@pytest.mark.uses_tuple_args +@pytest.mark.xfail( + strict=True, + reason="The consistency of the lengths of variable-length tuple arguments and the " + "'out' argument is not checked anywhere; the mismatch only surfaces as an internal " + "error (e.g. a bare 'AssertionError' from 'utils.tree_map') deep in the toolchain.", +) +def test_var_len_tuple_length_mismatch_rejected(cartesian_case): + # The return length equals the length of `a` or `b` depending on a runtime value, so + # with `len(a) != len(b)` no single length can be consistent with `out` — this should + # be rejected with a proper diagnostic when the concrete argument types are known. + @gtx.field_operator + def testee( + a: tuple[cases.IField, ...], b: tuple[cases.IField, ...], c: bool + ) -> tuple[cases.IField, ...]: + return a if c else b + + a = cases.allocate(cartesian_case, testee, "a", tuple_length=2)() + b = cases.allocate(cartesian_case, testee, "b", tuple_length=3)() + out = cases.allocate(cartesian_case, testee, cases.RETURN, tuple_length=2)() + + with pytest.raises(errors.DSLError): + cases.run(cartesian_case, testee, a, b, False, out=out) + + +@pytest.mark.uses_tuple_args +def test_tuple_comprehension_other_fo(cartesian_case): + @gtx.field_operator + def inner(tracer: cases.IField, factor: int32) -> cases.IField: + return tracer * factor + + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(inner(tracer, factor) for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(el * f for el in t), + ) + + +@gtx.field_operator +def _module_level_inner(tracer: cases.IField, factor: int32) -> cases.IField: + return tracer * factor + + +@pytest.mark.uses_tuple_args +def test_tuple_comprehension_module_level_fo(cartesian_case): + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(_module_level_inner(tracer, factor) for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, testee, ref=lambda t, f: tuple(el * f for el in t) + ) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_unstructured_shift +def test_tuple_comprehension_module_level_builtin(unstructured_case): + @gtx.field_operator + def testee(tracers: tuple[cases.EField, ...]) -> tuple[cases.VField, ...]: + return tuple(neighbor_sum(tracer(V2E), axis=V2EDim) for tracer in tracers) + + v2e = unstructured_case.offset_provider["V2E"].asnumpy() + valid = v2e != common._DEFAULT_SKIP_VALUE + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda t: tuple(np.sum(el[v2e], axis=1, where=valid) for el in t), + ) + + +@pytest.mark.uses_tuple_args +def test_tuple_comprehension_module_level_fo_used_outside_too(cartesian_case): + @gtx.field_operator + def testee( + tracers: tuple[cases.IField, ...], a: cases.IField, factor: int32 + ) -> tuple[cases.IField, ...]: + unrelated = _module_level_inner(a, factor) + return tuple(_module_level_inner(tracer, factor) for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, testee, ref=lambda t, a, f: tuple(el * f for el in t) + ) + + +@pytest.mark.uses_tuple_args +def test_nested_tuple_comprehension(cartesian_case): + @gtx.field_operator + def testee( + vals: tuple[tuple[cases.IField, ...], ...], factor: int32 + ) -> tuple[tuple[cases.IField, ...], ...]: + return tuple(tuple(grand_child * factor for grand_child in child) for child in vals) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(tuple(grand_child * f for grand_child in child) for child in t), + ) + + +@pytest.mark.uses_tuple_args +def test_nested_tuple_comprehension_shadowing_names(cartesian_case): + @gtx.field_operator + def testee( + vals: tuple[tuple[cases.IField, ...], ...], factor: int32 + ) -> tuple[tuple[cases.IField, ...], ...]: + return tuple(tuple(child * factor for child in child) for child in vals) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(tuple(child * f for child in child) for child in t), + ) + + +@pytest.mark.uses_tuple_args +def test_multi_target_tuple_comprehension(cartesian_case): + @gtx.field_operator + def testee(nested_tuple: tuple[tuple[int32, cases.IField], ...]) -> tuple[cases.IField, ...]: + return tuple(factor * tracer for factor, tracer in nested_tuple) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t: tuple(f * el for f, el in t), + ) + + @pytest.mark.uses_tuple_args @pytest.mark.uses_unstructured_shift @pytest.mark.xfail(reason="Iterator of tuple approach in lowering does not allow this.") @@ -280,5 +634,5 @@ def test_tuple_unpacking_too_few_values(cartesian_case): @gtx.field_operator(backend=cartesian_case.backend) def _invalid_unpack() -> tuple[int32, float64, int32]: - a, b, c = 1 + a, _b, _c = 1 return a diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py index 7ea11a7b69..7b93d7892e 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py @@ -7,16 +7,27 @@ # SPDX-License-Identifier: BSD-3-Clause import functools + import numpy as np -from typing import Tuple import pytest -from next_tests.integration_tests.cases import IDim, JDim, KDim, cartesian_case + from gt4py import next as gtx -from gt4py.next import float64, int32 -from gt4py.next.ffront.fbuiltins import where, broadcast +from gt4py.next import common, float64, int32, neighbor_sum +from gt4py.next.ffront.fbuiltins import broadcast, where + from next_tests.integration_tests import cases +from next_tests.integration_tests.cases import ( + E2V, + E2VDim, + IDim, + JDim, + KDim, + cartesian_case, + unstructured_case, +) from next_tests.integration_tests.cases_utils import ( exec_alloc_descriptor, + mesh_descriptor, ) @@ -113,6 +124,41 @@ def testee( ) +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +def test_with_tuples_and_local_condition(unstructured_case): + @gtx.field_operator + def testee( + a: cases.VField, + b: cases.VField, + c: cases.VField, + d: cases.VField, + ) -> tuple[cases.EField, cases.EField]: + cond = a(E2V) > c(E2V) + t = where(cond, (a(E2V), b(E2V)), (c(E2V), d(E2V))) + return neighbor_sum(t[0], axis=E2VDim), neighbor_sum(t[1], axis=E2VDim) + + e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b, c, d: ( + np.sum( + np.where(a[e2v_table] > c[e2v_table], a[e2v_table], c[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + np.sum( + np.where(a[e2v_table] > c[e2v_table], b[e2v_table], d[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + ), + ) + + @pytest.mark.uses_tuple_returns def test_conditional_nested_tuple(cartesian_case): @gtx.field_operator @@ -124,7 +170,9 @@ def conditional_nested_tuple( return where(mask, ((a, b), (b, a)), ((5.0, 7.0), (7.0, 5.0))) size = cartesian_case.default_sizes[IDim] - mask = cartesian_case.as_field([IDim], np.random.choice(a=[False, True], size=size)) + mask = cartesian_case.as_field( + [IDim], np.random.default_rng().choice(a=[False, True], size=size) + ) a = cases.allocate(cartesian_case, conditional_nested_tuple, "a")() b = cases.allocate(cartesian_case, conditional_nested_tuple, "b")() @@ -158,7 +206,9 @@ def conditional( return where(mask, a, b) size = cartesian_case.default_sizes[IDim] - mask = cartesian_case.as_field([IDim], np.random.choice(a=[False, True], size=(size))) + mask = cartesian_case.as_field( + [IDim], np.random.default_rng().choice(a=[False, True], size=size) + ) a = cases.allocate(cartesian_case, conditional, "a")() b = cases.allocate(cartesian_case, conditional, "b")() out = cases.allocate(cartesian_case, conditional, cases.RETURN)() @@ -180,7 +230,9 @@ def conditional_promotion(mask: cases.IBoolField, a: cases.IFloatField) -> cases return where(mask, a, 10.0) size = cartesian_case.default_sizes[IDim] - mask = cartesian_case.as_field([IDim], np.random.choice(a=[False, True], size=(size))) + mask = cartesian_case.as_field( + [IDim], np.random.default_rng().choice(a=[False, True], size=size) + ) a = cases.allocate(cartesian_case, conditional_promotion, "a")() out = cases.allocate(cartesian_case, conditional_promotion, cases.RETURN)() ref = np.where(mask.asnumpy(), a.asnumpy(), 10.0) @@ -214,7 +266,9 @@ def conditional_program( conditional_shifted(mask, a, b, out=out) size = cartesian_case.default_sizes[IDim] + 1 - mask = cartesian_case.as_field([IDim], np.random.choice(a=[False, True], size=(size))) + mask = cartesian_case.as_field( + [IDim], np.random.default_rng().choice(a=[False, True], size=size) + ) a = cases.allocate(cartesian_case, conditional_program, "a").extend({IDim: (0, 1)})() b = cases.allocate(cartesian_case, conditional_program, "b").extend({IDim: (0, 1)})() out = cases.allocate(cartesian_case, conditional_shifted, cases.RETURN)() diff --git a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py index 7e2d4fed6e..4ace4f8913 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py @@ -359,6 +359,41 @@ def broken(a: BrokenFieldAlias) -> gtx.Field[[IDim], float64]: assert re.search(r"\| +\^{19}(?!\^)", str(err)), str(err) +def test_bare_generator_expression_points_to_tuple_constructor(): + def with_genexp(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + b = (x for x in (a, a)) # noqa: F841 [unused-variable] + return a + + err = parse_error(with_genexp) + + assert isinstance(err, errors.UnsupportedPythonFeatureError) + assert err.message == "Unsupported Python syntax: generator expression." + assert any("'tuple(...)'" in hint for hint in err.hints) + + +def test_nested_generator_expression_suggests_single_for_clause(): + def with_nested_genexp(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + b = tuple(x for x in (a, a) for y in (a, a)) # noqa: F841 [unused-variable] + return a + + err = parse_error(with_nested_genexp) + + assert err.message == "Nested generator expressions are not supported." + assert any("single 'for' clause" in hint for hint in err.hints) + + +def test_conditional_in_generator_expression_explains_static_length(): + def with_filtered_genexp(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + b = tuple(x for x in (a, a) if True) # noqa: F841 [unused-variable] + return a + + err = parse_error(with_filtered_genexp) + + assert err.message == "Conditionals are not supported in generator expressions." + assert any("known at compile time" in note for note in err.notes) + assert "if True" in str(err) + + def test_unindexed_cartesian_offset_names_the_offset_as_written(): # The tag of 'IOff' is 'Ioff'; the message has to quote what the user wrote. def unindexed(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py index eb1927a776..c1738146d9 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py @@ -42,7 +42,6 @@ from gt4py.next.type_system import type_specifications as ts, type_translation from gt4py.next.iterator import ir as itir - Edge = gtx.Dimension("Edge") Vertex = gtx.Dimension("Vertex") V2EDim = gtx.Dimension("V2E", gtx.DimensionKind.LOCAL) @@ -386,6 +385,184 @@ def foo( assert lowered_inlined.expr == reference +def test_fixed_len_tuple_comprehension(): + def foo(a: tuple[gtx.Field[[TDim], float64], gtx.Field[[TDim], float64]], factor: float64): + return tuple(el * factor for el in a) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + lowered_inlined = inline_lambdas.InlineLambdas.apply(lowered) + + iterable = "__tuple_comprh_0" + reference = im.let(iterable, "a")( + im.make_tuple( + im.let("el", im.tuple_get(0, iterable))(im.op_as_fieldop("multiplies")("el", "factor")), + im.let("el", im.tuple_get(1, iterable))(im.op_as_fieldop("multiplies")("el", "factor")), + ) + ) + reference_inlined = inline_lambdas.InlineLambdas.apply(reference) + + assert lowered_inlined.expr == reference_inlined + + +def test_var_len_tuple_comprehension_scalar(): + def foo(a: tuple[float64, ...]): + return tuple(2.0 * el for el in a) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + + reference = im.call( + im.call("map_tuple")(im.lambda_("el")(im.multiplies_(im.literal("2.0", "float64"), "el"))) + )("a") + + assert lowered.expr == reference + + +def test_var_len_tuple_comprehension_field(): + def foo(a: tuple[gtx.Field[[TDim], float64], ...], factor: float64): + return tuple(el * factor for el in a) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + + reference = im.call( + im.call("map_tuple")(im.lambda_("el")(im.op_as_fieldop("multiplies")("el", "factor"))) + )("a") + + assert lowered.expr == reference + + +def test_var_len_tuple_comprehension_field_with_local_dim(): + def foo(a: tuple[gtx.Field[[Vertex, V2EDim], float64], ...]): + return tuple(2.0 * el for el in a) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + + two = im.literal("2.0", "float64") + reference = im.call( + im.call("map_tuple")( + im.lambda_("el")( + im.op_as_fieldop(im.map_list("multiplies"))( + im.op_as_fieldop("make_const_list")(two), "el" + ) + ) + ) + )("a") + + assert lowered.expr == reference + + +def test_fixed_len_tuple_comprehension_local_field(): + def foo(a: gtx.Field[[Edge], float64]): + return tuple(2.0 * el for el in (a(V2E), a(V2E))) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + lowered_inlined = inline_lambdas.InlineLambdas.apply(lowered) + + iterable = "__tuple_comprh_0" + two = im.literal("2.0", "float64") + reference = im.let( + iterable, + im.make_tuple(im.as_fieldop_neighbors("V2E", "a"), im.as_fieldop_neighbors("V2E", "a")), + )( + im.make_tuple( + im.let("el", im.tuple_get(0, iterable))( + im.op_as_fieldop(im.map_list("multiplies"))( + im.op_as_fieldop("make_const_list")(two), "el" + ) + ), + im.let("el", im.tuple_get(1, iterable))( + im.op_as_fieldop(im.map_list("multiplies"))( + im.op_as_fieldop("make_const_list")(two), "el" + ) + ), + ) + ) + reference_inlined = inline_lambdas.InlineLambdas.apply(reference) + + assert lowered_inlined.expr == reference_inlined + + +def test_fixed_len_tuple_comprehension_mixed_local_field_tuple_target(): + def foo( + a: gtx.Field[[Edge], float64], + b: gtx.Field[[Vertex], float64], + c: gtx.Field[[Edge], float64], + d: gtx.Field[[Vertex], float64], + ): + return tuple( + 2.0 * local_el + scalar_el for local_el, scalar_el in ((a(V2E), b), (c(V2E), d)) + ) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + lowered_inlined = inline_lambdas.InlineLambdas.apply(lowered) + + iterable = "__tuple_comprh_0" + two = im.literal("2.0", "float64") + local_term = im.op_as_fieldop(im.map_list("multiplies"))( + im.op_as_fieldop("make_const_list")(two), "local_el" + ) + element_expr = im.op_as_fieldop(im.map_list("plus"))( + local_term, im.op_as_fieldop("make_const_list")("scalar_el") + ) + tuple_el_0 = im.tuple_get(0, iterable) + tuple_el_1 = im.tuple_get(1, iterable) + mapped_tuple1 = im.make_tuple(im.as_fieldop_neighbors("V2E", "a"), "b") + mapped_tuple2 = im.make_tuple(im.as_fieldop_neighbors("V2E", "c"), "d") + reference = im.let(iterable, im.make_tuple(mapped_tuple1, mapped_tuple2))( + im.make_tuple( + im.let( + ("local_el", im.tuple_get(0, tuple_el_0)), + ("scalar_el", im.tuple_get(1, tuple_el_0)), + )(element_expr), + im.let( + ("local_el", im.tuple_get(0, tuple_el_1)), + ("scalar_el", im.tuple_get(1, tuple_el_1)), + )(element_expr), + ) + ) + reference_inlined = inline_lambdas.InlineLambdas.apply(reference) + + assert lowered_inlined.expr == reference_inlined + + +def test_fixed_len_tuple_comprehension_mixed_local_field(): + def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): + return tuple(2.0 * el for el in (a(V2E), b)) + + with pytest.raises( + errors.DSLError, + match="fixed-length tuples with differently typed iterable elements are not implemented", + ): + FieldOperatorParser.apply_to_function(foo) + + +def test_fixed_len_tuple_comprehension_mixed_field_domains(): + def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): + return tuple(2.0 * el for el in (a, b)) + + with pytest.raises( + errors.DSLError, + match="fixed-length tuples with differently typed iterable elements are not implemented", + ): + FieldOperatorParser.apply_to_function(foo) + + +def test_fixed_len_tuple_comprehension_tuple_target_mixed_element_types(): + def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): + return tuple(2.0 * left + right for left, right in ((a(V2E), b), (b, b))) + + with pytest.raises( + errors.DSLError, + match="fixed-length tuples with differently typed iterable elements are not implemented", + ): + FieldOperatorParser.apply_to_function(foo) + + def test_unary_minus(): def foo(inp: gtx.Field[[TDim], float64]): return -inp diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py index d72efaaa22..7847444535 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py @@ -269,6 +269,20 @@ def cast_scalar_temp(): _ = FieldOperatorParser.apply_to_function(cast_scalar_temp) +def test_scalar_cast_wrong_arity(): + def cast_scalar_two_args(): + return int32(1, 2) + + with pytest.raises(errors.DSLError, match=r"takes exactly one argument, got 2"): + _ = FieldOperatorParser.apply_to_function(cast_scalar_two_args) + + def cast_scalar_no_args(): + return int32() + + with pytest.raises(errors.DSLError, match=r"takes exactly one argument, got 0"): + _ = FieldOperatorParser.apply_to_function(cast_scalar_no_args) + + def test_conditional_wrong_mask_type(): def conditional_wrong_mask_type(a: gtx.Field[[TDim], float64]) -> gtx.Field[[TDim], float64]: return where(a, a, a) @@ -493,3 +507,68 @@ def tuple_index_failure( with pytest.raises(errors.DSLError, match=r"need .* literal"): _ = FieldOperatorParser.apply_to_function(tuple_index_failure) + + +def test_tuple_compr_list_target_failure(): + def testee(arg: tuple[tuple[int32, int32], ...]): + return tuple(p + q for [p, q] in arg) + + with pytest.raises(errors.UnsupportedPythonFeatureError, match=r"generator expression targets"): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_starred_target_failure(): + def testee(arg: tuple[tuple[int32, int32], ...]): + return tuple(p for p, *q in arg) + + with pytest.raises(errors.UnsupportedPythonFeatureError, match=r"generator expression targets"): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_subscript_target_failure(): + def testee(arg: tuple[int32, ...], x: tuple[int32, int32]): + return tuple(x[0] for x[0] in arg) + + with pytest.raises(errors.UnsupportedPythonFeatureError, match=r"generator expression targets"): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_duplicate_target_name_failure(): + def testee(arg: tuple[tuple[int32, int32], ...]): + return tuple(p for p, p in arg) + + with pytest.raises(errors.DSLError, match=r"Duplicate name 'p'"): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_non_tuple_iterable_failure(): + def testee(arg: float): + return tuple(_ for _ in arg) + + with pytest.raises( + errors.DSLError, + match=re.escape("Iterable in generator expression must be a tuple, got 'float64'."), + ): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_nested_tuple_compr_failure(): + def testee(nested_tuple: tuple[tuple[gtx.Field[[TDim], float64], ...], ...], factor: int32): + return tuple(grandchild * factor for child in nested_tuple for grandchild in child) + + with pytest.raises( + errors.DSLError, + match=re.escape("Nested generator expressions are not supported."), + ): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_unpacking_failure(): + def testee(arg: tuple[int32, ...]): + return tuple(a * b for a, b in arg) + + with pytest.raises( + errors.DSLError, + match=re.escape("Cannot unpack non-iterable 'int32' object."), + ): + _ = FieldOperatorParser.apply_to_function(testee) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py index 8486885076..2c895cc9e2 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py @@ -152,6 +152,100 @@ def test_copy_restrict_parsing(copy_restrict_program_def): pattern_node.match(past_node, raise_exception=True) +@pytest.fixture +def var_len_tuple_scale_op(): + @gtx.field_operator + def scale( + tracers: tuple[gtx.Field[[IDim], float64], ...], factor: float64 + ) -> tuple[gtx.Field[[IDim], float64], ...]: + return tuple(tracer * factor for tracer in tracers) + + return scale + + +def test_var_len_tuple_subscript_out_exception(var_len_tuple_scale_op): + scale = var_len_tuple_scale_op + + def var_len_tuple_slice_program( + tracers: tuple[gtx.Field[[IDim], float64], ...], + factor: float64, + out: tuple[gtx.Field[[IDim], float64], ...], + ): + scale(tracers, factor, out=out[0:2]) + + with pytest.raises(errors.DSLError, match="Subscripting a tuple"): + ProgramParser.apply_to_function(var_len_tuple_slice_program) + + +def test_var_len_tuple_subscript_arg_exception(var_len_tuple_scale_op): + scale = var_len_tuple_scale_op + + def var_len_tuple_slice_program( + tracers: tuple[gtx.Field[[IDim], float64], ...], + factor: float64, + out: tuple[gtx.Field[[IDim], float64], ...], + ): + scale(tracers[0:2], factor, out=out) + + with pytest.raises(errors.DSLError, match="Subscripting a tuple"): + ProgramParser.apply_to_function(var_len_tuple_slice_program) + + +def test_tuple_subscript_out_exception(): + @gtx.field_operator + def pair( + a: gtx.Field[[IDim], float64], + ) -> tuple[gtx.Field[[IDim], float64], gtx.Field[[IDim], float64]]: + return a, a + + def tuple_slice_program( + a: gtx.Field[[IDim], float64], + out: tuple[gtx.Field[[IDim], float64], gtx.Field[[IDim], float64]], + ): + pair(a, out=out[0:2]) + + with pytest.raises(errors.DSLError, match="Subscripting a tuple"): + ProgramParser.apply_to_function(tuple_slice_program) + + +def test_var_len_tuple_per_element_domain_exception(var_len_tuple_scale_op): + scale = var_len_tuple_scale_op + + def var_len_tuple_domain_program( + tracers: tuple[gtx.Field[[IDim], float64], ...], + factor: float64, + out: tuple[gtx.Field[[IDim], float64], ...], + ): + scale(tracers, factor, out=out, domain=({IDim: (0, 2)}, {IDim: (0, 2)})) + + with pytest.raises(errors.DSLError) as exc_info: + ProgramParser.apply_to_function(var_len_tuple_domain_program) + + assert exc_info.match("Invalid call to 'scale'") + assert ( + re.search("Per-element domain tuples are not supported", exc_info.value.__cause__.args[0]) + is not None + ) + + +def test_var_len_tuple_scan_arg(): + # A variable-length tuple program parameter forwarded to a scan operator: the scan + # signature check runs against the still generic (declared) argument type. + KDim = gtx.Dimension("KDim", kind=gtx.DimensionKind.VERTICAL) + + @gtx.scan_operator(axis=KDim, forward=True, init=float64(0.0)) + def scan_tracers(carry: float64, tracers: tuple[float64, ...]) -> float64: + return carry + + def scan_program( + tracers: tuple[gtx.Field[[KDim], float64], ...], + out: gtx.Field[[KDim], float64], + ): + scan_tracers(tracers, out=out) + + ProgramParser.apply_to_function(scan_program) + + def test_domain_exception_1(identity_def): domain_format_1 = gtx.field_operator(identity_def) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_source_utils.py b/tests/next_tests/unit_tests/ffront_tests/test_source_utils.py index 5735b4c25b..9a08544d91 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_source_utils.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_source_utils.py @@ -10,17 +10,19 @@ A generator expression, lambda or nested function compiles to its own code object, so names referenced only inside it appear in that object's `co_names` and never in -the enclosing function's. Collecting closure variables from the enclosing function -alone therefore misses them; `get_closure_vars_from_function` takes the global names -from the compiler's symbol table of the source instead, which covers nested scopes. +the enclosing function's. `get_closure_vars_from_function` therefore takes the global +names from the compiler's symbol table of the source, which covers nested scopes. Free variables need no such treatment: the enclosing code object carries a cell for them. """ +import pytest + import gt4py.next as gtx from gt4py.next import Dims, Dimension, float64, neighbor_sum from gt4py.next.ffront import source_utils +from gt4py.next.ffront.func_to_foast import FieldOperatorParser from gt4py.next.ffront.source_utils import get_closure_vars_from_function @@ -46,6 +48,27 @@ def _module_level_operator(tracers: tuple[CField, ...], factor: float64) -> tupl return tuple(scale(t, factor) for t in tracers) +@pytest.mark.parametrize( + "func", [_builtin_and_offset, _module_level_operator], ids=["builtin_and_offset", "operator"] +) +def test_module_level_name_used_only_in_comprehension(func): + """Module-level names are reachable from inside a comprehension body.""" + parsed = FieldOperatorParser.apply_to_function(func) + assert parsed.type is not None + + +def test_free_variables_still_resolve_through_the_parser(): + @gtx.field_operator + def local_scale(f: CField, factor: float64) -> CField: + return f * factor + + def uses_freevar(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]: + return tuple(local_scale(t, factor) for t in tracers) + + parsed = FieldOperatorParser.apply_to_function(uses_freevar) + assert parsed.type is not None + + def test_names_are_collected_from_the_nested_code_object(): """The names live in the generator expression's code object, not the function's.""" assert "neighbor_sum" not in _builtin_and_offset.__code__.co_names @@ -62,8 +85,6 @@ def test_comprehension_target_is_not_collected(): def test_free_variables_still_resolve(): - """The path that already worked, kept as a guard.""" - @gtx.field_operator def local_scale(f: CField, factor: float64) -> CField: return f * factor diff --git a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py index b6ae674312..1ddd7d6d1d 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py @@ -74,6 +74,34 @@ def temp_tuple(a: Field[[TDim], float64], b: Field[[TDim], int64]): ) +def test_tuple_comprehension_unpack_non_tuple(): + def foo(it: tuple[Field[[TDim], float64], ...]): + return tuple(a + b for a, b in it) + + with pytest.raises(errors.DSLError, match=r"Cannot unpack non-iterable"): + _ = FieldOperatorParser.apply_to_function(foo) + + +def test_tuple_comprehension_unpack_too_many_values(): + def foo( + it: tuple[ + tuple[Field[[TDim], float64], Field[[TDim], float64], Field[[TDim], float64]], ... + ], + ): + return tuple(a + b for a, b in it) + + with pytest.raises(errors.DSLError, match=r"Too many values to unpack \(expected 2\)"): + _ = FieldOperatorParser.apply_to_function(foo) + + +def test_tuple_comprehension_unpack_too_short_tuple(): + def foo(it: tuple[tuple[Field[[TDim], float64], Field[[TDim], float64]], ...]): + return tuple(a + b + c for a, b, c in it) + + with pytest.raises(errors.DSLError, match=r"Not enough values to unpack \(expected 3, got 2\)"): + _ = FieldOperatorParser.apply_to_function(foo) + + def test_adding_bool(): """Expect an error when using arithmetic on bools.""" diff --git a/tests/next_tests/unit_tests/test_containers.py b/tests/next_tests/unit_tests/test_containers.py index 8704174b09..4cd3e53e24 100644 --- a/tests/next_tests/unit_tests/test_containers.py +++ b/tests/next_tests/unit_tests/test_containers.py @@ -72,6 +72,50 @@ def test_make_container_extractor( ) +@dataclasses.dataclass +class PointDataclass: + x: int + y: int + + +@dataclasses.dataclass +class VarLenPointsDataclass: + points: tuple[PointDataclass, ...] + offset: int + + +def test_make_container_extractor_var_len_tuple(): + # the length of the variable-length tuple member is only known from the value, the + # same extractor must handle all lengths + extractor = named_collections.make_named_collection_extractor(VarLenPointsDataclass) + value = VarLenPointsDataclass( + points=(PointDataclass(1, 2), PointDataclass(3, 4), PointDataclass(5, 6)), offset=7 + ) + assert extractor(value) == (((1, 2), (3, 4), (5, 6)), 7) + assert extractor(VarLenPointsDataclass(points=(PointDataclass(1, 2),), offset=0)) == ( + ((1, 2),), + 0, + ) + + +@dataclasses.dataclass +class NestedVarLenDataclass: + groups: tuple[VarLenPointsDataclass, ...] + + +def test_make_container_extractor_nested_var_len_tuple(): + # nested comprehensions in the generated extractor reuse the variable name, + # relying on comprehension scoping + extractor = named_collections.make_named_collection_extractor(NestedVarLenDataclass) + value = NestedVarLenDataclass( + groups=( + VarLenPointsDataclass(points=(PointDataclass(1, 2), PointDataclass(3, 4)), offset=5), + VarLenPointsDataclass(points=(PointDataclass(6, 7),), offset=8), + ) + ) + assert extractor(value) == (((((1, 2), (3, 4)), 5), (((6, 7),), 8)),) + + @pytest.mark.parametrize( ["nested_tuple", "expected_container"], [(value, cnc.from_nested_tuple(cls, value)) for cls, value in cnc.PYCONTAINERS_SAMPLES.items()], diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py index 14a90409bd..be7f3aa659 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py @@ -117,6 +117,18 @@ def callable_type_info_cases(): returns=float_type, ), ) + var_len_scanop_type = ts_ffront.ScanOperatorType( + axis=KDim, + definition=ts.FunctionType( + pos_only_args=[], + pos_or_kw_args={ + "carry": float_type, + "a": ts.VarLenTupleType(element_type=int_type), + }, + kw_only_args={}, + returns=float_type, + ), + ) return [ # func_type, pos_only_args, kwargs, expected incompatibilities, return type @@ -348,6 +360,56 @@ def callable_type_info_cases(): ], ts.FieldType(dims=[IDim, JDim, KDim], dtype=float_type), ), + # scan operator with a variable-length tuple parameter + ( + # concrete, homogeneous argument + var_len_scanop_type, + [ + ts.TupleType( + types=[ + ts.FieldType(dims=[IDim, JDim], dtype=int_type), + ts.FieldType(dims=[IDim, JDim], dtype=int_type), + ] + ) + ], + {}, + [], + ts.FieldType(dims=[IDim, JDim, KDim], dtype=float_type), + ), + ( + # still generic argument, e.g. a variable-length tuple parameter of a program + var_len_scanop_type, + [ts.VarLenTupleType(element_type=ts.FieldType(dims=[IDim, JDim], dtype=int_type))], + {}, + [], + ts.FieldType(dims=[IDim, JDim, KDim], dtype=float_type), + ), + ( + # heterogeneous argument + var_len_scanop_type, + [ + ts.TupleType( + types=[ + ts.FieldType(dims=[IDim], dtype=int_type), + ts.FieldType(dims=[JDim], dtype=int_type), + ] + ) + ], + {}, + [ + r"Expected argument 'a' to be of type 'tuple\[Field\[\[\.\.\.\], int64\], " + r"\.\.\.\]', got 'tuple\[Field\[\[I\], int64\], Field\[\[J\], int64\]\]'." + ], + ts.FieldType(dims=[IDim, JDim, KDim], dtype=float_type), + ), + ( + # the empty tuple is a valid concretization of a variable-length tuple + var_len_scanop_type, + [ts.TupleType(types=[])], + {}, + [], + ts.FieldType(dims=[KDim], dtype=float_type), + ), ( ts.FunctionType( pos_only_args=[ @@ -464,6 +526,18 @@ def test_needs_value_extraction(type_spec: ts.TypeSpec, expected: bool): assert type_info.needs_value_extraction(type_spec) is expected +def test_is_concretizable_vararg_to_tuple(): + float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) + vararg = ts.VarLenTupleType(element_type=float_type) + + assert type_info.is_concretizable(vararg, ts.TupleType(types=[float_type, float_type])) + assert not type_info.is_concretizable(vararg, ts.TupleType(types=[int_type])) + # The empty tuple is a valid concretization of any variable-length tuple: there are + # no elements whose type could conflict with the element type. + assert type_info.is_concretizable(vararg, ts.TupleType(types=[])) + + def test_promote_lists(): float64 = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) int32 = ts.ScalarType(kind=ts.ScalarKind.INT32) @@ -479,3 +553,17 @@ def test_promote_lists(): type_info.promote(v2e_list, ts.ListType(element_type=int32, offset_type=V2EDim)) with pytest.raises(ValueError, match="non-lists"): type_info.promote(v2e_list, float64) + + +def test_is_concretizable_tuple_with_nested_vararg(): + float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + field_type = ts.FieldType(dims=[], dtype=float_type) + nested = ts.TupleType(types=[ts.VarLenTupleType(element_type=field_type), field_type]) + + assert type_info.is_concretizable( + nested, ts.TupleType(types=[ts.TupleType(types=[field_type, field_type]), field_type]) + ) + assert not type_info.is_concretizable( + nested, ts.TupleType(types=[ts.TupleType(types=[float_type]), field_type]) + ) + assert not type_info.is_concretizable(nested, ts.TupleType(types=[field_type])) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index 08783b6efd..7244545587 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -6,6 +6,7 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import dataclasses import enum import types import typing @@ -210,21 +211,52 @@ def test_make_symbol_type_from_typing(value, expected): assert type_translation.from_type_hint(value) == expected +def test_generic_tuple_constructor_type(): + # The type of the `tuple` builtin itself: `type[tuple]` (e.g. the symbol `tuple` + # referenced inside a field operator) maps to a constructor for a generic tuple. + expected = ts.ConstructorType( + definition=ts.FunctionType( + pos_only_args=[ts.DeferredType(constraint=None)], + pos_or_kw_args={}, + kw_only_args={}, + returns=ts.DeferredType(constraint=ts.VarLenTupleType), + ) + ) + assert ( + type_translation.make_constructor_type(ts.DeferredType(constraint=ts.TupleType)) == expected + ) + assert type_translation.from_type_hint(type[tuple]) == expected + + def test_invalid_symbol_types(): # Forward references with pytest.raises(ValueError, match="undefined forward references"): type_translation.from_type_hint("foo") # Tuples - with pytest.raises(ValueError, match="least one argument"): - type_translation.from_type_hint(typing.Tuple) - with pytest.raises(ValueError, match="least one argument"): - type_translation.from_type_hint(tuple) - - with pytest.raises(ValueError, match="Unbound tuples"): - type_translation.from_type_hint(tuple[int, ...]) - with pytest.raises(ValueError, match="Unbound tuples"): - type_translation.from_type_hint(typing.Tuple["float", ...]) + # Bare `tuple` and `typing.Tuple` (unparameterized) both return a DeferredType. + assert type_translation.from_type_hint(typing.Tuple) == ts.DeferredType(constraint=ts.TupleType) + assert type_translation.from_type_hint(tuple) == ts.DeferredType(constraint=ts.TupleType) + + # The empty-tuple annotation is not a valid type annotation. + with pytest.raises(ValueError, match="Tuple annotation"): + type_translation.from_type_hint(tuple[()]) + with pytest.raises(ValueError, match="Tuple annotation"): + type_translation.from_type_hint(typing.Tuple[()]) + + # Invalid element annotations are reported as ValueError, not internal errors. + with pytest.raises(ValueError, match="invalid element type"): + type_translation.from_type_hint(tuple[tuple, int]) + with pytest.raises(ValueError, match="invalid element type"): + type_translation.from_type_hint(tuple[tuple, ...]) + + # Variadic tuples (`tuple[T, ...]`) are now valid — returns a VarLenTupleType. + assert type_translation.from_type_hint(tuple[int, ...]) == ts.VarLenTupleType( + element_type=ts.ScalarType(kind=ts.ScalarKind.INT64) + ) + assert type_translation.from_type_hint(typing.Tuple["float", ...]) == ts.VarLenTupleType( + element_type=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ) # Fields with pytest.raises(ValueError, match="Field type requires two arguments"): @@ -282,6 +314,23 @@ def test_as_from_dtype(dtype): _FIELD = constructors.empty({cnc.TDim: 1}, gtx.float32) +_FIELD_TYPE = ts.FieldType(dims=[cnc.TDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT32)) + + +class VarLenMemberNamedTuple(typing.NamedTuple): + tracers: tuple[gtx.Field[gtx.Dims[cnc.TDim], gtx.float32], ...] + rho: gtx.Field[gtx.Dims[cnc.TDim], gtx.float32] + + +@dataclasses.dataclass +class VarLenMemberDataclass: + tracers: tuple[gtx.Field[gtx.Dims[cnc.TDim], gtx.float32], ...] + rho: gtx.Field[gtx.Dims[cnc.TDim], gtx.float32] + + +@dataclasses.dataclass +class ScalarMemberDataclass: + value: gtx.float32 @pytest.mark.parametrize( @@ -397,6 +446,57 @@ def test_from_value(value, expected): assert type_translation.from_value(value) == expected +@pytest.mark.parametrize("container", [VarLenMemberNamedTuple, VarLenMemberDataclass]) +def test_from_value_named_collection_var_len_tuple_member(container): + # the declared type keeps the variable-length tuple ... + assert type_translation.from_type_hint(container) == ts.NamedCollectionType( + types=[ts.VarLenTupleType(element_type=_FIELD_TYPE), _FIELD_TYPE], + keys=["tracers", "rho"], + original_python_type=_make_type_string_for_container(container), + ) + # ... while the member types of a value are taken from the runtime members, i.e. + # the tuple has a concrete length + value = container(tracers=(_FIELD, _FIELD), rho=_FIELD) + assert type_translation.from_value(value) == ts.NamedCollectionType( + types=[ts.TupleType(types=[_FIELD_TYPE, _FIELD_TYPE]), _FIELD_TYPE], + keys=["tracers", "rho"], + original_python_type=_make_type_string_for_container(container), + ) + + +def test_from_value_var_len_tuple_of_named_collections(): + # the declared type is a variable-length tuple of named collections ... + assert type_translation.from_type_hint( + tuple[cnc.NamedTupleNamedCollection, ...] + ) == ts.VarLenTupleType(element_type=NamedTupleNamedCollectionTypeSpec) + # ... while the type of a value has a concrete length + value = ( + cnc.NamedTupleNamedCollection(x=_FIELD, y=_FIELD), + cnc.NamedTupleNamedCollection(x=_FIELD, y=_FIELD), + ) + assert type_translation.from_value(value) == ts.TupleType( + types=[NamedTupleNamedCollectionTypeSpec, NamedTupleNamedCollectionTypeSpec] + ) + + +def test_from_value_named_collection_scalar_member(): + assert type_translation.from_value( + ScalarMemberDataclass(value=np.float32(0.0)) + ) == ts.NamedCollectionType( + types=[ts.ScalarType(kind=ts.ScalarKind.FLOAT32)], + keys=["value"], + original_python_type=_make_type_string_for_container(ScalarMemberDataclass), + ) + + +@pytest.mark.skipif(not __debug__, reason="the member compatibility is only checked in debug mode") +def test_from_value_named_collection_incompatible_member(): + # a plain Python 'float' is typed 'float64' by convention, which is incompatible + # with the member annotation 'float32'; note that this is only checked in debug mode + with pytest.raises(ValueError, match="Member 'value' of 'ScalarMemberDataclass'"): + type_translation.from_value(ScalarMemberDataclass(value=0.0)) + + def test_from_value_namespace_proxy(): import next_tests.artifacts.dummy_package as dummy_package