# Type Inference for the Functional IR ###### tags: `study` --- ## Recap: Constraint-based Type Inference ---- Constraint-based type inference is a two-pass approach: 1. Collect types where unkowns are filled with type variables; enforce type relations using constraints. 2. Resolve all collected constraints. This guarantees to derive the _most general type_ of an expression if it exists and fail otherwise. ---- ### What’s the Most General Type? ---- For the identity function $λ(x).x$ the most general type is $T → T$, where $T$ is a _type variable_. Type variables can be replaced by _any_ valid type and the expression will still be valid. ---- ### What Are the Tricky Parts in Our IR? ---- First, we have overloads: arithmetic operations work on integers, floats, columns of integers, and columns of floats. Second, some overloads are pretty strange: for example, `lift` takes any function that takes an arbitrary number of iterators (with any value type) and returns a value. Note that the accepted (stencil) function itself can be polymorphic and either work on scalars or columns. Third, there is `tuple_get`, which is overloaded for tuples of arbitrary sizes. So what’s the type of $λ(x).\mathrm{tuple\_get}(1, x)$? --- ## Examples of the New Type Inference Implementation ---- The identity function: $λ(x).x :: (T_0) → T_0$ Functions as arguments: $λ(x, y).x(y) :: ((T_0) → T_1, T_0) → T_1$ ---- Let’s look at a function (stencil!) calling _deref_: $λ(x).\mathrm{deref}(x) :: (\mathrm{It}[T_0^1]) → T_0^1$ So it takes an iterator of some type and returns a value of the same. But what’s the superscript ‘1’ telling us? It’s a variable for the polymorphic size! This is the first IR-specific feature. ---- And what about just _deref_? It’s the same as above (η-reduction): $\mathrm{deref} :: (\mathrm{It}[T_0^1]) → T_0^1$ So, all builtins are handled as real functions (passable to other functions) apart from _make_tuple_ and _tuple_get_ (which are part of the grammar). ---- Now let’s check a more complex thing: _lift_. This is what the inference algorithm gives: $\mathrm{lift} :: ((\mathrm{It}[T^1], …)_0 → T_2^1) → (\mathrm{It}[T^1], …)_0 → \mathrm{It}[T_2^1]$ Note the special argument type: it’s a tuple of an arbitrary number of iterators with arbitrary value types but all of equal size (column or scalar). The subscript ‘0’ shows that the argument types of the lifted stencil have to match the argument types of the input stencil. ---- A closer look at the superscripts on the example of scan: $$ (T_0^s:(\mathrm{It}[T^s], …)_1 → T_0^s, \mathrm{bool}^s, T_0^s) \\ → (\mathrm{It}[T^c], …)_1 → T_0^c $$ ---- Also quite fancy: $$ λ(x).\mathrm{tuple\_get}(1, x) :: \\ (\mathrm{ItOrVal}_0[(…, T_1, …)_2^3]) → \mathrm{ItOrVal}_0[T_1^3] $$ --- ## Probably Important Details ---- The basis of everything is a type (represented by `Val` in the source) that has three parameters: - A _kind_: either _iterator_, _value_, or unknown (a variable). - A _dtype_: a primitive data type (_float_, etc.) or unknown. - A _size_: either _column_, _scalar_, or unknown. ---- All builtins work on arguments of this type and return arguments of this type. The $\mathrm{It}[T_i^j]$, $T_i^j$, and $\mathrm{ItOrVal}_k[T_i^j]$ in all previous examples actually refer to instances of this type with more or less restrictions (e.g. take only iterators). ---- Also, _make_tuple_ and _tuple_get_ work on this type. This has some implications: for examples, you can not put functions into tuples as functions are of a different type. ---- This has some clear advantages. For example, one can call `some_stencil(make_tuple(a, b))`, where `a` and `b` are iterators. Additionally, on the callee-side `deref(tuple_get(0, x))` is equivalent to `tuple_get(0, deref(x))`. So basically there is no trouble with tuple-of-iterators vs. iterators-of-tuples. --- ## Limitations ---- Error reporting using constraint-based typing is a bit tricky. As the unification happens _after_ collecting all constraints, there is some effort required to track down the source of the error. And of course the current implementation does not invest any effort. ---- The current implementation interprets the function definitions as Scheme/Lisp `let*`-style statements; that is, functions can only call other functions that are defined before. This avoids recursion naturally, but requires us to either standardize that order in the IR (which probably makes sense) or just sort the definitions before type checking. ---- The current implementation has non-optimal run-time complexity ($O(n^2)$). There seem to be fancy tricks to make this faster. But not idea at which size of a program this starts to matter. Good implementations seem to be close to $O(n)$ in code size in practice. ---- Further, the current implementation does internally not distinguish variable types. That is, for example, unkown primitive types are represented by the same type variables as unknown kind or size values. This leads to more compact code and simpler type unification but of course allows for more general types than we might actually want. However, this could easily be changed.