# Colored Unstructured ## Idea Let us assume: - We have a fencil where unstructured connectivities are used to define the stencils. I.e. in terms of iterator model `shift` builtin is called with the tags like `E2V` or `V2E` and such. - The mesh (the set of vertices, edges and cells) has a regular (icosahedral in particular case) structure and it is possible to map it to a carthesizan grid. In genaral it would be different numbers of vertices, edges and cells per carthesizian point. In icosahedral case one "point" contains one vertex, three edges and two cells. This dictates us to introduce an extra dimension that is called `color`. The data on vertices will have one color, on edges three colors and on cells two. We would like to design a transformation that takes the problem that is formulated in terms of unstructured connectivities and returns the same semanticaly problem that is defined in terms of colored cartezian offsets. After the problem is transformed we are going to solve it with unmodified cartezian backend. ## Naive approach fails OK, let us try to keep it simple. Connectivites and cathesian offsets in the iterator model are just arguments of the `shift` builtin. We are going to just substitute `shift` argument sequences that contain connectivites. For example: `shift(C2E, 0, K, 3)` transforms to smth. like `shift(I, 1, C, -1, K, 3)`. This does not work because of colors -- it turns out that transformation itself depends on the color in the point where the iterator stays. I.e. for the previous example `shift(C2E, 0, K, 3)(it)` tranforms to `shift(I, 1, C, -1, K, 3)(i)` if `it` is at color `1` and say `shift(I, -1, C, 1, K, 3)(i)` if `it` is at color `0`. This problem reflects the fact that the `color` is not realy a dimension. Even though the original problem was a stencil (in the sense that in each point of the domain exact the same computation was performed), the transformed computation will not be the same for different colors. ## Treat each color as a separate field Say we have original stencil `foo`. We apply it on the domain with 2 colors. it takes two iterators that correspond to the fields with 3 and 2 colors: ```python foo(a : It3, b : It2) -> V2 ``` We can transform it to a couple of stencils. One per each domain color. Each function takes all input colors: ```python foo_0(a_0 : It, a_1 : It, a_2 : It, b_0 : It, b_1 : It) -> V foo_1(a_0 : It, a_1 : It, a_2 : It, b_0 : It, b_1 : It) -> V ``` or we can express it directly as a tuple of functions: ```python foo = (foo_0, foo_1) ``` Because `foo_0` and `foo_1` are stencils, `lift(foo)` is transformed `(lift(foo_0), lift(foo_1))` or `tuple(lift(f) for f in foo)`. In this approach function composition looks a bit ugly: `bar(lift(foo)(a, b), c)` is transformed to smth. like: `bar_0(lift(foo_0)(a_0, a_1, a_2, b_0, b_1), lift(foo_1)(a_0, a_1, a_2, b_0, b_1), c_0, c_1)` ## Steps - Deduce the number of colors for each iterator and value in all function parameters. - Deduce the number of colors for all returns for all functions. - Expand all `reduce` calls. I.e. lower them to `shift`'s. - Decompose all `shift`'s into argument pairs. Example (`shift(E2V, 3, V2E, 2, K, -1)(it)` -> `shift(E2V, 3)(shift(V2E, 2)(shift(K, -1)(it)))`). - Each function is transformed to a tuple of functions. One per color within return value. - Function paramters of the user defined functions are transformed to a concatenated sequence of colorized parameters (`f(a, b)` -> `f(a0, ... aN, b0, ... bM)`). - Each expression (function or builtin call, value reference) should be colorized. Colorization goes upside down: - We start with the body of some variant of colorized function (colorized function is a tuple of functions). It has a color that is an index within the tuple. - If this is a user function call we are choosing the function ref from the tuple and processing the arguments colorization according to the colors of parameters. - Else if it is a builtin that is not a connectivity `shift` and the color is propagated to the arguments. - Function and variable refs colorization is obvious -- just choose the item from the tuple. - Connectivity `shift` is special. It is driven by compile time lookup tables. Say we have `shift(Conn, Off)(<it_expr>)` and the color of that (outer) expression is `OuterColor`. We transform it to `shift(I, IOffset, J, JOffset)(<transformed_it_expr>)`, where `transformed_it_expr` has color `InnerColor`. `IOffset`, `JOffset` and `InnerColor` are taken from the lookup table. The index for lookup is `(Conn, Off, OuterColor)`.