# Calling stencils from stencils

###### tags: `archive`

We are used to this

BACKEND = "gtcuda"
dtype = np.float_

@gtscript.function
def laplacian(phi, alpha):
    return -alpha * phi[0, 0, 0] + 
          phi[-1, 0, 0] + phi[1, 0, 0] + phi[0, -1, 0] + phi[0, 1, 0]

@gtscript.stencil(backend=BACKEND, externals={"LIM": 0.01})
def diffusion_defs(in_phi: Field[dtype], bi_lap: Field[dtype],
                   *, alpha: float):

    from __externals__ import LIM

    with computation(PARALLEL), interval(...):
        # compute the laplacian-of-laplacian
        lap = laplacian(in_phi, alpha)
        bilap = laplacian(lap, alpha)

But a simple laplacian stencil would have to be rewritten as `stencil` instead of `function` causing code replication. Moreover, calling stencils can enable control flow.

There are two different mechanisms, with very different semantics.

### First Option

Calling a stencil outside of a `with computation` region.

BACKEND = "gtcuda"
dtype = np.float_

@gtscript.function
def lapf(phi, alpha):
    return -alpha * phi[0, 0, 0] + 
          phi[-1, 0, 0] + phi[1, 0, 0] + phi[0, -1, 0] + phi[0, 1, 0]

xyz = {"XYZ": 0.01}

@gtscript.stencil(backend=BACKEND, externals=xyz)
def laplacian(phi: Field[dtype], lap: Field[dtype], *, alpha: float):
    with computation(PARALLEL), interval(...):
        lap = -alpha * phi[0, 0, 0] + 
            phi[-1, 0, 0] + phi[1, 0, 0] + phi[0, -1, 0] + phi[0, 1, 0]


xyz = {"XYZ": 0.03}

@gtscript.stencil(backend=BACKEND, externals={"LIM": 0.01}, domain=[...])
def diffusion_defs(in_phi: Field[dtype], bi_lap: Field[dtype],
                   *, alpha: float):

    lap : Field[dtype] = create_field() # type hint is optional
    laplacian[this.domain, this.origin, {**xyz}](in_phi, lap, alpha)
    
    with computation(PARALLEL), interval(...):
        bilap = lapf(lap, alpha)

Observation: `backend` in the first stencil must be the same as the second, or lternatively the `backend` of the callee is ignored and the one from the caller is picked up.

The call to `laplacian` is in fact an **inlining**. To be able to do that we need to define the lap field, and it's dimensions have to be inferred from the implicit grid being passed. The exact syntax for defining the field can be taken from one of the syntax workshop proposals, but it not essential to be finalized, as far as the functionality works. **Alternatively, the grid, as the mesh in the unstructured case, could be explicitly passed to the stencil all the times.** `lap` is a three dimensional field in this case.

In this case the computation is well defined for any type of computation (parallel, forward, backward).

### Second Option

Another option is to treat the laplacian stencil as a regular function, by ignoring the annotations and by matching the computation to the vertical interval the caller is in when calling the stencil.

BACKEND = "gtcuda"
dtype = np.float_

@gtscript.function
def lapf(phi, alpha):
    return -alpha * phi[0, 0, 0] + 
          phi[-1, 0, 0] + phi[1, 0, 0] + phi[0, -1, 0] + phi[0, 1, 0]

@gtscript.stencil(backend=BACKEND)
def laplacian(phi, lap, alpha):
    with computation(PARALLEL), interval(...):
        lap = -alpha * phi[0, 0, 0] + 
            phi[-1, 0, 0] + phi[1, 0, 0] + phi[0, -1, 0] + phi[0, 1, 0]

@gtscript.stencil(backend=BACKEND, externals={"LIM": 0.01})
def diffusion_defs(in_phi: Field[dtype], bi_lap: Field[dtype],
                   *, alpha: float):

    with computation(PARALLEL), interval(...):
        lap : Field[dtype]
        laplacian(in_phi, lap)
        bilap = lapf(lap, alpha)

Now the `lap` field is a temporary and could be a 2D horizontal field. The callee is executed only in the vertical level where the caller is evaluated. The `with interval` statement in the callee is treated as a kind of switch-case statement, in which the brach is selected based on where the input values fall. The `with computation` have problems if there are vertical dependencies, especially if the caller has no dependencies. The `with computation` statement is basically ignored here, only the interval would have an effect.

### Additional ideas for future consideration

To avoid code replication we could think of *extracting* computations as functions from a stencil, like here, where `lapf` symbol will be available as a `@gtscript.function` with the provided signature:

@gtscript.stencil(backend=BACKEND)
def laplacian(phi, lap, alpha):
    with computation(PARALLEL, name=lapf(phi, alpha)->lap), interval(...):
        lap = -alpha * phi[0, 0, 0] + 
            phi[-1, 0, 0] + phi[1, 0, 0] + phi[0, -1, 0] + phi[0, 1, 0]

- Enhance `gtscript.stencil` lazy behavior:

`gtscript.stencil` could be used without arguments or some arguments but without a backend (no compilation, only syntax check as `gtscript.function`)

@gtscript.stencil
def diffusion_defs(in_phi: Field[dtype], bi_lap: Field[dtype],
                   *, alpha: float):

    from __externals__ import LIM

    with computation(PARALLEL), interval(...):
        # compute the laplacian-of-laplacian
        lap = laplacian(in_phi, alpha)
        bilap = laplacian(lap, alpha)

This can be later compiled as:

diffustion_stencils = diffusion_defs.compile(backend=..., externals=....)

- Enhance `gtscript.stencil` decorator accepting function definitions and a signature

@gtscript.function
def laplacian(phi, alpha):
    return -alpha * phi[0, 0, 0] + 
          phi[-1, 0, 0] + phi[1, 0, 0] + phi[0, -1, 0] + phi[0, 1, 0]

lap_stencil = gtscript.stencil(
    laplacian,
    signature="(phi: Field['dtype'], output<result>: Field['dtype'], *, alpha: float)",
    externals=...,)

## Goal

Implement a version of the first option (the second is independent and can be provided in a second time, as we can see a value for that use case, too).

The use cases to be tested should be fairly standard, for instance an HD or VA with the basic stencils to be composed as sequence of calls.

Provide the extent analysis for the fields defined outside of the `with computation` regions.

## No Goals

- Definitive syntax fo field definitions
- Use cases that falls outside the test cases of HD/VA. Other use cases discovered during the development shall be documented and discussed in order to determine if they can be supported or should be prohibited.

## Appetite

Cannot really tell. Feels to me like 6 weeks.

## Possible Solution

- Implement a caller-callee pair with not additional fields required (the fields are taken form the input parameters of the stencil)
- Front-end parser
- Decide if GTIR needs to be aware of the calling semantics or if the inlining should happen before the IR is produced
- In the latter case, the code should already work, otherwise there's need of intervention in GTC to lower GTIR into an expanded version
- Implement a field to be defined outside of the `with computation`
- This requires the front end to introduce the field as a node in GTIR
- Implement the extent analysis to determine its dimensions
- Proceed down to the code generation
- Implement a HD and or VA using this feature and vaidate the results