# Iterator IR improvements ## Overview - Tuple of fields on different locations ```python if some_condition: field_a: Field[Vertex, ...] = ... field_b: Field[Edge, ...] = ... else: field_a: Field[Vertex, ...] = ... field_b: Field[Edge, ...] = ... ``` Lowering options: Possible: ``` make_tuple(if_(...)(...)(...), if_(...)(...)(...), ) ``` Not possible (as tuple of it not allowed): ``` if_(...)(make_tuple(it_a_tb, it_b_tb))(make_tuple(it_a_tb, it_b_tb)) ``` - Multiple returns on different location types _Very similar to tuple problem above._ ```python @field_operator def field_op(inp_a: Field[[Vertex], ...], inp_b: Field[[Edge], ...]): return inp_a, inp_b ``` Felix: Computation is representable on ITIR so not needed before new version of IR. - It maybe doesn't need to be represented on ITIR but in lower level performance IRs (e.g. gtfn IR?) - Another option: moving both iterators to a fake location (`V_or_E`) with the bigger domain, and shifting both iterators to it with a fake "identity" connectivity and a mask/if - Masked output ```python @field_operator def field_op(inp: Field[[Vertex], ...]): return where(mask, inp+1, inp) ``` - Maybe the compiler can figure it out if there are not aliases for the variable. Check if it is possible to tell the compiler that there are not aliases or other similar hints. - Reduce/partial shifts Option 1: `reduce` takes a list of values and returns a value. `get_neighbor_values(...)` returns a (variable length) list. Option 2: `reduce` takes an iterator. `neighbor_shift` returns an iterator of lists. - Lowering of nested reduction (offsets in wrong order) - Broadcast ```python @field_operator def field_op(...): return broadcast(neighbor_sum(field, V2E), (V, V2E)) ``` Proposed solutions: ```python def stencil_like(*shifted_its) -> Value: return make_tuple(deref(shifted_its)...) def neighbor_sum(iterators): return reduce(sum)(lift(stencil_like)(...)) ``` Open question: - What happens if deref invalid (optional values) - Can we get size of tuple? Is this a list (just for reductions) instead of a tuple? - How to distinguish between user asking for tuple argument in reductor. - Direct access to sparse field neighbor. - Does this solve the broadcast problem? ```python @field_operator def nb_sum_broadcast(field: Field[[C, C2V]]) -> Field[[C, C2V]]: return broadcast(neighbor_sum(field, C2V), (C, C2V)) @field_operator def some_op(field: Field[[C, C2V], ...]): -> Field[[C, C2V]] return neighbor_sum(field, C2V) @field_operator def user_of_broadcasted(field): return some_op(field_op(field)) ``` ```python def nb_sum_broadcast(it): # it return a tuple of values on neighboring vertices return reduce(plus)(it) def some_op(nb_sum): def foo_stencil(*its): val = user_of_broadcasted(lift(neighbor_sum)(*its)) return deref(val) ``` - _ ```python @field_operator def number_of_neighbors(): return neighbor_sum(broadcast(1, (C, C2V)), C2V) # Possible solution ?? @field_operator def number_of_neighbors(): auto_f = magic_ITIR_builtin_with_index_in_domain return auto_f_as_ones * neighbor_sum(broadcast(1, (C, C2V)), C2V) ``` - Move static offset information to IR - (learning type inference pass) ### ITIR Type Inference Type inferrence on ITIR [Core](https://github.com/GridTools/gt4py/blob/functional/src/functional/iterator/type_inference.py)