# Formal Semantics for the Parallel Model

###### tags: `Felix’s remains`

**IMPORTANT NOTE: This document is about semantics only. It does not try to define any syntax. All code examples are for demonstration purposes only and not for showing real DSL syntax.**

## Goals

A clear understanding of the semantics of the current parallel model and its limitations and undefined behavior. Finally, a new well-defined parallel model with clear semantics, well-defined behavior and good composability should be derived.

## Definitions

* **Accessor**: a function taking offsets and returning a value.
* **Stencil**: a function taking zero or more accessor arguments and returning a value.
* **Grid**: set of points.
* **Compute domain**: a subset of a grid.
* **Fencil**: one or more pairs of a stencil closures (stencils with bound field arguments and return values) with associated compute domains.
* **Lifting** (of a stencil): transformation of a stencil function to a function that returns an accessor instead of a value.

## Observations on the Current Parallel Model

### Limitations

* Intended behavior of the backends is not clearly defined, often different backends behave differently.
* Horizontal regions are required, but very hard to incorporate into current model (with extended compute domains etc.).
* Composability is very limited.
* Many advanced features like automatic halo exchanges and control flow require a better understanding and definition of the underlying behavior.

### Behavior

* The current GTScript “stencils” according to the definitions above are actually fencils
* Temporary assignment behaves as lifting. Note the type annotations:

@gtscript.stencil(...)
def d2dx2(inp, out):
    with computation(PARALLEL), interval(...):
        # we define actually a new accessor function here!
        tmp: accessor = inp[I + 1] - inp[I - 1]
        out: float = tmp[I + 1] - tmp[I - 1]

* Alternatively, the assignments can be interpreted as assignments to temporary arrays (this requires some notion of implicit loop for each statement). Note that this is one possible implementation of lifting. The alternative implementation is inlining/on-the-fly computation.
* The current GTScript functions can be used to implement stencils (in the sense of the definition above).
* GT4Py inherited some features from GridTools that increase difficulty of optimization (input, output parameters, multiple definitions of same symbol/multiple writes to same field).

## Considerations for the New Parallel Model

* Arguments should be input only or output only.
* Symbols should be defined exactly once.
* Lifting allows for great composability of stencils:

@stencil
def ddx(y: accessor) -> value:
    return y[I + 1] - y[I - 1]
    
# explicit lifting
@stencil
def d2dx2(y: accessor) -> value:
    return ddx(lift(ddx)(y))
    
# implicit lifting by temporary assignment
@stencil
def d2dx2(y: accessor) -> value:
    tmp = ddx(y)
    return ddx(tmp)

* A typical example of applying boundary conditions to a field in numpy syntax:

    # Approximation of first derivative in inner domain
    dydx[1:-1] = (y[1:] - y[:-1]) / (2 * dx)
    # Apply different approximation or BC at the boundaries
    dydx[0] = (y[1] - y[0]) / dx
    dydx[-1] = (y[-1] - y[-2]) / dx
    # Continue using dydx in next computation
    ... = dydx[...]

* This can be handled directly within stencils:

@stencil
def ddx(y: accessor,
        on_left_boundary: accessor,
        on_right_boundary: accessor) -> value:
    if on_left_boundary:
        return y[I + 1] - y
    if on_right_boundary:
        return y - y[I - 1]
    return (y[I + 1] - y[I - 1]) / 2

* The composed can then be applied to an array easily using a fencil:

@fencil
def fency(y: Input, dydx: Output, output_domain: ComputeDomain):
    dydx = ddx(y) @ output_domain

* Handling of vertical computations can be implemented in this framework by choosing columns as the granularity for all operations.

## Appendix: Examples

**Note again: this is not a syntax proposal. Some arbitrary syntax is chosen. The examples are just there to provide an idea of how slightly more complex examples can be implemented using this model.**

### Heat Equation with FTCS Scheme (1D)

This is a typical example where boundary conditions are written after each stencil application.

@stencil
def d2dx2(y, dx):
    return (y[I + 1] - 2 * y + y[I - 1]) / dx**2
    
@stencil
def step(y, dx, dt, diff):
    d2ydx2 = d2dx2(y)
    return y + dt * diff * d2ydx2
    
@stencil
def step_with_bcs(y, dx, dt, diff, on_left_domain, on_right_domain):
    yp = step(y, dx, dt, diff)
    if on_left_domain:
        # zero-gradient BC
        return yp[I + 1]
    if on_right_domain:
        # Dirichlet BC (setting this every step is obviouly not necessary,
        # but done here for demonstration purposes)
        return 1
    return yp
    
@fencil
def apply_step(y0: Input,
               on_left_boundary: Input,
               on_right_boundary: Input,
               y1: Output,
               dx: float,
               dt: float,
               diff: float,
               domain: ComputeDomain):
    y1 = step_with_bcs(y0,
                       dx,
                       dt,
                       diff,
                       on_left_boundary,
                       on_right_boundary)

### Tridiagonal Solver

Having stencils taking whole columns as inputs and returning whole columns, allows to handle vertical solvers as pure stencil functions. All state required for the vertical sweeps can be encapsulated within the stencil.

# note: all stencil arguments are read-only (whole columns)
# read-write fields are hidden within a single stencil
@stencil
def forward(a, b, c, d):
    cp = column()
    dp = column()
    for k in krange[:1]:
        cp[k] = c[k] / b[k]
        dp[k] = d[k] / b[k]
    for k in krange[1:]:
        # note: cp and dp are read-write!
        cp[k] = c[k] / (b[k] - a[k] * cp[k - 1])
        dp[k] = (d[k] - a[k] * dp[k - 1]) / (b[k] - a[k] * cp[k - 1])
    # the return values are whole columns again
    return cp, dp
        
@stencil
def backward(cp, dp):
    for k in reversed(krange[-1:]):
        x[k] = dp[k]
    for k in reversed(krange[:-1]):
        x[k] = dp[k] - cp[k] * x[k + 1]
    return x
    
@stencil
def solve_tridiag(a, b, c, d):
    # the vertical solver functions can be composed as usual
    cp, dp = forward(a, b, c, d)
    x = backward(cp, dp)
    return x 

### Testability of Components

Unit testing of single components is an important feature and can be easily performed using the proposed model. Here an example of how to test the heat equation solver from above.

def test_d2dx2():
    # trivial fencil for applying the stencil to be tested
    @fencil
    def apply_d2dx2(y: Input, d2ydx2: Output, dx: float, domain: ComputeDomain):
        d2ydx2 = d2dx2(y, dx)
        
    # input field: default_origin is set to avoid out-of bounds accesses
    # this is user responsibility!
    y = storage(..., default_origin=(1, 0))
    np_y = np.asarray(y)
    d2dx2 = storage(...)
    apply_d2dx2(y, d2ydx2, domain=...)
    assert np.allclose(np.asarray(d2ydx2), (np_y[1:] - np_y[:-1]) / dx)
    
def test_step():
    @fencil
    def apply_d2dx2(y0: Input, y1: Output, dx: float, dt: float, diff: float, domain: ComputeDomain):
        y1 = step(y0, dx, dt, diff) @ domain
    ...

### Staggering

@stencil
def flux(c, u):
    # u is a staggered field, this returns a staggered flux, that is,
    # corresponds to f_{i + 0.5}
    return c if u >= 0 else c[I + 1]
    # could also be defined as follows, then we would return f_{i - 0.5}
    # it is user responsibility to use a consistent convention
    # return c[I - 1] if u >= 0 else c
    
@stencil
def dfdx(c, u, dx):
    f = flux(c, u)
    # user responsibility to take correct offsets for f_{i + 0.5}, f_{i - 0.5}
    # (the return value is non-staggered)
    return (f - f[I - 1]) / dx
    
@stencil
def step(c, u, dt, dx):
    return c + dt * dfdx(c, u, dx)
    
@fencil
def step_with_flux_output(c0: Input, u: Input, c1: Output, f: Output,
                          dx: float, dt: float, domain: FullDomain, flux_domain: FluxDomain):
    # backends can decide how they handle the redundant flux computation,
    # e.g. compute the flux and put it into a temporary buffer,
    # compute the flux twice in two separate loops (for each output),
    # do only one loop with OTF-computation of the flux and masked writes…
    f = flux(c0, u) @ FluxDomain
    c1 = step(c, u, dt, dx) @ FullDomain

## Appendix: Notes

* Handling of tuple return values in lifting:

@stencil
def foo(x: accessor) -> value:
    y, z = bar(x)
    return baz(y, z)
    
@stencil
def foo(x: accessor) -> value:
    y = lift(bar, return_index=0)(x)
    z = lift(bar, return_index=1)(x)
    return baz(y, z)