Hybrid indexing (absolute index in K)

[FlorianDeconinck]

Hybrid indexing: I/J remains relative and K is absolute

The pattern to solves are as follows.

do L=1,LM
  do J=1,JM
   do I=1,IM
      ...
      Field(i,j,LM)
      Field(i,j,0)
      ...

do L=1,LM
  do J=1,JM
   do I=1,IM
      ...
      computed_k_index = FindK(...)
      Field(i,j,computed_k_index)
      ...

Current syntax: #64 (comment)

---

Working branch: https://github.com/FlorianDeconinck/gt4py/tree/cartesian/feature/absolute_k_indexation
ToDo:

• Move syntax from field[K_at[...]] to field.at(K=...)
• gt:X & dace:gpu backends
• Variable access
• Hard check on frontend for I/J on center
• Hard check on frontend for read only
• Deal with mitigation (K dynamic indexing mitigation: check for out of bounds #65)

[FlorianDeconinck]

Solution proposal:

• Relative index

  • Summary: default indexation, read in all scheduling, no I/J write, K write in Forward/Backward
  • Legacy subscript: field[0, 0, 1] or field[0, 0, var]
  • Axis subscript: field[K+1] or field[K+var]

• Absolute index

  • Summary: forbidden in I/J (at first), read in K depending on interval
  • Axis subscript: field[K[-1]] (last K of interval) or field[K[var]] (var value for index)

  Question:

  • do we expose K as a readable variable? Might be a different ticket.

[FlorianDeconinck]

Using explicit function.

def stencil(KLCL: IntField):
    with computation(PARALLEL), interval(...):
        from __externals__ import KM
        
        #v = field[0, 0, K_at(KLCL)]  # INVALID
        v = field[K_at(KLCL)]        # Axis subscript
        v = field[K_at(KM)]          # Top layer, with KM an external given by NDSL

[FlorianDeconinck]

Dropping support for field[0, 0, K_at(KLCL)] that ways it's made extra clear that I/J is can't be used with offset when doing absolute K indexation.

[romanc]

I was thinking a bit about how this should look like and came up with one other option:

def stencil(KLCL: IntField):
    with computation(PARALLEL), interval(...):
        from __externals__ import KM
        
        v = field.at(K=KLCL)   # where KLCL is an integer field
        v = field.at(K=0)      # Bottom layer
        v = field.at(K=KM)     # Top layer, with KM an external given by NDSL

This is inspired by other languages' .at() function on arrays, e.g. C++, Js/Ts.

• ➖ Probably more effort needed in parsing (on the NDSL-side)
• ➕ Makes it abundantly clear that you can't write
• Still no way to use I / J with offset when doing absolute K indexing

[FlorianDeconinck]

I like it. Enrique from CSCS also proposed something like field.K.at[] or field.K[] feels like a similar approach, but I'd rather do away with [] and have a function.

[FlorianDeconinck]

Working branch with the field[K_at[2]] frontend for backend numpy and dace:cpu

[FlorianDeconinck]

Part 1 PR: FlorianDeconinck/gt4py#3, PR'ed into unstable/develop

[FlorianDeconinck]

Logged the mitigation task: #65

[FlorianDeconinck]

Syntax is now field.at(K=...)

[FlorianDeconinck]

PR up for NDSL to expose the bounds of K: NOAA-GFDL#78

[FlorianDeconinck]

Further debug on the dace:cpu side.

• dcir.IndexAccess whenever absolute K indexing comes into play
• parsing of the ast.Call when in a sub-ast component (e.g.OP)

Evap still not generating properly