From 212293f5279db609b77e61b8f3790a0c4be90c4d Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 22 Jun 2026 13:33:57 +0200 Subject: [PATCH] =?UTF-8?q?refactor[next]:=20(prototype)=20Approach=203=20?= =?UTF-8?q?=E2=80=94=20PartialTypeSpec=20lazy=20generic=20types?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Experiment for design comparison vs Stage-0 (#2660). NOT for merge. Reverts FieldType.dtype to the concrete ScalarType | ListType and expresses all dtype genericity through a separate PartialTypeSpec wrapper (target type-tag + partial field values, some of which may be TypeVarType). Specialization substitutes the bound vars and materializes a real concrete FieldType, so the IR/backend never observes a TypeVarType and every Stage-0 narrowing assert/cast is removed. partial_type_info.py: bind / substitute / specialize / is_generic over PartialTypeSpec. _approach3_deduction_sketch.py: illustrative deduction-impact sketch only (not imported/wired). --- .../ffront/foast_passes/type_deduction.py | 2 +- .../next/iterator/transforms/global_tmps.py | 12 +- .../runners/dace/lowering/gtir_dataflow.py | 4 +- .../runners/dace/lowering/gtir_to_sdfg.py | 1 - .../lowering/gtir_to_sdfg_concat_where.py | 1 - .../dace/lowering/gtir_to_sdfg_scan.py | 4 +- .../dace/lowering/gtir_to_sdfg_types.py | 8 +- .../_approach3_deduction_sketch.py | 102 ++++++++++ .../next/type_system/partial_type_info.py | 156 ++++++++++++++ src/gt4py/next/type_system/type_info.py | 192 ++---------------- .../next/type_system/type_specifications.py | 33 ++- 11 files changed, 319 insertions(+), 196 deletions(-) create mode 100644 src/gt4py/next/type_system/_approach3_deduction_sketch.py create mode 100644 src/gt4py/next/type_system/partial_type_info.py diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 2b9693cc1f..1b090489ff 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -831,7 +831,7 @@ def _visit_math_built_in(self, node: foast.Call, **kwargs: Any) -> foast.Call: print(f"Warning: return type of '{func_name}' might be inconsistent (not implemented).") # deduce return type - return_type: Optional[ts.FieldType | ts.ScalarType | ts.TypeVarType] = None + return_type: Optional[ts.FieldType | ts.ScalarType] = None if ( func_name in fbuiltins.UNARY_MATH_NUMBER_BUILTIN_NAMES + fbuiltins.UNARY_MATH_FP_BUILTIN_NAMES diff --git a/src/gt4py/next/iterator/transforms/global_tmps.py b/src/gt4py/next/iterator/transforms/global_tmps.py index 7aa8aa6cd5..49d2f3d5b4 100644 --- a/src/gt4py/next/iterator/transforms/global_tmps.py +++ b/src/gt4py/next/iterator/transforms/global_tmps.py @@ -181,16 +181,12 @@ def _transform_by_pattern( lambda x: next(uids["__tmp"]), result_collection_constructor=as_tuple, )(tmp_expr.type) - # the lowered IR is concrete, so `extract_dtype` never yields a `TypeVarType` here tmp_dtypes: ( ts.ScalarType | ts.ListType | tuple[ts.ScalarType | ts.ListType | tuple, ...] - ) = cast( - "ts.ScalarType | ts.ListType | tuple[ts.ScalarType | ts.ListType | tuple, ...]", - type_info.tree_map_type( - type_info.extract_dtype, - result_collection_constructor=as_tuple, - )(tmp_expr.type), - ) + ) = type_info.tree_map_type( + type_info.extract_dtype, + result_collection_constructor=as_tuple, + )(tmp_expr.type) tmp_domains: SymbolicDomain | tuple[SymbolicDomain | tuple, ...] = tmp_expr.annex.domain diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py index d8498b45c3..2e8983ccb8 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py +++ b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py @@ -104,9 +104,7 @@ class MemletExpr: @property def gt_dtype(self) -> ts.ScalarType | ts.ListType: - dtype = self.gt_field.dtype - assert isinstance(dtype, (ts.ScalarType, ts.ListType)) - return dtype + return self.gt_field.dtype def __post_init__(self) -> None: if isinstance(self.gt_dtype, ts.ListType): 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 aacc0ed7bc..223eff2d79 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 @@ -896,7 +896,6 @@ def _add_storage( dc_dtype = gtx_dace_args.as_dace_type(gt_type.dtype) all_dims = gt_type.dims else: # for 'ts.ListType' use 'offset_type' as local dimension - assert isinstance(gt_type.dtype, ts.ListType) assert gt_type.dtype.offset_type is not None assert gt_type.dtype.offset_type.kind == gtx_common.DimensionKind.LOCAL assert isinstance(gt_type.dtype.element_type, ts.ScalarType) diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_concat_where.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_concat_where.py index fde144dedc..ded61af77b 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_concat_where.py +++ b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_concat_where.py @@ -143,7 +143,6 @@ def testee(interior: cases.IJKField, boundary: cases.IJField) -> cases.IJKField: if isinstance(output_type.dtype, ts.ScalarType): all_dims = gtx_common.order_dimensions(output_type.dims) else: - assert isinstance(output_type.dtype, ts.ListType) assert output_type.dtype.offset_type all_dims = gtx_common.order_dimensions([*output_type.dims, output_type.dtype.offset_type]) diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py index da1a813f38..fddcb080b8 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py +++ b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py @@ -622,9 +622,7 @@ def _handle_dataflow_result_of_nested_sdfg( None, dace.Memlet.from_array(outer_dataname, outer_desc), ) - output_dtype = inner_data.gt_type.dtype - assert isinstance(output_dtype, (ts.ScalarType, ts.ListType)) - output_expr = gtir_dataflow.ValueExpr(outer_node, output_dtype) + output_expr = gtir_dataflow.ValueExpr(outer_node, inner_data.gt_type.dtype) return gtir_dataflow.DataflowOutputEdge(outer_ctx.state, output_expr) diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_types.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_types.py index 7f6fc45df8..83cd7660d8 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_types.py +++ b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_types.py @@ -86,11 +86,9 @@ def get_local_view( (dim, dace.symbolic.SymExpr(0) if self.origin is None else self.origin[i]) for i, dim in enumerate(self.gt_type.dims) ] - dtype = self.gt_type.dtype - # `FieldType.dtype` is widened to include `TypeVarType`, but generic operators are - # monomorphized before lowering, so only concrete dtypes reach here. - assert isinstance(dtype, (ts.ScalarType, ts.ListType)) - return gtir_dataflow.IteratorExpr(self.dc_node, dtype, field_origin, it_indices) + return gtir_dataflow.IteratorExpr( + self.dc_node, self.gt_type.dtype, field_origin, it_indices + ) raise NotImplementedError(f"Node type {type(self.gt_type)} not supported.") diff --git a/src/gt4py/next/type_system/_approach3_deduction_sketch.py b/src/gt4py/next/type_system/_approach3_deduction_sketch.py new file mode 100644 index 0000000000..f55e7db146 --- /dev/null +++ b/src/gt4py/next/type_system/_approach3_deduction_sketch.py @@ -0,0 +1,102 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +"""PROTOTYPE-ONLY: the FOAST type-deduction impact of Approach 3 (`PartialTypeSpec`). + +This is *not* wired into the toolchain. It is a faithful, compilable sketch of how the three +representative deduction operations in `ffront/foast_passes/type_deduction.py` would have to grow +a `PartialTypeSpec` branch, so the frontend invasiveness can be judged from real code instead of +prose. + +The contrast with Stage-0 (TypeVar-inside-FieldType): there, a generic field IS a `FieldType` +(with `dtype=TypeVarType`), so `case ts.FieldType(dims, dtype)` and `ts.FieldType(...)` +construction "just work" and the *predicates* (`promote`, `extract_dtype`, ...) were taught about +`TypeVarType` once, centrally. With `PartialTypeSpec` a generic field is a different object, so +every site that *matches* or *constructs* a `FieldType` in a path a generic type can reach needs a +parallel branch. +""" + +from __future__ import annotations + +from typing import Optional + +from gt4py.next.type_system import partial_type_info as pti, type_specifications as ts + + +# --- representative site 1: `with_altered_scalar_kind` (ffront/foast_passes/type_deduction.py) --- +# Stage-0 today (works unchanged on a generic field, because dtype just happens to be a TypeVarType): +# +# if isinstance(type_spec, ts.FieldType): +# return ts.FieldType(dims=type_spec.dims, +# dtype=with_altered_scalar_kind(type_spec.dtype, new_scalar_kind)) +# +# Approach 3: the predicate flips the dtype to bool, which is concrete, so the result is a concrete +# field even though the input was generic. But the *input* is now a PartialTypeSpec, so we must +# match it explicitly and reach into `fields`: +def with_altered_scalar_kind( + type_spec: ts.TypeSpec, new_scalar_kind: ts.ScalarKind +) -> ts.ScalarType | ts.FieldType: + if isinstance(type_spec, ts.PartialTypeSpec) and type_spec.target is ts.FieldType: + fields = dict(type_spec.fields) + # a generic dtype becomes a concrete `bool` scalar here -> a concrete FieldType + return ts.FieldType(dims=fields["dims"], dtype=ts.ScalarType(kind=new_scalar_kind)) + if isinstance(type_spec, ts.FieldType): + return ts.FieldType(dims=type_spec.dims, dtype=ts.ScalarType(kind=new_scalar_kind)) + if isinstance(type_spec, ts.ScalarType): + return ts.ScalarType(kind=new_scalar_kind, shape=type_spec.shape) + raise ValueError(f"Expected field or scalar type, got '{type_spec}'.") + + +# --- representative site 2: binary-op promotion (`_deduce_binop_type` -> type_info.promote) --- +# Stage-0 today: `type_info.promote(left.type, right.type)` — promote() was taught that +# `promote(T, T) = T` and that mixing T with a concrete dtype is the D3 strict-no-promotion error. +# +# Approach 3: `type_info.promote` is concrete-only again. Generic promotion is a *separate* code +# path keyed on "either operand is a PartialTypeSpec". The D3 rule lives here now, not in promote. +def promote_for_binop(left: ts.TypeSpec, right: ts.TypeSpec) -> ts.TypeSpec: + l_gen, r_gen = isinstance(left, ts.PartialTypeSpec), isinstance(right, ts.PartialTypeSpec) + if l_gen or r_gen: + if not (l_gen and r_gen): + # D3: a generic dtype combined with a concrete one is a (decoration-time) error. + raise ValueError( + f"Could not promote '{left}' and '{right}': a generic dtype can only be" + " combined with the same type variable." + ) + l_tvars, r_tvars = pti.type_vars(left), pti.type_vars(right) + if set(l_tvars) != set(r_tvars): + raise ValueError(f"Could not promote distinct type variables '{left}' and '{right}'.") + # `promote(T, T) = T`: dims still promote like concrete fields; here we just return one. + # (A real impl would promote the `dims` field and keep the shared TypeVar dtype.) + return left + # both concrete -> delegate to the untouched concrete promote + from gt4py.next.type_system import type_info + + return type_info.promote(left, right) # type: ignore[arg-type] # concrete branch + + +# --- representative site 3: subscript / dimension drop (`visit_Subscript`) --- +# Stage-0 today: +# case ts.FieldType(dims=dims, dtype=dtype): +# new_type = ts.FieldType(dims=[d for d in dims if d != idx.dim], dtype=dtype) +# +# Approach 3: add a parallel `PartialTypeSpec` arm that rebuilds the partial with the same dtype +# (still a TypeVar) but reduced dims: +def subscript_drop_dim(value_type: ts.TypeSpec, dropped: object) -> Optional[ts.TypeSpec]: + match value_type: + case ts.PartialTypeSpec(target=target, fields=fields) if target is ts.FieldType: + f = dict(fields) + # `PartialField = object`, so the field value is untyped and must be re-narrowed. + dims = f["dims"] + assert isinstance(dims, list) + return ts.PartialTypeSpec( + target=ts.FieldType, + fields=(("dims", [d for d in dims if d != dropped]), ("dtype", f["dtype"])), + ) + case ts.FieldType(dims=dims, dtype=dtype): + return ts.FieldType(dims=[d for d in dims if d != dropped], dtype=dtype) + return None diff --git a/src/gt4py/next/type_system/partial_type_info.py b/src/gt4py/next/type_system/partial_type_info.py new file mode 100644 index 0000000000..3f358ad34e --- /dev/null +++ b/src/gt4py/next/type_system/partial_type_info.py @@ -0,0 +1,156 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +"""Frontend-only utilities for `ts.PartialTypeSpec` (Approach 3 prototype). + +Genericity is expressed by wrapping a target `TypeSpec` class together with its constructor +fields, where some fields are `TypeVarType` placeholders. None of this is visible downstream of +specialization, so the IR/backend never has to narrow `FieldType.dtype`. +""" + +from __future__ import annotations + +from typing import Any, Sequence + +from gt4py.eve import extended_typing as xtyping +from gt4py.next.type_system import type_specifications as ts + + +def _walk_typevars(value: Any) -> xtyping.Iterator[ts.TypeVarType]: + """Yield every `TypeVarType` reachable inside a (possibly partial) field value.""" + if isinstance(value, ts.TypeVarType): + yield value + elif isinstance(value, ts.PartialTypeSpec): + for _, v in value.fields: + yield from _walk_typevars(v) + elif isinstance(value, (list, tuple)): + for v in value: + yield from _walk_typevars(v) + elif isinstance(value, ts.TupleType): + for v in value.types: + yield from _walk_typevars(v) + elif isinstance(value, ts.FieldType): + yield from _walk_typevars(value.dtype) + elif isinstance(value, ts.ListType): + yield from _walk_typevars(value.element_type) + + +def is_generic(type_: Any) -> bool: + """A `PartialTypeSpec` (or a structure containing one / a bare `TypeVarType`) is generic.""" + return next(_walk_typevars(type_), None) is not None + + +def type_vars(type_: Any) -> dict[str, ts.TypeVarType]: + """Collect the type variables occurring in ``type_``, keyed by name.""" + return {var.name: var for var in _walk_typevars(type_)} + + +def _substitute_value(value: Any, binding: xtyping.Mapping[str, ts.ScalarType]) -> Any: + """Replace bound type variables inside a single field value, recursing into nested partials.""" + if isinstance(value, ts.TypeVarType): + return binding.get(value.name, value) + if isinstance(value, ts.PartialTypeSpec): + return substitute(value, binding) + if isinstance(value, list): + return [_substitute_value(v, binding) for v in value] + if isinstance(value, tuple): + return tuple(_substitute_value(v, binding) for v in value) + return value + + +def substitute( + partial: ts.PartialTypeSpec, binding: xtyping.Mapping[str, ts.ScalarType] +) -> ts.PartialTypeSpec: + """Return a `PartialTypeSpec` with the bound type variables substituted (still possibly partial).""" + return ts.PartialTypeSpec( + target=partial.target, + fields=tuple((name, _substitute_value(value, binding)) for name, value in partial.fields), + ) + + +def _materialize(value: Any) -> Any: + """Turn any nested `PartialTypeSpec` into its concrete `target`, recursively. + + Precondition: ``value`` contains no unbound type variables. + """ + if isinstance(value, ts.PartialTypeSpec): + return value.target(**{name: _materialize(v) for name, v in value.fields}) + if isinstance(value, list): + return [_materialize(v) for v in value] + if isinstance(value, tuple): + return tuple(_materialize(v) for v in value) + return value + + +def specialize(partial: ts.PartialTypeSpec, binding: xtyping.Mapping[str, ts.ScalarType]) -> ts.TypeSpec: + """Substitute all type variables and instantiate the concrete ``target`` `TypeSpec`. + + Raises if a type variable remains unbound after substitution (the partial is still generic). + Nested `PartialTypeSpec`s are materialized too, so the result is fully concrete. + """ + substituted = substitute(partial, binding) + if is_generic(substituted): + unbound = ", ".join(sorted(type_vars(substituted))) + raise ValueError(f"Cannot specialize '{partial}': unbound type variable(s) {unbound}.") + result = _materialize(substituted) + assert isinstance(result, ts.TypeSpec) + return result + + +def _bind_var(var: ts.TypeVarType, dtype: ts.TypeSpec) -> dict[str, ts.ScalarType]: + if not isinstance(dtype, ts.ScalarType): + return {} + if dtype not in var.constraints: + raise ValueError(f"'{dtype}' does not satisfy the constraints of type variable '{var}'.") + return {var.name: dtype} + + +def _merge(parts: xtyping.Iterable[dict[str, ts.ScalarType]]) -> dict[str, ts.ScalarType]: + binding: dict[str, ts.ScalarType] = {} + for part in parts: + for name, dtype in part.items(): + if (prev := binding.get(name)) is not None and prev != dtype: + raise ValueError( + f"Type variable '{name}' is bound inconsistently: '{prev}' and '{dtype}'." + ) + binding[name] = dtype + return binding + + +def _bind_value(param_value: Any, arg: ts.TypeSpec) -> dict[str, ts.ScalarType]: + """Bind type variables in one partial field value against a concrete argument type.""" + if isinstance(param_value, ts.TypeVarType): + return _bind_var(param_value, arg) + if isinstance(param_value, ts.PartialTypeSpec): + return bind(param_value, arg) + return {} + + +def bind(partial: ts.PartialTypeSpec, arg: ts.TypeSpec) -> dict[str, ts.ScalarType]: + """Bind the type variables in ``partial`` by structurally matching the concrete ``arg``. + + Only the dtype field participates today (dtype-scoped, per ADR 0024); generalizing to other + fields (e.g. dims) is a matter of widening the structural match below. + """ + if partial.target is ts.FieldType: + dtype_value = dict(partial.fields).get("dtype") + # scalar arguments are promoted to zero-dimensional fields + arg_dtype = arg.dtype if isinstance(arg, ts.FieldType) else arg + return _bind_value(dtype_value, arg_dtype) + return {} + + +def bind_many( + params: Sequence[ts.TypeSpec], args: Sequence[ts.TypeSpec] +) -> dict[str, ts.ScalarType]: + """Compute the binding for a full parameter/argument list (mixed partial and concrete).""" + return _merge( + bind(param, arg) + for param, arg in zip(params, args) + if isinstance(param, ts.PartialTypeSpec) + ) diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index ab6c61fe2c..5ef3a49e48 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -85,8 +85,9 @@ def is_generic(symbol_type: ts.TypeSpec) -> bool: Figure out if a type contains parts that are only known when concrete arguments are given. Recurses into composite types, reporting ``True`` if any nested part is a `DeferredType` or - `TypeVarType`. Unlike :func:`is_concrete` (a shallow top-level check), this is deep, so a - tuple with a nested `DeferredType` is both concrete and generic. + a `PartialTypeSpec` (a dtype-generic frontend type). Unlike :func:`is_concrete` (a shallow + top-level check), this is deep, so a tuple with a nested `DeferredType` is both concrete and + generic. Note: this returns ``True`` for a bare ``astype`` constructor type, whose ``definition`` carries a ``DeferredType`` by design; callers that only care about *data* arguments must @@ -103,7 +104,7 @@ def is_generic(symbol_type: ts.TypeSpec) -> bool: >>> is_generic(ts.TupleType(types=[bool_type, ts.DeferredType(constraint=None)])) True """ - if isinstance(symbol_type, (ts.DeferredType, ts.TypeVarType)): + if isinstance(symbol_type, (ts.DeferredType, ts.PartialTypeSpec)): return True return any(is_generic(p) for p in _type_params(symbol_type)) @@ -235,9 +236,9 @@ def tree_map_type( ) -def extract_dtype(symbol_type: ts.TypeSpec) -> ts.ScalarType | ts.ListType | ts.TypeVarType: +def extract_dtype(symbol_type: ts.TypeSpec) -> ts.ScalarType | ts.ListType: """ - Extract the data type from ``symbol_type`` if it is `FieldType`, `ScalarType` or `TypeVarType`. + Extract the data type from ``symbol_type`` if it is a `FieldType` or `ScalarType`. Raise an error if no dtype can be found or the result would be ambiguous. @@ -255,8 +256,6 @@ def extract_dtype(symbol_type: ts.TypeSpec) -> ts.ScalarType | ts.ListType | ts. return dtype case ts.ScalarType() as dtype: return dtype - case ts.TypeVarType() as dtype: - return dtype raise ValueError(f"Can not unambiguosly extract data type from '{symbol_type}'.") @@ -271,17 +270,10 @@ def _scalar_kinds(scalar_types: tuple[type, ...]) -> frozenset[ts.ScalarKind]: def _is_field_or_scalar_of_kind(symbol_type: ts.TypeSpec, kinds: Collection[ts.ScalarKind]) -> bool: - """Check if ``symbol_type`` is a scalar or a field whose dtype kind is in ``kinds``. - - A type variable has the property iff all of its constraints have it. - """ - if isinstance(symbol_type, ts.TypeVarType): - return all(_is_field_or_scalar_of_kind(c, kinds) for c in symbol_type.constraints) + """Check if ``symbol_type`` is a scalar or a field whose dtype kind is in ``kinds``.""" if not isinstance(symbol_type, (ts.ScalarType, ts.FieldType)): return False dtype = extract_dtype(symbol_type) - if isinstance(dtype, ts.TypeVarType): - return all(_is_field_or_scalar_of_kind(c, kinds) for c in dtype.constraints) return isinstance(dtype, ts.ScalarType) and dtype.kind in kinds @@ -352,7 +344,7 @@ def is_arithmetic_scalar(symbol_type: ts.TypeSpec) -> bool: ... ) False """ - if not isinstance(symbol_type, (ts.ScalarType, ts.TypeVarType)): + if not isinstance(symbol_type, ts.ScalarType): return False return is_arithmetic(symbol_type) @@ -386,15 +378,6 @@ def is_arithmetic(symbol_type: ts.TypeSpec) -> bool: >>> is_arithmetic(ts.FieldType(dims=[], dtype=ts.ScalarType(kind=ts.ScalarKind.INT32))) True """ - # `is_arithmetic` cannot reuse `_is_field_or_scalar_of_kind`'s "all constraints - # share the kind" rule: a type variable is arithmetic if every constraint is - # arithmetic, even when the constraints mix floating point and integral kinds. - if isinstance(symbol_type, ts.TypeVarType): - return all(is_arithmetic(c) for c in symbol_type.constraints) - if isinstance(symbol_type, (ts.ScalarType, ts.FieldType)) and isinstance( - dtype := extract_dtype(symbol_type), ts.TypeVarType - ): - return is_arithmetic(dtype) return is_floating_point(symbol_type) or is_integral(symbol_type) @@ -470,7 +453,7 @@ def extract_dims(symbol_type: ts.TypeSpec) -> list[common.Dimension]: >>> extract_dims(ts.FieldType(dims=[I, J], dtype=ts.ScalarType(kind=ts.ScalarKind.INT64))) [Dimension(value='I', kind=), Dimension(value='J', kind=)] """ - if isinstance(symbol_type, (ts.ScalarType, ts.TypeVarType)): + if isinstance(symbol_type, ts.ScalarType): return [] if isinstance(symbol_type, ts.FieldType): return symbol_type.dims @@ -630,139 +613,14 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool: return False -def _bind_var(var: ts.TypeVarType, dtype: ts.TypeSpec) -> dict[str, ts.ScalarType]: - if not isinstance(dtype, ts.ScalarType): - # not a concrete scalar to bind to -- e.g. a `TypeVarType` (operator-from-operator - # call), a `DeferredType` (scan), or a `ListType` (local field). Leave it unbound; - # the caller is responsible for checking that no type variable remained unbound. - return {} - if dtype not in var.constraints: - raise ValueError(f"'{dtype}' does not satisfy the constraints of type variable '{var}'.") - return {var.name: dtype} - - -def _merge_bindings(parts: Iterable[dict[str, ts.ScalarType]]) -> dict[str, ts.ScalarType]: - binding: dict[str, ts.ScalarType] = {} - for part in parts: - for name, dtype in part.items(): - if (previous := binding.get(name)) is not None and previous != dtype: - raise ValueError( - f"Type variable '{name}' is bound inconsistently:" - f" '{previous}' and '{dtype}' (all arguments using '{name}'" - " must have the same dtype)." - ) - binding[name] = dtype - return binding - - -def _bind(param: ts.TypeSpec, arg: ts.TypeSpec) -> dict[str, ts.ScalarType]: - match param: - case ts.TypeVarType() as var: - return _bind_var(var, arg) - case ts.FieldType(dtype=ts.TypeVarType() as var): - # scalar arguments are promoted to zero-dimensional fields - return _bind_var(var, arg.dtype if isinstance(arg, ts.FieldType) else arg) - case ts.ListType(element_type=element_type) if isinstance(arg, ts.ListType): - return _bind(element_type, arg.element_type) - case ts.TupleType() | ts.NamedCollectionType() if isinstance( - arg, (ts.TupleType, ts.NamedCollectionType) - ): - # tolerant by design: a structural mismatch (e.g. tuple vs scalar) binds nothing - # here and is reported by the regular signature checks instead. - return _merge_bindings(_bind(p, a) for p, a in zip(param.types, arg.types)) - return {} - - -def bind_type_vars( - params: Sequence[ts.TypeSpec], args: Sequence[ts.TypeSpec] -) -> dict[str, ts.ScalarType]: - """ - Compute a binding of all type variables in ``params`` by structurally matching ``args``. - - Concrete (non-generic) parts of the parameters are ignored; a type variable position binds - only if the corresponding argument provides a concrete scalar dtype. The caller is - responsible for checking that no type variable remained unbound. - - Raises: - ValueError: If a type variable would be bound inconsistently or to a dtype that is - not one of its constraints. - - Examples: - >>> var = ts.TypeVarType(name="T", constraints=(ts.ScalarType(kind=ts.ScalarKind.FLOAT64),)) - >>> I = common.Dimension(value="I") - >>> binding = bind_type_vars( - ... [ts.FieldType(dims=[I], dtype=var)], - ... [ts.FieldType(dims=[I], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))], - ... ) - >>> print(binding["T"]) - float64 - """ - return _merge_bindings(_bind(param, arg) for param, arg in zip(params, args)) - - -def tree_map_type_params( - fun: Callable[[ts.TypeSpec], ts.TypeSpec], symbol_type: ts.TypeSpec -) -> ts.TypeSpec: - """Rebuild ``symbol_type`` applying ``fun`` to each immediate type-parameter sub-type. - - Leaf types are returned unchanged. - """ - match symbol_type: - case ts.FieldType(dims=dims, dtype=dtype): - new_dtype = fun(dtype) - assert isinstance(new_dtype, (ts.ScalarType, ts.ListType, ts.TypeVarType)) - return ts.FieldType(dims=dims, dtype=new_dtype) - case ts.ListType(element_type=element_type, offset_type=offset_type): - new_element_type = fun(element_type) - assert isinstance(new_element_type, ts.DataType) - return ts.ListType(element_type=new_element_type, offset_type=offset_type) - case ts.TupleType(types=types): - return ts.TupleType(types=[fun(t) for t in types]) - case ts.NamedCollectionType(types=types): - return ts.NamedCollectionType( - types=[fun(t) for t in types], - keys=symbol_type.keys, - original_python_type=symbol_type.original_python_type, - ) - case ts.FunctionType(): - return ts.FunctionType( - pos_only_args=[fun(t) for t in symbol_type.pos_only_args], - pos_or_kw_args={name: fun(t) for name, t in symbol_type.pos_or_kw_args.items()}, - kw_only_args={name: fun(t) for name, t in symbol_type.kw_only_args.items()}, - returns=fun(symbol_type.returns), - ) - return symbol_type - - -def substitute_type_vars( - type_: ts.TypeSpec, binding: xtyping.Mapping[str, ts.ScalarType] -) -> ts.TypeSpec: - """ - Replace all type variables in ``type_`` that are bound in ``binding``. - - Unbound type variables and all other generic parts (e.g. `DeferredType`) are kept as-is. - - Examples: - >>> var = ts.TypeVarType(name="T", constraints=(ts.ScalarType(kind=ts.ScalarKind.FLOAT64),)) - >>> I = common.Dimension(value="I") - >>> print( - ... substitute_type_vars( - ... ts.FieldType(dims=[I], dtype=var), - ... {"T": ts.ScalarType(kind=ts.ScalarKind.FLOAT64)}, - ... ) - ... ) - Field[[I], float64] - """ - if not binding or not is_generic(type_): - return type_ - if isinstance(type_, ts.TypeVarType): - return binding.get(type_.name, type_) - return tree_map_type_params(lambda t: substitute_type_vars(t, binding), type_) +# Binding / substitution / specialization of dtype-generic types lives in +# `gt4py.next.type_system.partial_type_info` (Approach 3): they operate on `ts.PartialTypeSpec`, +# never on a `FieldType` with an embedded type variable. def promote( - *types: ts.FieldType | ts.ScalarType | ts.TypeVarType, always_field: bool = False -) -> ts.FieldType | ts.ScalarType | ts.TypeVarType: + *types: ts.FieldType | ts.ScalarType, always_field: bool = False +) -> ts.FieldType | ts.ScalarType: """ Promote a set of field or scalar types to a common type. @@ -784,29 +642,17 @@ def promote( >>> promoted.dims == [I, J, K] and promoted.dtype == dtype True """ - if not always_field and all( - isinstance(type_, (ts.ScalarType, ts.TypeVarType)) for type_ in types - ): + if not always_field and all(isinstance(type_, ts.ScalarType) for type_ in types): if not all(type_ == types[0] for type_ in types): - if any(isinstance(type_, ts.TypeVarType) for type_ in types): - distinct_types = "', '".join(str(t) for t in dict.fromkeys(types)) - raise ValueError( - f"Could not promote '{distinct_types}': a generic dtype (type variable)" - " can only be combined with values of the same type variable," - " not with other dtypes." - ) raise ValueError("Could not promote scalars of different dtype (not implemented).") - if not all(type_.shape is None for type_ in types if isinstance(type_, ts.ScalarType)): + if not all(type_.shape is None for type_ in types): # type: ignore[union-attr] raise NotImplementedError("Shape promotion not implemented.") return types[0] - elif all(isinstance(type_, (ts.ScalarType, ts.FieldType, ts.TypeVarType)) for type_ in types): + elif all(isinstance(type_, (ts.ScalarType, ts.FieldType)) for type_ in types): dims = common.promote_dims(*(extract_dims(type_) for type_ in types)) extracted_dtypes = [extract_dtype(type_) for type_ in types] - assert all(isinstance(dtype, (ts.ScalarType, ts.TypeVarType)) for dtype in extracted_dtypes) - dtype = cast( # type variables promote like scalars (only with themselves) - ts.ScalarType | ts.TypeVarType, - promote(*extracted_dtypes), # type: ignore[arg-type] # checked above - ) + assert all(isinstance(dtype, ts.ScalarType) for dtype in extracted_dtypes) + dtype = cast(ts.ScalarType, promote(*extracted_dtypes)) # type: ignore[arg-type] # checked is `ScalarType` return ts.FieldType(dims=dims, dtype=dtype) raise TypeError("Expected a 'FieldType' or 'ScalarType'.") diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 8450bbb7d6..46453f1b68 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -151,7 +151,7 @@ class ListType(DataType): class FieldType(DataType, CallableType): dims: list[common.Dimension] - dtype: ScalarType | ListType | TypeVarType + dtype: ScalarType | ListType def __str__(self) -> str: dims = "..." if self.dims is Ellipsis else f"[{', '.join(dim.value for dim in self.dims)}]" @@ -260,3 +260,34 @@ def constructed_type(self) -> TypeSpec: class DomainType(DataType): dims: list[common.Dimension] + + +#: A field value of a `PartialTypeSpec`: either a concrete type-system value (the same kinds of +#: values that occur as constructor arguments of the wrapped `TypeSpec`) or a `TypeVarType` +#: standing in for a not-yet-bound dtype. Modeled as ``object`` because the constructor fields +#: of the wrapped classes are heterogeneous (e.g. ``dims: list[Dimension]``, ``dtype: ScalarType``). +PartialField = object + + +class PartialTypeSpec(TypeSpec): + """A type-system value with some constructor fields left as type variables. + + Frontend-only representation of a generic type. It records *which* concrete `TypeSpec` + class it will become (`target`) and the constructor field values (`fields`), some of which + may be `TypeVarType` placeholders instead of concrete values. `specialize` substitutes the + bound type variables and instantiates the real concrete `target`, so nothing downstream of + the frontend ever sees a `PartialTypeSpec`. + + As a frozen eve `DataModel` it inherits structural ``__eq__`` and ``content_hash`` (the same + machinery that keys the compilation cache for `FunctionType`), so it is a stable cache key + despite holding a mapping. + """ + + target: type[TypeSpec] + #: ``(field_name, value)`` pairs, in constructor order. A value is either concrete or a + #: `TypeVarType`. A tuple (not a dict) keeps ordering deterministic for hashing/equality. + fields: tuple[tuple[str, PartialField], ...] + + def __str__(self) -> str: + body = ", ".join(f"{name}={value}" for name, value in self.fields) + return f"Partial[{self.target.__name__}]({body})"