Refactor of GT4Py-DaCe bridge to expose all control-flow to Dace

[FlorianDeconinck]

The current bridge between gt4py.cartesian and DaCe uses a concept of expansion for stencils. This means that stencils, once they got through the entire GT4Py workflow, are pickled and shoved down a library node. Then the expansion mechanism turns those into a single map/tasklet per Horizontal Computation. This means that a lot of control-flow (mask, while loop, etc.) is hidden within from DaCe, leading to some optimization power being hidden.

We need to undo this and describe everything back to DaCe.

Expansion code lives around src/gt4py/cartesian/gtc/dace/expansion/*.py in GT4Py

Parent: GEOS-ESM/SMT-Nebulae#31

NB: This is a requirement for DaCe AD future features to kick in

---

• Expose more code to DaCe by exposing the mask/while to the SDFG instead of keeping them in the Tasklet
• ALTERNATIVELY: subtask further if the refactor work is bigger

[FlorianDeconinck]

Below is a small example showing the above problem

from gt4py.cartesian.gtscript import (
    computation,
    interval,
    PARALLEL,
)
from ndsl.boilerplate import get_factories_single_tile_orchestrated_cpu
from ndsl.constants import X_DIM, Y_DIM, Z_DIM
from ndsl.dsl.typing import FloatField
from ndsl import StencilFactory, orchestrate
import numpy as np

domain = (3, 3, 4)

stcil_fctry, ijk_qty_fctry = get_factories_single_tile_orchestrated_cpu(domain[0], domain[1], domain[2], 0)

def small_conditional(
    in_field: FloatField,
    out_field: FloatField
):
    with computation(PARALLEL), interval(...):
        if in_field > 5.0 and in_field < 20:
            out_field = in_field


class DaCeGT4Py_Bridge:
    def __init__(self, stencil_factory: StencilFactory):
        orchestrate(
            obj=self,
            config=stencil_factory.config.dace_config)
        self.stencil = stcil_fctry.from_dims_halo(
            func=small_conditional,
            compute_dims=[X_DIM, Y_DIM, Z_DIM],
        )

    def __call__(self, in_field: FloatField, out_field: FloatField):
        self.stencil(in_field, out_field)



if __name__ == "__main__":
    I = np.arange(domain[0]*domain[1]*domain[2], dtype=np.float64).reshape(domain)
    O = np.zeros(domain)

    bridge = DaCeGT4Py_Bridge(stcil_fctry)
    bridge(I, O)

    print(f"Input:\n{I}\n")
    print(f"Output:\n{O}\n")

This generate the following SDFG (under dace:cpu)

where the tasklet code carries the conditional:

mask_140660478452496_gen_0: dace.bool_
mask_140660478452496_gen_0 = ((gtIN__in_field > dace.float64(5.0)) and (gtIN__in_field < dace.float64(dace.int64(20))))
if mask_140660478452496_gen_0:
    gtOUT__out_field = gtIN__in_field

We want this conditional to be extracted before the tasklet. See dace.sdfg.SDFG:add_loop code to see how to dynamically create an if condition using the guard mechanism

[FlorianDeconinck]

Known native construct that are folded inside the Tasklet instead of described to DaCe:

• conditional
• while loops

GT4Py-DaCe uses an "expansion" mechanism. The stencils are dace.library_nodes while the code is being parsed (orchestration). Then expansion of those nodes is triggered and unrolls the stencil into an SDFG compatible description. It's within this code that the decision to write all code has tasklet is finalized. Refactor might need to impact pre-expansion as well in the DaceIRBuilder

[romanc]

Latest work is in https://github.com/romanc/gt4py/romanc/bridge-on-top-of-cleanups, which is re-based on top of the cleanups in PR GridTools/gt4py#1724 and should only contain changes necessary for the new bridge and no cleanups anymore.

The new bridge is currently pending on the following two DaCe issues:

• Dead Dataflow Elimination type hints do not work with all types spcl/dace#1710
• Scalar to symbol promotion fails if scalar is read in another state spcl/dace#1727

Working around those two, we still (only sometimes) have an issue with writes (in k) inside a k-loop (i.e. with computation FORWARD / BACKWARD). To be investigated further at a later point.

[romanc]

Latest work is still in https://github.com/romanc/gt4py/romanc/bridge-on-top-of-cleanups, which can be rebased on top of mainline gt4py once GridTools/gt4py#1752 to get a clearer picture.

Main ideas in that branch are the following

1. Introduce oir.CodeBlock to (recursively) break down oir.HorizontalExecutions into smaller pieces, which are then evaluated in a dcir.Tasklet.
2. Replace dcir.MaskStmt / dcir.While with dcir.Condition and dcir.WhileLoop (names aren't perfect and could be renamed again)
3. Add support for if statements and while loops in the state machine of sdfg_builder, i.e. add_condition and add_while on SDFGContext together with visit_Condition and visit_WhileLoop in the visitor (for the new dcir nodes).
4. Tasklets might define temporaries (scalars and/or tmp fields) in one tasklet and re-use them in others. We thus opt to "export" all temporaries and then "re-import" them again if needed. A DaCe pass called "DeadDataflowElimination" will take care of removing unnecessary "exports".
5. Memlets will need to be generated per tasklet at a more fine grained level compared to now where memlets are calculated just once per horizontal execution. Commit GridTools/gt4py@19f59a3 outlines the necessary changes. One thing I forgot in this PR is the fact that both, Conditions and WhileLoops generate tasklets on the fly in sdfg_builder. They will also need to calculate their memlets at this point.

[romanc]

Main chunk of this work was merged with PR GridTools/gt4py#1894. The PR pull this version of GT4Py into NDSL is NOAA-GFDL#119, which is expected to be merged in the coming days.

Follow-up issues include

• cartesian: re-evaluate InlineThreadLocalTransients in gt4py/dace bridge GridTools/gt4py#1896
• cartesian: refactor oportunities after dace bridge refactor GridTools/gt4py#1898
• cartesian: expose HorizontalRegions to DaCe GridTools/gt4py#1900

Given that, let's close this issue and treat potential fall-out (nothing big is expected since pace and pyFV3 test suites both pass) in follow-up issues.