# [GT4Py] Prepare ITIR type inference for extensions: replace constraint-based with type propagation

- Shaped by: Till
- Appetite (FTEs, weeks): full cycle
- Developers:

This project has already been proposed in cycle 17 and is just copied here with minor changes.

## Problem

The current type inference algorithm employed on ITIR is constraint based. Fundamentally, constraints are formulated between the type of the arguments of a function and its result. For example for the `plus` function the type of its two arguments must be equal to its return values. If additionally the types of the parameters of all top-level (i.e. not nested in another function) functions is specified, the types of all expressions can be fully inferred. The advantage of constraint based type inference algorithms is that types can propagate bi-directionally and expressions can naturally be partially typed (i.e. generic). In the forward direction: If the types of the arguments of a function are constrained (e.g. fixed to be a specific type) so is its return type. In the backward direction: If the return type of a function is constrained so are its arguments. Note however that while types can propagate both-ways the constraints don't have any directionality. A "simple" and also the employed algorithm in GT4Py traverses an ITIR expression, collects all constraints between subexpressions and then unifies them. This separation into two phases has the big disadvantage that in the unification phase we don't know anymore were the constraints came from. A small bug in the inference rules used in the first phase results in a cryptic error in the unification phase at which point the faulty rule is not known anymore. As a result debugging the type inference in its current form is far from practical and has lead to many long "bug hunts".

## Appetite

1 cycle

This task can be tackled by someone with good knowledge of the FOAST type inference which serves as a reference for this project. Knowledge of ITIR is not necessarily required, with few exceptions (e.g. columns, see below) the typing rules are relatively easy.

## Solution

Implement a new type inference that completely replaces the constraint based type inference.

<!--### Current state

- When [FOAST is lowered to ITIR](https://github.com/GridTools/gt4py/blob/8d91a7b047f5d3af57ac3c1407a412e69f6cea89/src/gt4py/next/ffront/foast_to_itir.py#L211) the types of all operator arguments is passed to the respective parameters (i.e. `itir.Sym` nodes) of all function arguments.
TODO: Example
- When starting from the embedded implementation of ITIR the type of all arguments are deduced from the arguments and-->

### Steps

- Introduce a new type specification for Iterators <!--`Iterator[current_location, defined_location, dtype]`-->. Sketch:

  class IteratorType(...):
    current_position: list[Dimension]
    defined_location: list[Dimension]
    dtype: ScalarType | TupleType  # what about column; check that tuple consists of scalars

- Modify `itir.Sym`, `itir.Expr` node to store a type specification, e.g. `Iterator` or `ScalarType`.
- Analogous to the [FOAST type inference](https://github.com/GridTools/gt4py/blob/main/src/gt4py/next/ffront/foast_passes/type_deduction.py) develop a pass that given an ITIR node deduces the type of all `itir.Expr` child nodes and stores the deduced type in the node.
- Replace all usages of the constraint based type inference and use the information stored in `itir.Sym`. There are two places were the type inferrence is used right now:
- [CollapseTuple pass](https://github.com/GridTools/gt4py/blob/8d91a7b047f5d3af57ac3c1407a412e69f6cea89/src/gt4py/next/iterator/transforms/collapse_tuple.py#L22)
- [Temporary extraction pass](https://github.com/GridTools/gt4py/blob/43eecfbea3aa1864176aa4204c584b73c408fae2/src/gt4py/next/iterator/transforms/global_tmps.py#L538)
- Develop unit tests starting from ITIR
- Wire up the field view and ITIR frontends
- Pass type information in [lowering from FOAST to ITIR](https://github.com/GridTools/gt4py/blob/8d91a7b047f5d3af57ac3c1407a412e69f6cea89/src/gt4py/next/ffront/foast_to_itir.py#L211)
- Pass type information when calling a traced stencil (TODO: link)

<!--Requirements:

- Supports getting tuple size
- Supports dtype on expr-->

Approach: Start from scratch and don't port the existing type inference.

## Rabbit holes
<!-- Details about the solution worth calling out to avoid problems -->

The typing rules for the scan involves columns and are not well documented. It might be possible to avoid this complication, but otherwise someone with good knowledge of ITIR might needed to figure out the exact rules.

## No-gos

- We are already aware of design aspects of the FOAST type inference and type system to improve. This is not part of the project, we explicitly want something that is analogous and can be improved in conjunction with the FOAST parts later on.
- We don't need to support every feature of the ITIR type inference. Only `itir.Expr`, `itir.Sym` should be supported, we don't need a type for other nodes, e.g. an `itir.FencilDefinition`. Note that this doesn't mean that that other nodes are ignored. For example the parameters of a fencil are certainly important. Nonetheless we don't need to introduce a new type for a fencil, since this is not used anywhere.

## Scratchpad

as_fieldop(λ(it) → ·⟪V2Eₒ, 0ₒ⟫(it), u⟨ Vertex: [0, 1), KDim: [0, 1) ⟩)(inp)

apply_stencil( λ(it)->deref(shift(V2E, 0)(it)), [V,K] ) ==> premap(V2E[0], f:[E,K])

apply_stencil( λ(it)->(λ()->deref(shift(V2E, 0)(it)))(), [V,K] ) ==> premap(V2E[0], f:[E,K])