---
# System prepended metadata

title: Iterator IR improvements
tags: [IR, scratchpad]

---

# 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)